From 23d4321be48af79e709af156471267d3e96ee80c Mon Sep 17 00:00:00 2001
From: Artem Andreenko <miolini@users.noreply.github.com>
Date: Wed, 22 May 2024 15:41:40 -0700
Subject: [PATCH] Implemented UNIX socket support for SOCKS5 server in yggstack
 command and updated README.md with usage instructions.

---
 README.md            |  6 +++++
 cmd/yggstack/main.go | 62 ++++++++++++++++++++++++++++++++++----------
 2 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/README.md b/README.md
index 5965bec..e52d91b 100644
--- a/README.md
+++ b/README.md
@@ -78,6 +78,12 @@ To run SOCKS proxy server listening on local port 1080 using generated configura
 ./yggstack -useconffile /path/to/yggdrasil.conf -socks 127.0.0.1:1080
 ```
 
+To run SOCKS proxy server listening on UNIX socket file `/tmp/yggstack.sock`:
+
+```
+./yggstack -useconffile /path/to/yggdrasil.conf -socks /tmp/yggstack.sock
+```
+
 To expose network services (like a Web server) listening on local port 8080 to Yggdrasil
 network address at port 80:
 
diff --git a/cmd/yggstack/main.go b/cmd/yggstack/main.go
index 41e029a..aff7ae0 100644
--- a/cmd/yggstack/main.go
+++ b/cmd/yggstack/main.go
@@ -5,6 +5,7 @@ import (
 	"crypto/ed25519"
 	"encoding/hex"
 	"encoding/json"
+	"errors"
 	"flag"
 	"fmt"
 	"net"
@@ -31,9 +32,10 @@ import (
 )
 
 type node struct {
-	core      *core.Core
-	multicast *multicast.Multicast
-	admin     *admin.AdminSocket
+	core           *core.Core
+	multicast      *multicast.Multicast
+	admin          *admin.AdminSocket
+	socks5Listener net.Listener
 }
 
 // The main function is responsible for configuring and starting Yggdrasil.
@@ -52,7 +54,7 @@ func main() {
 	getsnet := flag.Bool("subnet", false, "use in combination with either -useconf or -useconffile, outputs your IPv6 subnet")
 	getpkey := flag.Bool("publickey", false, "use in combination with either -useconf or -useconffile, outputs your public key")
 	loglevel := flag.String("loglevel", "info", "loglevel to enable")
-	socks := flag.String("socks", "", "address to listen on for SOCKS, i.e. :1080")
+	socks := flag.String("socks", "", "address to listen on for SOCKS, i.e. :1080; or UNIX socket file path, i.e. /tmp/yggstack.sock")
 	nameserver := flag.String("nameserver", "", "the Yggdrasil IPv6 address to use as a DNS server for SOCKS")
 	flag.Var(&expose, "exposetcp", "TCP ports to expose to the network, e.g. 22, 2022:22, 22:192.168.1.1:2022")
 	flag.Parse()
@@ -271,17 +273,45 @@ func main() {
 
 	// Create SOCKS server
 	{
-		if socks != nil && nameserver != nil && *socks != "" {
-			resolver := types.NewNameResolver(s, *nameserver)
-			socksOptions := []socks5.Option{
-				socks5.WithDial(s.DialContext),
-				socks5.WithResolver(resolver),
+		if socks != nil {
+			if nameserver != nil {
+				if strings.Contains(*socks, ":") {
+					logger.Infof("Starting SOCKS server on %s", *socks)
+					resolver := types.NewNameResolver(s, *nameserver)
+					socksOptions := []socks5.Option{
+						socks5.WithDial(s.DialContext),
+						socks5.WithResolver(resolver),
+					}
+					if logger.GetLevel("debug") {
+						socksOptions = append(socksOptions, socks5.WithLogger(logger))
+					}
+					server := socks5.NewServer(socksOptions...)
+					go server.ListenAndServe("tcp", *socks) // nolint:errcheck
+				} else {
+					logger.Infof("Starting SOCKS server with socket file %s", *socks)
+					_, err := os.Stat(*socks)
+					if os.IsNotExist(err) {
+						n.socks5Listener, err = net.Listen("unix", *socks)
+						if err != nil {
+							panic(err)
+						}
+						resolver := types.NewNameResolver(s, *nameserver)
+						socksOptions := []socks5.Option{
+							socks5.WithDial(s.DialContext),
+							socks5.WithResolver(resolver),
+						}
+						if logger.GetLevel("debug") {
+							socksOptions = append(socksOptions, socks5.WithLogger(logger))
+						}
+						server := socks5.NewServer(socksOptions...)
+						go server.Serve(n.socks5Listener) // nolint:errcheck
+					} else if err != nil {
+						logger.Errorf("Cannot create socket file %s: %s", *socks, err)
+					} else {
+						panic(errors.New(fmt.Sprintf("Socket file %s already exists", *socks)))
+					}
+				}
 			}
-			if logger.GetLevel("debug") {
-				socksOptions = append(socksOptions, socks5.WithLogger(logger))
-			}
-			server := socks5.NewServer(socksOptions...)
-			go server.ListenAndServe("tcp", *socks) // nolint:errcheck
 		}
 	}
 
@@ -317,6 +347,10 @@ func main() {
 	// Shut down the node.
 	_ = n.admin.Stop()
 	_ = n.multicast.Stop()
+	if n.socks5Listener != nil {
+		_ = n.socks5Listener.Close()
+		logger.Infof("Stopped UNIX socket listener")
+	}
 	n.core.Stop()
 }