fix

SagerNet · Apr 10, 2024 · 0467888 · 0467888
1 parent 40190f9
commit 0467888
Showing 1 changed file with 52 additions and 43 deletions.
diff --git a/common/bufio/nat.go b/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,
-		destination:   destination,
+		origin:        socksaddrWithoutPort(origin),
+		destination:   socksaddrWithoutPort(destination),
 	}
 }
 
 func NewNATPacketConn(conn N.NetPacketConn, origin M.Socksaddr, destination M.Socksaddr) NATPacketConn {
 	return &bidirectionalNATPacketConn{
 		NetPacketConn: conn,
-		origin:        origin,
-		destination:   destination,
+		origin:        socksaddrWithoutPort(origin),
+		destination:   socksaddrWithoutPort(destination),
 	}
 }
 
@@ -38,24 +38,23 @@ type unidirectionalNATPacketConn struct {
 
 func (c *unidirectionalNATPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err error) {
 	destination := M.SocksaddrFromNet(addr)
-	if destination == c.destination {
-		destination = c.origin
+	if socksaddrWithoutPort(destination) == c.destination {
+		destination = M.Socksaddr{
+			Addr: c.origin.Addr,
+			Fqdn: c.origin.Fqdn,
+			Port: destination.Port,
+		}
 	}
-	if c.destination.Port == c.origin.Port && destination.AddrString() == c.destination.AddrString() {
-		destination.Addr = c.origin.Addr
-		destination.Fqdn = c.origin.Fqdn
-	}
-	addr = destination.UDPAddr()
-	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 {
-		destination = c.origin
-	}
-	if c.destination.Port == c.origin.Port && destination.AddrString() == c.destination.AddrString() {
-		destination.Addr = c.origin.Addr
-		destination.Fqdn = c.origin.Fqdn
+	if socksaddrWithoutPort(destination) == c.destination {
+		destination = M.Socksaddr{
+			Addr: c.origin.Addr,
+			Fqdn: c.origin.Fqdn,
+			Port: destination.Port,
+		}
 	}
 	return c.NetPacketConn.WritePacket(buffer, destination)
 }
@@ -76,50 +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)
-	destination := M.SocksaddrFromNet(addr)
-	if err == nil && destination == c.origin {
-		destination = c.destination
+	if err != nil {
+		return
 	}
-	if c.destination.Port == c.origin.Port && destination.AddrString() == c.origin.AddrString() {
-		destination.Addr = c.destination.Addr
-		destination.Fqdn = c.destination.Fqdn
+	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) {
 	destination := M.SocksaddrFromNet(addr)
-	if destination == c.destination {
-		destination = c.origin
+	if socksaddrWithoutPort(destination) == c.destination {
+		destination = M.Socksaddr{
+			Addr: c.origin.Addr,
+			Fqdn: c.origin.Fqdn,
+			Port: destination.Port,
+		}
 	}
-	if c.destination.Port == c.origin.Port && destination.AddrString() == c.destination.AddrString() {
-		destination.Addr = c.origin.Addr
-		destination.Fqdn = c.origin.Fqdn
-	}
-	addr = destination.UDPAddr()
-	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 {
-		destination = c.destination
+	if err != nil {
+		return
 	}
-	if c.destination.Port == c.origin.Port && destination.AddrString() == c.origin.AddrString() {
-		destination.Addr = c.destination.Addr
-		destination.Fqdn = c.destination.Fqdn
+	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 {
-		destination = c.origin
-	}
-	if c.destination.Port == c.origin.Port && destination.AddrString() == c.destination.AddrString() {
-		destination.Addr = c.origin.Addr
-		destination.Fqdn = c.origin.Fqdn
+	if socksaddrWithoutPort(destination) == c.destination {
+		destination = M.Socksaddr{
+			Addr: c.origin.Addr,
+			Fqdn: c.origin.Fqdn,
+			Port: destination.Port,
+		}
 	}
 	return c.NetPacketConn.WritePacket(buffer, destination)
 }
@@ -131,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
+}