ProxyPool/internal/tester/http_tester.go

package tester

import (
	"context"
	"crypto/tls"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"golang.org/x/net/proxy"

	"proxyrotator/internal/model"
)

// HTTPTester HTTP 代理测试器
type HTTPTester struct {
	maxBodySize int64
}

// NewHTTPTester 创建测试器
func NewHTTPTester() *HTTPTester {
	return &HTTPTester{
		maxBodySize: 1024 * 1024, // 1MB
	}
}

// TestOne 测试单个代理
func (t *HTTPTester) TestOne(ctx context.Context, p model.Proxy, spec model.TestSpec) model.TestResult {
	result := model.TestResult{
		ProxyID:   p.ID,
		CheckedAt: time.Now(),
	}

	start := time.Now()

	// 创建 HTTP 客户端
	client, err := t.createClient(p, spec.Timeout)
	if err != nil {
		result.ErrorText = err.Error()
		result.LatencyMs = time.Since(start).Milliseconds()
		return result
	}

	// 创建请求
	method := spec.Method
	if method == "" {
		method = "GET"
	}

	req, err := http.NewRequestWithContext(ctx, method, spec.URL, nil)
	if err != nil {
		result.ErrorText = fmt.Sprintf("create request failed: %v", err)
		result.LatencyMs = time.Since(start).Milliseconds()
		return result
	}

	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36")

	// 发起请求
	resp, err := client.Do(req)
	if err != nil {
		result.ErrorText = err.Error()
		result.LatencyMs = time.Since(start).Milliseconds()
		return result
	}
	defer resp.Body.Close()

	result.LatencyMs = time.Since(start).Milliseconds()

	// 检查状态码
	if len(spec.ExpectStatus) > 0 {
		found := false
		for _, s := range spec.ExpectStatus {
			if resp.StatusCode == s {
				found = true
				break
			}
		}
		if !found {
			result.ErrorText = fmt.Sprintf("unexpected status: %d", resp.StatusCode)
			return result
		}
	}

	// 检查响应体关键字
	if spec.ExpectContains != "" {
		body, err := io.ReadAll(io.LimitReader(resp.Body, t.maxBodySize))
		if err != nil {
			result.ErrorText = fmt.Sprintf("read body failed: %v", err)
			return result
		}
		if !strings.Contains(string(body), spec.ExpectContains) {
			result.ErrorText = "expected content not found"
			return result
		}
	}

	result.OK = true
	return result
}

// TestBatch 并发测试多个代理
func (t *HTTPTester) TestBatch(ctx context.Context, proxies []model.Proxy, spec model.TestSpec, concurrency int) []model.TestResult {
	if concurrency <= 0 {
		concurrency = 10
	}

	jobs := make(chan model.Proxy, len(proxies))
	results := make(chan model.TestResult, len(proxies))

	// 启动 worker
	var wg sync.WaitGroup
	for i := 0; i < concurrency; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for p := range jobs {
				select {
				case <-ctx.Done():
					results <- model.TestResult{
						ProxyID:   p.ID,
						ErrorText: "context cancelled",
						CheckedAt: time.Now(),
					}
				default:
					results <- t.TestOne(ctx, p, spec)
				}
			}
		}()
	}

	// 发送任务
	go func() {
		for _, p := range proxies {
			jobs <- p
		}
		close(jobs)
	}()

	// 等待完成
	go func() {
		wg.Wait()
		close(results)
	}()

	// 收集结果
	out := make([]model.TestResult, 0, len(proxies))
	for r := range results {
		out = append(out, r)
	}

	return out
}

// createClient 根据代理类型创建 HTTP 客户端
func (t *HTTPTester) createClient(p model.Proxy, timeout time.Duration) (*http.Client, error) {
	if timeout <= 0 {
		timeout = 10 * time.Second
	}

	transport := &http.Transport{
		TLSClientConfig: &tls.Config{
			InsecureSkipVerify: true,
		},
		DisableKeepAlives:     true,
		TLSHandshakeTimeout:   timeout,
		ResponseHeaderTimeout: timeout,
		ExpectContinueTimeout: 1 * time.Second,
		// 代理连接设置
		ProxyConnectHeader: http.Header{},
	}

	switch p.Protocol {
	case model.ProtoHTTP, model.ProtoHTTPS:
		proxyURL := t.buildProxyURL(p)
		transport.Proxy = http.ProxyURL(proxyURL)

	case model.ProtoSOCKS5:
		dialer, err := t.createSOCKS5Dialer(p, timeout)
		if err != nil {
			return nil, err
		}
		transport.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
			return dialer.Dial(network, addr)
		}

	default:
		return nil, fmt.Errorf("unsupported protocol: %s", p.Protocol)
	}

	return &http.Client{
		Transport: transport,
		Timeout:   timeout,
		// 不自动跟随重定向，让我们检查原始响应
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			if len(via) >= 10 {
				return fmt.Errorf("too many redirects")
			}
			return nil
		},
	}, nil
}

// buildProxyURL 构建代理 URL
func (t *HTTPTester) buildProxyURL(p model.Proxy) *url.URL {
	scheme := "http"
	if p.Protocol == model.ProtoHTTPS {
		scheme = "https"
	}

	u := &url.URL{
		Scheme: scheme,
		Host:   fmt.Sprintf("%s:%d", p.Host, p.Port),
	}

	if p.Username != "" {
		u.User = url.UserPassword(p.Username, p.Password)
	}

	return u
}

// createSOCKS5Dialer 创建 SOCKS5 拨号器
func (t *HTTPTester) createSOCKS5Dialer(p model.Proxy, timeout time.Duration) (proxy.Dialer, error) {
	addr := fmt.Sprintf("%s:%d", p.Host, p.Port)

	var auth *proxy.Auth
	if p.Username != "" {
		auth = &proxy.Auth{
			User:     p.Username,
			Password: p.Password,
		}
	}

	// 创建基础拨号器带超时
	baseDialer := &net.Dialer{
		Timeout: timeout,
	}

	return proxy.SOCKS5("tcp", addr, auth, baseDialer)
}