refactor: udp

tunnel/connection.go

@@ -1,19 +1,20 @@
 package tunnel
 
 import (
-	"errors"
 	"net"
 	"time"
 
-	N "github.com/Dreamacro/clash/common/net"
-	"github.com/Dreamacro/clash/common/pool"
+	cN "github.com/Dreamacro/clash/common/net"
 	"github.com/Dreamacro/clash/component/resolver"
 	C "github.com/Dreamacro/clash/constant"
+	"github.com/sagernet/sing/common"
+	"github.com/sagernet/sing/common/buf"
+	M "github.com/sagernet/sing/common/metadata"
+	N "github.com/sagernet/sing/common/network"
 )
 
 func handleUDPToRemote(packet C.UDPPacket, pc C.PacketConn, metadata *C.Metadata) error {
-	defer packet.Drop()
-
+	defer packet.Data().Release()
 	// local resolve UDP dns
 	if !metadata.Resolved() {
 		ip, err := resolver.ResolveIP(metadata.Host)
@@ -23,12 +24,7 @@ func handleUDPToRemote(packet C.UDPPacket, pc C.PacketConn, metadata *C.Metadata
 		metadata.DstIP = ip
 	}
 
-	addr := metadata.UDPAddr()
-	if addr == nil {
-		return errors.New("udp addr invalid")
-	}
-
-	if _, err := pc.WriteTo(packet.Data(), addr); err != nil {
+	if err := pc.WritePacket(packet.Data(), metadata.Socksaddr()); err != nil {
 		return err
 	}
 	// reset timeout
@@ -37,30 +33,45 @@ func handleUDPToRemote(packet C.UDPPacket, pc C.PacketConn, metadata *C.Metadata
 	return nil
 }
 
-func handleUDPToLocal(packet C.UDPPacket, pc net.PacketConn, key string, fAddr net.Addr) {
-	buf := pool.Get(pool.UDPBufferSize)
-	defer pool.Put(buf)
+func handleUDPToLocal(packet C.UDPPacket, pc C.PacketConn, key string, fAddr net.Addr) {
 	defer natTable.Delete(key)
 	defer pc.Close()
-
-	for {
-		pc.SetReadDeadline(time.Now().Add(udpTimeout))
-		n, from, err := pc.ReadFrom(buf)
-		if err != nil {
-			return
-		}
-
-		if fAddr != nil {
-			from = fAddr
-		}
-
-		_, err = packet.WriteBack(buf[:n], from)
-		if err != nil {
-			return
-		}
+	var fSocksaddr M.Socksaddr
+	if fAddr != nil {
+		fSocksaddr = M.SocksaddrFromNet(fAddr)
 	}
+	_, _ = copyPacketTimeout(packet, pc, udpTimeout, fSocksaddr, fAddr != nil)
 }
 
 func handleSocket(ctx C.ConnContext, outbound net.Conn) {
-	N.Relay(ctx.Conn(), outbound)
+	cN.Relay(ctx.Conn(), outbound)
+}
+
+func copyPacketTimeout(dst N.PacketWriter, src N.TimeoutPacketReader, timeout time.Duration, fAddr M.Socksaddr, fOverride bool) (n int64, err error) {
+	_buffer := buf.StackNewPacket()
+	defer common.KeepAlive(_buffer)
+	buffer := common.Dup(_buffer)
+	buffer.IncRef()
+	defer buffer.DecRef()
+	var destination M.Socksaddr
+	for {
+		buffer.Reset()
+		err = src.SetReadDeadline(time.Now().Add(timeout))
+		if err != nil {
+			return
+		}
+		destination, err = src.ReadPacket(buffer)
+		if err != nil {
+			return
+		}
+		if fOverride {
+			destination = fAddr
+		}
+		dataLen := buffer.Len()
+		err = dst.WritePacket(buffer, destination)
+		if err != nil {
+			return
+		}
+		n += int64(dataLen)
+	}
 }

tunnel/statistic/tracker.go

@@ -6,6 +6,8 @@ import (
 
 	C "github.com/Dreamacro/clash/constant"
 	"github.com/gofrs/uuid"
+	"github.com/sagernet/sing/common/buf"
+	M "github.com/sagernet/sing/common/metadata"
 	"go.uber.org/atomic"
 )
 
@@ -115,6 +117,25 @@ func (ut *udpTracker) WriteTo(b []byte, addr net.Addr) (int, error) {
 	return n, err
 }
 
+func (ut *udpTracker) ReadPacket(buffer *buf.Buffer) (addr M.Socksaddr, err error) {
+	addr, err = ut.PacketConn.ReadPacket(buffer)
+	download := int64(buffer.Len())
+	ut.manager.PushDownloaded(download)
+	ut.DownloadTotal.Add(download)
+	return
+}
+
+func (ut *udpTracker) WritePacket(buffer *buf.Buffer, addr M.Socksaddr) error {
+	dataLen := buffer.Len()
+	err := ut.PacketConn.WritePacket(buffer, addr)
+	if err == nil {
+		upload := int64(dataLen)
+		ut.manager.PushUploaded(upload)
+		ut.UploadTotal.Add(upload)
+	}
+	return err
+}
+
 func (ut *udpTracker) Close() error {
 	ut.manager.Leave(ut)
 	return ut.PacketConn.Close()