Skip to content

Feat: per-request proxy #186

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
AltayAkkus wants to merge 3 commits into from
Closed

Feat: per-request proxy #186

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a3f5218
add proxy ContextKey
AltayAkkus Mar 24, 2026
547318a
add: socks5 server with request counter & tests for per-request proxies
AltayAkkus Mar 24, 2026
c7a4d3c
include per-request proxy docs in README
AltayAkkus Mar 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 16 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
280 changes: 254 additions & 26 deletions client_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ import (
"strconv"
"strings"
"sync"
"sync/atomic"
"testing"
"time"

Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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)
}
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand Down
Loading