add shadowsocks UDP relay support

gost/examples/bench/srv.go

@@ -41,6 +41,7 @@ func main() {
 	go sshTunnelServer()
 	// go http2Server()
 	go quicServer()
+	go shadowUDPServer()
 	select {}
 }
 
@@ -238,6 +239,16 @@ func quicServer() {
 	log.Fatal(s.Serve(ln, h))
 }
 
+func shadowUDPServer() {
+	s := &gost.Server{}
+	ln, err := gost.ShadowUDPListener(":18338", url.UserPassword("chacha20", "123456"), 30*time.Second)
+	if err != nil {
+		log.Fatal(err)
+	}
+	h := gost.ShadowUDPdHandler()
+	log.Fatal(s.Serve(ln, h))
+}
+
 var (
 	rawCert = []byte(`-----BEGIN CERTIFICATE-----
 MIIC+jCCAeKgAwIBAgIRAMlREhz8Miu1FQozsxbeqyMwDQYJKoZIhvcNAQELBQAw

gost/examples/ssu/ssu.go

+package main
+
+import (
+	"bytes"
+	"log"
+	"net"
+	"strconv"
+
+	"github.com/ginuerzh/gosocks5"
+	ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
+)
+
+func main() {
+	ssuClient()
+}
+
+func ssuClient() {
+	addr, err := net.ResolveUDPAddr("udp", ":18338")
+	if err != nil {
+		log.Fatal(err)
+	}
+	conn, err := net.ListenUDP("udp", nil)
+	if err != nil {
+		log.Fatal(err)
+	}
+	cp, err := ss.NewCipher("chacha20", "123456")
+	if err != nil {
+		log.Fatal(err)
+	}
+	cc := ss.NewSecurePacketConn(conn, cp, false)
+
+	raddr, _ := net.ResolveTCPAddr("udp", ":8080")
+	msg := []byte(`abcdefghijklmnopqrstuvwxyz`)
+	dgram := gosocks5.NewUDPDatagram(gosocks5.NewUDPHeader(0, 0, toSocksAddr(raddr)), msg)
+	buf := bytes.Buffer{}
+	dgram.Write(&buf)
+	if _, err := cc.WriteTo(buf.Bytes()[3:], addr); err != nil {
+		log.Fatal(err)
+	}
+}
+
+func toSocksAddr(addr net.Addr) *gosocks5.Addr {
+	host := "0.0.0.0"
+	port := 0
+	if addr != nil {
+		h, p, _ := net.SplitHostPort(addr.String())
+		host = h
+		port, _ = strconv.Atoi(p)
+	}
+	return &gosocks5.Addr{
+		Type: gosocks5.AddrIPv4,
+		Host: host,
+		Port: uint16(port),
+	}
+}

gost/forward.go

@@ -48,7 +48,6 @@ func (h *tcpDirectForwardHandler) Handle(conn net.Conn) {
 
 type udpDirectForwardHandler struct {
 	raddr   string
-	ttl     time.Duration
 	options *HandlerOptions
 }
 
@@ -168,7 +167,7 @@ func (h *udpRemoteForwardHandler) Handle(conn net.Conn) {
 }
 
 type udpDirectForwardListener struct {
-	ln       *net.UDPConn
+	ln       net.PacketConn
 	conns    map[string]*udpServerConn
 	connChan chan net.Conn
 	errChan  chan error
@@ -199,7 +198,7 @@ func UDPDirectForwardListener(addr string, ttl time.Duration) (Listener, error)
 func (l *udpDirectForwardListener) listenLoop() {
 	for {
 		b := make([]byte, mediumBufferSize)
-		n, raddr, err := l.ln.ReadFromUDP(b)
+		n, raddr, err := l.ln.ReadFrom(b)
 		if err != nil {
 			log.Logf("[udp] peer -> %s : %s", l.Addr(), err)
 			l.ln.Close()
@@ -226,7 +225,7 @@ func (l *udpDirectForwardListener) listenLoop() {
 		select {
 		case conn.rChan <- b[:n]:
 		default:
-			log.Logf("[udp] %s -> %s : write queue is full", raddr, l.Addr())
+			log.Logf("[udp] %s -> %s : read queue is full", raddr, l.Addr())
 		}
 	}
 }

gost/ss.go

 package gost
 
 import (
+	"bytes"
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"io"
 	"net"
@@ -9,6 +11,7 @@ import (
 	"strconv"
 	"time"
 
+	"github.com/ginuerzh/gosocks5"
 	"github.com/go-log/log"
 	ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
 )
@@ -207,3 +210,199 @@ func (h *shadowHandler) getRequest(conn net.Conn) (host string, err error) {
 	host = net.JoinHostPort(host, strconv.Itoa(int(port)))
 	return
 }
+
+type shadowUDPListener struct {
+	ln       net.PacketConn
+	conns    map[string]*udpServerConn
+	connChan chan net.Conn
+	errChan  chan error
+	ttl      time.Duration
+}
+
+// ShadowUDPListener creates a Listener for shadowsocks UDP relay server.
+func ShadowUDPListener(addr string, cipher *url.Userinfo, ttl time.Duration) (Listener, error) {
+	laddr, err := net.ResolveUDPAddr("udp", addr)
+	if err != nil {
+		return nil, err
+	}
+	ln, err := net.ListenUDP("udp", laddr)
+	if err != nil {
+		return nil, err
+	}
+
+	var method, password string
+	if cipher != nil {
+		method = cipher.Username()
+		password, _ = cipher.Password()
+	}
+	cp, err := ss.NewCipher(method, password)
+	if err != nil {
+		ln.Close()
+		return nil, err
+	}
+	l := &udpDirectForwardListener{
+		ln:       ss.NewSecurePacketConn(ln, cp, false),
+		conns:    make(map[string]*udpServerConn),
+		connChan: make(chan net.Conn, 1024),
+		errChan:  make(chan error, 1),
+		ttl:      ttl,
+	}
+	go l.listenLoop()
+	return l, nil
+}
+
+func (l *shadowUDPListener) listenLoop() {
+	for {
+		b := make([]byte, mediumBufferSize)
+		n, raddr, err := l.ln.ReadFrom(b[3:]) // add rsv and frag fields to make it the standard SOCKS5 UDP datagram
+		if err != nil {
+			log.Logf("[ssu] peer -> %s : %s", l.Addr(), err)
+			l.ln.Close()
+			l.errChan <- err
+			close(l.errChan)
+			return
+		}
+		if Debug {
+			log.Logf("[ssu] %s >>> %s : length %d", raddr, l.Addr(), n)
+		}
+
+		b[3] &= ss.AddrMask // remove OTA flag
+		conn, ok := l.conns[raddr.String()]
+		if !ok || conn.Closed() {
+			conn = newUDPServerConn(l.ln, raddr, l.ttl)
+			l.conns[raddr.String()] = conn
+
+			select {
+			case l.connChan <- conn:
+			default:
+				conn.Close()
+				log.Logf("[ssu] %s - %s: connection queue is full", raddr, l.Addr())
+			}
+		}
+
+		select {
+		case conn.rChan <- b[:n+3]: // we keep the addr info so that the handler can identify the destination.
+		default:
+			log.Logf("[ssu] %s -> %s : read queue is full", raddr, l.Addr())
+		}
+	}
+}
+
+func (l *shadowUDPListener) Accept() (conn net.Conn, err error) {
+	var ok bool
+	select {
+	case conn = <-l.connChan:
+	case err, ok = <-l.errChan:
+		if !ok {
+			err = errors.New("accpet on closed listener")
+		}
+	}
+	return
+}
+
+func (l *shadowUDPListener) Addr() net.Addr {
+	return l.ln.LocalAddr()
+}
+
+func (l *shadowUDPListener) Close() error {
+	return l.ln.Close()
+}
+
+type shadowUDPdHandler struct {
+	ttl     time.Duration
+	options *HandlerOptions
+}
+
+// ShadowUDPdHandler creates a server Handler for shadowsocks UDP relay server.
+func ShadowUDPdHandler(opts ...HandlerOption) Handler {
+	h := &udpDirectForwardHandler{
+		options: &HandlerOptions{},
+	}
+	for _, opt := range opts {
+		opt(h.options)
+	}
+	return h
+}
+
+func (h *shadowUDPdHandler) Handle(conn net.Conn) {
+	defer conn.Close()
+
+	var err error
+	var cc net.PacketConn
+	if h.options.Chain.IsEmpty() {
+		cc, err = net.ListenUDP("udp", nil)
+		if err != nil {
+			log.Logf("[udp] %s - : %s", conn.LocalAddr(), err)
+			return
+		}
+	} else {
+		var c net.Conn
+		c, err = getSOCKS5UDPTunnel(h.options.Chain, nil)
+		if err != nil {
+			log.Logf("[udp] %s - : %s", conn.LocalAddr(), err)
+			return
+		}
+		cc = &udpTunnelConn{Conn: c}
+	}
+	defer cc.Close()
+
+	log.Logf("[udp] %s <-> %s", conn.RemoteAddr(), conn.LocalAddr())
+	transportUDP(conn, cc)
+	log.Logf("[udp] %s >-< %s", conn.RemoteAddr(), conn.LocalAddr())
+}
+
+func transportUDP(sc net.Conn, cc net.PacketConn) error {
+	errc := make(chan error, 1)
+	go func() {
+		for {
+			dgram, err := gosocks5.ReadUDPDatagram(sc)
+			if err != nil {
+				errc <- err
+				return
+			}
+			if Debug {
+				log.Logf("[ssu] %s >>> %s length: %d", sc.RemoteAddr(), dgram.Header.Addr.String(), len(dgram.Data))
+			}
+			addr, err := net.ResolveUDPAddr("udp", dgram.Header.Addr.String())
+			if err != nil {
+				errc <- err
+				return
+			}
+			if _, err := cc.WriteTo(dgram.Data, addr); err != nil {
+				errc <- err
+				return
+			}
+		}
+	}()
+
+	go func() {
+		for {
+			b := make([]byte, mediumBufferSize)
+			n, addr, err := cc.ReadFrom(b)
+			if err != nil {
+				errc <- err
+				return
+			}
+			if Debug {
+				log.Logf("[ssu] %s <<< %s length: %d", sc.RemoteAddr(), addr, n)
+			}
+			dgram := gosocks5.NewUDPDatagram(gosocks5.NewUDPHeader(0, 0, toSocksAddr(addr)), b[:n])
+			buf := bytes.Buffer{}
+			dgram.Write(&buf)
+			if buf.Len() < 10 {
+				log.Logf("[ssu] %s <- %s : invalid udp datagram", sc.RemoteAddr(), addr)
+				continue
+			}
+			if _, err := sc.Write(buf.Bytes()[3:]); err != nil {
+				errc <- err
+				return
+			}
+		}
+	}()
+
+	err := <-errc
+	if err != nil && err == io.EOF {
+		err = nil
+	}
+	return err
+}