Fix(socks5): fully udp associate support (#233)

adapters/outbound/socks5.go

@@ -3,6 +3,8 @@ package adapters
 import (
 	"crypto/tls"
 	"fmt"
+	"io"
+	"io/ioutil"
 	"net"
 	"strconv"
 
@@ -51,24 +53,31 @@ func (ss *Socks5) Dial(metadata *C.Metadata) (net.Conn, error) {
 			Password: ss.pass,
 		}
 	}
-	if err := socks5.ClientHandshake(c, serializesSocksAddr(metadata), socks5.CmdConnect, user); err != nil {
+	if _, err := socks5.ClientHandshake(c, serializesSocksAddr(metadata), socks5.CmdConnect, user); err != nil {
 		return nil, err
 	}
 	return c, nil
 }
 
-func (ss *Socks5) DialUDP(metadata *C.Metadata) (net.PacketConn, net.Addr, error) {
+func (ss *Socks5) DialUDP(metadata *C.Metadata) (_ net.PacketConn, _ net.Addr, err error) {
 	c, err := dialTimeout("tcp", ss.addr, tcpTimeout)
+	if err != nil {
+		err = fmt.Errorf("%s connect error", ss.addr)
+		return
+	}
 
-	if err == nil && ss.tls {
+	if ss.tls {
 		cc := tls.Client(c, ss.tlsConfig)
 		err = cc.Handshake()
 		c = cc
 	}
 
-	if err != nil {
-		return nil, nil, fmt.Errorf("%s connect error", ss.addr)
-	}
+	defer func() {
+		if err != nil {
+			c.Close()
+		}
+	}()
+
 	tcpKeepAlive(c)
 	var user *socks5.User
 	if ss.user != "" {
@@ -78,10 +87,36 @@ func (ss *Socks5) DialUDP(metadata *C.Metadata) (net.PacketConn, net.Addr, error
 		}
 	}
 
-	if err := socks5.ClientHandshake(c, serializesSocksAddr(metadata), socks5.CmdUDPAssociate, user); err != nil {
-		return nil, nil, err
+	bindAddr, err := socks5.ClientHandshake(c, serializesSocksAddr(metadata), socks5.CmdUDPAssociate, user)
+	if err != nil {
+		err = fmt.Errorf("%v client hanshake error", err)
+		return
 	}
-	return &fakeUDPConn{Conn: c}, c.LocalAddr(), nil
+
+	addr, err := net.ResolveUDPAddr("udp", bindAddr.String())
+	if err != nil {
+		return
+	}
+
+	targetAddr, err := net.ResolveUDPAddr("udp", net.JoinHostPort(metadata.String(), metadata.DstPort))
+	if err != nil {
+		return
+	}
+
+	pc, err := net.ListenPacket("udp", "")
+	if err != nil {
+		return
+	}
+
+	go func() {
+		io.Copy(ioutil.Discard, c)
+		c.Close()
+		// A UDP association terminates when the TCP connection that the UDP
+		// ASSOCIATE request arrived on terminates. RFC1928
+		pc.Close()
+	}()
+
+	return &socksUDPConn{PacketConn: pc, rAddr: targetAddr}, addr, nil
 }
 
 func NewSocks5(option Socks5Option) *Socks5 {
@@ -108,3 +143,26 @@ func NewSocks5(option Socks5Option) *Socks5 {
 		tlsConfig:      tlsConfig,
 	}
 }
+
+type socksUDPConn struct {
+	net.PacketConn
+	rAddr net.Addr
+}
+
+func (uc *socksUDPConn) WriteTo(b []byte, addr net.Addr) (n int, err error) {
+	packet, err := socks5.EncodeUDPPacket(uc.rAddr.String(), b)
+	if err != nil {
+		return
+	}
+	return uc.PacketConn.WriteTo(packet, addr)
+}
+
+func (uc *socksUDPConn) ReadFrom(b []byte) (int, net.Addr, error) {
+	n, a, e := uc.PacketConn.ReadFrom(b)
+	addr, payload, err := socks5.DecodeUDPPacket(b)
+	if err != nil {
+		return 0, nil, err
+	}
+	copy(b, payload)
+	return n - len(addr) - 3, a, e
+}