Fix: some UDP issues (#265)

proxy/socks/udp.go

@@ -1,17 +1,13 @@
 package socks
 
 import (
+	"bytes"
 	"net"
 
 	adapters "github.com/Dreamacro/clash/adapters/inbound"
 	"github.com/Dreamacro/clash/common/pool"
 	"github.com/Dreamacro/clash/component/socks5"
 	C "github.com/Dreamacro/clash/constant"
-	"github.com/Dreamacro/clash/tunnel"
-)
-
-var (
-	_ = tunnel.NATInstance()
 )
 
 type SockUDPListener struct {
@@ -28,17 +24,17 @@ func NewSocksUDPProxy(addr string) (*SockUDPListener, error) {
 
 	sl := &SockUDPListener{l, addr, false}
 	go func() {
-		buf := pool.BufPool.Get().([]byte)
-		defer pool.BufPool.Put(buf[:cap(buf)])
 		for {
+			buf := pool.BufPool.Get().([]byte)
 			n, remoteAddr, err := l.ReadFrom(buf)
 			if err != nil {
+				pool.BufPool.Put(buf[:cap(buf)])
 				if sl.closed {
 					break
 				}
 				continue
 			}
-			go handleSocksUDP(l, buf[:n], remoteAddr)
+			handleSocksUDP(l, buf[:n], remoteAddr)
 		}
 	}()
 
@@ -54,12 +50,19 @@ func (l *SockUDPListener) Address() string {
 	return l.address
 }
 
-func handleSocksUDP(c net.PacketConn, packet []byte, remoteAddr net.Addr) {
-	target, payload, err := socks5.DecodeUDPPacket(packet)
+func handleSocksUDP(pc net.PacketConn, buf []byte, addr net.Addr) {
+	target, payload, err := socks5.DecodeUDPPacket(buf)
 	if err != nil {
-		// Unresolved UDP packet, do nothing
+		// Unresolved UDP packet, return buffer to the pool
+		pool.BufPool.Put(buf[:cap(buf)])
 		return
 	}
-	conn := newfakeConn(c, target.String(), remoteAddr, payload)
+	conn := &fakeConn{
+		PacketConn: pc,
+		remoteAddr: addr,
+		targetAddr: target,
+		buffer:     bytes.NewBuffer(payload),
+		bufRef:     buf,
+	}
 	tun.Add(adapters.NewSocket(target, conn, C.SOCKS, C.UDP))
 }

proxy/socks/utils.go

@@ -4,24 +4,16 @@ import (
 	"bytes"
 	"net"
 
+	"github.com/Dreamacro/clash/common/pool"
 	"github.com/Dreamacro/clash/component/socks5"
 )
 
 type fakeConn struct {
 	net.PacketConn
-	target     string
 	remoteAddr net.Addr
+	targetAddr socks5.Addr
 	buffer     *bytes.Buffer
-}
-
-func newfakeConn(conn net.PacketConn, target string, remoteAddr net.Addr, buf []byte) *fakeConn {
-	buffer := bytes.NewBuffer(buf)
-	return &fakeConn{
-		PacketConn: conn,
-		target:     target,
-		buffer:     buffer,
-		remoteAddr: remoteAddr,
-	}
+	bufRef     []byte
 }
 
 func (c *fakeConn) Read(b []byte) (n int, err error) {
@@ -29,7 +21,7 @@ func (c *fakeConn) Read(b []byte) (n int, err error) {
 }
 
 func (c *fakeConn) Write(b []byte) (n int, err error) {
-	packet, err := socks5.EncodeUDPPacket(c.target, b)
+	packet, err := socks5.EncodeUDPPacket(c.targetAddr, b)
 	if err != nil {
 		return
 	}
@@ -39,3 +31,9 @@ func (c *fakeConn) Write(b []byte) (n int, err error) {
 func (c *fakeConn) RemoteAddr() net.Addr {
 	return c.remoteAddr
 }
+
+func (c *fakeConn) Close() error {
+	err := c.PacketConn.Close()
+	pool.BufPool.Put(c.bufRef[:cap(c.bufRef)])
+	return err
+}