UA3F/src/internal/server/socks5/socks5.go
2025-11-06 13:24:03 +08:00

368 lines
10 KiB
Go

package socks5
import (
"encoding/binary"
"errors"
"fmt"
"io"
"net"
"strings"
"time"
"github.com/sirupsen/logrus"
"github.com/sunbk201/ua3f/internal/config"
"github.com/sunbk201/ua3f/internal/rewrite"
"github.com/sunbk201/ua3f/internal/server/utils"
"github.com/sunbk201/ua3f/internal/statistics"
)
// SOCKS5 constants
const (
socksVer5 = 0x05
socksNoAuth = 0x00
socksCmdConn = 0x01
socksCmdUDP = 0x03
socksATYPv4 = 0x01
socksATYDomain = 0x03
socksATYPv6 = 0x04
)
var (
ErrInvalidSocksVersion = errors.New("invalid socks version")
ErrInvalidSocksCmd = errors.New("invalid socks cmd")
)
// Server is a minimal SOCKS5 server that delegates HTTP UA rewriting to Rewriter.
type Server struct {
cfg *config.Config
rw *rewrite.Rewriter
listener net.Listener
}
// New returns a new Server with given config, rewriter, and version string.
func New(cfg *config.Config, rw *rewrite.Rewriter) *Server {
return &Server{
cfg: cfg,
rw: rw,
}
}
// Start begins listening for SOCKS5 clients.
func (s *Server) Start() (err error) {
if s.listener, err = net.Listen("tcp", s.cfg.ListenAddr); err != nil {
return fmt.Errorf("listen failed: %w", err)
}
// Start statistics worker
go statistics.StartRecorder()
var client net.Conn
for {
if client, err = s.listener.Accept(); err != nil {
logrus.Error("Accept failed: ", err)
continue
}
logrus.Debugf("Accept connection from %s", client.RemoteAddr().String())
go s.HandleClient(client)
}
}
// handleClient performs SOCKS5 negotiation and dispatches TCP/UDP handling.
func (s *Server) HandleClient(client net.Conn) {
// Handshake (no auth)
if err := s.socks5Auth(client); err != nil {
_ = client.Close()
return
}
destAddrPort, cmd, err := s.parseSocks5Request(client)
if err != nil {
if cmd == socksCmdUDP {
// UDP Associate
s.handleUDPAssociate(client)
_ = client.Close()
return
}
logrus.Debugf("[%s][%s] ParseSocks5Request failed: %s",
client.RemoteAddr().String(), destAddrPort, err.Error())
_ = client.Close()
return
}
// TCP CONNECT
target, err := s.socks5Connect(client, destAddrPort)
if err != nil {
logrus.Warn("Connect failed: ", err)
_ = client.Close()
return
}
s.ForwardTCP(client, target, destAddrPort)
}
// socks5Auth performs a minimal "no-auth" negotiation.
func (s *Server) socks5Auth(client net.Conn) error {
buf := make([]byte, 256)
// Read VER, NMETHODS
n, err := io.ReadFull(client, buf[:2])
if n != 2 {
if errors.Is(err, io.EOF) {
logrus.Warnf("[%s][Auth] read EOF", client.RemoteAddr().String())
} else {
logrus.Errorf("[%s][Auth] read header: %v", client.RemoteAddr().String(), err)
}
return fmt.Errorf("reading header: %w", err)
}
ver, nMethods := int(buf[0]), int(buf[1])
if ver != socksVer5 {
logrus.Errorf("[%s][Auth] invalid ver", client.RemoteAddr().String())
return ErrInvalidSocksVersion
}
// Read METHODS
n, err = io.ReadFull(client, buf[:nMethods])
if n != nMethods {
logrus.Errorf("[%s][Auth] read methods: %v", client.RemoteAddr().String(), err)
return fmt.Errorf("read methods: %w", err)
}
// Reply: no-auth
n, err = client.Write([]byte{socksVer5, socksNoAuth})
if n != 2 || err != nil {
logrus.Errorf("[%s][Auth] write rsp: %v", client.RemoteAddr().String(), err)
return fmt.Errorf("write rsp: %w", err)
}
return nil
}
// parseSocks5Request reads a single SOCKS5 request. Returns dest, cmd, and error.
func (s *Server) parseSocks5Request(client net.Conn) (string, byte, error) {
buf := make([]byte, 256)
// VER, CMD, RSV, ATYP
if _, err := io.ReadFull(client, buf[:4]); err != nil {
return "", 0, fmt.Errorf("read header: %w", err)
}
ver, cmd, atyp := buf[0], buf[1], buf[3]
if ver != socksVer5 {
return "", cmd, ErrInvalidSocksVersion
}
// UDP associate: let caller handle
if cmd == socksCmdUDP {
return "", socksCmdUDP, errors.New("UDP Associate")
}
if cmd != socksCmdConn {
return "", cmd, ErrInvalidSocksCmd
}
var addr string
switch atyp {
case socksATYPv4:
if _, err := io.ReadFull(client, buf[:4]); err != nil {
return "", cmd, fmt.Errorf("invalid IPv4: %w", err)
}
addr = fmt.Sprintf("%d.%d.%d.%d", buf[0], buf[1], buf[2], buf[3])
case socksATYDomain:
if _, err := io.ReadFull(client, buf[:1]); err != nil {
return "", cmd, fmt.Errorf("invalid hostname(len): %w", err)
}
addrLen := int(buf[0])
if _, err := io.ReadFull(client, buf[:addrLen]); err != nil {
return "", cmd, fmt.Errorf("invalid hostname: %w", err)
}
addr = string(buf[:addrLen])
case socksATYPv6:
return "", cmd, errors.New("IPv6: not supported yet")
default:
return "", cmd, errors.New("invalid atyp")
}
if _, err := io.ReadFull(client, buf[:2]); err != nil {
return "", cmd, fmt.Errorf("read port: %w", err)
}
port := binary.BigEndian.Uint16(buf[:2])
return fmt.Sprintf("%s:%d", addr, port), cmd, nil
}
// socks5Connect dials the target and responds success to the client.
func (s *Server) socks5Connect(client net.Conn, destAddrPort string) (net.Conn, error) {
target, err := utils.Connect(destAddrPort)
if err != nil {
// Reply failure
_, _ = client.Write([]byte{socksVer5, 0x01, 0x00, socksATYPv4, 0, 0, 0, 0, 0, 0})
return nil, fmt.Errorf("dial target %s: %w", destAddrPort, err)
}
// Reply success (bind set to 0.0.0.0:0)
if _, err = client.Write([]byte{socksVer5, 0x00, 0x00, socksATYPv4, 0, 0, 0, 0, 0, 0}); err != nil {
_ = target.Close()
return nil, err
}
return target, nil
}
// ForwardTCP proxies traffic in both directions.
// target->client uses raw copy.
// client->target is processed by the rewriter (or raw if cached).
func (s *Server) ForwardTCP(client, target net.Conn, destAddr string) {
// Server -> Client (raw)
go utils.CopyHalf(client, target)
if s.cfg.DirectForward {
// Client -> Server (raw)
go utils.CopyHalf(target, client)
return
}
// Client -> Server (rewriter)
go utils.ProxyHalf(target, client, s.rw, destAddr)
}
// handleUDPAssociate handles a UDP ASSOCIATE request by creating a UDP relay socket.
// Only IPv4 and domain ATYP are supported (no IPv6).
func (s *Server) handleUDPAssociate(client net.Conn) {
udpServer, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4zero, Port: 0})
if err != nil {
logrus.Errorf("[%s][UDP] ListenUDP failed: %v", client.RemoteAddr().String(), err)
return
}
defer udpServer.Close()
_, portStr, _ := net.SplitHostPort(udpServer.LocalAddr().String())
logrus.Debugf("[%s][UDP] ListenUDP on %s", client.RemoteAddr().String(), portStr)
portInt, _ := net.LookupPort("udp", portStr)
portBytes := make([]byte, 2)
binary.BigEndian.PutUint16(portBytes, uint16(portInt))
// Reply with chosen UDP port (bind addr set to 0.0.0.0)
if _, err = client.Write([]byte{socksVer5, 0x00, 0x00, socksATYPv4, 0, 0, 0, 0, portBytes[0], portBytes[1]}); err != nil {
logrus.Errorf("[%s][UDP] Write rsp failed: %v", client.RemoteAddr().String(), err)
return
}
buf := make([]byte, 65535)
udpPortMap := make(map[string][]byte)
var clientAddr *net.UDPAddr
isDomain := false
for {
_ = udpServer.SetReadDeadline(time.Now().Add(10 * time.Second))
n, fromAddr, err := udpServer.ReadFromUDP(buf)
if err != nil {
if strings.Contains(err.Error(), "i/o timeout") {
logrus.Debugf("[%s][UDP] ReadFromUDP timeout: %v", client.RemoteAddr().String(), err)
if !isAlive(client) {
logrus.Debugf("[%s][UDP] client is not alive", client.RemoteAddr().String())
break
}
} else {
logrus.Errorf("[%s][UDP] ReadFromUDP failed: %v", client.RemoteAddr().String(), err)
}
continue
}
if clientAddr == nil {
clientAddr = fromAddr
}
if clientAddr.IP.Equal(fromAddr.IP) && clientAddr.Port == fromAddr.Port {
// Packet from client -> forward to remote
atyp := buf[3]
var (
targetAddr string
targetPort uint16
payload []byte
header []byte
targetIP net.IP
)
switch atyp {
case socksATYPv4:
isDomain = false
targetAddr = fmt.Sprintf("%d.%d.%d.%d", buf[4], buf[5], buf[6], buf[7])
targetIP = net.ParseIP(targetAddr)
targetPort = binary.BigEndian.Uint16(buf[8:10])
payload = buf[10:n]
header = buf[0:10]
case socksATYDomain:
isDomain = true
addrLen := int(buf[4])
targetAddr = string(buf[5 : 5+addrLen])
targetIPAddr, err := net.ResolveIPAddr("ip", targetAddr)
if err != nil {
logrus.Errorf("[%s][UDP] ResolveIPAddr failed: %v", client.RemoteAddr().String(), err)
break
}
targetIP = targetIPAddr.IP
targetPort = binary.BigEndian.Uint16(buf[5+addrLen : 5+addrLen+2])
payload = buf[5+addrLen+2 : n]
header = buf[0 : 5+addrLen+2]
case socksATYPv6:
logrus.Errorf("[%s][UDP] IPv6: not supported yet", client.RemoteAddr().String())
return
default:
logrus.Errorf("[%s][UDP] invalid atyp", client.RemoteAddr().String())
return
}
remoteAddr := &net.UDPAddr{IP: targetIP, Port: int(targetPort)}
udpPortMap[remoteAddr.String()] = make([]byte, len(header))
copy(udpPortMap[remoteAddr.String()], header)
_ = udpServer.SetWriteDeadline(time.Now().Add(10 * time.Second))
if _, err = udpServer.WriteToUDP(payload, remoteAddr); err != nil {
logrus.Debugf("[%s][UDP] WriteToUDP to remote failed: %v", client.RemoteAddr().String(), err)
continue
}
} else {
// Packet from remote -> forward to client (rebuild header)
header := udpPortMap[fromAddr.String()]
if header == nil {
logrus.Errorf("[%s][UDP] udpPortMap invalid header", client.RemoteAddr().String())
continue
}
// For domain ATYP, preserve original head section size
if isDomain {
header = header[0:4]
}
body := append(header, buf[:n]...)
if _, err = udpServer.WriteToUDP(body, clientAddr); err != nil {
logrus.Debugf("[%s][UDP] WriteToUDP to client failed: %v", client.RemoteAddr().String(), err)
continue
}
}
}
}
// isAlive checks if a connection is still alive using a short read deadline.
func isAlive(conn net.Conn) bool {
one := make([]byte, 1)
_ = conn.SetReadDeadline(time.Now().Add(5 * time.Second))
_, err := conn.Read(one)
if err != nil {
switch {
case errors.Is(err, io.EOF):
logrus.Debugf("[%s] isAlive: EOF", conn.RemoteAddr().String())
return false
case strings.Contains(err.Error(), "use of closed network connection"):
logrus.Debugf("[%s] isAlive: closed", conn.RemoteAddr().String())
return false
case strings.Contains(err.Error(), "i/o timeout"):
logrus.Debugf("[%s] isAlive: timeout", conn.RemoteAddr().String())
return true
default:
logrus.Debugf("[%s] isAlive: %s", conn.RemoteAddr().String(), err.Error())
return false
}
}
_ = conn.SetReadDeadline(time.Time{})
return true
}