Improve bufio.NATPacketConn

common/bufio/nat.go

@@ -17,16 +17,16 @@ type NATPacketConn interface {
 func NewUnidirectionalNATPacketConn(conn N.NetPacketConn, origin M.Socksaddr, destination M.Socksaddr) NATPacketConn {
 	return &unidirectionalNATPacketConn{
 		NetPacketConn: conn,
-		origin:        origin,
+		origin:        socksaddrWithoutPort(origin),
-		destination:   destination,
+		destination:   socksaddrWithoutPort(destination),
 	}
 }
 
 func NewNATPacketConn(conn N.NetPacketConn, origin M.Socksaddr, destination M.Socksaddr) NATPacketConn {
 	return &bidirectionalNATPacketConn{
 		NetPacketConn: conn,
-		origin:        origin,
+		origin:        socksaddrWithoutPort(origin),
-		destination:   destination,
+		destination:   socksaddrWithoutPort(destination),
 	}
 }
 
@@ -37,15 +37,24 @@ type unidirectionalNATPacketConn struct {
 }
 
 func (c *unidirectionalNATPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
-	if M.SocksaddrFromNet(addr) == c.destination {
+	destination := M.SocksaddrFromNet(addr)
-		addr = c.origin.UDPAddr()
+	if socksaddrWithoutPort(destination) == c.destination {
+		destination = M.Socksaddr{
+			Addr: c.origin.Addr,
+			Fqdn: c.origin.Fqdn,
+			Port: destination.Port,
 		}
-	return c.NetPacketConn.WriteTo(p, addr)
+	}
+	return c.NetPacketConn.WriteTo(p, destination.UDPAddr())
 }
 
 func (c *unidirectionalNATPacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
-	if destination == c.destination {
+	if socksaddrWithoutPort(destination) == c.destination {
-		destination = c.origin
+		destination = M.Socksaddr{
+			Addr: c.origin.Addr,
+			Fqdn: c.origin.Fqdn,
+			Port: destination.Port,
+		}
 	}
 	return c.NetPacketConn.WritePacket(buffer, destination)
 }
@@ -66,30 +75,55 @@ type bidirectionalNATPacketConn struct {
 
 func (c *bidirectionalNATPacketConn) ReadFrom(p []byte) (n int, addr net.Addr, err error) {
 	n, addr, err = c.NetPacketConn.ReadFrom(p)
-	if err == nil && M.SocksaddrFromNet(addr) == c.origin {
+	if err != nil {
-		addr = c.destination.UDPAddr()
+		return
 	}
+	destination := M.SocksaddrFromNet(addr)
+	if socksaddrWithoutPort(destination) == c.origin {
+		destination = M.Socksaddr{
+			Addr: c.destination.Addr,
+			Fqdn: c.destination.Fqdn,
+			Port: destination.Port,
+		}
+	}
+	addr = destination.UDPAddr()
 	return
 }
 
 func (c *bidirectionalNATPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
-	if M.SocksaddrFromNet(addr) == c.destination {
+	destination := M.SocksaddrFromNet(addr)
-		addr = c.origin.UDPAddr()
+	if socksaddrWithoutPort(destination) == c.destination {
+		destination = M.Socksaddr{
+			Addr: c.origin.Addr,
+			Fqdn: c.origin.Fqdn,
+			Port: destination.Port,
 		}
-	return c.NetPacketConn.WriteTo(p, addr)
+	}
+	return c.NetPacketConn.WriteTo(p, destination.UDPAddr())
 }
 
 func (c *bidirectionalNATPacketConn) ReadPacket(buffer *buf.Buffer) (destination M.Socksaddr, err error) {
 	destination, err = c.NetPacketConn.ReadPacket(buffer)
-	if destination == c.origin {
+	if err != nil {
-		destination = c.destination
+		return
+	}
+	if socksaddrWithoutPort(destination) == c.origin {
+		destination = M.Socksaddr{
+			Addr: c.destination.Addr,
+			Fqdn: c.destination.Fqdn,
+			Port: destination.Port,
+		}
 	}
 	return
 }
 
 func (c *bidirectionalNATPacketConn) WritePacket(buffer *buf.Buffer, destination M.Socksaddr) error {
-	if destination == c.destination {
+	if socksaddrWithoutPort(destination) == c.destination {
-		destination = c.origin
+		destination = M.Socksaddr{
+			Addr: c.origin.Addr,
+			Fqdn: c.origin.Fqdn,
+			Port: destination.Port,
+		}
 	}
 	return c.NetPacketConn.WritePacket(buffer, destination)
 }
@@ -101,3 +135,8 @@ func (c *bidirectionalNATPacketConn) UpdateDestination(destinationAddress netip.
 func (c *bidirectionalNATPacketConn) Upstream() any {
 	return c.NetPacketConn
 }
+
+func socksaddrWithoutPort(destination M.Socksaddr) M.Socksaddr {
+	destination.Port = 0
+	return destination
+}

common/bufio/nat_wait.go

@@ -0,0 +1,39 @@
+package bufio
+
+import (
+	"github.com/sagernet/sing/common/buf"
+	M "github.com/sagernet/sing/common/metadata"
+	N "github.com/sagernet/sing/common/network"
+)
+
+func (c *bidirectionalNATPacketConn) CreatePacketReadWaiter() (N.PacketReadWaiter, bool) {
+	waiter, created := CreatePacketReadWaiter(c.NetPacketConn)
+	if !created {
+		return nil, false
+	}
+	return &waitBidirectionalNATPacketConn{c, waiter}, true
+}
+
+type waitBidirectionalNATPacketConn struct {
+	*bidirectionalNATPacketConn
+	readWaiter N.PacketReadWaiter
+}
+
+func (c *waitBidirectionalNATPacketConn) InitializeReadWaiter(options N.ReadWaitOptions) (needCopy bool) {
+	return c.readWaiter.InitializeReadWaiter(options)
+}
+
+func (c *waitBidirectionalNATPacketConn) WaitReadPacket() (buffer *buf.Buffer, destination M.Socksaddr, err error) {
+	buffer, destination, err = c.readWaiter.WaitReadPacket()
+	if err != nil {
+		return
+	}
+	if socksaddrWithoutPort(destination) == c.origin {
+		destination = M.Socksaddr{
+			Addr: c.destination.Addr,
+			Fqdn: c.destination.Fqdn,
+			Port: destination.Port,
+		}
+	}
+	return
+}