⚙️ Fix client UDP

Author: database64128

client/client.go

@@ -202,7 +202,7 @@ type ShadowsocksPacketConn interface {
 	net.PacketConn
 
 	// ReadFromZeroCopy eliminates copying by requiring that a big enough buffer is passed for reading.
-	ReadFromZeroCopy(b []byte) (payload []byte, address string, err error)
+	ReadFromZeroCopy(b []byte) (socksAddrStart, payloadStart, payloadLength int, err error)
 
 	// WriteToZeroCopy minimizes copying by requiring that enough space is reserved in b.
 	// The socks address is still being copied into the buffer.
@@ -368,7 +368,12 @@ func (c *packetConn) ReadFrom(b []byte) (int, net.Addr, error) {
 	}
 
 	// Decrypt in-place.
-	payload, address, err := ss.UnpackAndValidatePacket(c.cipher, c.aead, c.filter, c.sid, nil, cipherBuf[:n])
+	_, socksAddr, payload, err := ss.UnpackAndValidatePacket(c.cipher, c.aead, c.filter, c.sid, nil, cipherBuf[:n])
+	if err != nil {
+		return 0, nil, err
+	}
+
+	addr, err := socksAddr.Addr("udp")
 	if err != nil {
 		return 0, nil, err
 	}
@@ -378,37 +383,17 @@ func (c *packetConn) ReadFrom(b []byte) (int, net.Addr, error) {
 		err = io.ErrShortBuffer
 	}
 
-	return n, NewAddr(address, "udp"), err
+	return n, addr, err
 }
 
-func (c *packetConn) ReadFromZeroCopy(b []byte) (payload []byte, address string, err error) {
+func (c *packetConn) ReadFromZeroCopy(b []byte) (socksAddrStart, payloadStart, payloadLength int, err error) {
 	n, err := c.UDPConn.Read(b)
 	if err != nil {
 		return
 	}
 
-	payload, address, err = ss.UnpackAndValidatePacket(c.cipher, c.aead, c.filter, c.sid, nil, b[:n])
+	socksAddrStart, socksAddr, payload, err := ss.UnpackAndValidatePacket(c.cipher, c.aead, c.filter, c.sid, nil, b[:n])
+	payloadStart = socksAddrStart + len(socksAddr)
+	payloadLength = len(payload)
 	return
 }
-
-type addr struct {
-	address string
-	network string
-}
-
-func (a *addr) String() string {
-	return a.address
-}
-
-func (a *addr) Network() string {
-	return a.network
-}
-
-// NewAddr returns a net.Addr that holds an address of the form `host:port` with a domain name or IP as host.
-// Used for SOCKS addressing.
-func NewAddr(address, network string) net.Addr {
-	return &addr{
-		address: address,
-		network: network,
-	}
-}

client/client_test.go

@@ -121,7 +121,7 @@ func TestShadowsocksClient_ListenUDP(t *testing.T) {
 	}
 	defer conn.Close()
 	conn.SetReadDeadline(time.Now().Add(time.Second * 5))
-	pcrw := &packetConnReadWriter{PacketConn: conn, targetAddr: NewAddr(testTargetAddr, "udp")}
+	pcrw := &packetConnReadWriter{PacketConn: conn, targetAddr: onet.NewAddr(testTargetAddr, "udp")}
 	expectEchoPayload(pcrw, ss.MakeTestPayload(1024), make([]byte, 1024), t)
 
 	proxy.Close()
@@ -173,7 +173,7 @@ func BenchmarkShadowsocksClient_ListenUDP(b *testing.B) {
 	buf := make([]byte, service.UDPPacketBufferSize)
 	for n := 0; n < b.N; n++ {
 		payload := ss.MakeTestPayload(1024)
-		pcrw := &packetConnReadWriter{PacketConn: conn, targetAddr: NewAddr(testTargetAddr, "udp")}
+		pcrw := &packetConnReadWriter{PacketConn: conn, targetAddr: onet.NewAddr(testTargetAddr, "udp")}
 		b.StartTimer()
 		expectEchoPayload(pcrw, payload, buf, b)
 		b.StopTimer()
@@ -252,7 +252,7 @@ func startShadowsocksUDPEchoServer(expectedTgtAddr string, t testing.TB) (net.Co
 				t.Logf("Failed to read from UDP conn: %v", err)
 				return
 			}
-			buf, err := ss.Unpack(clientBuf, cipherBuf[:n], cipher)
+			_, buf, err := ss.Unpack(clientBuf, cipherBuf[:n], cipher)
 			if err != nil {
 				t.Fatalf("Failed to decrypt: %v", err)
 			}

client/nat.go

@@ -36,6 +36,9 @@ type natconn struct {
 
 	// The UDPConn where the last client packet was received from.
 	clientConn onet.UDPPacketConn
+
+	// packetAdapter translates packets between clientConn and proxyConn.
+	packetAdapter PacketAdapter
 }
 
 func isDNS(addr *net.UDPAddr) bool {
@@ -79,8 +82,8 @@ func (c *natconn) WriteToZeroCopy(b []byte, start, length int, socksAddr []byte)
 	return c.proxyConn.WriteToZeroCopy(b, start, length, socksAddr)
 }
 
-func (c *natconn) ReadFromZeroCopy(b []byte) (payload []byte, address string, err error) {
-	payload, address, err = c.proxyConn.ReadFromZeroCopy(b)
+func (c *natconn) ReadFromZeroCopy(b []byte) (socksAddrStart, payloadStart, payloadLength int, err error) {
+	socksAddrStart, payloadStart, payloadLength, err = c.proxyConn.ReadFromZeroCopy(b)
 	if err == nil {
 		c.onRead(nil)
 	}
@@ -104,7 +107,7 @@ func (c *natconn) timedCopy() {
 	packetBuf := make([]byte, service.UDPPacketBufferSize)
 
 	for {
-		payload, _, err := c.ReadFromZeroCopy(packetBuf)
+		socksAddrStart, payloadStart, payloadLength, err := c.ReadFromZeroCopy(packetBuf)
 		if err != nil {
 			if errors.Is(err, os.ErrDeadlineExceeded) {
 				return
@@ -113,7 +116,13 @@ func (c *natconn) timedCopy() {
 			continue
 		}
 
-		_, _, err = c.clientConn.WriteMsgUDP(payload, c.oobCache, c.clientAddr)
+		writeBuf, err := c.packetAdapter.EncapsulatePacket(packetBuf, socksAddrStart, payloadStart, payloadLength)
+		if err != nil {
+			log.Print(err)
+			continue
+		}
+
+		_, _, err = c.clientConn.WriteMsgUDP(writeBuf, c.oobCache, c.clientAddr)
 		if err != nil {
 			log.Print(err)
 		}
@@ -140,13 +149,14 @@ func (m *natmap) GetByClientAddress(key string) *natconn {
 	return m.keyConn[key]
 }
 
-func (m *natmap) set(clientAddr *net.UDPAddr, clientConn onet.UDPPacketConn, proxyConn ShadowsocksPacketConn, oobCache []byte) *natconn {
+func (m *natmap) set(clientAddr *net.UDPAddr, clientConn onet.UDPPacketConn, proxyConn ShadowsocksPacketConn, oobCache []byte, packetAdapter PacketAdapter) *natconn {
 	entry := &natconn{
 		proxyConn:      proxyConn,
 		defaultTimeout: m.timeout,
 		oobCache:       oobCache,
 		clientAddr:     clientAddr,
 		clientConn:     clientConn,
+		packetAdapter:  packetAdapter,
 	}
 
 	m.Lock()
@@ -168,8 +178,8 @@ func (m *natmap) del(key string) *natconn {
 	return nil
 }
 
-func (m *natmap) Add(clientAddr *net.UDPAddr, clientConn onet.UDPPacketConn, oobCache []byte, proxyConn ShadowsocksPacketConn) *natconn {
-	entry := m.set(clientAddr, clientConn, proxyConn, oobCache)
+func (m *natmap) Add(clientAddr *net.UDPAddr, clientConn onet.UDPPacketConn, oobCache []byte, proxyConn ShadowsocksPacketConn, packetAdapter PacketAdapter) *natconn {
+	entry := m.set(clientAddr, clientConn, proxyConn, oobCache, packetAdapter)
 	go func() {
 		entry.timedCopy()
 		if pc := m.del(clientAddr.String()); pc != nil {

client/tcp.go

@@ -62,18 +62,19 @@ func (s *TCPTunnel) listen() {
 
 			socksaddr, err := s.handshaker.Handshake(clientConn)
 			if err != nil {
-				if socksaddr == nil {
-					// Keep the connection open until EOF.
-					// Example use case: SOCKS5 UDP ASSOCIATE command.
-					b := make([]byte, 1)
-					_, err = io.ReadFull(clientConn, b)
-					if err == nil || err == io.ErrUnexpectedEOF {
-						return
-					}
-				}
 				log.Print(err)
 				return
 			}
+			if socksaddr == nil {
+				// Keep the connection open until EOF.
+				// Example use case: SOCKS5 UDP ASSOCIATE command.
+				b := make([]byte, 1)
+				_, err = io.ReadFull(clientConn, b)
+				if err != nil && err != io.ErrUnexpectedEOF {
+					log.Print(err)
+				}
+				return
+			}
 
 			proxyConn, err := s.client.DialTCP(nil, socksaddr, s.dialerTFO)
 			if err != nil {

client/udp.go

@@ -17,13 +17,13 @@ type UDPTunnel struct {
 	multiplexUDP  bool
 	natTimeout    time.Duration
 
-	client       Client
-	conn         *net.UDPConn
-	packetParser PacketParser
+	client        Client
+	conn          *net.UDPConn
+	packetAdapter PacketAdapter
 }
 
 func (s *UDPTunnel) Name() string {
-	return fmt.Sprint("UDP ", s.packetParser.Name(), " service")
+	return fmt.Sprint("UDP ", s.packetAdapter.Name(), " service")
 }
 
 func (s *UDPTunnel) Start() error {
@@ -61,7 +61,7 @@ func (s *UDPTunnel) listen() {
 			continue
 		}
 
-		payloadStart, payloadLength, detachedSocksAddr, err := s.packetParser.ParsePacket(packetBuf, ShadowsocksPacketConnFrontReserve, n)
+		payloadStart, payloadLength, detachedSocksAddr, err := s.packetAdapter.ParsePacket(packetBuf, ShadowsocksPacketConnFrontReserve, n)
 		if err != nil {
 			log.Print(err)
 			continue
@@ -76,7 +76,7 @@ func (s *UDPTunnel) listen() {
 				continue
 			}
 
-			proxyConn = nm.Add(clientAddr, s.conn, oobCache, spc)
+			proxyConn = nm.Add(clientAddr, s.conn, oobCache, spc, s.packetAdapter)
 		} else {
 			proxyConn.oobCache = oobCache
 		}
@@ -95,43 +95,54 @@ func (s *UDPTunnel) Stop() error {
 	return nil
 }
 
-// PacketParser parses an incoming packet and returns payload start index, payload length,
-// a detached socks address (if applicable), or an error.
-//
-// The detached socks address is only returned when the payload does not start with a socks address.
-type PacketParser interface {
+// PacketAdapter translates packets between a local interface and the proxy interface.
+type PacketAdapter interface {
 	Name() string
+
+	// ParsePacket parses an incoming packet and returns payload start index, payload length,
+	// a detached socks address (if applicable), or an error.
+	//
+	// The detached socks address is only returned when the payload does not start with a socks address.
 	ParsePacket(pkt []byte, start, length int) (payloadStart, payloadLength int, detachedSocksAddr []byte, err error)
+
+	// EncapsulatePacket encapsulates the decrypted packet from proxy
+	// into a new form so it's ready to be sent on the local interface.
+	// The encapsulation must not extend beyond the range of the full decrypted packet.
+	EncapsulatePacket(decryptedFullPacket []byte, socksAddrStart, payloadStart, payloadLength int) (pkt []byte, err error)
 }
 
-// SimpleTunnelPacketParser simply relays packets between clientConn and proxyConn.
-type SimpleTunnelPacketParser struct {
+// SimpleTunnelPacketAdapter simply relays packets between clientConn and proxyConn.
+type SimpleTunnelPacketAdapter struct {
 	remoteSocksAddr socks.Addr
 }
 
-func NewSimpleTunnelPacketParser(remoteSocksAddr socks.Addr) *SimpleTunnelPacketParser {
-	return &SimpleTunnelPacketParser{
+func NewSimpleTunnelPacketAdapter(remoteSocksAddr socks.Addr) *SimpleTunnelPacketAdapter {
+	return &SimpleTunnelPacketAdapter{
 		remoteSocksAddr: remoteSocksAddr,
 	}
 }
 
-func (p *SimpleTunnelPacketParser) Name() string {
+func (p *SimpleTunnelPacketAdapter) Name() string {
 	return "simple tunnel"
 }
 
-func (p *SimpleTunnelPacketParser) ParsePacket(_ []byte, start, length int) (payloadStart, payloadLength int, detachedSocksAddr []byte, err error) {
+func (p *SimpleTunnelPacketAdapter) ParsePacket(_ []byte, start, length int) (payloadStart, payloadLength int, detachedSocksAddr []byte, err error) {
 	return start, length, p.remoteSocksAddr, nil
 }
 
-// SimpleSocks5PacketParser is a minimal implementation of SOCKS5 UDP server.
+func (p *SimpleTunnelPacketAdapter) EncapsulatePacket(decryptedFullPacket []byte, _, payloadStart, payloadLength int) (pkt []byte, err error) {
+	return decryptedFullPacket[payloadStart : payloadStart+payloadLength], nil
+}
+
+// SimpleSocks5PacketAdapter is a minimal implementation of SOCKS5 UDP server.
 // It unconditionally accepts SOCKS5 UDP packets, no matter a corresponding UDP association exists or not.
-type SimpleSocks5PacketParser struct{}
+type SimpleSocks5PacketAdapter struct{}
 
-func (p *SimpleSocks5PacketParser) Name() string {
+func (p *SimpleSocks5PacketAdapter) Name() string {
 	return "simple SOCKS5"
 }
 
-func (p *SimpleSocks5PacketParser) ParsePacket(pkt []byte, start, length int) (payloadStart, payloadLength int, detachedSocksAddr []byte, err error) {
+func (p *SimpleSocks5PacketAdapter) ParsePacket(pkt []byte, start, length int) (payloadStart, payloadLength int, detachedSocksAddr []byte, err error) {
 	payloadStart = start + 3
 	if len(pkt) <= payloadStart {
 		return 0, 0, nil, ss.ErrShortPacket
@@ -152,14 +163,24 @@ func (p *SimpleSocks5PacketParser) ParsePacket(pkt []byte, start, length int) (p
 	return
 }
 
-// ShadowsocksNonePacketParser implements the 'none' mode of Shadowsocks.
-type ShadowsocksNonePacketParser struct{}
+func (p *SimpleSocks5PacketAdapter) EncapsulatePacket(decryptedFullPacket []byte, socksAddrStart, payloadStart, payloadLength int) (pkt []byte, err error) {
+	start := socksAddrStart - 3
+	// RSV
+	decryptedFullPacket[start] = 0
+	decryptedFullPacket[start+1] = 0
+	// FRAG
+	decryptedFullPacket[start+2] = 0
+	return decryptedFullPacket[start : payloadStart+payloadLength], nil
+}
+
+// ShadowsocksNonePacketAdapter implements the 'none' mode of Shadowsocks.
+type ShadowsocksNonePacketAdapter struct{}
 
-func (p *ShadowsocksNonePacketParser) Name() string {
+func (p *ShadowsocksNonePacketAdapter) Name() string {
 	return "Shadowsocks none"
 }
 
-func (p *ShadowsocksNonePacketParser) ParsePacket(pkt []byte, start, length int) (payloadStart, payloadLength int, detachedSocksAddr []byte, err error) {
+func (p *ShadowsocksNonePacketAdapter) ParsePacket(pkt []byte, start, length int) (payloadStart, payloadLength int, detachedSocksAddr []byte, err error) {
 	// Validate socks address.
 	_, err = socks.SplitAddr(pkt[start:])
 	if err != nil {
@@ -171,13 +192,17 @@ func (p *ShadowsocksNonePacketParser) ParsePacket(pkt []byte, start, length int)
 	return
 }
 
+func (p *ShadowsocksNonePacketAdapter) EncapsulatePacket(decryptedFullPacket []byte, socksAddrStart, payloadStart, payloadLength int) (pkt []byte, err error) {
+	return decryptedFullPacket[socksAddrStart : payloadStart+payloadLength], nil
+}
+
 func NewUDPSimpleTunnelService(tunnelListenAddress string, tunnelRemoteSocksAddr socks.Addr, multiplexUDP bool, natTimeout time.Duration, client Client) Service {
 	return &UDPTunnel{
 		listenAddress: tunnelListenAddress,
 		multiplexUDP:  multiplexUDP,
 		natTimeout:    natTimeout,
 		client:        client,
-		packetParser:  NewSimpleTunnelPacketParser(tunnelRemoteSocksAddr),
+		packetAdapter: NewSimpleTunnelPacketAdapter(tunnelRemoteSocksAddr),
 	}
 }
 
@@ -187,7 +212,7 @@ func NewUDPSimpleSocks5Service(socks5ListenAddress string, multiplexUDP bool, na
 		multiplexUDP:  multiplexUDP,
 		natTimeout:    natTimeout,
 		client:        client,
-		packetParser:  &SimpleSocks5PacketParser{},
+		packetAdapter: &SimpleSocks5PacketAdapter{},
 	}
 }
 
@@ -197,6 +222,6 @@ func NewUDPShadowsocksNoneService(ssNoneListenAddress string, multiplexUDP bool,
 		multiplexUDP:  multiplexUDP,
 		natTimeout:    natTimeout,
 		client:        client,
-		packetParser:  &ShadowsocksNonePacketParser{},
+		packetAdapter: &ShadowsocksNonePacketAdapter{},
 	}
 }

net/net.go

@@ -110,3 +110,25 @@ type ConnectionError struct {
 func NewConnectionError(status, message string, cause error) *ConnectionError {
 	return &ConnectionError{Status: status, Message: message, Cause: cause}
 }
+
+type addr struct {
+	address string
+	network string
+}
+
+func (a *addr) String() string {
+	return a.address
+}
+
+func (a *addr) Network() string {
+	return a.network
+}
+
+// NewAddr returns a net.Addr that holds an address of the form `host:port` with a domain name or IP as host.
+// Used for SOCKS addressing.
+func NewAddr(address, network string) net.Addr {
+	return &addr{
+		address: address,
+		network: network,
+	}
+}

service/nat.go

@@ -194,7 +194,7 @@ func (c *natconn) timedCopy(ses *session, sm metrics.ShadowsocksMetrics) {
 				ses.pid++
 			default:
 				headerStart = bodyStart - socksAddrLen
-				ss.WriteUDPAddrToSocksAddr(pkt[headerStart:], raddr)
+				socks.WriteUDPAddrAsSocksAddr(pkt[headerStart:], raddr)
 			}
 
 			// `plainTextBuf` concatenates the SOCKS address and body: