Improve: UDP relay refactor (#441)

Co-authored-by: Dreamacro <Dreamacro@vip.qq.com>

component/fakeip/pool.go

@@ -62,12 +62,26 @@ func (p *Pool) LookBack(ip net.IP) (string, bool) {
 	return "", false
 }
 
-// LookupHost return if host in host
-func (p *Pool) LookupHost(host string) bool {
+// LookupHost return if domain in host
+func (p *Pool) LookupHost(domain string) bool {
 	if p.host == nil {
 		return false
 	}
-	return p.host.Search(host) != nil
+	return p.host.Search(domain) != nil
+}
+
+// Exist returns if given ip exists in fake-ip pool
+func (p *Pool) Exist(ip net.IP) bool {
+	p.mux.Lock()
+	defer p.mux.Unlock()
+
+	if ip = ip.To4(); ip == nil {
+		return false
+	}
+
+	n := ipToUint(ip.To4())
+	offset := n - p.min + 1
+	return p.cache.Exist(offset)
 }
 
 // Gateway return gateway ip

component/socks5/socks5.go

@@ -2,6 +2,7 @@ package socks5
 
 import (
 	"bytes"
+	"encoding/binary"
 	"errors"
 	"io"
 	"net"
@@ -62,6 +63,25 @@ func (a Addr) String() string {
 	return net.JoinHostPort(host, port)
 }
 
+// UDPAddr converts a socks5.Addr to *net.UDPAddr
+func (a Addr) UDPAddr() *net.UDPAddr {
+	if len(a) == 0 {
+		return nil
+	}
+	switch a[0] {
+	case AtypIPv4:
+		var ip [net.IPv4len]byte
+		copy(ip[0:], a[1:1+net.IPv4len])
+		return &net.UDPAddr{IP: net.IP(ip[:]), Port: int(binary.BigEndian.Uint16(a[1+net.IPv4len : 1+net.IPv4len+2]))}
+	case AtypIPv6:
+		var ip [net.IPv6len]byte
+		copy(ip[0:], a[1:1+net.IPv6len])
+		return &net.UDPAddr{IP: net.IP(ip[:]), Port: int(binary.BigEndian.Uint16(a[1+net.IPv6len : 1+net.IPv6len+2]))}
+	}
+	// Other Atyp
+	return nil
+}
+
 // SOCKS errors as defined in RFC 1928 section 6.
 const (
 	ErrGeneralFailure       = Error(1)
@@ -338,6 +358,40 @@ func ParseAddr(s string) Addr {
 	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 hostip net.IP
+	var port int
+	if udpaddr, ok := addr.(*net.UDPAddr); ok {
+		hostip = udpaddr.IP
+		port = udpaddr.Port
+	} else if tcpaddr, ok := addr.(*net.TCPAddr); ok {
+		hostip = tcpaddr.IP
+		port = tcpaddr.Port
+	}
+
+	// fallback parse
+	if hostip == nil {
+		return ParseAddr(addr.String())
+	}
+
+	var parsed Addr
+	if ip4 := hostip.To4(); ip4.DefaultMask() != 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:], hostip)
+		binary.BigEndian.PutUint16(parsed[1+net.IPv6len:], uint16(port))
+	}
+	return parsed
+}
+
 // DecodeUDPPacket split `packet` to addr payload, and this function is mutable with `packet`
 func DecodeUDPPacket(packet []byte) (addr Addr, payload []byte, err error) {
 	if len(packet) < 5 {