quarkslab-proxyblob/pkg/proxy/socks/socks.go

// Package proxy implements SOCKS5 proxy functionality.
// It provides a complete SOCKS5 implementation following RFC 1928, supporting
// CONNECT and UDP ASSOCIATE (BIND not implemented) commands with NoAuth authentication.
package proxy

import (
	"context"
	"io"
	"net"
	"slices"

	"proxyblob/pkg/protocol"

	"github.com/google/uuid"
)

// SocksHandler implements a SOCKS5 protocol handler.
// It processes authentication, commands, and data transfer between clients
// and remote targets. The handler is safe for concurrent use.
type SocksHandler struct {
	*protocol.BaseHandler
}

// NewSocksHandler creates a SOCKS5 handler with the given connection.
// The connection is used for sending and receiving protocol messages.
func NewSocksHandler(ctx context.Context, conn net.Conn) *SocksHandler {
	handler := &SocksHandler{}
	handler.BaseHandler = protocol.NewBaseHandler(ctx, conn)
	handler.PacketHandler = handler
	return handler
}

// Start begins processing SOCKS5 requests. The address parameter is ignored
// as there are no listeners on the agent side.
func (h *SocksHandler) Start(address string) {
	go h.ReceiveLoop()
}

// Stop gracefully terminates the handler, closing all active connections
// and canceling the context.
func (h *SocksHandler) Stop() {
	h.CloseAllConnections()
	h.Cancel()
}

// OnNew handles new connection requests by initializing it and
// starting the SOCKS5 protocol flow.
func (h *SocksHandler) OnNew(connectionID uuid.UUID, data []byte) byte {
	// Check if the connection already exists
	if _, ok := h.Connections.Load(connectionID); ok {
		return protocol.ErrConnectionExists
	}

	// Create new connection
	conn := protocol.NewConnection(connectionID)
	h.Connections.Store(conn.ID, conn)

	// Create the virtual protocol connection
	conn.ProtocolConn = protocol.NewProtocolConn(h.Ctx, connectionID, h.BaseHandler)
	conn.StartDelivery()

	// Send ACK and process in a goroutine so ReceiveLoop never blocks on aznet writes
	go func() {
		errCode := h.SendConnAck(connectionID)
		if errCode != protocol.ErrNone {
			h.SendClose(connectionID, protocol.ErrConnectionClosed)
			return
		}
		h.processConnection(conn)
	}()
	return protocol.ErrNone
}

// OnAck reports ErrUnexpectedPacket as the agent only accepts incoming
// connections and does not initiate them.
func (h *SocksHandler) OnAck(connectionID uuid.UUID, data []byte) byte {
	return protocol.ErrUnexpectedPacket
}

// OnData processes incoming data for a connection.
// Data is forwarded directly (handled by transport layer).
func (h *SocksHandler) OnData(connectionID uuid.UUID, data []byte) byte {
	value, ok := h.Connections.Load(connectionID)
	if !ok {
		return protocol.ErrConnectionNotFound
	}
	conn := value.(*protocol.Connection)

	// Non-blocking delivery to per-connection goroutine
	if conn.ProtocolConn != nil {
		if !conn.Deliver(data) {
			return protocol.ErrBufferFull
		}
	}
	return protocol.ErrNone
}

// OnClose cleans up resources associated with a connection.
// It is safe to call multiple times.
func (h *SocksHandler) OnClose(connectionID uuid.UUID, errorCode byte) byte {
	value, ok := h.Connections.Load(connectionID)
	if !ok {
		return protocol.ErrNone // Connection already removed, nothing to do
	}
	conn := value.(*protocol.Connection)
	conn.Close()

	h.Connections.Delete(connectionID)
	return protocol.ErrNone
}

// processConnection handles the SOCKS5 protocol flow for a single connection.
// The flow consists of three phases:
//
//  1. Authentication method negotiation
//  2. Command processing (CONNECT, UDP ASSOCIATE)
//  3. Data transfer between client and target
func (h *SocksHandler) processConnection(conn *protocol.Connection) {
	// SOCKS protocol has 3 sequential phases
	errCode := h.handleAuthNegotiation(conn)
	if errCode != protocol.ErrNone {
		h.SendClose(conn.ID, errCode)
		return
	}

	errCode = h.handleCommand(conn)
	if errCode != protocol.ErrNone {
		h.SendClose(conn.ID, errCode)
		return
	}

	errCode = h.handleDataTransfer(conn)
	if errCode != protocol.ErrNone {
		h.SendClose(conn.ID, errCode)
		return
	}
}

// SendError sends a SOCKS5 error reply to the client.
// It maps internal error codes to SOCKS5 reply codes as defined in RFC 1928.
func (h *SocksHandler) SendError(conn *protocol.Connection, errCode byte) {
	// Default to general failure
	socksReplyCode := GeneralFailure

	// Map internal error codes to SOCKS reply codes
	switch errCode {
	case protocol.ErrNone:
		socksReplyCode = Succeeded
	case protocol.ErrNetworkUnreachable:
		socksReplyCode = NetworkUnreachable
	case protocol.ErrHostUnreachable:
		socksReplyCode = HostUnreachable
	case protocol.ErrConnectionRefused:
		socksReplyCode = ConnectionRefused
	case protocol.ErrTTLExpired:
		socksReplyCode = TTLExpired
	case protocol.ErrUnsupportedCommand:
		socksReplyCode = CommandNotSupported
	case protocol.ErrAddressNotSupported:
		socksReplyCode = AddressTypeNotSupported
	case protocol.ErrAuthFailed:
		socksReplyCode = NoAcceptableMethods
	}

	// Build and send error response
	response := []byte{Version5, socksReplyCode, 0x00, IPv4, 0, 0, 0, 0, 0, 0}
	h.SendData(conn.ID, response)
}

// handleAuthNegotiation processes the client's authentication method selection.
// Currently only the NO AUTHENTICATION REQUIRED (0x00) method is supported.
func (h *SocksHandler) handleAuthNegotiation(conn *protocol.Connection) byte {
	// Read SOCKS5 auth packet: [version(1)][nmethods(1)][methods(nmethods)]
	// Use stack allocation for small fixed-size header
	var headerBuf [2]byte
	header := headerBuf[:]
	if _, err := io.ReadFull(conn.ProtocolConn, header); err != nil {
		return protocol.ErrConnectionClosed
	}

	// Check version
	if header[0] != Version5 {
		return protocol.ErrInvalidSocksVersion
	}

	// Read methods
	nmethods := int(header[1])
	if nmethods == 0 {
		return protocol.ErrInvalidPacket
	}

	// Use stack buffer for typical case (most clients send 1-2 methods)
	// Max nmethods is 255, but allocate reasonable stack space
	var methodsBuf [4]byte
	var methods []byte
	if nmethods <= 4 {
		methods = methodsBuf[:nmethods]
	} else {
		methods = make([]byte, nmethods)
	}
	if _, err := io.ReadFull(conn.ProtocolConn, methods); err != nil {
		return protocol.ErrConnectionClosed
	}

	// Currently we only support NoAuth (0x00)
	if !slices.Contains(methods, NoAuth) {
		h.SendError(conn, protocol.ErrAuthFailed)
		return protocol.ErrAuthFailed
	}

	// Send response
	response := []byte{Version5, NoAuth}
	errCode := h.SendData(conn.ID, response)
	if errCode != protocol.ErrNone {
		return protocol.ErrConnectionClosed
	}

	return protocol.ErrNone
}

// handleCommand processes SOCKS5 commands from the client.
// Supported commands are:
//
//   - CONNECT (0x01): Establish TCP/IP stream connection
//   - UDP ASSOCIATE (0x03): UDP relay
//
// Unsupported commands are:
//
//   - BIND (0x02): TCP/IP port binding
func (h *SocksHandler) handleCommand(conn *protocol.Connection) byte {
	// Read SOCKS5 command header: [version(1)][cmd(1)][rsv(1)][atyp(1)]
	// Use stack allocation for fixed-size header
	var headerBuf [4]byte
	header := headerBuf[:]
	if _, err := io.ReadFull(conn.ProtocolConn, header); err != nil {
		return protocol.ErrConnectionClosed
	}

	// Check SOCKS version
	if header[0] != Version5 {
		h.SendError(conn, protocol.ErrInvalidSocksVersion)
		return protocol.ErrInvalidSocksVersion
	}

	cmd := header[1]
	atyp := header[3]

	// Read address based on address type
	var addr []byte
	switch atyp {
	case IPv4:
		// Use stack allocation for IPv4 (6 bytes: 4 IP + 2 port)
		var addrBuf [6]byte
		addr = addrBuf[:]
	case IPv6:
		// Use stack allocation for IPv6 (18 bytes: 16 IP + 2 port)
		var addrBuf [18]byte
		addr = addrBuf[:]
	case Domain:
		// Read domain length first
		var lenBuf [1]byte
		if _, err := io.ReadFull(conn.ProtocolConn, lenBuf[:]); err != nil {
			return protocol.ErrConnectionClosed
		}
		domainLen := int(lenBuf[0])
		// Domain name can be up to 255 bytes, use heap allocation
		addr = make([]byte, 1+domainLen+2) // length + domain + port
		addr[0] = lenBuf[0]
		if _, err := io.ReadFull(conn.ProtocolConn, addr[1:]); err != nil {
			return protocol.ErrConnectionClosed
		}
	default:
		h.SendError(conn, protocol.ErrAddressNotSupported)
		return protocol.ErrAddressNotSupported
	}

	// Read the address and port
	if atyp != Domain {
		if _, err := io.ReadFull(conn.ProtocolConn, addr); err != nil {
			return protocol.ErrConnectionClosed
		}
	}

	// Build complete command data for handlers
	cmdData := append(header, addr...)

	var errCode byte
	switch cmd {
	case Connect:
		errCode = h.handleConnect(conn, cmdData)
	case Bind:
		errCode = h.handleBind(conn, cmdData)
	case UDPAssociate:
		errCode = h.handleUDPAssociate(conn)
	default:
		h.SendError(conn, protocol.ErrUnsupportedCommand)
		return protocol.ErrUnsupportedCommand
	}

	return errCode
}

// handleDataTransfer manages the flow of data between client and target.
// Each command handler implements its own data transfer mechanism.
func (h *SocksHandler) handleDataTransfer(conn *protocol.Connection) byte {
	// Each command handler takes care of data transfer
	// Just wait for connection to be closed
	<-conn.Closed
	return protocol.ErrNone
}