socks5,ss: improve udp server

proxy/socks5/server.go

@@ -14,6 +14,8 @@ import (
 	"github.com/nadoo/glider/proxy/protocol/socks"
 )
 
+var nm sync.Map
+
 // NewSocks5Server returns a socks5 proxy server.
 func NewSocks5Server(s string, p proxy.Proxy) (proxy.Server, error) {
 	return NewSocks5(s, nil, p)
@@ -104,56 +106,74 @@ func (s *Socks5) ListenAndServeUDP() {
 
 	log.F("[socks5] listening UDP on %s", s.addr)
 
-	var nm sync.Map
-	buf := make([]byte, proxy.UDPBufSize)
-
 	for {
 		c := NewPktConn(lc, nil, nil, true, nil)
+		buf := pool.GetBuffer(proxy.UDPBufSize)
 
-		n, raddr, err := c.ReadFrom(buf)
+		n, srcAddr, err := c.ReadFrom(buf)
 		if err != nil {
 			log.F("[socks5u] remote read error: %v", err)
 			continue
 		}
 
-		var pc *PktConn
+		var session *Session
-		v, ok := nm.Load(raddr.String())
+		sessionKey := srcAddr.String()
-		if !ok && v == nil {
-			if c.tgtAddr == nil {
-				log.F("[socks5u] can not get target address, not a valid request")
-				continue
-			}
-
-			lpc, dialer, nextHop, err := s.proxy.DialUDP("udp", c.tgtAddr.String())
-			if err != nil {
-				log.F("[socks5u] remote dial error: %v", err)
-				continue
-			}
-
-			pc = NewPktConn(lpc, nextHop, nil, false, nil)
-			nm.Store(raddr.String(), pc)
-
-			go func() {
-				proxy.RelayUDP(c, raddr, pc, 2*time.Minute)
-				pc.Close()
-				nm.Delete(raddr.String())
-			}()
-
-			log.F("[socks5u] %s <-> %s via %s", raddr, c.tgtAddr, dialer.Addr())
 
+		v, ok := nm.Load(sessionKey)
+		if !ok || v == nil {
+			session = newSession(sessionKey, srcAddr, c)
+			nm.Store(sessionKey, session)
+			go s.serveSession(session)
 		} else {
-			pc = v.(*PktConn)
+			session = v.(*Session)
 		}
 
-		_, err = pc.WriteTo(buf[:n], pc.writeAddr)
+		session.msgCh <- buf[:n]
-		if err != nil {
-			log.F("[socks5u] remote write error: %v", err)
-			continue
-		}
-
-		// log.F("[socks5u] %s <-> %s", raddr, c.tgtAddr)
 	}
+}
 
+func (s *Socks5) serveSession(session *Session) {
+	dstC, dialer, writeTo, err := s.proxy.DialUDP("udp", session.srcPC.tgtAddr.String())
+	if err != nil {
+		log.F("[socks5u] remote dial error: %v", err)
+		return
+	}
+	dstPC := NewPktConn(dstC, writeTo, nil, false, nil)
+	defer dstPC.Close()
+
+	go func() {
+		proxy.RelayUDP(session.srcPC, session.src, dstPC, 2*time.Minute)
+		nm.Delete(session.key)
+		close(session.finCh)
+	}()
+
+	log.F("[socks5u] %s <-> %s via %s", session.src, session.srcPC.tgtAddr, dialer.Addr())
+
+	for {
+		select {
+		case p := <-session.msgCh:
+			_, err = dstPC.WriteTo(p, writeTo)
+			if err != nil {
+				log.F("[socks5u] writeTo %s error: %v", writeTo, err)
+			}
+			pool.PutBuffer(p)
+		case <-session.finCh:
+			return
+		}
+	}
+}
+
+// Session is a udp session
+type Session struct {
+	key   string
+	src   net.Addr
+	srcPC *PktConn
+	msgCh chan []byte
+	finCh chan struct{}
+}
+
+func newSession(key string, src net.Addr, srcPC *PktConn) *Session {
+	return &Session{key, src, srcPC, make(chan []byte, 32), make(chan struct{})}
 }
 
 // Handshake fast-tracks SOCKS initialization to get target address to connect.

proxy/ss/server.go

@@ -8,10 +8,13 @@ import (
 	"time"
 
 	"github.com/nadoo/glider/log"
+	"github.com/nadoo/glider/pool"
 	"github.com/nadoo/glider/proxy"
 	"github.com/nadoo/glider/proxy/protocol/socks"
 )
 
+var nm sync.Map
+
 // NewSSServer returns a ss proxy server.
 func NewSSServer(s string, p proxy.Proxy) (proxy.Server, error) {
 	return NewSS(s, nil, p)
@@ -82,7 +85,7 @@ func (s *SS) Serve(c net.Conn) {
 	}
 }
 
-// ListenAndServeUDP serves udp ss requests.
+// ListenAndServe listens on server's addr and serves connections.
 func (s *SS) ListenAndServeUDP() {
 	lc, err := net.ListenPacket("udp", s.addr)
 	if err != nil {
@@ -91,52 +94,75 @@ func (s *SS) ListenAndServeUDP() {
 	}
 	defer lc.Close()
 
-	lc = s.PacketConn(lc)
-
 	log.F("[ss] listening UDP on %s", s.addr)
 
-	var nm sync.Map
+	lc = s.PacketConn(lc)
-	buf := make([]byte, proxy.UDPBufSize)
-
 	for {
 		c := NewPktConn(lc, nil, nil, true)
+		buf := pool.GetBuffer(proxy.UDPBufSize)
 
-		n, raddr, err := c.ReadFrom(buf)
+		n, srcAddr, err := c.ReadFrom(buf)
 		if err != nil {
 			log.F("[ssu] remote read error: %v", err)
 			continue
 		}
 
-		var pc *PktConn
+		var session *Session
-		v, ok := nm.Load(raddr.String())
+		sessionKey := srcAddr.String()
-		if !ok && v == nil {
-			lpc, dialer, nextHop, err := s.proxy.DialUDP("udp", c.tgtAddr.String())
-			if err != nil {
-				log.F("[ssu] remote dial error: %v", err)
-				continue
-			}
-
-			pc = NewPktConn(lpc, nextHop, nil, false)
-			nm.Store(raddr.String(), pc)
-
-			go func() {
-				proxy.RelayUDP(c, raddr, pc, 2*time.Minute)
-				pc.Close()
-				nm.Delete(raddr.String())
-			}()
-
-			log.F("[ssu] %s <-> %s via %s", raddr, c.tgtAddr, dialer.Addr())
 
+		v, ok := nm.Load(sessionKey)
+		if !ok || v == nil {
+			session = newSession(sessionKey, srcAddr, c)
+			nm.Store(sessionKey, session)
+			go s.serveSession(session)
 		} else {
-			pc = v.(*PktConn)
+			session = v.(*Session)
 		}
 
-		_, err = pc.WriteTo(buf[:n], pc.writeAddr)
+		session.msgCh <- buf[:n]
-		if err != nil {
-			log.F("[ssu] remote write error: %v", err)
-			continue
-		}
-
-		// log.F("[ssu] %s <-> %s", raddr, c.tgtAddr)
 	}
 }
+
+func (s *SS) serveSession(session *Session) {
+	dstC, dialer, writeTo, err := s.proxy.DialUDP("udp", session.srcPC.tgtAddr.String())
+	if err != nil {
+		log.F("[ssu] remote dial error: %v", err)
+		return
+	}
+	dstPC := NewPktConn(dstC, writeTo, nil, false)
+	defer dstPC.Close()
+
+	go func() {
+		proxy.RelayUDP(session.srcPC, session.src, dstPC, 2*time.Minute)
+		nm.Delete(session.key)
+		close(session.finCh)
+	}()
+
+	log.F("[ssu] %s <-> %s via %s", session.src, session.srcPC.tgtAddr, dialer.Addr())
+
+	for {
+		select {
+		case p := <-session.msgCh:
+			_, err = dstPC.WriteTo(p, writeTo)
+			if err != nil {
+				log.F("[ssu] writeTo %s error: %v", writeTo, err)
+			}
+			pool.PutBuffer(p)
+		case <-session.finCh:
+			return
+		}
+	}
+}
+
+// Session is a udp session
+type Session struct {
+	key   string
+	src   net.Addr
+	srcPC *PktConn
+	msgCh chan []byte
+	finCh chan struct{}
+}
+
+func newSession(key string, src net.Addr, srcPC *PktConn) *Session {
+	return &Session{key, src, srcPC, make(chan []byte, 32), make(chan struct{})}
+}

proxy/tproxy/server.go

@@ -11,18 +11,17 @@ import (
 	"github.com/nadoo/glider/proxy"
 )
 
-// TProxy struct.
+var nm sync.Map
-type TProxy struct {
-	proxy proxy.Proxy
-	addr  string
-}
 
 func init() {
 	proxy.RegisterServer("tproxy", NewTProxyServer)
 }
 
-// nat mapping
+// TProxy struct.
-var nm sync.Map
+type TProxy struct {
+	proxy proxy.Proxy
+	addr  string
+}
 
 // NewTProxy returns a tproxy.
 func NewTProxy(s string, p proxy.Proxy) (*TProxy, error) {
@@ -138,7 +137,6 @@ func (s *TProxy) serveSession(session *Session) {
 			return
 		}
 	}
-
 }
 
 // Session is a udp session

proxy/udp/udp.go

@@ -12,19 +12,19 @@ import (
 	"github.com/nadoo/glider/proxy"
 )
 
-// UDP struct.
+var nm sync.Map
-type UDP struct {
-	addr   string
-	dialer proxy.Dialer
-	proxy  proxy.Proxy
-}
 
 func init() {
 	proxy.RegisterDialer("udp", NewUDPDialer)
 	proxy.RegisterServer("udp", NewUDPServer)
 }
 
-var nm sync.Map
+// UDP struct.
+type UDP struct {
+	addr   string
+	dialer proxy.Dialer
+	proxy  proxy.Proxy
+}
 
 // NewUDP returns a udp struct.
 func NewUDP(s string, d proxy.Dialer, p proxy.Proxy) (*UDP, error) {
@@ -57,7 +57,7 @@ func NewUDPServer(s string, p proxy.Proxy) (proxy.Server, error) {
 func (s *UDP) ListenAndServe() {
 	c, err := net.ListenPacket("udp", s.addr)
 	if err != nil {
-		log.F("[udp] failed to listen on %s: %v", s.addr, err)
+		log.F("[udp] failed to listen on UDP %s: %v", s.addr, err)
 		return
 	}
 	defer c.Close()

proxy/unix/server.go

@@ -12,12 +12,12 @@ import (
 	"github.com/nadoo/glider/proxy"
 )
 
+var nm sync.Map
+
 func init() {
 	proxy.RegisterServer("unix", NewUnixServer)
 }
 
-var nm sync.Map
-
 // NewUnixServer returns a unix domain socket server.
 func NewUnixServer(s string, p proxy.Proxy) (proxy.Server, error) {
 	schemes := strings.SplitN(s, ",", 2)