Implement SOCKS5 UDP (ACL not yet integrated)

cmd/client.go

@@ -126,6 +126,23 @@ func client(config *clientConfig) {
 							"dst": reqAddr,
 						}).Debug("SOCKS5 TCP EOF")
 					}
+				},
+				func(addr net.Addr) {
+					logrus.WithFields(logrus.Fields{
+						"src": addr.String(),
+					}).Debug("SOCKS5 UDP associate")
+				},
+				func(addr net.Addr, err error) {
+					if err != io.EOF {
+						logrus.WithFields(logrus.Fields{
+							"error": err,
+							"src":   addr.String(),
+						}).Info("SOCKS5 UDP error")
+					} else {
+						logrus.WithFields(logrus.Fields{
+							"src": addr.String(),
+						}).Debug("SOCKS5 UDP EOF")
+					}
 				})
 			if err != nil {
 				logrus.WithField("error", err).Fatal("Failed to initialize SOCKS5 server")

docs/socks5/udpchk.py

@@ -0,0 +1,19 @@
+import socks
+import socket
+
+
+def main():
+    s = socks.socksocket(socket.AF_INET, socket.SOCK_DGRAM)
+    s.set_proxy(socks.SOCKS5, "127.0.0.1", 1080)
+    # Raw DNS request
+    req = b"\x12\x34\x01\x00\x00\x01\x00\x00\x00\x00\x00\x00\x05\x62\x61\x69\x64\x75\x03\x63\x6f\x6d\x00\x00\x01\x00\x01"
+    s.sendto(req, ("8.8.8.8", 53))
+    (rsp, address) = s.recvfrom(4096)
+    if rsp[0] == req[0] and rsp[1] == req[1]:
+        print("UDP check passed")
+    else:
+        print("Invalid response")
+
+
+if __name__ == "__main__":
+    main()

pkg/socks5/server.go

@@ -16,6 +16,8 @@ import (
 	"time"
 )
 
+const udpBufferSize = 65535
+
 var (
 	ErrUnsupportedCmd = errors.New("unsupported command")
 	ErrUserPassAuth   = errors.New("invalid username or password")
@@ -30,8 +32,10 @@ type Server struct {
 	ACLEngine  *acl.Engine
 	DisableUDP bool
 
-	TCPRequestFunc func(addr net.Addr, reqAddr string, action acl.Action, arg string)
+	TCPRequestFunc   func(addr net.Addr, reqAddr string, action acl.Action, arg string)
-	TCPErrorFunc   func(addr net.Addr, reqAddr string, err error)
+	TCPErrorFunc     func(addr net.Addr, reqAddr string, err error)
+	UDPAssociateFunc func(addr net.Addr)
+	UDPErrorFunc     func(addr net.Addr, err error)
 
 	tcpListener *net.TCPListener
 }
@@ -39,7 +43,8 @@ type Server struct {
 func NewServer(hyClient *core.Client, addr string, authFunc func(username, password string) bool, tcpTimeout time.Duration,
 	aclEngine *acl.Engine, disableUDP bool,
 	tcpReqFunc func(addr net.Addr, reqAddr string, action acl.Action, arg string),
-	tcpErrorFunc func(addr net.Addr, reqAddr string, err error)) (*Server, error) {
+	tcpErrorFunc func(addr net.Addr, reqAddr string, err error),
+	udpAssocFunc func(addr net.Addr), udpErrorFunc func(addr net.Addr, err error)) (*Server, error) {
 	tAddr, err := net.ResolveTCPAddr("tcp", addr)
 	if err != nil {
 		return nil, err
@@ -49,15 +54,17 @@ func NewServer(hyClient *core.Client, addr string, authFunc func(username, passw
 		m = socks5.MethodUsernamePassword
 	}
 	s := &Server{
-		HyClient:       hyClient,
+		HyClient:         hyClient,
-		AuthFunc:       authFunc,
+		AuthFunc:         authFunc,
-		Method:         m,
+		Method:           m,
-		TCPAddr:        tAddr,
+		TCPAddr:          tAddr,
-		TCPTimeout:     tcpTimeout,
+		TCPTimeout:       tcpTimeout,
-		ACLEngine:      aclEngine,
+		ACLEngine:        aclEngine,
-		DisableUDP:     disableUDP,
+		DisableUDP:       disableUDP,
-		TCPRequestFunc: tcpReqFunc,
+		TCPRequestFunc:   tcpReqFunc,
-		TCPErrorFunc:   tcpErrorFunc,
+		TCPErrorFunc:     tcpErrorFunc,
+		UDPAssociateFunc: udpAssocFunc,
+		UDPErrorFunc:     udpErrorFunc,
 	}
 	return s, nil
 }
@@ -142,8 +149,12 @@ func (s *Server) handle(c *net.TCPConn, r *socks5.Request) error {
 		return s.handleTCP(c, r)
 	} else if r.Cmd == socks5.CmdUDP {
 		// UDP
-		_ = sendReply(c, socks5.RepCommandNotSupported)
+		if !s.DisableUDP {
-		return ErrUnsupportedCmd
+			return s.handleUDP(c, r)
+		} else {
+			_ = sendReply(c, socks5.RepCommandNotSupported)
+			return ErrUnsupportedCmd
+		}
 	} else {
 		_ = sendReply(c, socks5.RepCommandNotSupported)
 		return ErrUnsupportedCmd
@@ -207,6 +218,107 @@ func (s *Server) handleTCP(c *net.TCPConn, r *socks5.Request) error {
 	}
 }
 
+func (s *Server) handleUDP(c *net.TCPConn, r *socks5.Request) error {
+	s.UDPAssociateFunc(c.RemoteAddr())
+	var closeErr error
+	defer func() {
+		s.UDPErrorFunc(c.RemoteAddr(), closeErr)
+	}()
+	// Start local UDP server
+	udpConn, err := net.ListenUDP("udp", &net.UDPAddr{
+		IP:   s.TCPAddr.IP,
+		Zone: s.TCPAddr.Zone,
+	})
+	if err != nil {
+		_ = sendReply(c, socks5.RepServerFailure)
+		closeErr = err
+		return err
+	}
+	defer udpConn.Close()
+	// HyClient UDP session
+	hyUDP, err := s.HyClient.DialUDP()
+	if err != nil {
+		_ = sendReply(c, socks5.RepServerFailure)
+		closeErr = err
+		return err
+	}
+	defer hyUDP.Close()
+	// Send UDP server addr to the client
+	atyp, addr, port, err := socks5.ParseAddress(udpConn.LocalAddr().String())
+	if err != nil {
+		_ = sendReply(c, socks5.RepServerFailure)
+		closeErr = err
+		return err
+	}
+	_, _ = socks5.NewReply(socks5.RepSuccess, atyp, addr, port).WriteTo(c)
+	// Let UDP server do its job, we hold the TCP connection here
+	go s.udpServer(udpConn, hyUDP)
+	buf := make([]byte, 1024)
+	for {
+		if s.TCPTimeout != 0 {
+			_ = c.SetDeadline(time.Now().Add(s.TCPTimeout))
+		}
+		_, err := c.Read(buf)
+		if err != nil {
+			closeErr = err
+			break
+		}
+	}
+	// As the TCP connection closes, so does the UDP server & HyClient session
+	return nil
+}
+
+func (s *Server) udpServer(c *net.UDPConn, hyUDP core.UDPConn) {
+	var clientAddr *net.UDPAddr
+	buf := make([]byte, udpBufferSize)
+	// Local to remote
+	for {
+		n, cAddr, err := c.ReadFromUDP(buf)
+		if err != nil {
+			break
+		}
+		d, err := socks5.NewDatagramFromBytes(buf[:n])
+		if err != nil || d.Frag != 0 {
+			// Ignore bad packets
+			continue
+		}
+		if clientAddr == nil {
+			// Whoever sends the first valid packet is our client
+			clientAddr = cAddr
+			go func() {
+				// Start remote to local
+				for {
+					bs, _, err := hyUDP.ReadFrom()
+					if err != nil {
+						break
+					}
+					// RFC 1928 is very ambiguous on how to properly use DST.ADDR and DST.PORT in reply packets
+					// So we just fill in zeros for now. Works fine for all the SOCKS5 clients I tested
+					d := socks5.NewDatagram(socks5.ATYPIPv4, []byte{0x00, 0x00, 0x00, 0x00}, []byte{0x00, 0x00}, bs)
+					_, err = c.WriteTo(d.Bytes(), clientAddr)
+					if err != nil {
+						break
+					}
+				}
+			}()
+		} else if cAddr.String() != clientAddr.String() {
+			// Not our client, bye
+			continue
+		}
+		_, _, _, addr := parseDatagramRequestAddress(d)
+		/*
+			action, arg := acl.ActionProxy, ""
+			if s.ACLEngine != nil {
+				action, arg = s.ACLEngine.Lookup(domain, ip)
+			}
+		*/
+		err = hyUDP.WriteTo(d.Data, addr)
+		if err != nil {
+			break
+		}
+	}
+}
+
 func sendReply(conn *net.TCPConn, rep byte) error {
 	p := socks5.NewReply(rep, socks5.ATYPIPv4, []byte{0x00, 0x00, 0x00, 0x00}, []byte{0x00, 0x00})
 	_, err := p.WriteTo(conn)
@@ -222,3 +334,13 @@ func parseRequestAddress(r *socks5.Request) (domain string, ip net.IP, port stri
 		return "", r.DstAddr, p, net.JoinHostPort(net.IP(r.DstAddr).String(), p)
 	}
 }
+
+func parseDatagramRequestAddress(r *socks5.Datagram) (domain string, ip net.IP, port string, addr string) {
+	p := strconv.Itoa(int(binary.BigEndian.Uint16(r.DstPort)))
+	if r.Atyp == socks5.ATYPDomain {
+		d := string(r.DstAddr[1:])
+		return d, nil, p, net.JoinHostPort(d, p)
+	} else {
+		return "", r.DstAddr, p, net.JoinHostPort(net.IP(r.DstAddr).String(), p)
+	}
+}