From b5416c9f4208d0cc5556c502f2bb3eeab527e1a2 Mon Sep 17 00:00:00 2001
From: V2Ray <admin@v2ray.com>
Date: Mon, 7 Sep 2015 22:26:37 +0200
Subject: [PATCH] more on socks proxy

---
 io/socks/socks.go      | 91 ++++++++++++++++++++++++++++++++++++++++--
 io/socks/socks_test.go | 52 ++++++++++++++++++++++++
 2 files changed, 140 insertions(+), 3 deletions(-)

diff --git a/io/socks/socks.go b/io/socks/socks.go
index bb84d73c9..20662d6bc 100644
--- a/io/socks/socks.go
+++ b/io/socks/socks.go
@@ -2,6 +2,7 @@
 package socks
 
 import (
+	"encoding/binary"
 	"fmt"
 	"io"
 )
@@ -13,6 +14,7 @@ const (
 // Authentication request header of Socks5 protocol
 type Socks5AuthenticationRequest struct {
 	version     byte
+	nMethods    byte
 	authMethods [256]byte
 }
 
@@ -33,14 +35,21 @@ func ReadAuthentication(reader io.Reader) (auth Socks5AuthenticationRequest, err
 		return
 	}
 
-	nMethods := buffer[1]
-	if nMethods <= 0 {
+	auth.nMethods = buffer[1]
+	if auth.nMethods <= 0 {
 		err = fmt.Errorf("Zero length of authentication methods")
 		return
 	}
 
-	buffer = make([]byte, nMethods)
+	buffer = make([]byte, auth.nMethods)
 	nBytes, err = reader.Read(buffer)
+	if err != nil {
+		return
+	}
+	if nBytes != int(auth.nMethods) {
+		err = fmt.Errorf("Unmatching number of auth methods, expecting %d, but got %d", auth.nMethods, nBytes)
+		return
+	}
 	copy(auth.authMethods[:nBytes], buffer)
 	return
 }
@@ -64,3 +73,79 @@ func WriteAuthentication(writer io.Writer, response Socks5AuthenticationResponse
 	}
 	return nil
 }
+
+type Socks5Request struct {
+	version  byte
+	command  byte
+	addrType byte
+	ipv4     [4]byte
+	domain   string
+	ipv6     [16]byte
+	port     uint16
+}
+
+func ReadRequest(reader io.Reader) (request *Socks5Request, err error) {
+	request = new(Socks5Request)
+	buffer := make([]byte, 4)
+	nBytes, err := reader.Read(buffer)
+	if err != nil {
+		return
+	}
+	if nBytes < len(buffer) {
+		err = fmt.Errorf("Unable to read request.")
+		return
+	}
+
+	request.version = buffer[0]
+	request.command = buffer[1]
+	// buffer[2] is a reserved field
+	request.addrType = buffer[3]
+	switch request.addrType {
+	case 0x01:
+		nBytes, err = reader.Read(request.ipv4[:])
+		if err != nil {
+			return
+		}
+		if nBytes != 4 {
+			err = fmt.Errorf("Unable to read IPv4 address.")
+			return
+		}
+	case 0x03:
+		buffer = make([]byte, 257)
+		nBytes, err = reader.Read(buffer)
+		if err != nil {
+			return
+		}
+		domainLength := buffer[0]
+		if nBytes != int(domainLength)+1 {
+			err = fmt.Errorf("Unable to read domain")
+			return
+		}
+		request.domain = string(buffer[1 : domainLength+1])
+	case 0x04:
+		nBytes, err = reader.Read(request.ipv6[:])
+		if err != nil {
+			return
+		}
+		if nBytes != 16 {
+			err = fmt.Errorf("Unable to read IPv4 address.")
+			return
+		}
+	default:
+		err = fmt.Errorf("Unexpected address type %d", request.addrType)
+		return
+	}
+
+	buffer = make([]byte, 2)
+	nBytes, err = reader.Read(buffer)
+	if err != nil {
+		return
+	}
+	if nBytes != 2 {
+		err = fmt.Errorf("Unable to read port.")
+		return
+	}
+
+	request.port = binary.BigEndian.Uint16(buffer)
+	return
+}
diff --git a/io/socks/socks_test.go b/io/socks/socks_test.go
index e96c5668f..eea2af905 100644
--- a/io/socks/socks_test.go
+++ b/io/socks/socks_test.go
@@ -1,9 +1,31 @@
 package socks
 
 import (
+	"bytes"
 	"testing"
 )
 
+func TestAuthenticationRequestRead(t *testing.T) {
+	rawRequest := []byte{
+		0x05, // version
+		0x01, // nMethods
+		0x02, // methods
+	}
+	request, err := ReadAuthentication(bytes.NewReader(rawRequest))
+	if err != nil {
+		t.Errorf("Unexpected error %v", err)
+	}
+	if request.version != 0x05 {
+		t.Errorf("Expected version 5, but got %d", request.version)
+	}
+	if request.nMethods != 0x01 {
+		t.Errorf("Expected nMethod 1, but got %d", request.nMethods)
+	}
+	if request.authMethods[0] != 0x02 {
+		t.Errorf("Expected method 2, but got %d", request.authMethods[0])
+	}
+}
+
 func TestAuthenticationResponseToBytes(t *testing.T) {
 	socksVersion := uint8(5)
 	authMethod := uint8(1)
@@ -16,3 +38,33 @@ func TestAuthenticationResponseToBytes(t *testing.T) {
 		t.Errorf("Unexpected Socks auth method %d", bytes[1])
 	}
 }
+
+func TestRequestRead(t *testing.T) {
+	rawRequest := []byte{
+		0x05,                   // version
+		0x01,                   // cmd connect
+		0x00,                   // reserved
+		0x01,                   // ipv4 type
+		0x72, 0x72, 0x72, 0x72, // 114.114.114.114
+		0x00, 0x35, // port 53
+	}
+	request, err := ReadRequest(bytes.NewReader(rawRequest))
+	if err != nil {
+		t.Errorf("Unexpected error %v", err)
+	}
+	if request.version != 0x05 {
+		t.Errorf("Expected version 5, but got %d", request.version)
+	}
+	if request.command != 0x01 {
+		t.Errorf("Expected command 1, but got %d", request.command)
+	}
+	if request.addrType != 0x01 {
+		t.Errorf("Expected addresstype 1, but got %d", request.addrType)
+	}
+	if !bytes.Equal([]byte{0x72, 0x72, 0x72, 0x72}, request.ipv4[:]) {
+		t.Errorf("Expected IPv4 address 114.114.114.114, but got %v", request.ipv4[:])
+	}
+	if request.port != 53 {
+		t.Errorf("Expected port 53, but got %d", request.port)
+	}
+}