wormhole-client/internal/proxy/socks5.go

package proxy

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"time"

	"github.com/azoic/wormhole-client/pkg/logger"
)

// SOCKS5Proxy SOCKS5代理客户端
type SOCKS5Proxy struct {
	serverAddr string
	username   string
	password   string
	timeout    time.Duration
}

// NewSOCKS5Proxy 创建SOCKS5代理客户端
func NewSOCKS5Proxy(serverAddr, username, password string, timeout time.Duration) *SOCKS5Proxy {
	return &SOCKS5Proxy{
		serverAddr: serverAddr,
		username:   username,
		password:   password,
		timeout:    timeout,
	}
}

// CreateHTTPProxy 创建HTTP代理服务器
func (p *SOCKS5Proxy) CreateHTTPProxy(localPort int) *http.Server {
	proxyHandler := &httpProxyHandler{
		socks5Proxy: p,
	}

	server := &http.Server{
		Addr:    fmt.Sprintf(":%d", localPort),
		Handler: proxyHandler,
	}

	return server
}

// httpProxyHandler HTTP代理处理器
type httpProxyHandler struct {
	socks5Proxy *SOCKS5Proxy
}

func (h *httpProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	logger.Info("Processing request: %s %s", r.Method, r.URL.String())

	if r.Method == http.MethodConnect {
		h.handleHTTPSProxy(w, r)
	} else {
		h.handleHTTPProxy(w, r)
	}
}

// handleHTTPSProxy 处理HTTPS代理请求 (CONNECT方法)
func (h *httpProxyHandler) handleHTTPSProxy(w http.ResponseWriter, r *http.Request) {
	destConn, err := h.socks5Proxy.DialTCP(r.Host)
	if err != nil {
		logger.Error("Failed to connect via SOCKS5: %v", err)
		http.Error(w, "Bad Gateway", http.StatusBadGateway)
		return
	}
	defer destConn.Close()

	w.WriteHeader(http.StatusOK)
	hijacker, ok := w.(http.Hijacker)
	if !ok {
		logger.Error("Hijacking not supported")
		http.Error(w, "Internal Server Error", http.StatusInternalServerError)
		return
	}

	clientConn, _, err := hijacker.Hijack()
	if err != nil {
		logger.Error("Failed to hijack connection: %v", err)
		return
	}
	defer clientConn.Close()

	// 双向数据转发
	go h.copyData(clientConn, destConn)
	h.copyData(destConn, clientConn)
}

// handleHTTPProxy 处理HTTP代理请求
func (h *httpProxyHandler) handleHTTPProxy(w http.ResponseWriter, r *http.Request) {
	// 通过SOCKS5连接到目标服务器
	destConn, err := h.socks5Proxy.DialTCP(r.Host)
	if err != nil {
		logger.Error("Failed to connect via SOCKS5: %v", err)
		http.Error(w, "Bad Gateway", http.StatusBadGateway)
		return
	}
	defer destConn.Close()

	// 发送HTTP请求
	if err := r.Write(destConn); err != nil {
		logger.Error("Failed to write request: %v", err)
		http.Error(w, "Bad Gateway", http.StatusBadGateway)
		return
	}

	// 读取响应并返回给客户端
	if _, err := io.Copy(w, destConn); err != nil {
		logger.Error("Failed to copy response: %v", err)
	}
}

// DialTCP 通过SOCKS5连接到目标地址
func (p *SOCKS5Proxy) DialTCP(address string) (net.Conn, error) {
	// 连接到SOCKS5代理服务器
	conn, err := net.DialTimeout("tcp", p.serverAddr, p.timeout)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to SOCKS5 server: %v", err)
	}

	// 执行SOCKS5握手
	if err := p.performSOCKS5Handshake(conn, address); err != nil {
		conn.Close()
		return nil, fmt.Errorf("SOCKS5 handshake failed: %v", err)
	}

	logger.Debug("Successfully connected to %s via SOCKS5 proxy", address)
	return conn, nil
}

// performSOCKS5Handshake 执行SOCKS5握手协议
func (p *SOCKS5Proxy) performSOCKS5Handshake(conn net.Conn, targetAddr string) error {
	// 简化的SOCKS5握手实现
	// 实际项目中应该完整实现SOCKS5协议

	// 发送认证方法选择
	authMethods := []byte{0x05, 0x01, 0x02} // 版本5，1个方法，用户名密码认证
	if _, err := conn.Write(authMethods); err != nil {
		return fmt.Errorf("failed to send auth methods: %v", err)
	}

	// 读取服务器响应
	response := make([]byte, 2)
	if _, err := conn.Read(response); err != nil {
		return fmt.Errorf("failed to read auth response: %v", err)
	}

	if response[0] != 0x05 || response[1] != 0x02 {
		return fmt.Errorf("unsupported authentication method")
	}

	// 发送用户名密码
	authData := []byte{0x01} // 子协议版本
	authData = append(authData, byte(len(p.username)))
	authData = append(authData, []byte(p.username)...)
	authData = append(authData, byte(len(p.password)))
	authData = append(authData, []byte(p.password)...)

	if _, err := conn.Write(authData); err != nil {
		return fmt.Errorf("failed to send credentials: %v", err)
	}

	// 读取认证结果
	authResult := make([]byte, 2)
	if _, err := conn.Read(authResult); err != nil {
		return fmt.Errorf("failed to read auth result: %v", err)
	}

	if authResult[1] != 0x00 {
		return fmt.Errorf("authentication failed")
	}

	// 发送连接请求
	host, portStr, err := net.SplitHostPort(targetAddr)
	if err != nil {
		return fmt.Errorf("invalid target address: %v", err)
	}

	// 简化的连接请求（实际实现应该支持域名解析）
	connectReq := []byte{0x05, 0x01, 0x00, 0x03} // 版本，连接命令，保留字段，域名类型
	connectReq = append(connectReq, byte(len(host)))
	connectReq = append(connectReq, []byte(host)...)

	// 添加端口
	portNum := 80 // 默认HTTP端口
	if portStr != "" {
		// 简化处理：如果端口是443则用443，否则用80
		if portStr == "443" {
			portNum = 443
		}
	}
	connectReq = append(connectReq, byte(portNum>>8), byte(portNum&0xFF))

	if _, err := conn.Write(connectReq); err != nil {
		return fmt.Errorf("failed to send connect request: %v", err)
	}

	// 读取连接响应
	connectResp := make([]byte, 10) // 简化的响应读取
	if _, err := conn.Read(connectResp); err != nil {
		return fmt.Errorf("failed to read connect response: %v", err)
	}

	if connectResp[1] != 0x00 {
		return fmt.Errorf("connection failed, status: %d", connectResp[1])
	}

	return nil
}

// copyData 数据复制
func (h *httpProxyHandler) copyData(dst, src net.Conn) {
	defer dst.Close()
	defer src.Close()
	io.Copy(dst, src)
}