Perf(SOCKS5): optimize memory footprint with authentication (#315)

This commit is contained in:
cty
2023-11-13 06:28:42 +08:00
committed by GitHub
parent 010765138c
commit e86b3b7dc5
2 changed files with 51 additions and 9 deletions
+19 -9
View File
@@ -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")
}
+32
View File
@@ -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")
}