From 70d9b0bbce8e9564df2b5d59ecb60304911175ad Mon Sep 17 00:00:00 2001
From: Anton <sazan95@gmail.com>
Date: Mon, 6 Mar 2023 00:48:21 +0300
Subject: [PATCH] Added allowed IPs support

---
 README.md                                   |  6 ++++
 docker-compose.build.yml                    | 11 +++++++
 server.go                                   | 17 +++++++++--
 vendor/github.com/armon/go-socks5/socks5.go | 32 +++++++++++++++++++++
 4 files changed, 63 insertions(+), 3 deletions(-)
 create mode 100644 docker-compose.build.yml

diff --git a/README.md b/README.md
index 093110b..5008a13 100644
--- a/README.md
+++ b/README.md
@@ -19,6 +19,12 @@ Leave `PROXY_USER` and `PROXY_PASSWORD` empty for skip authentication options wh
 |PROXY_PASSWORD|String|EMPTY|Set proxy password for auth, used with PROXY_USER|
 |PROXY_PORT|String|1080|Set listen port for application inside docker container|
 |TZ|String|UTC|Set Timezone like in many common Operation Systems|
+|ALLOWED_IPS|String|Empty|Set allowed IP's that can connect to proxy, separator `,`|
+
+`ALLOWED_IPS` parameter is not included in `serjs/go-socks5-proxy` image.\
+You can build your image with:
+`docker-compose -f docker-compose.build.yml up -d`\
+Just don't forget to set parameters in the `.env` file.
 
 ## Test running service
 
diff --git a/docker-compose.build.yml b/docker-compose.build.yml
new file mode 100644
index 0000000..6e87be8
--- /dev/null
+++ b/docker-compose.build.yml
@@ -0,0 +1,11 @@
+version: '3'
+services:
+  socks5-build:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    container_name: socks5-build
+    env_file: .env
+    ports:
+      - "1080:1080"
+    restart: unless-stopped
diff --git a/server.go b/server.go
index 2a6e47f..41d28ae 100644
--- a/server.go
+++ b/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)
diff --git a/vendor/github.com/armon/go-socks5/socks5.go b/vendor/github.com/armon/go-socks5/socks5.go
index a17be68..2d630fb 100644
--- a/vendor/github.com/armon/go-socks5/socks5.go
+++ b/vendor/github.com/armon/go-socks5/socks5.go
@@ -55,6 +55,7 @@ type Config struct {
 type Server struct {
 	config      *Config
 	authMethods map[uint8]Authenticator
+	isIPAllowed func(net.IP) bool
 }
 
 // New creates a new Server and potentially returns an error
@@ -93,6 +94,11 @@ func New(conf *Config) (*Server, error) {
 		server.authMethods[a.GetCode()] = a
 	}
 
+	// Set default IP whitelist function
+	server.isIPAllowed = func(ip net.IP) bool {
+		return true // default allow all IPs
+	}
+
 	return server, nil
 }
 
@@ -117,11 +123,37 @@ func (s *Server) Serve(l net.Listener) error {
 	return nil
 }
 
+// SetIPWhitelist sets the function to check if a given IP is allowed
+func (s *Server) SetIPWhitelist(allowedIPs []net.IP) {
+	s.isIPAllowed = func(ip net.IP) bool {
+		for _, allowedIP := range allowedIPs {
+			if ip.Equal(allowedIP) {
+				return true
+			}
+		}
+		return false
+	}
+}
+
 // ServeConn is used to serve a single connection.
 func (s *Server) ServeConn(conn net.Conn) error {
 	defer conn.Close()
 	bufConn := bufio.NewReader(conn)
 
+	// Check client IP against whitelist
+	clientIP, _, err := net.SplitHostPort(conn.RemoteAddr().String())
+	if err != nil {
+		s.config.Logger.Printf("[ERR] socks: Failed to get client IP address: %v", err)
+		return err
+	}
+	ip := net.ParseIP(clientIP)
+	if s.isIPAllowed(ip) {
+		s.config.Logger.Printf("[INFO] socks: Connection from allowed IP address: %s", clientIP)
+	} else {
+		s.config.Logger.Printf("[WARN] socks: Connection from not allowed IP address: %s", clientIP)
+		return fmt.Errorf("connection from not allowed IP address")
+	}
+
 	// Read the version byte
 	version := []byte{0}
 	if _, err := bufConn.Read(version); err != nil {