Added allowed IPs support

server.go

@@ -2,6 +2,7 @@ package main
 
 import (
 	"log"
+	"net"
 	"os"
 
 	"github.com/armon/go-socks5"
@@ -9,9 +10,10 @@ import (
 )
 
 type params struct {
-	User     string `env:"PROXY_USER" envDefault:""`
-	Password string `env:"PROXY_PASSWORD" envDefault:""`
-	Port     string `env:"PROXY_PORT" envDefault:"1080"`
+	User       string   `env:"PROXY_USER" envDefault:""`
+	Password   string   `env:"PROXY_PASSWORD" envDefault:""`
+	Port       string   `env:"PROXY_PORT" envDefault:"1080"`
+	AllowedIPs []string `env:"ALLOWED_IPS" envSeparator:"," envDefault:""`
 }
 
 func main() {
@@ -40,6 +42,15 @@ func main() {
 		log.Fatal(err)
 	}
 
+	// Set IP whitelist
+	if len(cfg.AllowedIPs) > 0 {
+		whitelist := make([]net.IP, len(cfg.AllowedIPs))
+		for i, ip := range cfg.AllowedIPs {
+			whitelist[i] = net.ParseIP(ip)
+		}
+		server.SetIPWhitelist(whitelist)
+	}
+
 	log.Printf("Start listening proxy service on port %s\n", cfg.Port)
 	if err := server.ListenAndServe("tcp", ":"+cfg.Port); err != nil {
 		log.Fatal(err)