proxy/pkg/socks5/server.go

package socks5

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"log"
	"log/slog"
	"net"
	"os"
	"proxy-server/pkg/utils"
	"strconv"
)

const (
	SocksVersion = byte(5)
)

type Config struct {
	Name string

	Host string
	Port uint16

	// 认证方法
	AuthMethods []Authenticator

	// 域名解析
	Resolver NameResolver

	// 自定义认证规则
	Rules RuleSet

	// 地址重写
	Rewriter AddressRewriter

	// 用于 bind 和 associate
	BindIP net.IP

	// Logger
	Logger *log.Logger

	// 自定义连接流程
	Dial func(network, addr string) (net.Conn, error)
}

type Server struct {
	config *Config
	Name   string
	Port   uint16
	Conn   chan ProxyData
}

// New 创建服务器
func New(conf *Config) (*Server, error) {
	if len(conf.AuthMethods) == 0 {
		return nil, ConfigError("认证方法不能为空")
	}

	if conf.Resolver == nil {
		conf.Resolver = DNSResolver{}
	}

	if conf.Rules == nil {
		conf.Rules = PermitAll()
	}

	if conf.Logger == nil {
		conf.Logger = log.New(os.Stdout, "", log.LstdFlags)
	}

	if conf.Dial == nil {
		conf.Dial = func(network, addr string) (net.Conn, error) {
			return net.Dial(network, addr)
		}
	}

	return &Server{
		config: conf,
		Name:   conf.Name,
		Port:   conf.Port,
		Conn:   make(chan ProxyData, 100),
	}, nil
}

// Run 监听端口
func (server *Server) Run() error {
	host := server.config.Host
	port := server.config.Port
	addr := net.JoinHostPort(host, strconv.Itoa(int(port)))

	slog.Info("启动 socks5 代理服务")

	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	defer closeListener(listener)

	slog.Info("代理服务已启动，正在监听端口 " + addr)

	for {
		conn, err := listener.Accept()
		if err != nil {
			slog.Error("客户端连接失败", err)
			continue
		}

		go func() {
			err := server.serve(conn)
			if err != nil {
				slog.Error("连接异常退出", err)
			}
		}()
	}
}

// serve 建立连接
func (server *Server) serve(conn net.Conn) error {
	slog.Info("收到来自" + conn.RemoteAddr().String() + "的连接")

	reader := bufio.NewReader(conn)

	// 认证
	slog.Debug("开始认证流程")
	authContext, err := server.authenticate(reader, conn)
	if err != nil {
		conn.Close()
		slog.Error("认证失败", err)
		return err
	} else {
		slog.Debug("认证完成")
	}

	// 处理连接请求
	slog.Debug("处理连接请求")
	request, err := server.request(reader, conn)
	if err != nil {
		slog.Error("连接请求处理失败", err)
		return err
	} else {
		slog.Debug("连接请求处理完成")
	}

	request.AuthContext = authContext
	client, ok := conn.RemoteAddr().(*net.TCPAddr)
	if !ok {
		return fmt.Errorf("获取客户端地址失败")
	}

	request.RemoteAddr = &AddrSpec{
		IP:   client.IP,
		Port: client.Port,
	}

	// 处理请求
	slog.Debug("开始代理流量")
	err = server.handle(request, conn)
	if err != nil {
		return err
	}

	return nil
}

// checkVersion 检查客户端版本
func checkVersion(reader io.Reader) error {
	version, err := utils.ReadByte(reader)
	if err != nil {
		return err
	}

	slog.Debug("客户端请求版本", "version", version)

	if version != SocksVersion {
		return errors.New("客户端版本不兼容")
	}

	return nil
}

// closeListener 关闭监听并处理可能的错误
func closeListener(listener net.Listener) {
	err := listener.Close()
	if err != nil {
		slog.Info("结束监听端口")
	}
}

// closeConnection 关闭连接并处理可能的错误
func closeConnection(conn net.Conn) {
	err := conn.Close()
	if err != nil {
		slog.Error("连接异常关闭", err)
	} else {
		slog.Info("已关闭来自" + conn.RemoteAddr().String() + "的连接")
	}
}