From 29dd4bf9d16a461a34c6cc167931a5ca7501412e Mon Sep 17 00:00:00 2001
From: Mark Pashmfouroush <mark@markpash.me>
Date: Fri, 12 Apr 2024 12:47:47 +0100
Subject: [PATCH] proxy: udp tweaks

Signed-off-by: Mark Pashmfouroush <mark@markpash.me>
---
 proxy/pkg/socks5/common.go | 20 +++++++++++---------
 proxy/pkg/socks5/server.go |  3 +++
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/proxy/pkg/socks5/common.go b/proxy/pkg/socks5/common.go
index bffccf645..cc71fabee 100644
--- a/proxy/pkg/socks5/common.go
+++ b/proxy/pkg/socks5/common.go
@@ -6,7 +6,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"math"
 	"net"
 	"strconv"
 	"strings"
@@ -20,7 +19,7 @@ var (
 )
 
 const (
-	maxUdpPacket = math.MaxUint16 - 28
+	maxUdpPacket = 2048
 )
 
 const (
@@ -207,11 +206,14 @@ func readAddr(r io.Reader) (*address, error) {
 	default:
 		return nil, errUnrecognizedAddrType
 	}
+
 	var port [2]byte
 	if _, err := io.ReadFull(r, port[:]); err != nil {
 		return nil, err
 	}
+
 	address.Port = int(binary.BigEndian.Uint16(port[:]))
+
 	return address, nil
 }
 
@@ -321,8 +323,8 @@ func (cc *udpCustomConn) RemoteAddr() net.Addr {
 
 func (cc *udpCustomConn) asyncReadPackets() {
 	go func() {
+		tempBuf := make([]byte, maxUdpPacket)
 		for {
-			tempBuf := make([]byte, maxUdpPacket)
 			n, addr, err := cc.ReadFrom(tempBuf)
 			if err != nil {
 				cc.packetQueue <- &readStruct{
@@ -331,18 +333,18 @@ func (cc *udpCustomConn) asyncReadPackets() {
 				}
 				break
 			}
-			if cc.sourceAddr == nil {
-				cc.sourceAddr = addr
-			}
-			packetData := tempBuf[:n]
-			if len(packetData) < 3 {
+			if n < 3 {
 				cc.packetQueue <- &readStruct{
 					data: nil,
 					err:  err,
 				}
 				break
 			}
-			reader := bytes.NewBuffer(packetData[3:])
+			if cc.sourceAddr == nil {
+				cc.sourceAddr = addr
+			}
+
+			reader := bytes.NewBuffer(tempBuf[3:n])
 			targetAddr, err := readAddr(reader)
 
 			if err != nil {
diff --git a/proxy/pkg/socks5/server.go b/proxy/pkg/socks5/server.go
index 0b53270ff..ff6be811a 100644
--- a/proxy/pkg/socks5/server.go
+++ b/proxy/pkg/socks5/server.go
@@ -396,6 +396,7 @@ func (s *Server) embedHandleAssociate(req *request, udpConn net.PacketConn) erro
 				s.Logger.Debug(err.Error())
 				continue
 			}
+
 			if targetAddr == nil {
 				targetAddr = &net.UDPAddr{
 					IP:   addr.IP,
@@ -403,10 +404,12 @@ func (s *Server) embedHandleAssociate(req *request, udpConn net.PacketConn) erro
 				}
 				wantTarget = targetAddr.String()
 			}
+
 			if addr.String() != wantTarget {
 				s.Logger.Debug("ignore non-target addresses", "address", addr)
 				continue
 			}
+
 			_, err = udpConn.WriteTo(reader.Bytes(), targetAddr)
 			if err != nil {
 				return err