Support whitelist now

README.en.md

@@ -32,7 +32,7 @@ _✨ Easily forward your local port to the public network. ✨_
 + [x] Very easy to use
 + [x] Support TCP
 + [ ] Support UDP
-+ [ ] Support IP whitelist
++ [x] Support IP whitelist
 
 ## Usages
 

README.md

@@ -32,7 +32,7 @@ _✨ 基于 Go 的本地端口转发工具，开箱即用 ✨_
 + [x] 开箱即用
 + [x] 支持 TCP
 + [ ] 支持 UDP
-+ [ ] 支持 IP 白名单
++ [x] 支持 IP 白名单
 
 ## 用法
 

common/config.go

@@ -9,7 +9,7 @@ import (
 type serverConfig struct {
 	Port      int      `yaml:"port"`
 	Token     string   `yaml:"token"`
-	WhiteList []string `yaml:"white_list"`
+	Whitelist []string `yaml:"whitelist"`
 }
 
 type clientConfig struct {

handler/server.go

@@ -26,6 +26,11 @@ func ServeForever() {
 			println(err.Error())
 			continue
 		}
+		ip := conn.RemoteAddr().(*net.TCPAddr).IP.String()
+		if !isInWhitelist(ip) {
+			conn.Close()
+			continue
+		}
 		go handleClientConnection(conn)
 	}
 }

handler/whitelist.go

@@ -0,0 +1,52 @@
+package handler
+
+import (
+	"go-public/common"
+	"sync"
+)
+
+var ipStore = struct {
+	sync.RWMutex
+	ips map[string]bool
+}{
+	ips: make(map[string]bool),
+}
+
+var whitelistCacheThreshold = 4
+
+func isInWhitelist(ip string) bool {
+	// If whitelist is empty, just return true
+	if len(common.ServerConfig.Whitelist) == 0 {
+		return true
+	}
+	// If ip is localhost, just return true
+	if ip == "::1" {
+		return true
+	}
+	// If whitelist is too short, we don't cache it
+	if len(common.ServerConfig.Whitelist) < whitelistCacheThreshold {
+		for _, v := range common.ServerConfig.Whitelist {
+			if v == ip {
+				return true
+			}
+		}
+		return false
+	}
+	// Check the cache
+	ipStore.RLock()
+	if _, ok := ipStore.ips[ip]; ok {
+		ipStore.RUnlock()
+		return true
+	}
+	ipStore.RUnlock()
+	// Check the whitelist
+	for _, v := range common.ServerConfig.Whitelist {
+		if v == ip {
+			ipStore.Lock()
+			ipStore.ips[ip] = true
+			ipStore.Unlock()
+			return true
+		}
+	}
+	return false
+}