Improve(socks5): unify addr parsing

metadata/metadata.go

@@ -1,12 +1,8 @@
 package metadata
 
 import (
-	"bytes"
-	"encoding/binary"
 	"net"
 	"strconv"
-
-	"github.com/xjasonlyu/tun2socks/v2/transport/socks5"
 )
 
 // Metadata contains metadata of transport protocol sessions.
@@ -55,23 +51,6 @@ func (m *Metadata) UDPAddr() *net.UDPAddr {
 	}
 }
 
-func (m *Metadata) SerializeSocksAddr() socks5.Addr {
-	var (
-		buf  [][]byte
-		port [2]byte
-	)
-	binary.BigEndian.PutUint16(port[:], m.DstPort)
-
-	if m.DstIP.To4() != nil /* IPv4 */ {
-		aType := socks5.AtypIPv4
-		buf = [][]byte{{aType}, m.DstIP.To4(), port[:]}
-	} else /* IPv6 */ {
-		aType := socks5.AtypIPv6
-		buf = [][]byte{{aType}, m.DstIP.To16(), port[:]}
-	}
-	return bytes.Join(buf, nil)
-}
-
 // Addr implements the net.Addr interface.
 type Addr struct {
 	metadata *Metadata

proxy/shadowsocks.go

@@ -61,7 +61,7 @@ func (ss *Shadowsocks) DialContext(ctx context.Context, metadata *M.Metadata) (c
 	}
 
 	c = ss.cipher.StreamConn(c)
-	_, err = c.Write(metadata.SerializeSocksAddr())
+	_, err = c.Write(serializeSocksAddr(metadata))
 	return
 }
 
@@ -89,9 +89,9 @@ type ssPacketConn struct {
 func (pc *ssPacketConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
 	var packet []byte
 	if ma, ok := addr.(*M.Addr); ok {
-		packet, err = socks5.EncodeUDPPacket(ma.Metadata().SerializeSocksAddr(), b)
+		packet, err = socks5.EncodeUDPPacket(serializeSocksAddr(ma.Metadata()), b)
 	} else {
-		packet, err = socks5.EncodeUDPPacket(socks5.ParseAddrToSocksAddr(addr), b)
+		packet, err = socks5.EncodeUDPPacket(socks5.ParseAddr(addr), b)
 	}
 
 	if err != nil {

proxy/socks5.go

@@ -59,7 +59,7 @@ func (ss *Socks5) DialContext(ctx context.Context, metadata *M.Metadata) (c net.
 		}
 	}
 
-	_, err = socks5.ClientHandshake(c, metadata.SerializeSocksAddr(), socks5.CmdConnect, user)
+	_, err = socks5.ClientHandshake(c, serializeSocksAddr(metadata), socks5.CmdConnect, user)
 	return
 }
 
@@ -142,9 +142,9 @@ type socksPacketConn struct {
 func (pc *socksPacketConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
 	var packet []byte
 	if ma, ok := addr.(*M.Addr); ok {
-		packet, err = socks5.EncodeUDPPacket(ma.Metadata().SerializeSocksAddr(), b)
+		packet, err = socks5.EncodeUDPPacket(serializeSocksAddr(ma.Metadata()), b)
 	} else {
-		packet, err = socks5.EncodeUDPPacket(socks5.ParseAddrToSocksAddr(addr), b)
+		packet, err = socks5.EncodeUDPPacket(socks5.ParseAddr(addr), b)
 	}
 
 	if err != nil {
@@ -178,3 +178,7 @@ func (pc *socksPacketConn) Close() error {
 	pc.tcpConn.Close()
 	return pc.PacketConn.Close()
 }
+
+func serializeSocksAddr(m *M.Metadata) socks5.Addr {
+	return socks5.SerializeAddr("", m.DstIP, m.DstPort)
+}

transport/socks5/socks5.go

@@ -278,74 +278,55 @@ func SplitAddr(b []byte) Addr {
 	return b[:addrLen]
 }
 
-// ParseAddr parses the address in string s. Returns nil if failed.
-func ParseAddr(s string) Addr {
+// SerializeAddr serializes destination address and port to Addr.
+// If a domain name is provided, AtypDomainName would be used first.
+func SerializeAddr(domainName string, dstIP net.IP, dstPort uint16) Addr {
+	var (
+		buf  [][]byte
+		port [2]byte
+	)
+	binary.BigEndian.PutUint16(port[:], dstPort)
+
+	if domainName != "" /* Domain Name */ {
+		length := len(domainName)
+		buf = [][]byte{{AtypDomainName, uint8(length)}, []byte(domainName), port[:]}
+	} else if dstIP.To4() != nil /* IPv4 */ {
+		buf = [][]byte{{AtypIPv4}, dstIP.To4(), port[:]}
+	} else /* IPv6 */ {
+		buf = [][]byte{{AtypIPv6}, dstIP.To16(), port[:]}
+	}
+	return bytes.Join(buf, nil)
+}
+
+// ParseAddr parses a socks addr from net.Addr.
+// This is a fast path of ParseAddrString(addr.String())
+func ParseAddr(addr net.Addr) Addr {
+	switch v := addr.(type) {
+	case *net.TCPAddr:
+		return SerializeAddr("", v.IP, uint16(v.Port))
+	case *net.UDPAddr:
+		return SerializeAddr("", v.IP, uint16(v.Port))
+	default:
+		return ParseAddrString(addr.String())
+	}
+}
+
+// ParseAddrString parses the address in string s to Addr. Returns nil if failed.
+func ParseAddrString(s string) Addr {
 	host, port, err := net.SplitHostPort(s)
 	if err != nil {
 		return nil
 	}
 
-	var addr Addr
-	if ip := net.ParseIP(host); ip != nil {
-		if ip4 := ip.To4(); ip4 != nil {
-			addr = make([]byte, 1+net.IPv4len+2)
-			addr[0] = AtypIPv4
-			copy(addr[1:], ip4)
-		} else {
-			addr = make([]byte, 1+net.IPv6len+2)
-			addr[0] = AtypIPv6
-			copy(addr[1:], ip)
-		}
-	} else {
-		if len(host) > 255 {
-			return nil
-		}
-		addr = make([]byte, 1+1+len(host)+2)
-		addr[0] = AtypDomainName
-		addr[1] = byte(len(host))
-		copy(addr[2:], host)
-	}
-
-	p, err := strconv.ParseUint(port, 10, 16)
+	dstPort, err := strconv.ParseUint(port, 10, 16)
 	if err != nil {
 		return nil
 	}
-	binary.BigEndian.PutUint16(addr[len(addr)-2:], uint16(p))
 
-	return addr
-}
-
-// ParseAddrToSocksAddr parse a socks addr from net.addr
-// This is a fast path of ParseAddr(addr.String())
-func ParseAddrToSocksAddr(addr net.Addr) Addr {
-	var ip net.IP
-	var port int
-	if udpAddr, ok := addr.(*net.UDPAddr); ok {
-		ip = udpAddr.IP
-		port = udpAddr.Port
-	} else if tcpAddr, ok := addr.(*net.TCPAddr); ok {
-		ip = tcpAddr.IP
-		port = tcpAddr.Port
+	if ip := net.ParseIP(host); ip != nil {
+		return SerializeAddr("", ip, uint16(dstPort))
 	}
-
-	// fallback parse
-	if ip == nil {
-		return ParseAddr(addr.String())
-	}
-
-	var parsed Addr
-	if ip4 := ip.To4(); ip4 != nil {
-		parsed = make([]byte, 1+net.IPv4len+2)
-		parsed[0] = AtypIPv4
-		copy(parsed[1:], ip4)
-		binary.BigEndian.PutUint16(parsed[1+net.IPv4len:], uint16(port))
-	} else {
-		parsed = make([]byte, 1+net.IPv6len+2)
-		parsed[0] = AtypIPv6
-		copy(parsed[1:], ip)
-		binary.BigEndian.PutUint16(parsed[1+net.IPv6len:], uint16(port))
-	}
-	return parsed
+	return SerializeAddr(host, nil, uint16(dstPort))
 }
 
 // DecodeUDPPacket split `packet` to addr payload, and this function is mutable with `packet`