diff --git a/README.md b/README.md index 7c34de6..e0de2fc 100644 --- a/README.md +++ b/README.md @@ -115,6 +115,22 @@ func main() { <-feedbackChan } ``` +### Per-Request Proxy + +By default, the proxy configured in `HTTPClientSettings` applies to all requests. You can override it on a per-request basis using `WithProxy`: + +```go +req, err := http.NewRequest("GET", "https://example.com", nil) +if err != nil { + panic(err) +} + +// Use a different proxy for this specific request +req = req.WithContext(warc.WithProxy(req.Context(), "socks5://other-proxy:1080")) +resp, err := client.Do(req) +``` + +This follows the same context-based pattern as `WithFeedbackChannel`. Proxy dialers are cached internally, so reusing the same proxy URL across requests is efficient. ### DNS Resolution and Proxy Behavior diff --git a/client_test.go b/client_test.go index a77b0a5..48f0e09 100644 --- a/client_test.go +++ b/client_test.go @@ -20,6 +20,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "testing" "time" @@ -130,6 +131,40 @@ func newTestImageServer(t testing.TB, st int) *httptest.Server { })) } +// socks5.RuleSet that permits all connections and increments a counter on every request +type countingRuleSet struct { + count atomic.Int64 +} + +func (r *countingRuleSet) Allow(ctx context.Context, req *socks5.Request) (context.Context, bool) { + r.count.Add(1) + return ctx, true +} + +// starts a SOCKS5 proxy on a random port and returns its +// address, a connection counter, and a cleanup function that stops the server. +func startSOCKS5Server(t *testing.T) (addr string, counter *atomic.Int64, cleanup func()) { + t.Helper() + rule := &countingRuleSet{} + proxyServer := socks5.NewServer(socks5.WithRule(rule)) + listener, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("failed to listen for proxy: %v", err) + } + stopChan := make(chan struct{}) + go func() { + defer listener.Close() + go func() { + <-stopChan + listener.Close() + }() + if err := proxyServer.Serve(listener); err != nil && !strings.Contains(err.Error(), "use of closed network connection") { + panic(err) + } + }() + return listener.Addr().String(), &rule.count, func() { close(stopChan) } +} + func (e *errorReadCloser) Read(p []byte) (int, error) { if len(e.data) > 0 && e.readBefore > 0 { // Read up to min(len(p), readBefore, len(data)) @@ -649,46 +684,66 @@ func TestHTTPClientDNSFailure(t *testing.T) { } func TestHTTPClientWithProxy(t *testing.T) { - var ( - rotatorSettings = defaultRotatorSettings(t) - err error - ) + rotatorSettings := defaultRotatorSettings(t) - // init socks5 proxy server - proxyServer := socks5.NewServer() - listener, err := net.Listen("tcp", "127.0.0.1:0") + proxyAddr, proxyCounter, stopProxy := startSOCKS5Server(t) + defer stopProxy() + + server := newTestImageServer(t, http.StatusOK) + defer server.Close() + + httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ + RotatorSettings: rotatorSettings, + Proxy: fmt.Sprintf("socks5://%s", proxyAddr), + }) if err != nil { - t.Fatalf("failed to listen for proxy: %v", err) + t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } + waitForErrors := drainErrChan(t, httpClient.ErrChan) - // Create a channel to signal server stop - stopChan := make(chan struct{}) + req, err := http.NewRequest("GET", server.URL, nil) + if err != nil { + t.Fatal(err) + } - go func() { - defer listener.Close() + resp, err := httpClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() - go func() { - <-stopChan - listener.Close() - }() + io.Copy(io.Discard, resp.Body) - if err := proxyServer.Serve(listener); err != nil && !strings.Contains(err.Error(), "use of closed network connection") { - panic(err) - } - }() + httpClient.Close() + waitForErrors() - proxyAddr := listener.Addr().String() - // Defer sending the stop signal - defer close(stopChan) + if c := proxyCounter.Load(); c != 1 { + t.Fatalf("expected proxy to handle 1 connection, got %d", c) + } + + files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") + if err != nil { + t.Fatal(err) + } + + for _, path := range files { + testFileSingleHashCheck(t, path, "sha1:UIRWL5DFIPQ4MX3D3GFHM2HCVU3TZ6I3", []string{"26872"}, 1, server.URL+"/") + } +} + +func TestHTTPClientWithPerRequestProxy(t *testing.T) { + rotatorSettings := defaultRotatorSettings(t) + + proxyAddr, proxyCounter, stopProxy := startSOCKS5Server(t) + defer stopProxy() - // init test HTTP endpoint server := newTestImageServer(t, http.StatusOK) defer server.Close() - // init the HTTP client responsible for recording HTTP(s) requests / responses + // Client created with NO default proxy httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ RotatorSettings: rotatorSettings, - Proxy: fmt.Sprintf("socks5://%s", proxyAddr)}) + }) if err != nil { t.Fatalf("Unable to init WARC writing HTTP client: %s", err) } @@ -699,6 +754,9 @@ func TestHTTPClientWithProxy(t *testing.T) { t.Fatal(err) } + // Set proxy on this specific request via context + req = req.WithContext(WithProxy(req.Context(), fmt.Sprintf("socks5://%s", proxyAddr))) + resp, err := httpClient.Do(req) if err != nil { t.Fatal(err) @@ -710,6 +768,10 @@ func TestHTTPClientWithProxy(t *testing.T) { httpClient.Close() waitForErrors() + if c := proxyCounter.Load(); c != 1 { + t.Fatalf("expected proxy to handle 1 connection, got %d", c) + } + files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") if err != nil { t.Fatal(err) @@ -720,6 +782,172 @@ func TestHTTPClientWithProxy(t *testing.T) { } } +func TestHTTPClientPerRequestProxyOverridesDefault(t *testing.T) { + rotatorSettings := defaultRotatorSettings(t) + + // Proxy A: a live proxy set as the client default (should NOT be used) + proxyAAddr, proxyACounter, stopProxyA := startSOCKS5Server(t) + defer stopProxyA() + + // Proxy B: the per-request override (should be used) + proxyBAddr, proxyBCounter, stopProxyB := startSOCKS5Server(t) + defer stopProxyB() + + server := newTestImageServer(t, http.StatusOK) + defer server.Close() + + // Client created with proxy A as the default + httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ + RotatorSettings: rotatorSettings, + Proxy: fmt.Sprintf("socks5://%s", proxyAAddr), + }) + if err != nil { + t.Fatalf("Unable to init WARC writing HTTP client: %s", err) + } + waitForErrors := drainErrChan(t, httpClient.ErrChan) + + req, err := http.NewRequest("GET", server.URL, nil) + if err != nil { + t.Fatal(err) + } + + // Override with proxy B + req = req.WithContext(WithProxy(req.Context(), fmt.Sprintf("socks5://%s", proxyBAddr))) + + resp, err := httpClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + + io.Copy(io.Discard, resp.Body) + + httpClient.Close() + waitForErrors() + + if c := proxyACounter.Load(); c != 0 { + t.Fatalf("expected default proxy A to handle 0 connections, got %d", c) + } + if c := proxyBCounter.Load(); c != 1 { + t.Fatalf("expected per-request proxy B to handle 1 connection, got %d", c) + } + + files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") + if err != nil { + t.Fatal(err) + } + + for _, path := range files { + testFileSingleHashCheck(t, path, "sha1:UIRWL5DFIPQ4MX3D3GFHM2HCVU3TZ6I3", []string{"26872"}, 1, server.URL+"/") + } +} + +func TestHTTPClientPerRequestProxyBypassDefault(t *testing.T) { + rotatorSettings := defaultRotatorSettings(t) + + // Start a live proxy and set it as the client default + proxyAddr, proxyCounter, stopProxy := startSOCKS5Server(t) + defer stopProxy() + + server := newTestImageServer(t, http.StatusOK) + defer server.Close() + + httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ + RotatorSettings: rotatorSettings, + Proxy: fmt.Sprintf("socks5://%s", proxyAddr), + }) + if err != nil { + t.Fatalf("Unable to init WARC writing HTTP client: %s", err) + } + waitForErrors := drainErrChan(t, httpClient.ErrChan) + + req, err := http.NewRequest("GET", server.URL, nil) + if err != nil { + t.Fatal(err) + } + + // Force direct connection by passing empty string, bypassing the default proxy + req = req.WithContext(WithProxy(req.Context(), "")) + + resp, err := httpClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + + io.Copy(io.Discard, resp.Body) + + httpClient.Close() + waitForErrors() + + if c := proxyCounter.Load(); c != 0 { + t.Fatalf("expected default proxy to handle 0 connections (bypassed), got %d", c) + } + + files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") + if err != nil { + t.Fatal(err) + } + + for _, path := range files { + testFileSingleHashCheck(t, path, "sha1:UIRWL5DFIPQ4MX3D3GFHM2HCVU3TZ6I3", []string{"26872"}, 1, server.URL+"/") + } +} + +func TestHTTPClientPerRequestProxyCacheReuse(t *testing.T) { + rotatorSettings := defaultRotatorSettings(t) + + proxyAddr, proxyCounter, stopProxy := startSOCKS5Server(t) + defer stopProxy() + + server := newTestImageServer(t, http.StatusOK) + defer server.Close() + + httpClient, err := NewWARCWritingHTTPClient(HTTPClientSettings{ + RotatorSettings: rotatorSettings, + }) + if err != nil { + t.Fatalf("Unable to init WARC writing HTTP client: %s", err) + } + waitForErrors := drainErrChan(t, httpClient.ErrChan) + + proxyURL := fmt.Sprintf("socks5://%s", proxyAddr) + + // Make two requests through the same per-request proxy to exercise cache reuse + for i := 0; i < 2; i++ { + req, err := http.NewRequest("GET", server.URL, nil) + if err != nil { + t.Fatal(err) + } + + req = req.WithContext(WithProxy(req.Context(), proxyURL)) + + resp, err := httpClient.Do(req) + if err != nil { + t.Fatalf("request %d failed: %s", i, err) + } + + io.Copy(io.Discard, resp.Body) + resp.Body.Close() + } + + httpClient.Close() + waitForErrors() + + if c := proxyCounter.Load(); c != 2 { + t.Fatalf("expected proxy to handle 2 connections, got %d", c) + } + + files, err := filepath.Glob(rotatorSettings.OutputDirectory + "/*") + if err != nil { + t.Fatal(err) + } + + for _, path := range files { + testFileSingleHashCheck(t, path, "sha1:UIRWL5DFIPQ4MX3D3GFHM2HCVU3TZ6I3", []string{"26872"}, 2, server.URL+"/") + } +} + func TestHTTPClientConcurrent(t *testing.T) { var ( rotatorSettings = defaultRotatorSettings(t) diff --git a/dialer.go b/dialer.go index 623d5d0..5c0414b 100644 --- a/dialer.go +++ b/dialer.go @@ -39,6 +39,11 @@ const ( // This is used internally to retrieve the wrapped connection for advanced use cases. // Use WithWrappedConnection() helper function for convenience. ContextKeyWrappedConn contextKey = "wrappedConn" + + // ContextKeyProxy is the context key for per-request proxy override. + // When provided, the request will use this proxy instead of the client-level default. + // Use WithProxy() helper function for convenience. + ContextKeyProxy contextKey = "proxy" ) // WithFeedbackChannel adds a feedback channel to the request context. @@ -62,6 +67,19 @@ func WithWrappedConnection(ctx context.Context, wrappedConnChan chan *CustomConn return context.WithValue(ctx, ContextKeyWrappedConn, wrappedConnChan) } +// WithProxy adds a per-request proxy URL to the request context. +// When provided, this proxy will be used instead of the client-level default. +// Pass an empty string to force a direct connection, bypassing any default proxy. +func WithProxy(ctx context.Context, proxyURL string) context.Context { + return context.WithValue(ctx, ContextKeyProxy, proxyURL) +} + +// resolvedProxy holds the proxy dialer and its DNS behavior for a single dial operation. +type resolvedProxy struct { + dialer proxy.ContextDialer + needsHostname bool +} + // dnsExchanger is an interface for DNS clients that can exchange messages type dnsExchanger interface { ExchangeContext(ctx context.Context, m *dns.Msg, address string) (r *dns.Msg, rtt time.Duration, err error) @@ -69,7 +87,8 @@ type dnsExchanger interface { type customDialer struct { proxyDialer proxy.ContextDialer - proxyNeedsHostname bool // true if proxy requires hostname (socks5h, http), false if can use IP (socks5) + proxyNeedsHostname bool // true if proxy requires hostname (socks5h, http), false if can use IP (socks5) + proxyCache sync.Map // caches *resolvedProxy by proxy URL string client *CustomHTTPClient DNSConfig *dns.ClientConfig DNSClient dnsExchanger @@ -81,6 +100,43 @@ type customDialer struct { dnsRoundRobinIndex atomic.Uint32 } +// Returns the configured proxy for a given request context. +// The per-request proxy takes precedence over the client-level default. +// An empty string explicitly forces a direct connection (bypasses the default proxy). +// Returns nil if no proxy is configured. +func (d *customDialer) resolveProxy(ctx context.Context) (*resolvedProxy, error) { + if proxyURL, ok := ctx.Value(ContextKeyProxy).(string); ok { + if proxyURL == "" { + return nil, nil + } + if cached, ok := d.proxyCache.Load(proxyURL); ok { + return cached.(*resolvedProxy), nil + } + + u, err := url.Parse(proxyURL) + if err != nil { + return nil, err + } + + pd, err := proxy.FromURL(u, d) + if err != nil { + return nil, err + } + + rp := &resolvedProxy{ + dialer: pd.(proxy.ContextDialer), + needsHostname: u.Scheme == "socks5h" || u.Scheme == "socks4a" || u.Scheme == "http" || u.Scheme == "https", + } + d.proxyCache.Store(proxyURL, rp) + return rp, nil + } + + if d.proxyDialer != nil { + return &resolvedProxy{dialer: d.proxyDialer, needsHostname: d.proxyNeedsHostname}, nil + } + return nil, nil +} + var emptyPayloadDigests = []string{ "sha1:3I42H3S6NNFQ2MSVX7XZKYAYSCX5QBYJ", "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", @@ -104,16 +160,16 @@ type dialResult struct { // It returns the first established connection and closes the other. // Otherwise it returns an error from the primary address. // This implements Happy Eyeballs (RFC 8305). -func (d *customDialer) dialParallel(ctx context.Context, network string, primaryAddr, fallbackAddr string, primaryIP, fallbackIP net.IP) (net.Conn, net.IP, error) { +func (d *customDialer) dialParallel(ctx context.Context, network string, primaryAddr, fallbackAddr string, primaryIP, fallbackIP net.IP, rp *resolvedProxy) (net.Conn, net.IP, error) { if fallbackAddr == "" && primaryAddr == "" { return nil, nil, errors.New("no addresses available") } if fallbackAddr == "" { - conn, err := d.dialSingle(ctx, network+"6", primaryAddr, primaryIP) + conn, err := d.dialSingle(ctx, network+"6", primaryAddr, primaryIP, rp) return conn, primaryIP, err } if primaryAddr == "" { - conn, err := d.dialSingle(ctx, network+"4", fallbackAddr, fallbackIP) + conn, err := d.dialSingle(ctx, network+"4", fallbackAddr, fallbackIP, rp) return conn, fallbackIP, err } @@ -131,7 +187,7 @@ func (d *customDialer) dialParallel(ctx context.Context, network string, primary } else { addr, ip, netType = fallbackAddr, fallbackIP, network+"4" } - conn, err := d.dialSingle(ctx, netType, addr, ip) + conn, err := d.dialSingle(ctx, netType, addr, ip, rp) select { case results <- dialResult{conn: conn, err: err, primary: primary, done: true, ip: ip}: case <-returned: @@ -180,9 +236,9 @@ func (d *customDialer) dialParallel(ctx context.Context, network string, primary } // dialSingle performs a single dial attempt -func (d *customDialer) dialSingle(ctx context.Context, network, address string, resolvedIP net.IP) (net.Conn, error) { - if d.proxyDialer != nil { - return d.proxyDialer.DialContext(ctx, network, address) +func (d *customDialer) dialSingle(ctx context.Context, network, address string, resolvedIP net.IP, rp *resolvedProxy) (net.Conn, error) { + if rp != nil { + return rp.dialer.DialContext(ctx, network, address) } if d.client.randomLocalIP { @@ -348,10 +404,15 @@ func (d *customDialer) wrapConnection(ctx context.Context, c net.Conn, scheme st } func (d *customDialer) CustomDialContext(ctx context.Context, network, address string) (conn net.Conn, err error) { - if d.proxyDialer != nil && d.proxyNeedsHostname { + rp, err := d.resolveProxy(ctx) + if err != nil { + return nil, err + } + + if rp != nil && rp.needsHostname { // Remote DNS proxy (socks5h, socks4a, http, https) // Skip DNS archiving to avoid privacy leak and ensure accuracy. - conn, err = d.proxyDialer.DialContext(ctx, network, address) + conn, err = rp.dialer.DialContext(ctx, network, address) if err != nil { return nil, err @@ -381,7 +442,7 @@ func (d *customDialer) CustomDialContext(ctx context.Context, network, address s } // Use Happy Eyeballs: IPv6 primary, IPv4 fallback - conn, _, err = d.dialParallel(ctx, network, ipv6Addr, ipv4Addr, ipv6, ipv4) + conn, _, err = d.dialParallel(ctx, network, ipv6Addr, ipv4Addr, ipv6, ipv4, rp) if err != nil { return nil, err @@ -396,12 +457,16 @@ func (d *customDialer) CustomDial(network, address string) (net.Conn, error) { func (d *customDialer) CustomDialTLSContext(ctx context.Context, network, address string) (net.Conn, error) { var plainConn net.Conn - var err error - if d.proxyDialer != nil && d.proxyNeedsHostname { + rp, err := d.resolveProxy(ctx) + if err != nil { + return nil, err + } + + if rp != nil && rp.needsHostname { // Remote DNS proxy (socks5h, socks4a, http, https) // Skip DNS archiving to avoid privacy leak and ensure accuracy. - plainConn, err = d.proxyDialer.DialContext(ctx, network, address) + plainConn, err = rp.dialer.DialContext(ctx, network, address) if err != nil { return nil, err } @@ -430,7 +495,7 @@ func (d *customDialer) CustomDialTLSContext(ctx context.Context, network, addres } // Use Happy Eyeballs: IPv6 primary, IPv4 fallback - plainConn, _, err = d.dialParallel(ctx, network, ipv6Addr, ipv4Addr, ipv6, ipv4) + plainConn, _, err = d.dialParallel(ctx, network, ipv6Addr, ipv4Addr, ipv6, ipv4, rp) if err != nil { return nil, err } @@ -568,7 +633,8 @@ func (d *customDialer) writeWARCFromConnection(ctx context.Context, reqPipe, res case <-ctx.Done(): return default: - if d.proxyDialer == nil { + perRequestProxy, _ := ctx.Value(ContextKeyProxy).(string) + if d.proxyDialer == nil && perRequestProxy == "" { switch addr := conn.RemoteAddr().(type) { case *net.TCPAddr: IP := addr.IP.String()