add socks5 udp relay

pkg/handler/socks/v5/udp.go

@@ -0,0 +1,168 @@
+package v5
+
+import (
+	"bytes"
+	"context"
+	"io"
+	"io/ioutil"
+	"net"
+	"time"
+
+	"github.com/go-gost/gosocks5"
+	"github.com/go-gost/gost/pkg/internal/bufpool"
+	"github.com/go-gost/gost/pkg/logger"
+)
+
+func (h *socks5Handler) handleUDP(ctx context.Context, conn net.Conn, req *gosocks5.Request) {
+	h.logger = h.logger.WithFields(map[string]interface{}{
+		"cmd": "udp",
+	})
+
+	relay, err := net.ListenUDP("udp", nil)
+	if err != nil {
+		h.logger.Error(err)
+		reply := gosocks5.NewReply(gosocks5.Failure, nil)
+		reply.Write(conn)
+		if h.logger.IsLevelEnabled(logger.DebugLevel) {
+			h.logger.Debug(reply)
+		}
+		return
+	}
+	defer relay.Close()
+
+	saddr, _ := gosocks5.NewAddr(relay.LocalAddr().String())
+	if saddr == nil {
+		saddr = &gosocks5.Addr{}
+	}
+	saddr.Type = 0
+	saddr.Host, _, _ = net.SplitHostPort(conn.LocalAddr().String()) // replace the IP to the out-going interface's
+	reply := gosocks5.NewReply(gosocks5.Succeeded, saddr)
+	if err := reply.Write(conn); err != nil {
+		h.logger.Error(err)
+		return
+	}
+	if h.logger.IsLevelEnabled(logger.DebugLevel) {
+		h.logger.Debug(reply)
+	}
+
+	h.logger = h.logger.WithFields(map[string]interface{}{
+		"bind": saddr.String(),
+	})
+	h.logger.Infof("bind on %s OK", saddr.String())
+
+	if !h.chain.IsEmpty() {
+
+	}
+
+	peer, err := net.ListenUDP("udp", nil)
+	if err != nil {
+		h.logger.Error(err)
+		return
+	}
+	defer peer.Close()
+
+	go h.transportUDP(relay, peer)
+
+	t := time.Now()
+	h.logger.Infof("%s <-> %s", conn.RemoteAddr(), saddr)
+	io.Copy(ioutil.Discard, conn)
+	h.logger.
+		WithFields(map[string]interface{}{"duration": time.Since(t)}).
+		Infof("%s >-< %s", conn.RemoteAddr(), saddr)
+}
+
+func (h *socks5Handler) transportUDP(relay, peer net.PacketConn) (err error) {
+	const bufSize = 65 * 1024
+	errc := make(chan error, 2)
+
+	var clientAddr net.Addr
+
+	go func() {
+		b := bufpool.Get(bufSize)
+		defer bufpool.Put(b)
+
+		for {
+			n, laddr, err := relay.ReadFrom(b)
+			if err != nil {
+				errc <- err
+				return
+			}
+			if clientAddr == nil {
+				clientAddr = laddr
+			}
+
+			var addr gosocks5.Addr
+			header := gosocks5.UDPHeader{
+				Addr: &addr,
+			}
+			hlen, err := header.ReadFrom(bytes.NewReader(b[:n]))
+			if err != nil {
+				errc <- err
+				return
+			}
+
+			raddr, err := net.ResolveUDPAddr("udp", addr.String())
+			if err != nil {
+				continue // drop silently
+			}
+
+			if h.bypass != nil && h.bypass.Contains(raddr.String()) {
+				h.logger.Warn("bypass: ", raddr)
+				continue // bypass
+			}
+
+			data := b[hlen:n]
+			if _, err := peer.WriteTo(data, raddr); err != nil {
+				errc <- err
+				return
+			}
+			if h.logger.IsLevelEnabled(logger.DebugLevel) {
+				h.logger.Debugf("%s >>> %s: %v data: %d",
+					clientAddr, raddr, b[:hlen], len(data))
+			}
+		}
+	}()
+
+	go func() {
+		b := bufpool.Get(bufSize)
+		defer bufpool.Put(b)
+
+		const dataPos = 1024
+
+		for {
+			n, raddr, err := peer.ReadFrom(b[dataPos:])
+			if err != nil {
+				errc <- err
+				return
+			}
+			if clientAddr == nil {
+				continue
+			}
+			if h.bypass != nil && h.bypass.Contains(raddr.String()) {
+				h.logger.Warn("bypass: ", raddr)
+				continue // bypass
+			}
+
+			socksAddr, _ := gosocks5.NewAddr(raddr.String())
+			if socksAddr == nil {
+				socksAddr = &gosocks5.Addr{}
+			}
+			socksAddr.Type = 0
+			addrLen := socksAddr.Length()
+			socksAddr.Encode(b[dataPos-addrLen : dataPos])
+
+			hlen := addrLen + 3
+			if _, err := relay.WriteTo(b[dataPos-hlen:dataPos+n], clientAddr); err != nil {
+				errc <- err
+				return
+			}
+
+			if h.logger.IsLevelEnabled(logger.DebugLevel) {
+				h.logger.Debugf("%s <<< %s: %v data: %d",
+					clientAddr, raddr, b[dataPos-hlen:dataPos], n)
+			}
+		}
+	}()
+
+	return <-errc
+}