code cleanup

go.mod

@@ -5,15 +5,18 @@ go 1.18
 require (
 	github.com/MakeNowJust/heredoc/v2 v2.0.1
 	github.com/akamensky/argparse v1.3.1
-	github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5
 	github.com/go-ini/ini v1.66.4
+	github.com/txthinking/socks5 v0.0.0-20220615051428-39268faee3e6
 	golang.zx2c4.com/wireguard v0.0.0-20220829161405-d1d08426b27b
 	suah.dev/protect v1.2.0
 )
 
 require (
 	github.com/google/btree v1.0.1 // indirect
+	github.com/patrickmn/go-cache v2.1.0+incompatible // indirect
 	github.com/stretchr/testify v1.8.0 // indirect
+	github.com/txthinking/runnergroup v0.0.0-20210608031112-152c7c4432bf // indirect
+	github.com/txthinking/x v0.0.0-20210326105829-476fab902fbe // indirect
 	golang.org/x/crypto v0.0.0-20220315160706-3147a52a75dd // indirect
 	golang.org/x/net v0.0.0-20220225172249-27dd8689420f // indirect
 	golang.org/x/sys v0.0.0-20220315194320-039c03cc5b86 // indirect

go.sum

@@ -2,8 +2,6 @@ github.com/MakeNowJust/heredoc/v2 v2.0.1 h1:rlCHh70XXXv7toz95ajQWOWQnN4WNLt0TdpZ
 github.com/MakeNowJust/heredoc/v2 v2.0.1/go.mod h1:6/2Abh5s+hc3g9nbWLe9ObDIOhaRrqsyY9MWy+4JdRM=
 github.com/akamensky/argparse v1.3.1 h1:kP6+OyvR0fuBH6UhbE6yh/nskrDEIQgEA1SUXDPjx4g=
 github.com/akamensky/argparse v1.3.1/go.mod h1:S5kwC7IuDcEr5VeXtGPRVZ5o/FdhcMlQz4IZQuw64xA=
-github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5 h1:0CwZNZbxp69SHPdPJAN/hZIm0C4OItdklCFmMRWYpio=
-github.com/armon/go-socks5 v0.0.0-20160902184237-e75332964ef5/go.mod h1:wHh0iHkYZB8zMSxRWpUBQtwG5a7fFgvEO+odwuTv2gs=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
@@ -11,6 +9,8 @@ github.com/go-ini/ini v1.66.4 h1:dKjMqkcbkzfddhIhyglTPgMoJnkvmG+bSLrU9cTHc5M=
 github.com/go-ini/ini v1.66.4/go.mod h1:ByCAeIL28uOIIG0E3PJtZPDL8WnHpFKFOtgjp+3Ies8=
 github.com/google/btree v1.0.1 h1:gK4Kx5IaGY9CD5sPJ36FHiBJ6ZXl0kilRiiCj+jdYp4=
 github.com/google/btree v1.0.1/go.mod h1:xXMiIv4Fb/0kKde4SpL7qlzvu5cMJDRkFDxJfI9uaxA=
+github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
+github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
@@ -18,6 +18,12 @@ github.com/stretchr/objx v0.4.0/go.mod h1:YvHI0jy2hoMjB+UWwv71VJQ9isScKT/TqJzVSS
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.8.0 h1:pSgiaMZlXftHpm5L7V1+rVB+AZJydKsMxsQBIJw4PKk=
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
+github.com/txthinking/runnergroup v0.0.0-20210608031112-152c7c4432bf h1:7PflaKRtU4np/epFxRXlFhlzLXZzKFrH5/I4so5Ove0=
+github.com/txthinking/runnergroup v0.0.0-20210608031112-152c7c4432bf/go.mod h1:CLUSJbazqETbaR+i0YAhXBICV9TrKH93pziccMhmhpM=
+github.com/txthinking/socks5 v0.0.0-20220615051428-39268faee3e6 h1:8DkPbOq/EPxbD5VJajKuvssiYZJSrlpeetcGfrBoBVE=
+github.com/txthinking/socks5 v0.0.0-20220615051428-39268faee3e6/go.mod h1:7NloQcrxaZYKURWph5HLxVDlIwMHJXCPkeWPtpftsIg=
+github.com/txthinking/x v0.0.0-20210326105829-476fab902fbe h1:gMWxZxBFRAXqoGkwkYlPX2zvyyKNWJpxOxCrjqJkm5A=
+github.com/txthinking/x v0.0.0-20210326105829-476fab902fbe/go.mod h1:WgqbSEmUYSjEV3B1qmee/PpP2NYEz4bL9/+mF1ma+s4=
 golang.org/x/crypto v0.0.0-20220315160706-3147a52a75dd h1:XcWmESyNjXJMLahc3mqVQJcgSTDxFxhETVlfk9uGc38=
 golang.org/x/crypto v0.0.0-20220315160706-3147a52a75dd/go.mod h1:IxCIyHEi3zRg3s0A5j5BB6A9Jmi73HwBIUl50j+osU4=
 golang.org/x/net v0.0.0-20220225172249-27dd8689420f h1:oA4XRj0qtSt8Yo1Zms0CUlsT3KG69V2UGQWPBxujDmc=

routine.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/subtle"
 	"errors"
+	"github.com/txthinking/socks5"
 	"io"
 	"log"
 	"math/rand"
@@ -11,9 +12,6 @@ import (
 	"os"
 	"strconv"
 
-	"github.com/armon/go-socks5"
-
-	"golang.zx2c4.com/wireguard/tun/netstack"
 	"net/netip"
 )
 
@@ -26,12 +24,6 @@ type CredentialValidator struct {
 	password string
 }
 
-// VirtualTun stores a reference to netstack network and DNS configuration
-type VirtualTun struct {
-	tnet      *netstack.Net
-	systemDNS bool
-}
-
 // RoutineSpawner spawns a routine (e.g. socks5, tcp static routes) after the configuration is parsed
 type RoutineSpawner interface {
 	SpawnRoutine(vt *VirtualTun)
@@ -121,18 +113,12 @@ func (d VirtualTun) resolveToAddrPort(endpoint *addressPort) (*netip.AddrPort, e
 
 // SpawnRoutine spawns a socks5 server.
 func (config *Socks5Config) SpawnRoutine(vt *VirtualTun) {
-	conf := &socks5.Config{Dial: vt.tnet.DialContext, Resolver: vt}
-	if username := config.Username; username != "" {
-		validator := CredentialValidator{username: username}
-		validator.password = config.Password
-		conf.Credentials = validator
-	}
-	server, err := socks5.New(conf)
+	server, err := socks5.NewClassicServer(config.BindAddress, "0.0.0.0", config.Username, config.Password, 15, 15)
 	if err != nil {
 		log.Fatal(err)
 	}
-
-	if err := server.ListenAndServe("tcp", config.BindAddress); err != nil {
+	err = server.ListenAndServe(vt)
+	if err != nil {
 		log.Fatal(err)
 	}
 }

tunnel.go

@@ -0,0 +1,179 @@
+package wireproxy
+
+import (
+	"fmt"
+	"github.com/patrickmn/go-cache"
+	"github.com/txthinking/socks5"
+	"golang.zx2c4.com/wireguard/tun/netstack"
+	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
+	"io"
+	"log"
+	"net"
+	"time"
+)
+
+// VirtualTun stores a reference to netstack network and DNS configuration
+type VirtualTun struct {
+	tnet                 *netstack.Net
+	systemDNS            bool
+	mappedPortToNatEntry map[uint16]string
+	natEntryToMappedPort *cache.Cache
+}
+
+type NatEntry struct {
+	key        string
+	srcAddr    net.Addr
+	mappedPort uint16
+	conn       *gonet.UDPConn
+}
+
+func (d *VirtualTun) connect(w io.Writer, r *socks5.Request) (net.Conn, error) {
+	if socks5.Debug {
+		log.Println("Call:", r.Address())
+	}
+	tmp, err := d.tnet.Dial("tcp", r.Address())
+	if err != nil {
+		var p *socks5.Reply
+		if r.Atyp == socks5.ATYPIPv4 || r.Atyp == socks5.ATYPDomain {
+			p = socks5.NewReply(socks5.RepHostUnreachable, socks5.ATYPIPv4, []byte{0x00, 0x00, 0x00, 0x00}, []byte{0x00, 0x00})
+		} else {
+			p = socks5.NewReply(socks5.RepHostUnreachable, socks5.ATYPIPv6, []byte(net.IPv6zero), []byte{0x00, 0x00})
+		}
+		if _, err := p.WriteTo(w); err != nil {
+			return nil, err
+		}
+		return nil, err
+	}
+
+	a, addr, port, err := socks5.ParseAddress(tmp.LocalAddr().String())
+	if err != nil {
+		var p *socks5.Reply
+		if r.Atyp == socks5.ATYPIPv4 || r.Atyp == socks5.ATYPDomain {
+			p = socks5.NewReply(socks5.RepHostUnreachable, socks5.ATYPIPv4, []byte{0x00, 0x00, 0x00, 0x00}, []byte{0x00, 0x00})
+		} else {
+			p = socks5.NewReply(socks5.RepHostUnreachable, socks5.ATYPIPv6, []byte(net.IPv6zero), []byte{0x00, 0x00})
+		}
+		if _, err := p.WriteTo(w); err != nil {
+			return nil, err
+		}
+		return nil, err
+	}
+	p := socks5.NewReply(socks5.RepSuccess, a, addr, port)
+	if _, err := p.WriteTo(w); err != nil {
+		return nil, err
+	}
+
+	return tmp, nil
+}
+
+func (d *VirtualTun) TCPHandle(s *socks5.Server, c *net.TCPConn, r *socks5.Request) error {
+	if r.Cmd == socks5.CmdConnect {
+		rc, err := d.connect(c, r)
+		if err != nil {
+			return err
+		}
+		defer rc.Close()
+		go func() {
+			var bf [1024 * 2]byte
+			for {
+				if s.TCPTimeout != 0 {
+					if err := rc.SetDeadline(time.Now().Add(time.Duration(s.TCPTimeout) * time.Second)); err != nil {
+						return
+					}
+				}
+				i, err := rc.Read(bf[:])
+				if err != nil {
+					return
+				}
+				if _, err := c.Write(bf[0:i]); err != nil {
+					return
+				}
+			}
+		}()
+		var bf [1024 * 2]byte
+		for {
+			if s.TCPTimeout != 0 {
+				if err := c.SetDeadline(time.Now().Add(time.Duration(s.TCPTimeout) * time.Second)); err != nil {
+					return nil
+				}
+			}
+			i, err := c.Read(bf[:])
+			if err != nil {
+				return nil
+			}
+			if _, err := rc.Write(bf[0:i]); err != nil {
+				return nil
+			}
+		}
+	}
+	if r.Cmd == socks5.CmdUDP {
+		caddr, err := r.UDP(c, s.ServerAddr)
+		if err != nil {
+			return err
+		}
+		srcAddr := caddr.String()
+		mappedPort := uint16(caddr.Port)
+		tries := 0
+		for _, occupied := d.mappedPortToNatEntry[mappedPort]; occupied; mappedPort++ {
+			tries++
+			if tries > 65535 {
+				return fmt.Errorf("nat table is full")
+			}
+		}
+		laddr, err := net.ResolveUDPAddr("udp", fmt.Sprintf(":%d", mappedPort))
+		if err != nil {
+			return err
+		}
+		conn, err := d.tnet.ListenUDP(laddr)
+		if err != nil {
+			fmt.Println("fic")
+			return err
+		}
+		entry := &NatEntry{
+			key:        srcAddr,
+			srcAddr:    caddr,
+			conn:       conn,
+			mappedPort: mappedPort,
+		}
+		d.mappedPortToNatEntry[mappedPort] = srcAddr
+		d.natEntryToMappedPort.Set(srcAddr, entry, 0)
+		go func() {
+			buf := make([]byte, 65536)
+			for n, from, err := conn.ReadFrom(buf); err == nil; n, from, err = conn.ReadFrom(buf) {
+				a, addr, port, err := socks5.ParseAddress(from.String())
+				if err != nil {
+					log.Println(err)
+					break
+				}
+				d.natEntryToMappedPort.Set(srcAddr, entry, 0)
+				d1 := socks5.NewDatagram(a, addr, port, buf[0:n])
+				if _, err := s.UDPConn.WriteToUDP(d1.Bytes(), caddr); err != nil {
+					log.Println(err)
+					break
+				}
+			}
+			_ = conn.Close()
+			d.natEntryToMappedPort.Delete(srcAddr)
+		}()
+		fmt.Printf("%s udp mapped to port %d\n", srcAddr, mappedPort)
+		return nil
+	}
+	return socks5.ErrUnsupportCmd
+}
+
+func (d *VirtualTun) UDPHandle(server *socks5.Server, addr *net.UDPAddr, datagram *socks5.Datagram) error {
+	srcAddr := addr.String()
+	entry, ok := d.natEntryToMappedPort.Get(srcAddr)
+	if !ok {
+		return fmt.Errorf("this udp address %s is not associated", srcAddr)
+	}
+	natEntry := entry.(*NatEntry)
+	// refresh timeout
+	d.natEntryToMappedPort.Set(srcAddr, entry, 0)
+	raddr, err := net.ResolveUDPAddr("udp", datagram.Address())
+	if err != nil {
+		return err
+	}
+	_, err = natEntry.conn.WriteTo(datagram.Data, raddr)
+	return err
+}

wireguard.go

@@ -3,8 +3,9 @@ package wireproxy
 import (
 	"bytes"
 	"fmt"
-
+	"github.com/patrickmn/go-cache"
 	"net/netip"
+	"time"
 
 	"github.com/MakeNowJust/heredoc/v2"
 	"golang.zx2c4.com/wireguard/conn"
@@ -74,8 +75,17 @@ func StartWireguard(conf *DeviceConfig) (*VirtualTun, error) {
 		return nil, err
 	}
 
+	natTable := cache.New(30*time.Second, time.Minute)
+	natTableRev := make(map[uint16]string)
+	natTable.OnEvicted(func(srcAddr string, i interface{}) {
+		entry := i.(*NatEntry)
+		_ = entry.conn.Close()
+		delete(natTableRev, entry.mappedPort)
+	})
 	return &VirtualTun{
 		tnet:                 tnet,
 		systemDNS:            len(setting.dns) == 0,
+		mappedPortToNatEntry: natTableRev,
+		natEntryToMappedPort: natTable,
 	}, nil
 }