From ca40be560091dd069210eb49ffb7d4e43df9d233 Mon Sep 17 00:00:00 2001
From: Peter Huang <peter.huang@practiceevolve.com.au>
Date: Sat, 24 Jul 2021 21:22:03 +1000
Subject: [PATCH] Added Regex destination matching

---
 README.md  |  9 ++++++++-
 ruleset.go | 24 ++++++++++++++++++++++++
 server.go  | 17 +++++++++++------
 3 files changed, 43 insertions(+), 7 deletions(-)
 create mode 100644 ruleset.go

diff --git a/README.md b/README.md
index 093110b..bcc009f 100644
--- a/README.md
+++ b/README.md
@@ -7,7 +7,13 @@ Simple socks5 server using go-socks5 with authentication options
 
 ## Start container with proxy
 
-```docker run -d --name socks5 -p 1080:1080 -e PROXY_USER=<PROXY_USER> -e PROXY_PASSWORD=<PROXY_PASSWORD>  serjs/go-socks5-proxy```
+```
+docker run -d --name socks5 -p 1080:1080 \
+    -e PROXY_USER=<PROXY_USER> \
+    -e PROXY_PASSWORD=<PROXY_PASSWORD> \
+    -e ALLOWED_DEST_FQDN=<REGEX_PATTERN> \
+    serjs/go-socks5-proxy
+```
 
 Leave `PROXY_USER` and `PROXY_PASSWORD` empty for skip authentication options while running socks5 server.
 
@@ -18,6 +24,7 @@ Leave `PROXY_USER` and `PROXY_PASSWORD` empty for skip authentication options wh
 |PROXY_USER|String|EMPTY|Set proxy user (also required existed PROXY_PASS)|
 |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|
+|ALLOWED_DEST_FQDN|String|EMPTY|Allowed destination address regular expression pattern. Default allows all.|
 |TZ|String|UTC|Set Timezone like in many common Operation Systems|
 
 ## Test running service
diff --git a/ruleset.go b/ruleset.go
new file mode 100644
index 0000000..4b9f6d6
--- /dev/null
+++ b/ruleset.go
@@ -0,0 +1,24 @@
+package main
+
+import (
+	"regexp"
+
+	"github.com/armon/go-socks5"
+	"golang.org/x/net/context"
+)
+
+// PermitDestAddrPattern returns a RuleSet which selectively allows addresses
+func PermitDestAddrPattern(pattern string) socks5.RuleSet {
+	return &PermitDestAddrPatternRuleSet{pattern}
+}
+
+// PermitDestAddrPatternRuleSet is an implementation of the RuleSet which
+// enables filtering supported destination address
+type PermitDestAddrPatternRuleSet struct {
+	AllowedFqdnPattern string
+}
+
+func (p *PermitDestAddrPatternRuleSet) Allow(ctx context.Context, req *socks5.Request) (context.Context, bool) {
+	match, _ := regexp.MatchString(p.AllowedFqdnPattern, req.DestAddr.FQDN)
+	return ctx, match
+}
diff --git a/server.go b/server.go
index 21adf81..efdd533 100644
--- a/server.go
+++ b/server.go
@@ -9,9 +9,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"`
+	AllowedDestFqdn string `env:"ALLOWED_DEST_FQDN" envDefault:""`
 }
 
 func main() {
@@ -23,7 +24,7 @@ func main() {
 	}
 
 	//Initialize socks5 config
-	socsk5conf := &socks5.Config{
+	socks5conf := &socks5.Config{
 		Logger: log.New(os.Stdout, "", log.LstdFlags),
 	}
 
@@ -32,10 +33,14 @@ func main() {
 			os.Getenv("PROXY_USER"): os.Getenv("PROXY_PASSWORD"),
 		}
 		cator := socks5.UserPassAuthenticator{Credentials: creds}
-		socsk5conf.AuthMethods = []socks5.Authenticator{cator}
+		socks5conf.AuthMethods = []socks5.Authenticator{cator}
 	}
 
-	server, err := socks5.New(socsk5conf)
+	if cfg.AllowedDestFqdn != "" {
+		socks5conf.Rules = PermitDestAddrPattern(cfg.AllowedDestFqdn)
+	}
+
+	server, err := socks5.New(socks5conf)
 	if err != nil {
 		log.Fatal(err)
 	}