Merge pull request #71 from lixin9311/shadowsocks-org

UDP Associate support for socks5

README.md

@@ -7,7 +7,7 @@ GoDoc at https://godoc.org/github.com/shadowsocks/go-shadowsocks2/
 
 ## Features
 
-- SOCKS5 proxy 
+- SOCKS5 proxy (Including UDP Associate)
 - Support for Netfilter TCP redirect (IPv6 should work but not tested)
 - UDP tunneling (e.g. relay DNS packets)
 - TCP tunneling (e.g. benchmark with iperf3)
@@ -40,8 +40,8 @@ respectively.
 
 ```sh
 go-shadowsocks2 -c ss://AEAD_CHACHA20_POLY1305:your-password@[server_address]:8488 \
-     -verbose -socks :1080 -udptun :8053=8.8.8.8:53,:8054=8.8.4.4:53 \
-                           -tcptun :8053=8.8.8.8:53,:8054=8.8.4.4:53
+     -verbose -socks :1080 -u -udptun :8053=8.8.8.8:53,:8054=8.8.4.4:53 \
+                              -tcptun :8053=8.8.8.8:53,:8054=8.8.4.4:53
 ```
 
 Replace `[server_address]` with the server's public address.

main.go

@@ -15,6 +15,7 @@ import (
 	"time"
 
 	"github.com/shadowsocks/go-shadowsocks2/core"
+	"github.com/shadowsocks/go-shadowsocks2/socks"
 )
 
 var config struct {
@@ -42,6 +43,7 @@ func main() {
 		RedirTCP6 string
 		TCPTun    string
 		UDPTun    string
+		UDPSocks  bool
 	}
 
 	flag.BoolVar(&config.Verbose, "verbose", false, "verbose mode")
@@ -52,6 +54,7 @@ func main() {
 	flag.StringVar(&flags.Server, "s", "", "server listen address or url")
 	flag.StringVar(&flags.Client, "c", "", "client connect address or url")
 	flag.StringVar(&flags.Socks, "socks", "", "(client-only) SOCKS listen address")
+	flag.BoolVar(&flags.UDPSocks, "u", false, "(client-only) Enable UDP support for SOCKS")
 	flag.StringVar(&flags.RedirTCP, "redir", "", "(client-only) redirect TCP from this address")
 	flag.StringVar(&flags.RedirTCP6, "redir6", "", "(client-only) redirect TCP IPv6 from this address")
 	flag.StringVar(&flags.TCPTun, "tcptun", "", "(client-only) TCP tunnel (laddr1=raddr1,laddr2=raddr2,...)")
@@ -113,7 +116,11 @@ func main() {
 		}
 
 		if flags.Socks != "" {
+			socks.UDPEnabled = flags.UDPSocks
 			go socksLocal(flags.Socks, addr, ciph.StreamConn)
+			if flags.UDPSocks {
+				go udpSocksLocal(flags.Socks, addr, ciph.PacketConn)
+			}
 		}
 
 		if flags.RedirTCP != "" {

socks/socks.go

@@ -7,6 +7,9 @@ import (
 	"strconv"
 )
 
+// UDPEnabled is the toggle for UDP support
+var UDPEnabled = false
+
 // SOCKS request commands as defined in RFC 1928 section 4.
 const (
 	CmdConnect      = 1
@@ -38,6 +41,7 @@ const (
 	ErrTTLExpired           = Error(6)
 	ErrCommandNotSupported  = Error(7)
 	ErrAddressNotSupported  = Error(8)
+	InfoUDPAssociate        = Error(9)
 )
 
 // MaxAddrLen is the maximum size of SOCKS address in bytes.
@@ -184,14 +188,26 @@ func Handshake(rw io.ReadWriter) (Addr, error) {
 	if _, err := io.ReadFull(rw, buf[:3]); err != nil {
 		return nil, err
 	}
-	if buf[1] != CmdConnect {
-		return nil, ErrCommandNotSupported
-	}
 	addr, err := readAddr(rw, buf)
 	if err != nil {
 		return nil, err
 	}
-	// write VER REP RSV ATYP BND.ADDR BND.PORT
-	_, err = rw.Write([]byte{5, 0, 0, 1, 0, 0, 0, 0, 0, 0})
-	return addr, err
+	switch buf[1] {
+	case CmdConnect:
+		_, err = rw.Write([]byte{5, 0, 0, 1, 0, 0, 0, 0, 0, 0}) // SOCKS v5, reply succeeded
+	case CmdUDPAssociate:
+		if !UDPEnabled {
+			return nil, ErrCommandNotSupported
+		}
+		listenAddr := ParseAddr(rw.(net.Conn).LocalAddr().String())
+		_, err = rw.Write(append([]byte{5, 0, 0}, listenAddr...)) // SOCKS v5, reply succeeded
+		if err != nil {
+			return nil, ErrCommandNotSupported
+		}
+		err = InfoUDPAssociate
+	default:
+		return nil, ErrCommandNotSupported
+	}
+
+	return addr, err // skip VER, CMD, RSV fields
 }

tcp.go

@@ -43,9 +43,23 @@ func tcpLocal(addr, server string, shadow func(net.Conn) net.Conn, getAddr func(
 		go func() {
 			defer c.Close()
 			c.(*net.TCPConn).SetKeepAlive(true)
-
 			tgt, err := getAddr(c)
 			if err != nil {
+
+				// UDP: keep the connection until disconnect then free the UDP socket
+				if err == socks.InfoUDPAssociate {
+					buf := []byte{}
+					// block here
+					for {
+						_, err := c.Read(buf)
+						if err, ok := err.(net.Error); ok && err.Timeout() {
+							continue
+						}
+						logf("UDP Associate End.")
+						return
+					}
+				}
+
 				logf("failed to get target address: %v", err)
 				return
 			}

udp.go

@@ -10,6 +10,14 @@ import (
 	"github.com/shadowsocks/go-shadowsocks2/socks"
 )
 
+type mode int
+
+const (
+	remoteServer mode = iota
+	relayClient
+	socksClient
+)
+
 const udpBufSize = 64 * 1024
 
 // Listen on laddr for UDP packets, encrypt and send to server to reach target.
@@ -55,7 +63,7 @@ func udpLocal(laddr, server, target string, shadow func(net.PacketConn) net.Pack
 			}
 
 			pc = shadow(pc)
-			nm.Add(raddr, c, pc, false)
+			nm.Add(raddr, c, pc, relayClient)
 		}
 
 		_, err = pc.WriteTo(buf[:len(tgt)+n], srvAddr)
@@ -66,6 +74,51 @@ func udpLocal(laddr, server, target string, shadow func(net.PacketConn) net.Pack
 	}
 }
 
+// Listen on laddr for Socks5 UDP packets, encrypt and send to server to reach target.
+func udpSocksLocal(laddr, server string, shadow func(net.PacketConn) net.PacketConn) {
+	srvAddr, err := net.ResolveUDPAddr("udp", server)
+	if err != nil {
+		logf("UDP server address error: %v", err)
+		return
+	}
+
+	c, err := net.ListenPacket("udp", laddr)
+	if err != nil {
+		logf("UDP local listen error: %v", err)
+		return
+	}
+	defer c.Close()
+
+	nm := newNATmap(config.UDPTimeout)
+	buf := make([]byte, udpBufSize)
+
+	for {
+		n, raddr, err := c.ReadFrom(buf)
+		if err != nil {
+			logf("UDP local read error: %v", err)
+			continue
+		}
+
+		pc := nm.Get(raddr.String())
+		if pc == nil {
+			pc, err = net.ListenPacket("udp", "")
+			if err != nil {
+				logf("UDP local listen error: %v", err)
+				continue
+			}
+			logf("UDP socks tunnel %s <-> %s <-> %s", laddr, server, socks.Addr(buf[3:]))
+			pc = shadow(pc)
+			nm.Add(raddr, c, pc, socksClient)
+		}
+
+		_, err = pc.WriteTo(buf[3:n], srvAddr)
+		if err != nil {
+			logf("UDP local write error: %v", err)
+			continue
+		}
+	}
+}
+
 // Listen on addr for encrypted packets and basically do UDP NAT.
 func udpRemote(addr string, shadow func(net.PacketConn) net.PacketConn) {
 	c, err := net.ListenPacket("udp", addr)
@@ -109,7 +162,7 @@ func udpRemote(addr string, shadow func(net.PacketConn) net.PacketConn) {
 				continue
 			}
 
-			nm.Add(raddr, c, pc, true)
+			nm.Add(raddr, c, pc, remoteServer)
 		}
 
 		_, err = pc.WriteTo(payload, tgtUDPAddr) // accept only UDPAddr despite the signature
@@ -159,19 +212,19 @@ func (m *natmap) Del(key string) net.PacketConn {
 	return nil
 }
 
-func (m *natmap) Add(peer net.Addr, dst, src net.PacketConn, srcIncluded bool) {
+func (m *natmap) Add(peer net.Addr, dst, src net.PacketConn, role mode) {
 	m.Set(peer.String(), src)
 
 	go func() {
-		timedCopy(dst, peer, src, m.timeout, srcIncluded)
+		timedCopy(dst, peer, src, m.timeout, role)
 		if pc := m.Del(peer.String()); pc != nil {
 			pc.Close()
 		}
 	}()
 }
 
 // copy from src to dst at target with read timeout
-func timedCopy(dst net.PacketConn, target net.Addr, src net.PacketConn, timeout time.Duration, srcIncluded bool) error {
+func timedCopy(dst net.PacketConn, target net.Addr, src net.PacketConn, timeout time.Duration, role mode) error {
 	buf := make([]byte, udpBufSize)
 
 	for {
@@ -181,14 +234,17 @@ func timedCopy(dst net.PacketConn, target net.Addr, src net.PacketConn, timeout
 			return err
 		}
 
-		if srcIncluded { // server -> client: add original packet source
+		switch role {
+		case remoteServer: // server -> client: add original packet source
 			srcAddr := socks.ParseAddr(raddr.String())
 			copy(buf[len(srcAddr):], buf[:n])
 			copy(buf, srcAddr)
 			_, err = dst.WriteTo(buf[:len(srcAddr)+n], target)
-		} else { // client -> user: strip original packet source
+		case relayClient: // client -> user: strip original packet source
 			srcAddr := socks.SplitAddr(buf[:n])
 			_, err = dst.WriteTo(buf[len(srcAddr):n], target)
+		case socksClient: // client -> socks5 program: just set RSV and FRAG = 0
+			_, err = dst.WriteTo(append([]byte{0, 0, 0}, buf[:n]...), target)
 		}
 
 		if err != nil {