From 9eb6d4ea3990677b4120ab91974301c0d2fa4aa5 Mon Sep 17 00:00:00 2001
From: Mark Pashmfouroush <mark@markpash.me>
Date: Sun, 5 May 2024 17:20:20 +0100
Subject: [PATCH] proxy: some socks5 optimizations

Signed-off-by: Mark Pashmfouroush <mark@markpash.me>
---
 proxy/pkg/mixed/proxy.go   | 24 +++++++++---------------
 proxy/pkg/socks5/common.go | 17 ++++++++---------
 2 files changed, 17 insertions(+), 24 deletions(-)

diff --git a/proxy/pkg/mixed/proxy.go b/proxy/pkg/mixed/proxy.go
index c8c45db7f..28a9b57d3 100644
--- a/proxy/pkg/mixed/proxy.go
+++ b/proxy/pkg/mixed/proxy.go
@@ -63,20 +63,20 @@ type Option func(*Proxy)
 // SwitchConn wraps a net.Conn and a bufio.Reader
 type SwitchConn struct {
 	net.Conn
-	reader *bufio.Reader
+	*bufio.Reader
 }
 
 // NewSwitchConn creates a new SwitchConn
 func NewSwitchConn(conn net.Conn) *SwitchConn {
 	return &SwitchConn{
 		Conn:   conn,
-		reader: bufio.NewReader(conn),
+		Reader: bufio.NewReaderSize(conn, 2048),
 	}
 }
 
 // Read reads data into p, first from the bufio.Reader, then from the net.Conn
 func (c *SwitchConn) Read(p []byte) (n int, err error) {
-	return c.reader.Read(p)
+	return c.Reader.Read(p)
 }
 
 func (p *Proxy) ListenAndServe() error {
@@ -116,6 +116,7 @@ func (p *Proxy) ListenAndServe() error {
 			// Start a new goroutine to handle each connection
 			// This way, the server can handle multiple connections concurrently
 			go func() {
+				defer conn.Close()
 				err := p.handleConnection(conn)
 				if err != nil {
 					p.logger.Error(err.Error()) // Log errors from ServeConn
@@ -129,23 +130,16 @@ func (p *Proxy) handleConnection(conn net.Conn) error {
 	// Create a SwitchConn
 	switchConn := NewSwitchConn(conn)
 
-	// Read one byte to determine the protocol
-	buf := make([]byte, 1)
-	_, err := switchConn.Read(buf)
+	// Peek one byte to determine the protocol
+	buf, err := switchConn.Peek(1)
 	if err != nil {
 		return err
 	}
 
-	// Unread the byte so it's available for the next read
-	err = switchConn.reader.UnreadByte()
-	if err != nil {
-		return err
-	}
-
-	switch {
-	case buf[0] == 5:
+	switch buf[0] {
+	case 5:
 		err = p.socks5Proxy.ServeConn(switchConn)
-	case buf[0] == 4:
+	case 4:
 		err = p.socks4Proxy.ServeConn(switchConn)
 	default:
 		err = p.httpProxy.ServeConn(switchConn)
diff --git a/proxy/pkg/socks5/common.go b/proxy/pkg/socks5/common.go
index bffccf645..014ceee80 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 (
@@ -321,8 +320,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 +330,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 {