From 3cb301dea56967418aef05dca3d532c8816452df Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Wed, 18 Dec 2019 10:27:52 +0800
Subject: [PATCH] improve ssu data transport

---
 ss.go | 112 +++++++++++++++++++++++++++++++---------------------------
 1 file changed, 60 insertions(+), 52 deletions(-)

diff --git a/ss.go b/ss.go
index 9a62a0d..f779d3d 100644
--- a/ss.go
+++ b/ss.go
@@ -455,41 +455,45 @@ func (h *shadowUDPdHandler) Handle(conn net.Conn) {
 
 func (h *shadowUDPdHandler) transportUDP(sc net.Conn, cc net.PacketConn) error {
 	errc := make(chan error, 1)
+
 	go func() {
 		for {
-			b := mPool.Get().([]byte)
-			defer mPool.Put(b)
+			er := func() (err error) {
+				b := lPool.Get().([]byte)
+				defer lPool.Put(b)
 
-			b[0] = 0
-			b[1] = 0
-			b[2] = 0
+				b[0] = 0
+				b[1] = 0
+				b[2] = 0
 
-			n, err := sc.Read(b[3:]) // add rsv and frag fields to make it the standard SOCKS5 UDP datagram
-			if err != nil {
-				// log.Logf("[ssu] %s - %s : %s", sc.RemoteAddr(), sc.LocalAddr(), err)
-				errc <- err
+				// add rsv and frag fields to make it the standard SOCKS5 UDP datagram
+				n, err := sc.Read(b[3:])
+				if err != nil {
+					// log.Logf("[ssu] %s - %s : %s", sc.RemoteAddr(), sc.LocalAddr(), err)
+					return
+				}
+				dgram, err := gosocks5.ReadUDPDatagram(bytes.NewReader(b[:n+3]))
+				if err != nil {
+					log.Logf("[ssu] %s - %s : %s", sc.RemoteAddr(), sc.LocalAddr(), 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 {
+					return
+				}
+				if h.options.Bypass.Contains(addr.String()) {
+					log.Log("[ssu] bypass", addr)
+					return // bypass
+				}
+				_, err = cc.WriteTo(dgram.Data, addr)
 				return
-			}
-			dgram, err := gosocks5.ReadUDPDatagram(bytes.NewReader(b[:n+3]))
-			if err != nil {
-				log.Logf("[ssu] %s - %s : %s", sc.RemoteAddr(), sc.LocalAddr(), err)
-				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 h.options.Bypass.Contains(addr.String()) {
-				log.Log("[ssu] bypass", addr)
-				continue // bypass
-			}
-			if _, err := cc.WriteTo(dgram.Data, addr); err != nil {
-				errc <- err
+			}()
+
+			if er != nil {
+				errc <- er
 				return
 			}
 		}
@@ -497,30 +501,34 @@ func (h *shadowUDPdHandler) transportUDP(sc net.Conn, cc net.PacketConn) error {
 
 	go func() {
 		for {
-			b := mPool.Get().([]byte)
-			defer mPool.Put(b)
+			er := func() (err error) {
+				b := lPool.Get().([]byte)
+				defer lPool.Put(b)
 
-			n, addr, err := cc.ReadFrom(b)
-			if err != nil {
-				errc <- err
+				n, addr, err := cc.ReadFrom(b)
+				if err != nil {
+					return
+				}
+				if Debug {
+					log.Logf("[ssu] %s <<< %s length: %d", sc.RemoteAddr(), addr, n)
+				}
+				if h.options.Bypass.Contains(addr.String()) {
+					log.Log("[ssu] bypass", addr)
+					return // bypass
+				}
+				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)
+					return // ignore invalid datagram
+				}
+				_, err = sc.Write(buf.Bytes()[3:])
 				return
-			}
-			if Debug {
-				log.Logf("[ssu] %s <<< %s length: %d", sc.RemoteAddr(), addr, n)
-			}
-			if h.options.Bypass.Contains(addr.String()) {
-				log.Log("[ssu] bypass", addr)
-				continue // bypass
-			}
-			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
+			}()
+
+			if er != nil {
+				errc <- er
 				return
 			}
 		}