fix(proxy): surface chosen ProxyURL on Request and Response

colly.go

@@ -208,7 +208,15 @@ var collectorCounter uint32
 // other packages.
 type key int
 
-// ProxyURLKey is the context key for the request proxy address.
+// ProxyURLKey is the context key for the per-request proxy URL holder
+// (a *string shared across the request and any clone Go's http.Client may
+// fork during Do, so the value survives forkReq on the error path).
+//
+// ProxyFunc implementations should write the chosen proxy via this pattern:
+//
+//	if h, _ := req.Context().Value(colly.ProxyURLKey).(*string); h != nil {
+//	    *h = chosen.String()
+//	}
 const (
 	ProxyURLKey key = iota
 	CheckRevisitKey
@@ -673,7 +681,9 @@ func (c *Collector) scrape(u, method string, depth int, requestData io.Reader, c
 	}
 	// note: once 1.13 is minimum supported Go version,
 	// replace this with http.NewRequestWithContext
-	req = req.WithContext(context.WithValue(c.Context, CheckRevisitKey, checkRevisit))
+	req = req.WithContext(context.WithValue(
+		context.WithValue(c.Context, CheckRevisitKey, checkRevisit),
+		ProxyURLKey, new(string)))
 
 	if err := c.requestCheck(parsedURL, method, req.GetBody, depth, checkRevisit); err != nil {
 		return err
@@ -724,29 +734,38 @@ func (c *Collector) fetch(u, method string, depth int, requestData io.Reader, ct
 		req = hTrace.WithTrace(req)
 	}
 	origURL := req.URL
+	// Read the per-request *string holder set by the ProxyFunc. Called twice:
+	// inside checkResponseHeadersFunc so OnResponseHeaders sees it, and again
+	// after Cache returns so the error path (where the headers callback
+	// never fires) still surfaces the proxy URL.
+	syncProxyURL := func(r *http.Request) {
+		if h, _ := r.Context().Value(ProxyURLKey).(*string); h != nil && *h != "" {
+			request.ProxyURL = *h
+		}
+	}
 	checkResponseHeadersFunc := func(req *http.Request, statusCode int, headers http.Header) bool {
 		if req.URL != origURL {
 			request.URL = req.URL
 			request.Headers = &req.Header
 		}
-		c.handleOnResponseHeaders(&Response{Ctx: ctx, Request: request, StatusCode: statusCode, Headers: &headers})
+		syncProxyURL(req)
+		c.handleOnResponseHeaders(&Response{Ctx: ctx, Request: request, ProxyURL: request.ProxyURL, StatusCode: statusCode, Headers: &headers})
 		return !request.abort
 	}
 	checkRequestHeadersFunc := func(req *http.Request) bool {
 		c.handleOnRequestHeaders(request)
 		return !request.abort
 	}
 	response, err := c.backend.Cache(req, c.MaxBodySize, checkRequestHeadersFunc, checkResponseHeadersFunc, c.CacheDir, c.CacheExpiration)
-	if proxyURL, ok := req.Context().Value(ProxyURLKey).(string); ok {
-		request.ProxyURL = proxyURL
-	}
+	syncProxyURL(req)
 	if err := c.handleOnError(response, err, request, ctx); err != nil {
 		return err
 	}
 	c.responseCount.Add(1)
 	response.Ctx = ctx
 	response.Request = request
 	response.Trace = hTrace
+	response.ProxyURL = request.ProxyURL
 
 	err = response.fixCharset(c.DetectCharset, request.ResponseCharacterEncoding)
 	if err != nil {
@@ -1328,6 +1347,7 @@ func (c *Collector) handleOnError(response *Response, err error, request *Reques
 			Ctx:     ctx,
 		}
 	}
+	response.ProxyURL = request.ProxyURL
 	if c.debugger != nil {
 		c.debugger.Event(createEvent("error", request.ID, c.ID, map[string]string{
 			"url":    request.URL.String(),

proxy/proxy.go

@@ -15,7 +15,6 @@
 package proxy
 
 import (
-	"context"
 	"net/http"
 	"net/url"
 	"sync/atomic"
@@ -31,9 +30,13 @@ type roundRobinSwitcher struct {
 func (r *roundRobinSwitcher) GetProxy(pr *http.Request) (*url.URL, error) {
 	index := atomic.AddUint32(&r.index, 1) - 1
 	u := r.proxyURLs[index%uint32(len(r.proxyURLs))]
-
-	ctx := context.WithValue(pr.Context(), colly.ProxyURLKey, u.String())
-	*pr = *pr.WithContext(ctx)
+	// Write through the per-request *string holder colly placed in the
+	// context, so the chosen proxy is visible on Request.ProxyURL even when
+	// the request fails before response headers (forkReq isolates ctx field
+	// rewrites; pointer writes survive).
+	if h, _ := pr.Context().Value(colly.ProxyURLKey).(*string); h != nil {
+		*h = u.String()
+	}
 	return u, nil
 }
 

proxy/proxy_test.go

@@ -0,0 +1,94 @@
+// Copyright 2018 Adam Tauber
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package proxy
+
+import (
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/gocolly/colly/v2"
+)
+
+// TestRoundRobinProxySwitcher_PropagatesProxyURL is the minimal smoke test:
+// after a Visit through the switcher, the response must carry a non-empty
+// ProxyURL on both Request and Response.
+func TestRoundRobinProxySwitcher_PropagatesProxyURL(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		fmt.Fprintln(w, "ok")
+	}))
+	defer ts.Close()
+
+	rp, err := RoundRobinProxySwitcher(ts.URL)
+	if err != nil {
+		t.Fatalf("RoundRobinProxySwitcher: %v", err)
+	}
+
+	c := colly.NewCollector(colly.IgnoreRobotsTxt())
+	c.SetProxyFunc(rp)
+
+	var called bool
+	c.OnResponse(func(r *colly.Response) {
+		called = true
+		if r.Request.ProxyURL == "" {
+			t.Errorf("Request.ProxyURL is empty — ProxyURLKey not propagated")
+		}
+		if r.ProxyURL == "" {
+			t.Errorf("Response.ProxyURL is empty")
+		}
+	})
+
+	if err := c.Visit("http://example.com/"); err != nil {
+		t.Fatalf("Visit: %v", err)
+	}
+	if !called {
+		t.Fatal("OnResponse never fired")
+	}
+}
+
+// TestRoundRobinProxySwitcher_ProxyURLOnError ensures the chosen proxy URL
+// is still recorded when the request fails before any response headers
+// arrive (e.g. dial refused) — so OnError can report which proxy was tried.
+func TestRoundRobinProxySwitcher_ProxyURLOnError(t *testing.T) {
+	ln := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
+	dead := ln.URL
+	ln.Close() // guarantees dial refused on dead
+
+	rp, err := RoundRobinProxySwitcher(dead)
+	if err != nil {
+		t.Fatalf("RoundRobinProxySwitcher: %v", err)
+	}
+	c := colly.NewCollector(colly.IgnoreRobotsTxt())
+	c.SetProxyFunc(rp)
+
+	var called bool
+	c.OnError(func(r *colly.Response, _ error) {
+		called = true
+		if r.Request.ProxyURL != dead {
+			t.Errorf("Request.ProxyURL = %q, want %q", r.Request.ProxyURL, dead)
+		}
+		if r.ProxyURL != dead {
+			t.Errorf("Response.ProxyURL = %q, want %q", r.ProxyURL, dead)
+		}
+	})
+
+	if err := c.Visit("http://example.com/"); err == nil {
+		t.Fatal("expected Visit to fail")
+	}
+	if !called {
+		t.Fatal("OnError never fired")
+	}
+}

response.go

@@ -42,6 +42,9 @@ type Response struct {
 	// Trace contains the HTTPTrace for the request. Will only be set by the
 	// collector if Collector.TraceHTTP is set to true.
 	Trace *HTTPTrace
+	// ProxyURL is the proxy address that handled the request, mirrored from
+	// Request.ProxyURL for convenience.
+	ProxyURL string
 }
 
 // Save writes response body to disk