From f1cb45146f23b515c3f89be9c12300fef03af0e5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E6=B2=B3?= <shixi_liuhe1@staff.sina.com.cn>
Date: Tue, 3 Dec 2019 00:46:30 +0800
Subject: [PATCH] update socks5 udp

---
 client/client.go       | 7 ++++++-
 lib/common/pool.go     | 2 +-
 server/proxy/socks5.go | 9 +++++++--
 3 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/client/client.go b/client/client.go
index 3ae167b..a48907e 100755
--- 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
index 1f7a47e..3ed19dc 100644
--- 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
index 91fb451..603d1e4 100755
--- 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
 			}