refactor socks handshake

proxy/socks/protocol.go

@@ -1,13 +1,13 @@
 package socks
 
 import (
+	"encoding/binary"
 	"io"
 
 	"v2ray.com/core/common"
 	"v2ray.com/core/common/buf"
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/common/protocol"
-	"v2ray.com/core/common/serial"
 )
 
 const (
@@ -43,28 +43,26 @@ type ServerSession struct {
 	port   net.Port
 }
 
-func (s *ServerSession) Handshake(reader io.Reader, writer io.Writer) (*protocol.RequestHeader, error) {
+func (s *ServerSession) handshake4(cmd byte, reader io.Reader, writer io.Writer) (*protocol.RequestHeader, error) {
-	buffer := buf.New()
-	defer buffer.Release()
-
-	request := new(protocol.RequestHeader)
-
-	if _, err := buffer.ReadFullFrom(reader, 2); err != nil {
-		return nil, newError("insufficient header").Base(err)
-	}
-
-	version := buffer.Byte(0)
-	if version == socks4Version {
 	if s.config.AuthType == AuthType_PASSWORD {
 		writeSocks4Response(writer, socks4RequestRejected, net.AnyIP, net.Port(0)) // nolint: errcheck
 		return nil, newError("socks 4 is not allowed when auth is required.")
 	}
 
+	var port net.Port
+	var address net.Address
+
+	{
+		buffer := buf.StackNew()
 		if _, err := buffer.ReadFullFrom(reader, 6); err != nil {
+			buffer.Release()
 			return nil, newError("insufficient header").Base(err)
 		}
-		port := net.PortFromBytes(buffer.BytesRange(2, 4))
+		port = net.PortFromBytes(buffer.BytesRange(0, 2))
-		address := net.IPAddress(buffer.BytesRange(4, 8))
+		address = net.IPAddress(buffer.BytesRange(2, 6))
+		buffer.Release()
+	}
+
 	if _, err := ReadUntilNull(reader); /* user id */ err != nil {
 		return nil, err
 	}
@@ -76,26 +74,30 @@ func (s *ServerSession) Handshake(reader io.Reader, writer io.Writer) (*protocol
 		address = net.DomainAddress(domain)
 	}
 
-		switch buffer.Byte(1) {
+	switch cmd {
 	case cmdTCPConnect:
-			request.Command = protocol.RequestCommandTCP
+		request := &protocol.RequestHeader{
-			request.Address = address
+			Command: protocol.RequestCommandTCP,
-			request.Port = port
+			Address: address,
-			request.Version = socks4Version
+			Port:    port,
+			Version: socks4Version,
+		}
 		if err := writeSocks4Response(writer, socks4RequestGranted, net.AnyIP, net.Port(0)); err != nil {
 			return nil, err
 		}
 		return request, nil
 	default:
 		writeSocks4Response(writer, socks4RequestRejected, net.AnyIP, net.Port(0)) // nolint: errcheck
-			return nil, newError("unsupported command: ", buffer.Byte(1))
+		return nil, newError("unsupported command: ", cmd)
-		}
 	}
+}
 
-	if version == socks5Version {
+func (s *ServerSession) auth5(nMethod byte, reader io.Reader, writer io.Writer) error {
-		nMethod := int32(buffer.Byte(1))
+	buffer := buf.StackNew()
-		if _, err := buffer.ReadFullFrom(reader, nMethod); err != nil {
+	defer buffer.Release()
-			return nil, newError("failed to read auth methods").Base(err)
+
+	if _, err := buffer.ReadFullFrom(reader, int32(nMethod)); err != nil {
+		return newError("failed to read auth methods").Base(err)
 	}
 
 	var expectedAuth byte = authNotRequired
@@ -103,37 +105,51 @@ func (s *ServerSession) Handshake(reader io.Reader, writer io.Writer) (*protocol
 		expectedAuth = authPassword
 	}
 
-		if !hasAuthMethod(expectedAuth, buffer.BytesRange(2, 2+nMethod)) {
+	if !hasAuthMethod(expectedAuth, buffer.BytesRange(0, int32(nMethod))) {
 		writeSocks5AuthenticationResponse(writer, socks5Version, authNoMatchingMethod) // nolint: errcheck
-			return nil, newError("no matching auth method")
+		return newError("no matching auth method")
 	}
 
 	if err := writeSocks5AuthenticationResponse(writer, socks5Version, expectedAuth); err != nil {
-			return nil, newError("failed to write auth response").Base(err)
+		return newError("failed to write auth response").Base(err)
 	}
 
 	if expectedAuth == authPassword {
 		username, password, err := ReadUsernamePassword(reader)
 		if err != nil {
-				return nil, newError("failed to read username and password for authentication").Base(err)
+			return newError("failed to read username and password for authentication").Base(err)
 		}
 
 		if !s.config.HasAccount(username, password) {
 			writeSocks5AuthenticationResponse(writer, 0x01, 0xFF) // nolint: errcheck
-				return nil, newError("invalid username or password")
+			return newError("invalid username or password")
 		}
 
 		if err := writeSocks5AuthenticationResponse(writer, 0x01, 0x00); err != nil {
-				return nil, newError("failed to write auth response").Base(err)
+			return newError("failed to write auth response").Base(err)
 		}
 	}
 
-		buffer.Clear()
+	return nil
+}
+
+func (s *ServerSession) handshake5(nMethod byte, reader io.Reader, writer io.Writer) (*protocol.RequestHeader, error) {
+	if err := s.auth5(nMethod, reader, writer); err != nil {
+		return nil, err
+	}
+
+	var cmd byte
+	{
+		buffer := buf.StackNew()
 		if _, err := buffer.ReadFullFrom(reader, 3); err != nil {
+			buffer.Release()
 			return nil, newError("failed to read request").Base(err)
 		}
+		cmd = buffer.Byte(1)
+		buffer.Release()
+	}
 
-		cmd := buffer.Byte(1)
+	request := new(protocol.RequestHeader)
 	switch cmd {
 	case cmdTCPConnect, cmdTorResolve, cmdTorResolvePTR:
 		// We don't have a solution for Tor case now. Simply treat it as connect command.
@@ -152,11 +168,9 @@ func (s *ServerSession) Handshake(reader io.Reader, writer io.Writer) (*protocol
 		return nil, newError("unknown command ", cmd)
 	}
 
-		buffer.Clear()
-
 	request.Version = socks5Version
 
-		addr, port, err := addrParser.ReadAddressPort(buffer, reader)
+	addr, port, err := addrParser.ReadAddressPort(nil, reader)
 	if err != nil {
 		return nil, newError("failed to read address").Base(err)
 	}
@@ -178,9 +192,28 @@ func (s *ServerSession) Handshake(reader io.Reader, writer io.Writer) (*protocol
 	}
 
 	return request, nil
+}
+
+// Handshake performs a Socks4/4a/5 handshake.
+func (s *ServerSession) Handshake(reader io.Reader, writer io.Writer) (*protocol.RequestHeader, error) {
+	buffer := buf.StackNew()
+	if _, err := buffer.ReadFullFrom(reader, 2); err != nil {
+		buffer.Release()
+		return nil, newError("insufficient header").Base(err)
 	}
 
+	version := buffer.Byte(0)
+	cmd := buffer.Byte(1)
+	buffer.Release()
+
+	switch version {
+	case socks4Version:
+		return s.handshake4(cmd, reader, writer)
+	case socks5Version:
+		return s.handshake5(cmd, reader, writer)
+	default:
 		return nil, newError("unknown Socks version: ", version)
+	}
 }
 
 // ReadUsernamePassword reads Socks 5 username/password message from the given reader.
@@ -264,12 +297,13 @@ func writeSocks5Response(writer io.Writer, errCode byte, address net.Address, po
 }
 
 func writeSocks4Response(writer io.Writer, errCode byte, address net.Address, port net.Port) error {
-	buffer := buf.New()
+	buffer := buf.StackNew()
 	defer buffer.Release()
 
 	common.Must(buffer.WriteByte(0x00))
 	common.Must(buffer.WriteByte(errCode))
-	common.Must2(serial.WriteUint16(buffer, port.Value()))
+	portBytes := buffer.Extend(2)
+	binary.BigEndian.PutUint16(portBytes, port.Value())
 	common.Must2(buffer.Write(address.IP()))
 	return buf.WriteAllBytes(writer, buffer.Bytes())
 }