Eliminate per-connection 407 probe for known-auth proxies - #182
Eliminate per-connection 407 probe for known-auth proxies#182rurouni88 wants to merge 17 commits into
Conversation
In connectViaProxy, after a stale cache hit forces a cold probe on a new transport tr2, the function was falling through to return tr.hijack() -- hijacking the 407-dead tr connection instead of the live tr2 tunnel. Introduce activeTr (*transport) initialised to &tr; set to &tr2 in the stale-eviction branch. Return activeTr.hijack() at the end so the correct transport is always hijacked regardless of which path was taken. Also clear Proxy-Authorization before the cold probe on tr2. After retryConnectWithAuth returns a stale 407, req carries the auth header from the last attempt inside that helper. The bare cold probe on tr2 must not send credentials -- clearing the header here is consistent with how retryConnectWithAuth itself clears between attempts.
After a stale cache entry is evicted and a cold probe on tr2 receives a new 407, retryConnectWithAuth is called with the fresh schemes. On success the cache was left empty, so the next request would cold-probe again instead of using the warm entry. Store the fresh proxyAuthInfo before retryConnectWithAuth so the entry is populated regardless of whether the re-auth succeeds or not. If it fails, connectViaProxy returns an error and the entry is harmless.
parseProxyAuthenticateSchemes normalises scheme names to lowercase. The assertion was checking for "Basic" but the cached value is "basic". Aligns the assertion with the storage invariant.
Extends authCacheMockProxy with a stale407Once field: returns 407 only on the first bare CONNECT, then 200 on subsequent bare CONNECTs. Adds newStaleOnceTestProxy helper and a new test that exercises the full stale-evict-then-succeed path: stale cache hit -> 407 evicts entry -> cold probe succeeds -> non-nil connection returned.
…ession on second request
ProxyFinder gains an onPACUpdate func() hook called after every successful PAC re-download. ProxyHandler registers a sync.Map range-delete as the implementation, so auth cache entries are cleared whenever the PAC file changes and proxy topology may have shifted.
…me fixtures, tighten bare-probe assertions, assert cache repopulation
samuong
left a comment
samuong
left a comment
There was a problem hiding this comment.
thanks for putting this together! it's looking really good, but i have a few suggestions below
|
|
||
| var resp *http.Response | ||
| if cached, ok := authCache.Load(proxyURL.Host); ok && auth != nil { | ||
| // Phase 1: cache hit — skip unauthenticated probe. |
There was a problem hiding this comment.
It says "phase 1" here, which to me suggests that two things will happen in sequence, i.e. phase 1 and then phase 2. But this is an if-else statement, with phase 2 in the else block, so we'll do either one or the other?
| if err := tr2.dial(proxyURL); err != nil { | ||
| log.Printf("[%d] Error dialling proxy %s: %v", id, proxyURL.Host, err) | ||
| return nil, err | ||
| } | ||
| req.Header.Del("Proxy-Authorization") | ||
| resp2, err := tr2.RoundTrip(req) | ||
| if err != nil { | ||
| log.Printf("[%d] Error reading CONNECT response: %v", id, err) | ||
| return nil, err | ||
| } | ||
| if resp2.StatusCode == http.StatusProxyAuthRequired && auth != nil { | ||
| log.Printf("[%d] Got %q response, retrying with auth", id, resp2.Status) | ||
| schemes := parseProxyAuthenticateSchemes(resp2.Header) | ||
| _ = resp2.Body.Close() | ||
| authResp2, err := retryConnectWithAuth(req, proxyURL, auth, schemes, &tr2) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| if authResp2.StatusCode != http.StatusProxyAuthRequired { | ||
| authCache.Store(proxyURL.Host, proxyAuthInfo{schemes: schemes}) | ||
| } | ||
| log.Printf("[%d] Got %q response", id, authResp2.Status) | ||
| resp = authResp2 | ||
| } else { | ||
| resp = resp2 | ||
| } |
There was a problem hiding this comment.
I'm not sure if I've missed something, but this looks largely similar to "phase 2", and I think this duplication (the "stale cache entry" case above, and the "cold probe - no cached entry" case below) is a result of handling the cache hit and miss in two separate if-else branches.
In my mind, the logic would've been something like this:
- Try to get schemes from cache
- If no cached schemes, do cold probe:
a. dial, send CONNECT
b. if 407 then parse schemes, cache them, fall through
c. if 200 then hijack and return (no auth needed at all) - retryConnectWithAuth using cached or probed schemes
- if still 407 then evict cache and return error
- if 200 then hijack and return
Let me know if I've oversimplified this in my head, and I'm missing something?
There was a problem hiding this comment.
Almost. Correction on step 4: A stale 407 isn't a dead end. It evicts the cache entry and runs a cold probe on a fresh transport, which can actually succeed. Self-healing cycle, not an error return.
Good catch on the duplication. It's been extracted to a helper function.
The if-else has to stay, because a cache hit skips the probe entirely. Collapsing to a single linear flow would cold-probe every time and scrap the optimisation I'm asking for, bringing me back to Square 1 😆
There was a problem hiding this comment.
Sorry for taking so long to reply to this, I've had a lot on my plate over the last week.
In the case of a cache hit, we'll skip step 2 which is the cold probe - this is the optimisation we want?
In step 4, we either reach it after a cold probe + retry with auth (steps 2-3), or using a cached scheme. In both cases we should evict the cache entry and return an error response to the client (eg bad gateway). So I thought it would be a dead end, let me know if I'm missing something?
Everything else looks fine to me, if we can close this one item I think we're ready to merge.
There was a problem hiding this comment.
All good. I've been up watching World Cup. Argentina vs Egypt was wow.
Let me think about this one. I thought Step 4 should self-heal (evict + cold probe again). The scenario the self-heal was designed for is the proxy changing its auth scheme mid-session, but I just realised that scenario is already handled by the PAC flush being wired up in proxyfinder.go. When the network changes (VPN, Wi-Fi switch), the PAC re-downloads and onPACUpdate clears the cache.
The likely failure case for 407 after retryConnectWithAuth is bad credentials, not network change anyways.
TLDR; The self-heal adds complexity and an extra round-trip for an edge case that's not going to happen, whilst not handling the common failure case well. Good spot, let me go back to the drawing board on this one.
There was a problem hiding this comment.
Removed it completely, as extraneous logic. Been in World Cup mode, and forgot to mention that had been done.
Please review if you're happy with it. Thanks.
| func connectViaProxy(req *http.Request, proxyURL *url.URL, auth *authChain, | ||
| authCache *sync.Map) (net.Conn, error) { |
There was a problem hiding this comment.
this function requires both authChain and now authCache, maybe we should just turn it into a method?
| func connectViaProxy(req *http.Request, proxyURL *url.URL, auth *authChain, | |
| authCache *sync.Map) (net.Conn, error) { | |
| func (ph *ProxyHandler) connectViaProxy(req *http.Request, proxyURL *url.URL) (net.Conn, error) { |
| // Construction order is load-bearing: authCache must exist before | ||
| // onPACUpdate is captured, and blockProxy before proxyFinder is used. |
There was a problem hiding this comment.
The term "load-bearing" was a bit confusing to me, maybe we can reword it to something like this?
| // Construction order is load-bearing: authCache must exist before | |
| // onPACUpdate is captured, and blockProxy before proxyFinder is used. | |
| // Note that proxyHandler and proxyFinder are mutually dependent: | |
| // proxyHandler calls proxyFinder.blockProxy(), and proxyFinder calls | |
| // proxyHandler.authCache.Clear(). |
| t.Helper() | ||
| req, err := http.NewRequest(http.MethodConnect, "https://target.example.com:443", nil) | ||
| require.NoError(t, err) | ||
| req.Host = "target.example.com:443" |
There was a problem hiding this comment.
the call to http.NewRequest() takes a url, how come req.Host gets set again?
There was a problem hiding this comment.
http.NewRequest sets req.URL but leaves req.Host empty.
But rather than hardcoding the value twice, there is a more intuitive way to fix this, which I'll make happen.
There was a problem hiding this comment.
Oh you're right. For some reason I think I was able to remove this line and still have the tests pass, but I'm away from my computer now so I can't confirm. If you need to leave it, feel free to do so.
There was a problem hiding this comment.
The mock proxy doesn't care about the CONNECT target, it only looks at auth headers, so tests would pass with an empty req.Host.
But our req.Host = req.URL.Host fix is recommended, because it makes the test request match what a real CONNECT looks like, and documents why the assignment is needed.
| m.mu.Lock() | ||
| defer m.mu.Unlock() |
There was a problem hiding this comment.
iiuc all the tests run the request and then the assertions sequentially, do we need a mutex at all?
if not, maybe we could simplify it by getting rid of the mutex and just accessing the count fields directly in the assertions?
There was a problem hiding this comment.
I'd recommend keeping it.
The sequential assertion argument only holds if you trust net/http's internal sync chain to substitute for an explicit happens-before guarantee, which isn't what the Go memory model gives you. ServeHTTP runs on its own goroutine, and the increment fires before the response write, so the read in counts() is racing against a goroutine boundary net/http never promised to fence for you.
The mutex is cheap, race-detector clean, and the conventional pattern for any httptest.Server handler touching shared state, and removing it trades a clear guarantee for an implicit one.
|
|
||
| func (pf *ProxyFinder) checkForUpdates() { | ||
| pf.Lock() | ||
| defer pf.Unlock() |
There was a problem hiding this comment.
i think it's a lot easier to reason about concurrency when we unlock the mutex with a defer statement; do you need to strictly call pf.onPACUpdate outside of the lock?
if so, i wonder if we could do something like this?
func (pf *ProxyFinder) checkForUpdates() {
notify := func() {}
// notify gets called after the mutex is unlocked
defer func() { notify() }()
pf.Lock()
defer pf.Unlock()
// ...
}
alternatively, split it out into two functions: one is protected by the lock, the other is responsible for calling notify()
There was a problem hiding this comment.
Done. Went with your defer pattern.
Had to add a nil guard when assigning notify = pf.onPACUpdate since some call sites pass nil for the callback, which caused a panic in tests.
Removed hand-rolled authCache anonymous function. Replaced with Clear() Co-authored-by: Sam Uong <samuong@gmail.com>
- Rename "Phase 1/Phase 2" branch labels to "Cache hit" / "Cold probe" - Extract coldProbe helper to eliminate duplication between stale-eviction and cold-probe paths; store to cache only on auth success - Convert connectViaProxy to a pointer-receiver method on ProxyHandler, removing auth and authCache parameters - Rewrite main.go construction comment to describe mutual dependency between proxyHandler and proxyFinder explicitly - Add comment to proxyfinder.go explaining why onPACUpdate is called outside the lock - Fix makeConnectReq to derive req.Host from req.URL.Host
|
Screencap of traffic test after fixes from initial round of review. |
|
Ran a benchmark today. Anyways, these tests are a N=100 tests against our Corporate LLM Gateway via our Corporate Proxy. |
After retryConnectWithAuth returns a 407, evict the cache entry and return an error rather than falling back to a cold probe. The self-heal was intended to handle proxy auth scheme changes mid-session, but that scenario is already covered by the PAC flush wired in proxyfinder.go (onPACUpdate clears the cache on network change). The common 407-after- auth failure is bad credentials, which a cold probe retry cannot fix. Remove TestConnectAuthCache_EvictsOnStale407_ThenSucceeds and its newStaleOnceTestProxy helper. Update TestConnectAuthCache_EvictsOnStale407 to assert an error is returned rather than a bare probe being fired.
2 participants
Summary
Adds a per-session auth cache (keyed on proxy host) to ProxyHandler. After the first CONNECT tunnel to a proxy receives a 407, the advertised schemes are cached. Subsequent tunnels skip the unauthenticated probe and go straight to authentication.
Net result: One 407 on session warm-up, zero on every subsequent CONNECT to the same proxy host. A stale entry self-heals in one extra round-trip.
What's added
What's changed (with why)
How to validate
Manual: with a live NTLM proxy, the first CONNECT logs a 407 handshake; all subsequent CONNECTs go straight to Attempting NTLM authentication with no preceding 407. Verified on Windows against a corporate NTLM proxy. Screencap with redacted info below)
Notes for review