Search/proxy.go

package main

import (
	"fmt"
	"net/http"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/proxy"
)

// ProxyConfig holds configuration for a single proxy.
type ProxyConfig struct {
	Address  string
	Username string
	Password string
}

// ProxyClient provides an HTTP client pool for proxies.
type ProxyClient struct {
	clients []*http.Client
	lock    sync.Mutex
	index   int
}

// Package-level proxy clients
var (
	metaProxyClient    *ProxyClient
	crawlerProxyClient *ProxyClient
)

// NewProxyClientPool creates a pool of HTTP clients with proxies.
func NewProxyClientPool(proxies []ProxyConfig, timeout time.Duration) (*ProxyClient, error) {
	if len(proxies) == 0 {
		return nil, fmt.Errorf("no proxies provided")
	}

	clients := make([]*http.Client, len(proxies))

	for i, proxyConfig := range proxies {
		var auth *proxy.Auth
		if proxyConfig.Username != "" || proxyConfig.Password != "" {
			auth = &proxy.Auth{
				User:     proxyConfig.Username,
				Password: proxyConfig.Password,
			}
		}

		dialer, err := proxy.SOCKS5("tcp", proxyConfig.Address, auth, proxy.Direct)
		if err != nil {
			return nil, fmt.Errorf("failed to create SOCKS5 dialer for %s: %w", proxyConfig.Address, err)
		}

		transport := &http.Transport{Dial: dialer.Dial}
		clients[i] = &http.Client{
			Transport: transport,
			Timeout:   timeout,
		}
	}

	return &ProxyClient{clients: clients}, nil
}

// Do sends an HTTP request using the next proxy in the pool.
func (p *ProxyClient) Do(req *http.Request) (*http.Response, error) {
	p.lock.Lock()
	client := p.clients[p.index]
	p.index = (p.index + 1) % len(p.clients)
	p.lock.Unlock()
	return client.Do(req)
}

func (p *ProxyClient) GetProxy() string {
	p.lock.Lock()
	defer p.lock.Unlock()

	if len(p.clients) == 0 {
		return ""
	}

	// Round-robin proxy retrieval
	client := p.clients[p.index]
	p.index = (p.index + 1) % len(p.clients)

	// Assume each client has a proxy string saved
	// Example implementation depends on how your proxies are configured
	proxyTransport, ok := client.Transport.(*http.Transport)
	if ok && proxyTransport.Proxy != nil {
		proxyURL, _ := proxyTransport.Proxy(nil)
		if proxyURL != nil {
			return proxyURL.String()
		}
	}

	return ""
}

// ParseProxies parses the proxy strings in the format ADDRESS:PORT or ADDRESS:PORT:USER:PASSWORD.
func ParseProxies(proxyStrings []string) []ProxyConfig {
	var proxies []ProxyConfig
	for _, proxy := range proxyStrings {
		parts := strings.Split(proxy, ":")
		if len(parts) == 2 { // ADDRESS:PORT
			proxies = append(proxies, ProxyConfig{
				Address: fmt.Sprintf("%s:%s", parts[0], parts[1]),
			})
		} else if len(parts) == 4 { // ADDRESS:PORT:USER:PASSWORD
			proxies = append(proxies, ProxyConfig{
				Address:  fmt.Sprintf("%s:%s", parts[0], parts[1]),
				Username: parts[2],
				Password: parts[3],
			})
		} else {
			fmt.Printf("Invalid proxy format: %s\n", proxy)
		}
	}
	return proxies
}

// InitProxies initializes the proxy clients for Meta and Crawler proxies.
func InitProxies() {
	// Initialize Meta Proxy Client
	if config.MetaProxyEnabled {
		metaProxies := ParseProxies(config.MetaProxies)
		client, err := NewProxyClientPool(metaProxies, 30*time.Second)
		if err != nil {
			if config.MetaProxyStrict {
				panic(fmt.Sprintf("Failed to initialize Meta proxies: %v", err))
			}
			fmt.Printf("Warning: Meta proxy initialization failed: %v\n", err)
		}
		metaProxyClient = client
	}

	// Initialize Crawler Proxy Client
	if config.CrawlerProxyEnabled {
		crawlerProxies := ParseProxies(config.CrawlerProxies)
		client, err := NewProxyClientPool(crawlerProxies, 30*time.Second)
		if err != nil {
			if config.CrawlerProxyStrict {
				panic(fmt.Sprintf("Failed to initialize Crawler proxies: %v", err))
			}
			fmt.Printf("Warning: Crawler proxy initialization failed: %v\n", err)
		}
		crawlerProxyClient = client
	}
}

// func main() {
// 	config := loadConfig()

// 	// Initialize proxies if enabled
// 	if config.CrawlerProxyEnabled || config.MetaProxyEnabled {
// 		InitProxies()
// 	}

// 	// Example usage
// 	if metaProxyClient != nil {
// 		req, _ := http.NewRequest("GET", "https://example.com", nil)
// 		resp, err := metaProxyClient.Do(req)
// 		if err != nil {
// 			fmt.Printf("Error using MetaProxyClient: %v\n", err)
// 		} else {
// 			fmt.Printf("Meta Proxy Response Status: %s\n", resp.Status)
// 			resp.Body.Close()
// 		}
// 	}
// }