From e86b3b7dc56f03993786114c0e075af3a428b490 Mon Sep 17 00:00:00 2001
From: cty <ctychen2216@gmail.com>
Date: Sun, 12 Nov 2023 23:28:42 +0100
Subject: [PATCH] Perf(SOCKS5): optimize memory footprint with authentication
 (#315)

---
 transport/socks5/socks5.go      | 28 +++++++++++++++++++---------
 transport/socks5/socks5_test.go | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 9 deletions(-)
 create mode 100644 transport/socks5/socks5_test.go

diff --git a/transport/socks5/socks5.go b/transport/socks5/socks5.go
index 04ad4e0..1b4c8c9 100644
--- a/transport/socks5/socks5.go
+++ b/transport/socks5/socks5.go
@@ -11,6 +11,15 @@ import (
 	"strconv"
 )
 
+// AuthMethod is the authentication method as defined in RFC 1928 section 3.
+type AuthMethod = uint8
+
+// SOCKS authentication methods as defined in RFC 1928 section 3.
+const (
+	MethodNoAuth   AuthMethod = 0x00
+	MethodUserPass AuthMethod = 0x02
+)
+
 // Version is the protocol version as defined in RFC 1928 section 4.
 const Version = 0x05
 
@@ -162,9 +171,9 @@ func ClientHandshake(rw io.ReadWriter, addr Addr, command Command, user *User) (
 
 	var method uint8
 	if user != nil {
-		method = 0x02 /* USERNAME/PASSWORD */
+		method = MethodUserPass /* USERNAME/PASSWORD */
 	} else {
-		method = 0x00 /* NO AUTHENTICATION REQUIRED */
+		method = MethodNoAuth /* NO AUTHENTICATION REQUIRED */
 	}
 
 	// VER, NMETHODS, METHODS
@@ -181,23 +190,24 @@ func ClientHandshake(rw io.ReadWriter, addr Addr, command Command, user *User) (
 		return nil, errors.New("socks version mismatched")
 	}
 
-	if buf[1] == 0x02 /* USERNAME/PASSWORD */ {
+	if buf[1] == MethodUserPass /* USERNAME/PASSWORD */ {
 		if user == nil {
 			return nil, errors.New("auth required")
 		}
 
+		authMsgLen := 1 + 1 + len(user.Username) + 1 + len(user.Password)
+		if authMsgLen > MaxAuthLen {
+			return nil, errors.New("auth message too long")
+		}
+
 		// password protocol version
-		authMsg := &bytes.Buffer{}
+		authMsg := bytes.NewBuffer(make([]byte, 0, authMsgLen))
 		authMsg.WriteByte(0x01 /* VER */)
 		authMsg.WriteByte(byte(len(user.Username)) /* ULEN */)
 		authMsg.WriteString(user.Username /* UNAME */)
 		authMsg.WriteByte(byte(len(user.Password)) /* PLEN */)
 		authMsg.WriteString(user.Password /* PASSWD */)
 
-		if len(authMsg.Bytes()) > MaxAuthLen {
-			return nil, errors.New("auth message too long")
-		}
-
 		if _, err := rw.Write(authMsg.Bytes()); err != nil {
 			return nil, err
 		}
@@ -210,7 +220,7 @@ func ClientHandshake(rw io.ReadWriter, addr Addr, command Command, user *User) (
 			return nil, errors.New("rejected username/password")
 		}
 
-	} else if buf[1] != 0x00 /* NO AUTHENTICATION REQUIRED */ {
+	} else if buf[1] != MethodNoAuth /* NO AUTHENTICATION REQUIRED */ {
 		return nil, errors.New("unsupported method")
 	}
 
diff --git a/transport/socks5/socks5_test.go b/transport/socks5/socks5_test.go
new file mode 100644
index 0000000..575949e
--- /dev/null
+++ b/transport/socks5/socks5_test.go
@@ -0,0 +1,32 @@
+package socks5
+
+import (
+	"bufio"
+	"bytes"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSocks5ClientHandshake(t *testing.T) {
+	// Mock server responses
+	readBuffer := &bytes.Buffer{}
+	readBuffer.Write([]byte{Version, MethodUserPass})
+	readBuffer.Write([]byte{Version, 0x00 /* STATUS of SUCCESS */})
+	readBuffer.Write([]byte{Version, 0x00 /* STATUS of SUCCESS */, 0x00 /* RSV */})
+	readBuffer.Write([]byte{AtypIPv4, 0x1, 0x2, 0x3, 0x4, 0x0, 0x0 /* IPv4: 1.2.3.4:0 */})
+	reader := bufio.NewReader(bytes.NewReader(readBuffer.Bytes()))
+
+	writeBuffer := &bytes.Buffer{}
+	writer := bufio.NewWriter(writeBuffer)
+
+	io := bufio.NewReadWriter(reader, writer)
+
+	addr, err := ClientHandshake(io, []byte{0x0, 0x0, 0x0, 0x0, 0x0, 0x0}, CmdConnect, &User{
+		Username: "test",
+		Password: "6ab49d8b-a009-44e4-bd53-fbdb48fbe7eb",
+	})
+
+	assert.Nil(t, err, "Failed to perform SOCKS5 client handshake: %v", err)
+	assert.Equal(t, "1.2.3.4:0", addr.String(), "Incorrect address obtained from SOCKS5 client handshake")
+}