internal/runtime/executor/proxy_helpers.go

package executor

import (
	"context"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
	cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
	log "github.com/sirupsen/logrus"
	"golang.org/x/net/proxy"
)

var transportPool sync.Map

// newProxyAwareHTTPClient creates an HTTP client with proper proxy configuration priority:
// 1. Use auth.ProxyURL if configured (highest priority)
// 2. Use cfg.ProxyURL if auth proxy is not configured
// 3. Use RoundTripper from context if neither are configured
func newProxyAwareHTTPClient(ctx context.Context, cfg *config.Config, auth *cliproxyauth.Auth, timeout time.Duration) *http.Client {
	httpClient := &http.Client{}
	if timeout > 0 {
		httpClient.Timeout = timeout
	}

	// Priority 1: Use auth.ProxyURL if configured
	var proxyURL string
	if auth != nil {
		proxyURL = strings.TrimSpace(auth.ProxyURL)
	}

	// Priority 2: Use cfg.ProxyURL if auth proxy is not configured
	if proxyURL == "" && cfg != nil {
		proxyURL = strings.TrimSpace(cfg.ProxyURL)
	}

	keepAliveEnabled := true
	if cfg != nil {
		keepAliveEnabled = cfg.HTTPKeepAlive.Enabled
	}

	// If we have a proxy URL configured, set up the transport.
	if proxyURL != "" {
		if keepAliveEnabled {
			if pooled := getOrCreatePooledTransport(proxyURL, cfg); pooled != nil {
				httpClient.Transport = pooled
				return httpClient
			}
		} else {
			transport := buildProxyTransport(proxyURL, cfg, false)
			if transport != nil {
				httpClient.Transport = transport
				return httpClient
			}
		}
		// If proxy setup failed, log and fall through to context RoundTripper.
		log.Debugf("failed to setup proxy from URL: %s, falling back to context/default transport", proxyURL)
	}

	// Priority 3: Use RoundTripper from context (typically from RoundTripperFor)
	if rt, ok := ctx.Value("cliproxy.roundtripper").(http.RoundTripper); ok && rt != nil {
		httpClient.Transport = rt
		return httpClient
	}

	if keepAliveEnabled {
		httpClient.Transport = getOrCreatePooledTransport("", cfg)
	} else {
		httpClient.Transport = buildProxyTransport("", cfg, false)
	}

	return httpClient
}

func pooledTransportKey(proxyURL string, cfg *config.Config) string {
	idle := 120
	maxIdle := 512
	maxPerHost := 128
	if cfg != nil {
		if cfg.HTTPKeepAlive.IdleConnTimeoutSeconds > 0 {
			idle = cfg.HTTPKeepAlive.IdleConnTimeoutSeconds
		}
		if cfg.HTTPKeepAlive.MaxIdleConns > 0 {
			maxIdle = cfg.HTTPKeepAlive.MaxIdleConns
		}
		if cfg.HTTPKeepAlive.MaxIdleConnsPerHost > 0 {
			maxPerHost = cfg.HTTPKeepAlive.MaxIdleConnsPerHost
		}
	}
	return strings.Join([]string{proxyURL, strconv.Itoa(idle), strconv.Itoa(maxIdle), strconv.Itoa(maxPerHost)}, "|")
}

func getOrCreatePooledTransport(proxyURL string, cfg *config.Config) *http.Transport {
	key := pooledTransportKey(proxyURL, cfg)
	if existing, ok := transportPool.Load(key); ok {
		if tr, ok := existing.(*http.Transport); ok {
			return tr
		}
	}
	created := buildProxyTransport(proxyURL, cfg, true)
	if created == nil {
		return nil
	}
	actual, _ := transportPool.LoadOrStore(key, created)
	if tr, ok := actual.(*http.Transport); ok {
		return tr
	}
	return created
}

// buildProxyTransport creates an HTTP transport configured for the given proxy URL.
// It supports SOCKS5, HTTP, and HTTPS proxy protocols.
func buildProxyTransport(proxyURL string, cfg *config.Config, keepAliveEnabled bool) *http.Transport {
	parsedURL := &url.URL{}
	if proxyURL != "" {
		var errParse error
		parsedURL, errParse = url.Parse(proxyURL)
		if errParse != nil {
			log.Errorf("parse proxy URL failed: %v", errParse)
			return nil
		}
	}
	var transport *http.Transport
	maxIdle := 512
	maxPerHost := 128
	idleTimeout := 120 * time.Second
	if cfg != nil {
		if cfg.HTTPKeepAlive.MaxIdleConns > 0 {
			maxIdle = cfg.HTTPKeepAlive.MaxIdleConns
		}
		if cfg.HTTPKeepAlive.MaxIdleConnsPerHost > 0 {
			maxPerHost = cfg.HTTPKeepAlive.MaxIdleConnsPerHost
		}
		if cfg.HTTPKeepAlive.IdleConnTimeoutSeconds > 0 {
			idleTimeout = time.Duration(cfg.HTTPKeepAlive.IdleConnTimeoutSeconds) * time.Second
		}
	}

	baseTransport := &http.Transport{
		MaxIdleConns:          maxIdle,
		MaxIdleConnsPerHost:   maxPerHost,
		IdleConnTimeout:       idleTimeout,
		TLSHandshakeTimeout:   10 * time.Second,
		ExpectContinueTimeout: 1 * time.Second,
		ForceAttemptHTTP2:     true,
		DisableKeepAlives:     !keepAliveEnabled,
	}

	if parsedURL.Scheme == "" {
		return baseTransport
	}

	// Handle different proxy schemes
	if parsedURL.Scheme == "socks5" {
		// Configure SOCKS5 proxy with optional authentication
		var proxyAuth *proxy.Auth
		if parsedURL.User != nil {
			username := parsedURL.User.Username()
			password, _ := parsedURL.User.Password()
			proxyAuth = &proxy.Auth{User: username, Password: password}
		}
		dialer, errSOCKS5 := proxy.SOCKS5("tcp", parsedURL.Host, proxyAuth, proxy.Direct)
		if errSOCKS5 != nil {
			log.Errorf("create SOCKS5 dialer failed: %v", errSOCKS5)
			return nil
		}
		// Set up a custom transport using the SOCKS5 dialer
		transport = baseTransport.Clone()
		transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
			return dialer.Dial(network, addr)
		}
	} else if parsedURL.Scheme == "http" || parsedURL.Scheme == "https" {
		// Configure HTTP or HTTPS proxy
		transport = baseTransport.Clone()
		transport.Proxy = http.ProxyURL(parsedURL)
	} else {
		log.Errorf("unsupported proxy scheme: %s", parsedURL.Scheme)
		return nil
	}

	return transport
}