Fix(socks5): fully udp associate support (#233)

tunnel/connection.go

@@ -63,54 +63,43 @@ func (t *Tunnel) handleHTTP(request *adapters.HTTPAdapter, outbound net.Conn) {
 	}
 }
 
+func (t *Tunnel) handleUDPToRemote(conn net.Conn, pc net.PacketConn, addr net.Addr) {
+	buf := pool.BufPool.Get().([]byte)
+	defer pool.BufPool.Put(buf[:cap(buf)])
+
+	n, err := conn.Read(buf)
+	if err != nil {
+		return
+	}
+	if _, err = pc.WriteTo(buf[:n], addr); err != nil {
+		return
+	}
+	t.traffic.Up() <- int64(n)
+}
+
+func (t *Tunnel) handleUDPToLocal(conn net.Conn, pc net.PacketConn) {
+	buf := pool.BufPool.Get().([]byte)
+	defer pool.BufPool.Put(buf[:cap(buf)])
+
+	for {
+		n, _, err := pc.ReadFrom(buf)
+		if err != nil {
+			return
+		}
+
+		n, err = conn.Write(buf[:n])
+		if err != nil {
+			return
+		}
+		t.traffic.Down() <- int64(n)
+	}
+}
+
 func (t *Tunnel) handleSocket(request *adapters.SocketAdapter, outbound net.Conn) {
 	conn := newTrafficTrack(outbound, t.traffic)
 	relay(request, conn)
 }
 
-func (t *Tunnel) handleUDPOverTCP(conn net.Conn, pc net.PacketConn, addr net.Addr) error {
-	ch := make(chan error, 1)
-
-	go func() {
-		buf := pool.BufPool.Get().([]byte)
-		defer pool.BufPool.Put(buf)
-		for {
-			n, err := conn.Read(buf)
-			if err != nil {
-				ch <- err
-				return
-			}
-			pc.SetReadDeadline(time.Now().Add(120 * time.Second))
-			if _, err = pc.WriteTo(buf[:n], addr); err != nil {
-				ch <- err
-				return
-			}
-			t.traffic.Up() <- int64(n)
-			ch <- nil
-		}
-	}()
-
-	buf := pool.BufPool.Get().([]byte)
-	defer pool.BufPool.Put(buf)
-
-	for {
-		pc.SetReadDeadline(time.Now().Add(120 * time.Second))
-		n, _, err := pc.ReadFrom(buf)
-		if err != nil {
-			break
-		}
-
-		if _, err := conn.Write(buf[:n]); err != nil {
-			break
-		}
-
-		t.traffic.Down() <- int64(n)
-	}
-
-	<-ch
-	return nil
-}
-
 // relay copies between left and right bidirectionally.
 func relay(leftConn, rightConn net.Conn) {
 	ch := make(chan error)

tunnel/session.go

@@ -0,0 +1,22 @@
+package tunnel
+
+import (
+	"sync"
+	"time"
+
+	nat "github.com/Dreamacro/clash/component/nat-table"
+)
+
+var (
+	natTable *nat.Table
+	natOnce  sync.Once
+
+	natTimeout = 120 * time.Second
+)
+
+func NATInstance() *nat.Table {
+	natOnce.Do(func() {
+		natTable = nat.New(natTimeout)
+	})
+	return natTable
+}

tunnel/tunnel.go

@@ -103,9 +103,20 @@ func (t *Tunnel) needLookupIP(metadata *C.Metadata) bool {
 }
 
 func (t *Tunnel) handleConn(localConn C.ServerAdapter) {
-	defer localConn.Close()
-	metadata := localConn.Metadata()
+	defer func() {
+		var conn net.Conn
+		switch adapter := localConn.(type) {
+		case *InboundAdapter.HTTPAdapter:
+			conn = adapter.Conn
+		case *InboundAdapter.SocketAdapter:
+			conn = adapter.Conn
+		}
+		if _, ok := conn.(*net.TCPConn); ok {
+			localConn.Close()
+		}
+	}()
 
+	metadata := localConn.Metadata()
 	if !metadata.Valid() {
 		log.Warnln("[Metadata] not valid: %#v", metadata)
 		return
@@ -138,18 +149,32 @@ func (t *Tunnel) handleConn(localConn C.ServerAdapter) {
 		}
 	}
 
-	if metadata.NetWork == C.UDP {
-		pc, addr, err := proxy.DialUDP(metadata)
+	switch metadata.NetWork {
+	case C.TCP:
+		t.handleTCPConn(localConn, metadata, proxy)
+	case C.UDP:
+		t.handleUDPConn(localConn, metadata, proxy)
+	}
+}
+
+func (t *Tunnel) handleUDPConn(localConn C.ServerAdapter, metadata *C.Metadata, proxy C.Proxy) {
+	pc, addr := natTable.Get(localConn.RemoteAddr())
+	if pc == nil {
+		var err error
+		pc, addr, err = proxy.DialUDP(metadata)
 		if err != nil {
 			log.Warnln("Proxy[%s] connect [%s --> %s] error: %s", proxy.Name(), metadata.SrcIP.String(), metadata.String(), err.Error())
 			return
 		}
-		defer pc.Close()
 
-		t.handleUDPOverTCP(localConn, pc, addr)
-		return
+		natTable.Set(localConn.RemoteAddr(), pc, addr)
+		go t.handleUDPToLocal(localConn, pc)
 	}
 
+	t.handleUDPToRemote(localConn, pc, addr)
+}
+
+func (t *Tunnel) handleTCPConn(localConn C.ServerAdapter, metadata *C.Metadata, proxy C.Proxy) {
 	remoConn, err := proxy.Dial(metadata)
 	if err != nil {
 		log.Warnln("Proxy[%s] connect [%s --> %s] error: %s", proxy.Name(), metadata.SrcIP.String(), metadata.String(), err.Error())
@@ -196,6 +221,7 @@ func (t *Tunnel) match(metadata *C.Metadata) (C.Proxy, error) {
 			}
 
 			if metadata.NetWork == C.UDP && !adapter.SupportUDP() {
+				log.Debugln("%v UDP is not supported", adapter.Name())
 				continue
 			}