Eliminate per-connection 407 probe for known-auth proxies

main.go

@@ -200,8 +200,18 @@ func main() {
 
 func createServer(port int, pacurl string, auth *authChain, enableSocks bool) *http.Server {
 	pacWrapper := NewPACWrapper(PACData{Port: port})
-	proxyFinder := NewProxyFinder(pacurl, pacWrapper, enableSocks)
-	proxyHandler := NewProxyHandler(auth, getProxyFromContext, proxyFinder.blockProxy)
+	// Note that proxyHandler and proxyFinder are mutually dependent:
+	// proxyHandler calls proxyFinder.blockProxy(), and proxyFinder calls
+	// proxyHandler.authCache.Clear().
+	var blockProxy func(string)
+	proxyHandler := NewProxyHandler(auth, getProxyFromContext, func(host string) {
+		if blockProxy != nil {
+			blockProxy(host)
+		}
+	})
+	onPACUpdate := proxyHandler.authCache.Clear
+	proxyFinder := NewProxyFinder(pacurl, pacWrapper, enableSocks, onPACUpdate)
+	blockProxy = proxyFinder.blockProxy
 	mux := http.NewServeMux()
 	pacWrapper.SetupHandlers(mux)
 

multiauth_integration_test.go

@@ -405,7 +405,8 @@ func TestConnectViaProxy_FallsThroughOn407(t *testing.T) {
 	require.NoError(t, err)
 	req.Host = "example.com:443"
 
-	conn, err := connectViaProxy(req, proxy.URL(), chain)
+	ph := NewProxyHandler(chain, nil, func(string) {})
+	conn, err := ph.connectViaProxy(req, proxy.URL())
 	require.NoError(t, err)
 	defer conn.Close() //nolint:errcheck //nolint:errcheck
 
@@ -433,7 +434,8 @@ func TestConnectViaProxy_RefusesBasicDowngrade(t *testing.T) {
 	require.NoError(t, err)
 	req.Host = "example.com:443"
 
-	_, err = connectViaProxy(req, proxy.URL(), chain)
+	ph2 := NewProxyHandler(chain, nil, func(string) {})
+	_, err = ph2.connectViaProxy(req, proxy.URL())
 	require.Error(t, err)
 	assert.ErrorIs(t, err, errNoMatchingAuthMethod)
 }

proxy.go

@@ -26,6 +26,7 @@ import (
 	"net/http"
 	"net/url"
 	"strings"
+	"sync"
 
 	"golang.org/x/net/proxy"
 )
@@ -133,7 +134,7 @@ func isToken(s string) bool {
 
 // parseProxyAuthenticateSchemes returns the deduplicated, lower-cased
 // scheme names from a 407 response's Proxy-Authenticate header(s),
-// preserving the order in which they appeared. RFC 7235 §4.3 allows
+// preserving the order in which they appeared. RFC 9110 §11.3 allows
 // either multiple Proxy-Authenticate header fields OR a single header
 // field containing a comma-separated list of challenges (RFC 7230 list
 // extension), so we honour both. The portion after the scheme token
@@ -157,11 +158,11 @@ func parseProxyAuthenticateSchemes(header http.Header) []string {
 
 // splitChallengeNames extracts the scheme names from a single
 // Proxy-Authenticate header field value that may contain multiple
-// challenges joined by commas (RFC 7235 §4.3 + RFC 7230 §7 list
+// challenges joined by commas (RFC 9110 §11.3 + RFC 7230 §7 list
 // extension). It walks the value byte-by-byte so that commas inside
 // quoted-strings are not treated as challenge separators.
 //
-// RFC 7235's challenge grammar is:
+// RFC 9110's challenge grammar is:
 //
 //	challenge = auth-scheme [ 1*SP ( token68 /
 //	    [ ( "," / auth-param ) *( OWS "," [ OWS auth-param ] ) ] ) ]
@@ -178,7 +179,7 @@ func parseProxyAuthenticateSchemes(header http.Header) []string {
 // will cause `bar` to be misclassified as a scheme name. This is benign
 // because no real authenticator's scheme() returns `bar`, so the picker
 // silently ignores it. We accept this rather than implementing full
-// RFC 7235 grammar parsing for the tail of unparseable proxies.
+// RFC 9110 grammar parsing for the tail of unparseable proxies.
 //
 // The returned strings are the raw auth-scheme tokens; the caller
 // validates them via isToken.
@@ -220,17 +221,23 @@ func splitChallengeNames(value string) []string {
 	return names
 }
 
+type proxyAuthInfo struct {
+	// schemes are lower-cased, matching parseProxyAuthenticateSchemes output.
+	schemes []string
+}
+
 type ProxyHandler struct {
 	transport *http.Transport
 	auth      *authChain
 	block     func(string)
+	authCache *sync.Map
 }
 
 type proxyFunc func(*http.Request) (*url.URL, error)
 
 func NewProxyHandler(auth *authChain, proxy proxyFunc, block func(string)) ProxyHandler {
 	tr := &http.Transport{Proxy: proxy, TLSClientConfig: tlsClientConfig}
-	return ProxyHandler{tr, auth, block}
+	return ProxyHandler{tr, auth, block, &sync.Map{}}
 }
 
 func (ph ProxyHandler) WrapHandler(next http.Handler) http.Handler {
@@ -270,7 +277,7 @@ func (ph ProxyHandler) handleConnect(w http.ResponseWriter, req *http.Request) {
 	if proxyURL == nil {
 		server, err = connectDirect(req)
 	} else {
-		server, err = connectViaProxy(req, proxyURL, ph.auth)
+		server, err = ph.connectViaProxy(req, proxyURL)
 		var oe *net.OpError
 		if errors.As(err, &oe) && oe.Op == "proxyconnect" {
 			log.Printf("[%d] Temporarily blocking proxy: %q", id, proxyURL.Host)
@@ -345,7 +352,40 @@ func connectDirect(req *http.Request) (net.Conn, error) {
 	return server, err
 }
 
-func connectViaProxy(req *http.Request, proxyURL *url.URL, auth *authChain) (net.Conn, error) {
+// coldProbe dials proxyURL on tr, sends an unauthenticated CONNECT, and handles
+// the 407 response by caching the advertised schemes and retrying with auth.
+// The caller is responsible for deferring tr.Close().
+func coldProbe(req *http.Request, proxyURL *url.URL, auth *authChain,
+	tr *transport, authCache *sync.Map) (*http.Response, error) {
+	id := req.Context().Value(contextKeyID)
+	if err := tr.dial(proxyURL); err != nil {
+		log.Printf("[%d] Error dialling proxy %s: %v", id, proxyURL.Host, err)
+		return nil, err
+	}
+	req.Header.Del("Proxy-Authorization")
+	resp, err := tr.RoundTrip(req)
+	if err != nil {
+		log.Printf("[%d] Error reading CONNECT response: %v", id, err)
+		return nil, err
+	}
+	if resp.StatusCode != http.StatusProxyAuthRequired || auth == nil {
+		return resp, nil
+	}
+	log.Printf("[%d] Got %q response, retrying with auth", id, resp.Status)
+	schemes := parseProxyAuthenticateSchemes(resp.Header)
+	_ = resp.Body.Close()
+	authResp, err := retryConnectWithAuth(req, proxyURL, auth, schemes, tr)
+	if err != nil {
+		return nil, err
+	}
+	if authResp.StatusCode != http.StatusProxyAuthRequired {
+		authCache.Store(proxyURL.Host, proxyAuthInfo{schemes: schemes})
+	}
+	log.Printf("[%d] Got %q response", id, authResp.Status)
+	return authResp, nil
+}
+
+func (ph *ProxyHandler) connectViaProxy(req *http.Request, proxyURL *url.URL) (net.Conn, error) {
 	id := req.Context().Value(contextKeyID)
 
 	// SOCKS5 short-circuit: SOCKS5 has its own authentication model
@@ -377,26 +417,27 @@ func connectViaProxy(req *http.Request, proxyURL *url.URL, auth *authChain) (net
 
 	var tr transport
 	defer tr.Close() //nolint:errcheck
-	if err := tr.dial(proxyURL); err != nil {
-		log.Printf("[%d] Error dialling proxy %s: %v", id, proxyURL.Host, err)
-		return nil, err
-	}
-	resp, err := tr.RoundTrip(req)
-	if err != nil {
-		log.Printf("[%d] Error reading CONNECT response: %v", id, err)
-		return nil, err
-	}
-	if resp.StatusCode == http.StatusProxyAuthRequired && auth != nil {
-		log.Printf("[%d] Got %q response, retrying with auth", id, resp.Status)
-		schemes := parseProxyAuthenticateSchemes(resp.Header)
-		_ = resp.Body.Close()
-		// resp is now stale; the retry helper returns a fresh one.
-		authResp, err := retryConnectWithAuth(req, proxyURL, auth, schemes, &tr)
+	var resp *http.Response
+	if cached, ok := ph.authCache.Load(proxyURL.Host); ok && ph.auth != nil {
+		// Cache hit — skip unauthenticated probe.
+		info := cached.(proxyAuthInfo)
+		authResp, err := retryConnectWithAuth(req, proxyURL, ph.auth, info.schemes, &tr)
 		if err != nil {
 			return nil, err
 		}
+		if authResp.StatusCode == http.StatusProxyAuthRequired {
+			// Stale or invalid cache entry — evict and return error.
+			ph.authCache.Delete(proxyURL.Host)
+		}
 		log.Printf("[%d] Got %q response", id, authResp.Status)
 		resp = authResp
+	} else {
+		// Cold probe — no cached entry; populate cache on first 407.
+		var err error
+		resp, err = coldProbe(req, proxyURL, ph.auth, &tr, ph.authCache)
+		if err != nil {
+			return nil, err
+		}
 	}
 	_ = resp.Body.Close()
 	if resp.StatusCode == http.StatusProxyAuthRequired {

proxy_test.go

@@ -26,6 +26,7 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"strings"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -426,3 +427,146 @@ func TestSOCKS5Proxy(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, "/testpath", string(body))
 }
+
+// Auth-cache unit tests — covers proxyAuthInfo / authCache behaviour in connectViaProxy.
+
+type authCacheMockProxy struct {
+	mu             sync.Mutex
+	bareConnects   int
+	authedConnects int
+	respondWith407 bool
+}
+
+func (m *authCacheMockProxy) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	if req.Method != http.MethodConnect {
+		http.Error(w, "only CONNECT supported", http.StatusMethodNotAllowed)
+		return
+	}
+	hasAuth := req.Header.Get("Proxy-Authorization") != ""
+	m.mu.Lock()
+	if hasAuth {
+		m.authedConnects++
+	} else {
+		m.bareConnects++
+	}
+	m.mu.Unlock()
+	if m.respondWith407 || !hasAuth {
+		w.Header().Set("Proxy-Authenticate", "Basic realm=\"proxy\"")
+		w.WriteHeader(http.StatusProxyAuthRequired)
+		return
+	}
+	w.WriteHeader(http.StatusOK)
+}
+
+func (m *authCacheMockProxy) counts() (bare, authed int) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.bareConnects, m.authedConnects
+}
+
+func newAuthCacheTestProxy(t *testing.T, respondWith407 bool) (*httptest.Server, *authCacheMockProxy) {
+	t.Helper()
+	mock := &authCacheMockProxy{respondWith407: respondWith407}
+	srv := httptest.NewServer(mock)
+	t.Cleanup(srv.Close)
+	return srv, mock
+}
+
+
+func makeConnectReq(t *testing.T) *http.Request {
+	t.Helper()
+	req, err := http.NewRequest(http.MethodConnect, "https://target.example.com:443", nil)
+	require.NoError(t, err)
+	req.Host = req.URL.Host
+	return req
+}
+
+func TestConnectAuthCache_PopulatesOnFirst407(t *testing.T) {
+	proxySrv, mock := newAuthCacheTestProxy(t, false)
+	proxyURL, err := url.Parse(proxySrv.URL)
+	require.NoError(t, err)
+	auth := newAuthChain(newBasicAuthenticator("user:pass"))
+	cache := &sync.Map{}
+	req := makeConnectReq(t)
+	ph := ProxyHandler{auth: auth, authCache: cache}
+	conn, err := ph.connectViaProxy(req, proxyURL)
+	require.NoError(t, err)
+	if conn != nil {
+		conn.Close()
+	}
+	bare, _ := mock.counts()
+	assert.Equal(t, 1, bare, "expected exactly one unauthenticated probe")
+	val, ok := cache.Load(proxyURL.Host)
+	require.True(t, ok, "authCache must have an entry for the proxy host after a 407")
+	info, ok := val.(proxyAuthInfo)
+	require.True(t, ok, "cached value must be of type proxyAuthInfo")
+	assert.NotEmpty(t, info.schemes, "cached schemes must not be empty")
+	assert.Contains(t, info.schemes, "basic", "Basic scheme must be recorded in the cache")
+}
+
+func TestConnectAuthCache_SkipsProbeOnCacheHit(t *testing.T) {
+	proxySrv, mock := newAuthCacheTestProxy(t, false)
+	proxyURL, err := url.Parse(proxySrv.URL)
+	require.NoError(t, err)
+	auth := newAuthChain(newBasicAuthenticator("user:pass"))
+	cache := &sync.Map{}
+	cache.Store(proxyURL.Host, proxyAuthInfo{schemes: []string{"basic"}})
+	req := makeConnectReq(t)
+	ph := ProxyHandler{auth: auth, authCache: cache}
+	conn, err := ph.connectViaProxy(req, proxyURL)
+	require.NoError(t, err)
+	if conn != nil {
+		conn.Close()
+	}
+	bare, authed := mock.counts()
+	assert.Equal(t, 0, bare, "cache hit must suppress the unauthenticated probe")
+	assert.Equal(t, 1, authed, "expected exactly one authenticated CONNECT")
+}
+
+func TestConnectAuthCache_EvictsOnStale407(t *testing.T) {
+	proxySrv, _ := newAuthCacheTestProxy(t, true)
+	proxyURL, err := url.Parse(proxySrv.URL)
+	require.NoError(t, err)
+	auth := newAuthChain(newBasicAuthenticator("user:pass"))
+	cache := &sync.Map{}
+	cache.Store(proxyURL.Host, proxyAuthInfo{schemes: []string{"basic"}})
+	req := makeConnectReq(t)
+	ph := ProxyHandler{auth: auth, authCache: cache}
+	_, err = ph.connectViaProxy(req, proxyURL)
+	assert.Error(t, err, "stale 407 must return an error — no self-heal")
+	_, stillCached := cache.Load(proxyURL.Host)
+	assert.False(t, stillCached, "stale cache entry must be evicted on 407")
+}
+
+
+func TestConnectAuthCache_NoProbeOnSecondRequest(t *testing.T) {
+	proxySrv, mock := newAuthCacheTestProxy(t, false)
+	proxyURL, err := url.Parse(proxySrv.URL)
+	require.NoError(t, err)
+	auth := newAuthChain(newBasicAuthenticator("user:pass"))
+	cache := &sync.Map{}
+
+	// Pass 1: cold path — cache is empty, proxy will issue a 407 probe.
+	ph := ProxyHandler{auth: auth, authCache: cache}
+	req1 := makeConnectReq(t)
+	conn1, err := ph.connectViaProxy(req1, proxyURL)
+	require.NoError(t, err, "pass 1 must succeed")
+	require.NotNil(t, conn1, "pass 1 connection must be non-nil")
+	conn1.Close()
+
+	bare, _ := mock.counts()
+	assert.Equal(t, 1, bare, "exactly one bare probe expected after pass 1")
+	_, cached := cache.Load(proxyURL.Host)
+	assert.True(t, cached, "cache must have an entry for the proxy host after pass 1")
+
+	// Pass 2: warm path — cache hit must suppress the bare probe entirely.
+	req2 := makeConnectReq(t)
+	conn2, err := ph.connectViaProxy(req2, proxyURL)
+	require.NoError(t, err, "pass 2 must succeed")
+	require.NotNil(t, conn2, "pass 2 connection must be non-nil")
+	conn2.Close()
+
+	bare, authed := mock.counts()
+	assert.Equal(t, 1, bare, "bare probe count must remain 1 after pass 2 (no new probe)")
+	assert.Equal(t, 2, authed, "both passes must have used authenticated CONNECT")
+}