Implemented PT2 authentication mode

2 files changed, 36 insertions, 17 deletions

diff --git a/common/socks5/socks5.go b/common/socks5/socks5.go
index 42b19eb..74e1175 100644
--- a/common/socks5/socks5.go
+++ b/common/socks5/socks5.go
@@ -40,6 +40,7 @@ package socks5
 import (
 	"bufio"
 	"bytes"
+	"encoding/binary"
 	"fmt"
 	"io"
 	"net"
@@ -61,6 +62,7 @@ const (
 
 	authNoneRequired        = 0x00
 	authUsernamePassword    = 0x02
+	authPT2PrivateMethod    = 0x80
 	authNoAcceptableMethods = 0xff
 
 	requestTimeout = 5 * time.Second
@@ -124,7 +126,7 @@ type Request struct {
 // Handshake attempts to handle a incoming client handshake over the provided
 // connection and receive the SOCKS5 request.  The routine handles sending
 // appropriate errors if applicable, but will not close the connection.
-func Handshake(conn net.Conn) (*Request, error) {
+func Handshake(conn net.Conn, needOptions bool) (*Request, error) {
 	// Arm the handshake timeout.
 	var err error
 	if err = conn.SetDeadline(time.Now().Add(requestTimeout)); err != nil {
@@ -144,7 +146,7 @@ func Handshake(conn net.Conn) (*Request, error) {
 
 	// Negotiate the protocol version and authentication method.
 	var method byte
-	if method, err = req.negotiateAuth(); err != nil {
+	if method, err = req.negotiateAuth(needOptions); err != nil {
 		return nil, err
 	}
 
@@ -186,7 +188,7 @@ func (req *Request) Reply(code ReplyCode) error {
 	return req.flushBuffers()
 }
 
-func (req *Request) negotiateAuth() (byte, error) {
+func (req *Request) negotiateAuth(needOptions bool) (byte, error) {
 	// The client sends a version identifier/selection message.
 	//	uint8_t ver (0x05)
 	//  uint8_t nmethods (>= 1).
@@ -209,11 +211,19 @@ func (req *Request) negotiateAuth() (byte, error) {
 	}
 
 	// Pick the best authentication method, prioritizing authenticating
-	// over not if both options are present.
-	if bytes.IndexByte(methods, authUsernamePassword) != -1 {
-		method = authUsernamePassword
-	} else if bytes.IndexByte(methods, authNoneRequired) != -1 {
-		method = authNoneRequired
+	// over not if both options are present and SOCKS header options are needed.
+	if needOptions {
+		if bytes.IndexByte(methods, authPT2PrivateMethod) != -1 {
+			method = authPT2PrivateMethod
+		} else if bytes.IndexByte(methods, authNoneRequired) != -1 {
+			method = authNoneRequired
+		}
+	} else {
+		if bytes.IndexByte(methods, authNoneRequired) != -1 {
+			method = authNoneRequired
+		} else if bytes.IndexByte(methods, authPT2PrivateMethod) != -1 {
+			method = authPT2PrivateMethod
+		}
 	}
 
 	// The server sends a method selection message.
@@ -231,8 +241,8 @@ func (req *Request) authenticate(method byte) error {
 	switch method {
 	case authNoneRequired:
 		// No authentication required.
-	case authUsernamePassword:
-		if err := req.authRFC1929(); err != nil {
+	case authPT2PrivateMethod:
+		if err := req.authPT2(); err != nil {
 			return err
 		}
 	case authNoAcceptableMethods:
@@ -349,6 +359,15 @@ func (req *Request) readByteVerify(descr string, expected byte) error {
 	return nil
 }
 
+func (req *Request) readUint32() (uint32, error) {
+	buffer, err := req.readBytes(4)
+	if err != nil {
+		return 0, err
+	}
+
+	return binary.BigEndian.Uint32(buffer), nil
+}
+
 func (req *Request) readBytes(n int) ([]byte, error) {
 	b := make([]byte, n)
 	if _, err := io.ReadFull(req.rw, b); err != nil {
diff --git a/common/socks5/socks_test.go b/common/socks5/socks_test.go
index 720476f..0e7e1df 100644
--- a/common/socks5/socks_test.go
+++ b/common/socks5/socks_test.go
@@ -91,7 +91,7 @@ func TestAuthInvalidVersion(t *testing.T) {
 
 	// VER = 03, NMETHODS = 01, METHODS = [00]
 	c.writeHex("030100")
-	if _, err := req.negotiateAuth(); err == nil {
+	if _, err := req.negotiateAuth(false); err == nil {
 		t.Error("negotiateAuth(InvalidVersion) succeded")
 	}
 }
@@ -105,7 +105,7 @@ func TestAuthInvalidNMethods(t *testing.T) {
 
 	// VER = 05, NMETHODS = 00
 	c.writeHex("0500")
-	if method, err = req.negotiateAuth(); err != nil {
+	if method, err = req.negotiateAuth(false); err != nil {
 		t.Error("negotiateAuth(No Methods) failed:", err)
 	}
 	if method != authNoAcceptableMethods {
@@ -125,7 +125,7 @@ func TestAuthNoneRequired(t *testing.T) {
 
 	// VER = 05, NMETHODS = 01, METHODS = [00]
 	c.writeHex("050100")
-	if method, err = req.negotiateAuth(); err != nil {
+	if method, err = req.negotiateAuth(false); err != nil {
 		t.Error("negotiateAuth(None) failed:", err)
 	}
 	if method != authNoneRequired {
@@ -145,7 +145,7 @@ func TestAuthUsernamePassword(t *testing.T) {
 
 	// VER = 05, NMETHODS = 01, METHODS = [02]
 	c.writeHex("050102")
-	if method, err = req.negotiateAuth(); err != nil {
+	if method, err = req.negotiateAuth(false); err != nil {
 		t.Error("negotiateAuth(UsernamePassword) failed:", err)
 	}
 	if method != authUsernamePassword {
@@ -166,7 +166,7 @@ func TestAuthBoth(t *testing.T) {
 
 	// VER = 05, NMETHODS = 02, METHODS = [00, 02]
 	c.writeHex("05020002")
-	if method, err = req.negotiateAuth(); err != nil {
+	if method, err = req.negotiateAuth(false); err != nil {
 		t.Error("negotiateAuth(Both) failed:", err)
 	}
 	if method != authUsernamePassword {
@@ -186,7 +186,7 @@ func TestAuthUnsupported(t *testing.T) {
 
 	// VER = 05, NMETHODS = 01, METHODS = [01] (GSSAPI)
 	c.writeHex("050101")
-	if method, err = req.negotiateAuth(); err != nil {
+	if method, err = req.negotiateAuth(false); err != nil {
 		t.Error("negotiateAuth(Unknown) failed:", err)
 	}
 	if method != authNoAcceptableMethods {
@@ -207,7 +207,7 @@ func TestAuthUnsupported2(t *testing.T) {
 
 	// VER = 05, NMETHODS = 03, METHODS = [00,01,02]
 	c.writeHex("0503000102")
-	if method, err = req.negotiateAuth(); err != nil {
+	if method, err = req.negotiateAuth(false); err != nil {
 		t.Error("negotiateAuth(Unknown2) failed:", err)
 	}
 	if method != authUsernamePassword {