update socks5 udp

zhoucheng45 · Dec 2, 2019 · f1cb451 · f1cb451
1 parent 8d27a17
commit f1cb451
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 4 deletions.
diff --git a/client/client.go b/client/client.go
@@ -225,7 +225,12 @@ func (s *TRPClient) handleUdp(serverConn net.Conn) {
 			buf := bytes.Buffer{}
 			dgram := common.NewUDPDatagram(common.NewUDPHeader(0, 0, common.ToSocksAddr(raddr)), b[:n])
 			dgram.Write(&buf)
-			if _, err := serverConn.Write(buf.Bytes()); err != nil {
+			b, err := conn.GetLenBytes(buf.Bytes())
+			if err != nil {
+				logs.Warn("get len bytes error", err.Error())
+				continue
+			}
+			if _, err := serverConn.Write(b); err != nil {
 				logs.Error("write data to remote  error", err.Error())
 				return
 			}

diff --git a/lib/common/pool.go b/lib/common/pool.go
@@ -7,7 +7,7 @@ import (
 
 const PoolSize = 64 * 1024
 const PoolSizeSmall = 100
-const PoolSizeUdp = 1472
+const PoolSizeUdp = 1472 + 200
 const PoolSizeCopy = 32 << 10
 const PoolSizeBuffer = 4096
 const PoolSizeWindow = PoolSizeBuffer - 2 - 4 - 4 - 1

diff --git a/server/proxy/socks5.go b/server/proxy/socks5.go
@@ -247,16 +247,21 @@ func (s *Sock5ModeServer) handleUDP(c net.Conn) {
 	}()
 
 	go func() {
+		var l int32
 		b := common.BufPoolUdp.Get().([]byte)
 		defer common.BufPoolUdp.Put(b)
 		defer c.Close()
 		for {
-			n, err := target.Read(b)
+			if err := binary.Read(target, binary.LittleEndian, &l); err != nil || l >= common.PoolSizeUdp || l <= 0 {
+				logs.Warn("read len bytes error", err.Error())
+				return
+			}
+			binary.Read(target, binary.LittleEndian, b[:l])
 			if err != nil {
 				logs.Warn("read data form client error", err.Error())
 				return
 			}
-			if _, err := reply.WriteTo(b[:n], clientAddr); err != nil {
+			if _, err := reply.WriteTo(b[:l], clientAddr); err != nil {
 				logs.Warn("write data to user ", err.Error())
 				return
 			}