From d4de2070d2f0b7387f32be2ee27f746e6c9fbd2d Mon Sep 17 00:00:00 2001 From: Ben Browning <56071+bbrowning@users.noreply.github.com> Date: Thu, 13 Aug 2026 18:37:54 +0000 Subject: [PATCH] Retry once with a forced credential refresh on upstream 401 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Vertex AI requests can fail with Google's own 401 ("invalid authentication credentials") even though GCloudInjector.Inject considered its cached OAuth token locally valid — most likely because a long-running proxy process's cached token can look unexpired by wall clock while genuinely being expired server-side after a host suspend/resume cycle. Previously the only fix was a full `paude stop` + `start`, which restarts the proxy process and forces a fresh token. GCloudInjector gains ForceRefresh(), which discards the cached credentials and rebuilds them from the original ADC source so the next token fetch is a real network exchange rather than a reused token. Store gains MatchInjector() so callers can look up which injector would handle a request. The proxy's response pipeline now buffers request bodies for any request handled by a Refresher injector and, on an upstream 401, calls ForceRefresh and retries the request exactly once with the newly fetched token before giving up. This is root-cause-agnostic: whatever caused the local validity check to disagree with the server, trusting the server's rejection and retrying with a guaranteed-fresh credential recovers without requiring a full process restart. Co-Authored-By: Claude Opus 4.8 --- internal/credentials/gcloud.go | 132 +++++++++++++++------ internal/credentials/gcloud_test.go | 95 +++++++++++++++ internal/credentials/store.go | 83 +++++++++----- internal/proxy/integration_test.go | 132 +++++++++++++++++++++ internal/proxy/proxy.go | 156 ++++++++++++++++++++++++- internal/proxy/refresh_retry_test.go | 165 +++++++++++++++++++++++++++ 6 files changed, 697 insertions(+), 66 deletions(-) create mode 100644 internal/proxy/refresh_retry_test.go diff --git a/internal/credentials/gcloud.go b/internal/credentials/gcloud.go index 49fca23..e53cdf5 100644 --- a/internal/credentials/gcloud.go +++ b/internal/credentials/gcloud.go @@ -19,9 +19,10 @@ import ( // Google Application Default Credentials. It handles automatic // token refresh. Always overrides any existing Authorization header. type GCloudInjector struct { + mu sync.Mutex credentials *google.Credentials - initOnce sync.Once initErr error + initialized bool adcPath string adcJSON []byte scopes []string @@ -47,47 +48,102 @@ func NewGCloudInjectorFromJSON(data []byte) *GCloudInjector { } func (g *GCloudInjector) init() error { - g.initOnce.Do(func() { - var data []byte - if len(g.adcJSON) > 0 { - data = g.adcJSON - } else { - var err error - data, err = os.ReadFile(g.adcPath) - if err != nil { - g.initErr = fmt.Errorf("read ADC file %s: %w", g.adcPath, err) - return - } - } + g.mu.Lock() + defer g.mu.Unlock() + if g.initialized { + return g.initErr + } + return g.recordInit(g.doInit()) +} - // Custom HTTP client for OAuth2 token refresh. DisableKeepAlives forces - // fresh connections (~1/hour refresh rate, so no benefit to pooling). - httpClient := &http.Client{ - Timeout: timeouts.ResponseHeader, - Transport: &http.Transport{ - Proxy: nil, // token refresh must go directly to Google, never through HTTP_PROXY - DisableKeepAlives: true, - TLSClientConfig: &tls.Config{ - MinVersion: tls.VersionTLS12, - }, - TLSHandshakeTimeout: timeouts.TLSHandshake, - ResponseHeaderTimeout: timeouts.ResponseHeader, - }, - } +// ForceRefresh discards any cached credentials/token and rebuilds the +// underlying google.Credentials from the original ADC source. The rebuilt +// TokenSource starts from a token with no AccessToken/Expiry, so the next +// Inject call is guaranteed to perform a real network token exchange rather +// than reusing a cached token that Inject still considered locally valid. +// +// This exists because a long-running proxy process can end up with a +// cached token that looks unexpired by its local clock while actually +// being rejected by Google (e.g. after a host suspend/resume cycle) — the +// caller invokes this after seeing an upstream 401 to force a real refresh +// before retrying. +func (g *GCloudInjector) ForceRefresh() error { + g.mu.Lock() + defer g.mu.Unlock() + err := g.doInit() + if err != nil && g.credentials != nil { + // Keep the last-good credentials on a failed rebuild. Discarding them + // would poison the injector — Inject would fail with 502 forever, and a + // 502 (unlike a 401) never re-triggers this retry path, so it could + // never recover. Report the error so the caller skips the now-pointless + // retry, but leave initErr/initialized untouched so Inject keeps working. + log.Printf("WARN gcloud force refresh failed, keeping existing credentials: %v", err) + return err + } + return g.recordInit(err) +} - // Use context.Background() with custom HTTP client — this context is stored by - // the oauth2 library and reused for all token refresh HTTP calls. It must NOT - // be canceled or have a short timeout. - ctx := context.WithValue(context.Background(), oauth2.HTTPClient, httpClient) - creds, err := google.CredentialsFromJSON(ctx, data, g.scopes...) +// recordInit stores the outcome of a doInit attempt and returns it, so init() +// and ForceRefresh() share one place that marks the injector initialized. +// Callers must hold g.mu. +func (g *GCloudInjector) recordInit(err error) error { + g.initErr = err + g.initialized = true + return err +} + +// doInit rebuilds credentials from the original ADC source and, on success, +// swaps them into g.credentials. On failure it returns the error WITHOUT +// mutating any state, so a failed rebuild never discards previously-valid +// credentials. Callers must hold g.mu and are responsible for recording +// initialized/initErr. +func (g *GCloudInjector) doInit() error { + var data []byte + if len(g.adcJSON) > 0 { + data = g.adcJSON + } else { + var err error + data, err = os.ReadFile(g.adcPath) if err != nil { - g.initErr = fmt.Errorf("parse ADC credentials: %w", err) - return + return fmt.Errorf("read ADC file %s: %w", g.adcPath, err) } + } + + // Custom HTTP client for OAuth2 token refresh. DisableKeepAlives forces + // fresh connections (~1/hour refresh rate, so no benefit to pooling). + httpClient := &http.Client{ + Timeout: timeouts.ResponseHeader, + Transport: &http.Transport{ + Proxy: nil, // token refresh must go directly to Google, never through HTTP_PROXY + DisableKeepAlives: true, + TLSClientConfig: &tls.Config{ + MinVersion: tls.VersionTLS12, + }, + TLSHandshakeTimeout: timeouts.TLSHandshake, + ResponseHeaderTimeout: timeouts.ResponseHeader, + }, + } + + // Use context.Background() with custom HTTP client — this context is stored by + // the oauth2 library and reused for all token refresh HTTP calls. It must NOT + // be canceled or have a short timeout. + ctx := context.WithValue(context.Background(), oauth2.HTTPClient, httpClient) + creds, err := google.CredentialsFromJSON(ctx, data, g.scopes...) + if err != nil { + return fmt.Errorf("parse ADC credentials: %w", err) + } + + g.credentials = creds + return nil +} - g.credentials = creds - }) - return g.initErr +// tokenSource returns the current TokenSource under the lock, so a +// concurrent ForceRefresh swapping g.credentials can't race with Inject +// reading it. +func (g *GCloudInjector) tokenSource() oauth2.TokenSource { + g.mu.Lock() + defer g.mu.Unlock() + return g.credentials.TokenSource } // Inject sets the Authorization: Bearer header with a fresh OAuth2 token. @@ -103,7 +159,7 @@ func (g *GCloudInjector) Inject(req *http.Request) InjectResult { return InjectFailed } - token, err := g.credentials.TokenSource.Token() + token, err := g.tokenSource().Token() if err != nil { log.Printf("ERROR gcloud token refresh failed: %v", err) return InjectFailed diff --git a/internal/credentials/gcloud_test.go b/internal/credentials/gcloud_test.go index 6c7657d..51dc6b5 100644 --- a/internal/credentials/gcloud_test.go +++ b/internal/credentials/gcloud_test.go @@ -2,6 +2,8 @@ package credentials import ( "net/http" + "os" + "path/filepath" "testing" ) @@ -101,3 +103,96 @@ func TestGCloudInjector_NeverReturnsInjectAuthRequired(t *testing.T) { }) } } + +// validAuthorizedUserADC is a syntactically valid (but unusable — the +// refresh token isn't real) authorized_user ADC JSON blob, sufficient for +// google.CredentialsFromJSON to succeed at parsing/building a Credentials +// object without making any network call. +const validAuthorizedUserADC = `{ + "type": "authorized_user", + "client_id": "test-client-id", + "client_secret": "test-client-secret", + "refresh_token": "test-refresh-token" +}` + +func TestGCloudInjector_ForceRefresh_RebuildsCredentials(t *testing.T) { + inj := NewGCloudInjectorFromJSON([]byte(validAuthorizedUserADC)) + + if err := inj.init(); err != nil { + t.Fatalf("init: %v", err) + } + before := inj.credentials + + if err := inj.ForceRefresh(); err != nil { + t.Fatalf("ForceRefresh: %v", err) + } + after := inj.credentials + + if before == after { + t.Error("ForceRefresh should rebuild credentials from the original ADC source, not reuse the cached object") + } +} + +func TestGCloudInjector_ForceRefresh_BeforeInit(t *testing.T) { + // ForceRefresh must work even if Inject/init was never called yet. + inj := NewGCloudInjectorFromJSON([]byte(validAuthorizedUserADC)) + + if err := inj.ForceRefresh(); err != nil { + t.Fatalf("ForceRefresh: %v", err) + } + if inj.credentials == nil { + t.Error("ForceRefresh should populate credentials") + } +} + +func TestGCloudInjector_ForceRefresh_PropagatesInitFailure(t *testing.T) { + inj := NewGCloudInjector("/nonexistent/path/to/adc.json") + + if err := inj.ForceRefresh(); err == nil { + t.Error("ForceRefresh should fail when the ADC file doesn't exist") + } +} + +// A ForceRefresh that fails to rebuild (e.g. a transient ADC file read error +// during a host suspend/resume hiccup) must not poison an already-initialized +// injector. Before the fix, doInit set initErr and left initialized=true, so +// every subsequent Inject returned failure -> 502 forever, and since a 502 +// (not a 401) never re-triggers the retry path, the injector could never +// recover without a full process restart — the exact outage the retry feature +// exists to prevent. +func TestGCloudInjector_ForceRefresh_FailurePreservesCredentials(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "adc.json") + if err := os.WriteFile(path, []byte(validAuthorizedUserADC), 0600); err != nil { + t.Fatal(err) + } + + inj := NewGCloudInjector(path) + if err := inj.init(); err != nil { + t.Fatalf("init: %v", err) + } + good := inj.credentials + if good == nil { + t.Fatal("expected credentials after successful init") + } + + // Simulate the ADC source becoming momentarily unreadable. + if err := os.Remove(path); err != nil { + t.Fatal(err) + } + + // ForceRefresh can't rebuild, so it reports the error (so the caller skips + // the now-pointless retry)... + if err := inj.ForceRefresh(); err == nil { + t.Error("ForceRefresh should report an error when the ADC source is unreadable") + } + + // ...but it must NOT discard the last-good credentials or poison the + // injector: the object is unchanged and init() still succeeds. + if inj.credentials != good { + t.Error("failed ForceRefresh must preserve the last-good credentials") + } + if err := inj.init(); err != nil { + t.Errorf("init should still succeed after a failed ForceRefresh, got %v", err) + } +} diff --git a/internal/credentials/store.go b/internal/credentials/store.go index 6924437..a813115 100644 --- a/internal/credentials/store.go +++ b/internal/credentials/store.go @@ -28,6 +28,15 @@ type Injector interface { Inject(req *http.Request) InjectResult } +// Refresher is implemented by injectors that can discard cached credential +// state and force a real refresh on the next Inject call. Used to recover +// from a credential that Inject considered locally valid but that upstream +// rejected (e.g. a cached OAuth token whose local validity check disagreed +// with the server's). +type Refresher interface { + ForceRefresh() error +} + // Route maps a domain pattern to a credential injector. type Route struct { // DomainSuffix matches if the hostname ends with this suffix. @@ -76,45 +85,65 @@ func (s *Store) InjectCredentials(req *http.Request) InjectResult { s.mu.RLock() defer s.mu.RUnlock() - if req == nil || req.URL == nil { + route, matchedPattern, host := s.matchRoute(req) + if route == nil { return InjectNoMatch } - host := req.URL.Host + result := route.Injector.Inject(req) + switch result { + case InjectOK: + log.Printf("CREDENTIAL_INJECT host=%s pattern=%s method=%s path=%s", host, matchedPattern, req.Method, req.URL.Path) + case InjectFailed: + log.Printf("CREDENTIAL_INJECT_FAILED host=%s pattern=%s method=%s path=%s", host, matchedPattern, req.Method, req.URL.Path) + case InjectAuthRequired: + log.Printf("CREDENTIAL_AUTH_REQUIRED host=%s pattern=%s method=%s path=%s", host, matchedPattern, req.Method, req.URL.Path) + } + return result +} + +// MatchInjector returns the injector that would handle req, or nil if no +// route matches. Used by the response pipeline to retry with the same +// injector after a forced refresh. +func (s *Store) MatchInjector(req *http.Request) Injector { + s.mu.RLock() + defer s.mu.RUnlock() + + route, _, _ := s.matchRoute(req) + if route == nil { + return nil + } + return route.Injector +} + +// matchRoute finds the first route matching req. Callers must hold s.mu +// (read lock is sufficient). Returns the matched route, its display +// pattern, and the normalized host, or a nil route if nothing matched. +func (s *Store) matchRoute(req *http.Request) (route *Route, pattern string, host string) { + if req == nil || req.URL == nil { + return nil, "", "" + } + + host = req.URL.Host if idx := strings.LastIndex(host, ":"); idx != -1 { host = host[:idx] } host = strings.ToLower(host) - for _, route := range s.routes { - matched := false - matchedPattern := "" - pathMatched := route.PathPrefix == "" || pathMatchesPrefix(req.URL.Path, route.PathPrefix) - methodMatched := len(route.Methods) == 0 || route.Methods[strings.ToUpper(req.Method)] - - if pathMatched && methodMatched && route.ExactDomain != "" && host == route.ExactDomain { - matched = true - matchedPattern = route.ExactDomain - } else if pathMatched && methodMatched && route.DomainSuffix != "" && strings.HasSuffix(host, route.DomainSuffix) { - matched = true - matchedPattern = "*" + route.DomainSuffix - } + for i := range s.routes { + r := &s.routes[i] + pathMatched := r.PathPrefix == "" || pathMatchesPrefix(req.URL.Path, r.PathPrefix) + methodMatched := len(r.Methods) == 0 || r.Methods[strings.ToUpper(req.Method)] - if matched { - result := route.Injector.Inject(req) - switch result { - case InjectOK: - log.Printf("CREDENTIAL_INJECT host=%s pattern=%s method=%s path=%s", host, matchedPattern, req.Method, req.URL.Path) - case InjectFailed: - log.Printf("CREDENTIAL_INJECT_FAILED host=%s pattern=%s method=%s path=%s", host, matchedPattern, req.Method, req.URL.Path) - case InjectAuthRequired: - log.Printf("CREDENTIAL_AUTH_REQUIRED host=%s pattern=%s method=%s path=%s", host, matchedPattern, req.Method, req.URL.Path) - } - return result + if pathMatched && methodMatched && r.ExactDomain != "" && host == r.ExactDomain { + return r, r.ExactDomain, host + } + if pathMatched && methodMatched && r.DomainSuffix != "" && strings.HasSuffix(host, r.DomainSuffix) { + return r, "*" + r.DomainSuffix, host } } - return InjectNoMatch + return nil, "", host } func pathMatchesPrefix(path, prefix string) bool { diff --git a/internal/proxy/integration_test.go b/internal/proxy/integration_test.go index a264466..bc89fd6 100644 --- a/internal/proxy/integration_test.go +++ b/internal/proxy/integration_test.go @@ -14,6 +14,8 @@ import ( "os" "path/filepath" "strings" + "sync" + "sync/atomic" "testing" "time" @@ -878,6 +880,136 @@ func TestIntegration_CredentialInjectionFailure_Returns502(t *testing.T) { } } +// refreshableInjector is a test double implementing both credentials.Injector +// and credentials.Refresher. It injects "stale" until ForceRefresh is +// called, then injects "fresh" — simulating a credential whose local +// validity check disagreed with upstream until forced to refresh. +type refreshableInjector struct { + mu sync.Mutex + refreshed bool +} + +func (r *refreshableInjector) Inject(req *http.Request) credentials.InjectResult { + r.mu.Lock() + defer r.mu.Unlock() + if r.refreshed { + req.Header.Set("Authorization", "Bearer fresh") + } else { + req.Header.Set("Authorization", "Bearer stale") + } + return credentials.InjectOK +} + +func (r *refreshableInjector) ForceRefresh() error { + r.mu.Lock() + defer r.mu.Unlock() + r.refreshed = true + return nil +} + +func TestIntegration_GCloudTokenRefreshOnUpstream401(t *testing.T) { + skipIntegration(t) + + ca, err := GenerateCA() + if err != nil { + t.Fatalf("generate CA: %v", err) + } + + var hits int32 + upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&hits, 1) + if r.Header.Get("Authorization") == "Bearer fresh" { + w.WriteHeader(http.StatusOK) + return + } + w.WriteHeader(http.StatusUnauthorized) + _, _ = w.Write([]byte(`[{"error":{"code":401,"message":"Request had invalid authentication credentials."}}]`)) + })) + defer upstream.Close() + + upstreamURL, _ := url.Parse(upstream.URL) + upstreamHostname := upstreamURL.Hostname() + + df := filter.NewDomainFilter(upstreamHostname) + + store := credentials.NewStore() + store.AddRoute(credentials.Route{ + ExactDomain: upstreamHostname, + Injector: &refreshableInjector{}, + }) + + upstreamCAs := upstreamCertPool(t, upstream) + upstreamCert := upstream.TLS.Certificates[0] + upstreamCA, _ := x509.ParseCertificate(upstreamCert.Certificate[0]) + + proxyAddr, cleanup := startTestProxy(t, ca, df, store, nil, upstreamCAs) + defer cleanup() + + client := httpClientViaProxy(t, proxyAddr, ca.Certificate, upstreamCA) + + resp, err := client.Post(upstream.URL+"/v1/generate", "application/json", strings.NewReader(`{"prompt":"hello"}`)) + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusOK { + t.Errorf("expected 200 OK after retry, got %d", resp.StatusCode) + } + if got := atomic.LoadInt32(&hits); got != 2 { + t.Errorf("expected upstream to be hit exactly twice (original + one retry), got %d", got) + } +} + +func TestIntegration_NonRefreshableInjector_401PassesThroughWithoutRetry(t *testing.T) { + skipIntegration(t) + + ca, err := GenerateCA() + if err != nil { + t.Fatalf("generate CA: %v", err) + } + + var hits int32 + upstream := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&hits, 1) + w.WriteHeader(http.StatusUnauthorized) + })) + defer upstream.Close() + + upstreamURL, _ := url.Parse(upstream.URL) + upstreamHostname := upstreamURL.Hostname() + + df := filter.NewDomainFilter(upstreamHostname) + + store := credentials.NewStore() + store.AddRoute(credentials.Route{ + ExactDomain: upstreamHostname, + Injector: &credentials.BearerInjector{Token: "bad-token"}, + }) + + upstreamCAs := upstreamCertPool(t, upstream) + upstreamCert := upstream.TLS.Certificates[0] + upstreamCA, _ := x509.ParseCertificate(upstreamCert.Certificate[0]) + + proxyAddr, cleanup := startTestProxy(t, ca, df, store, nil, upstreamCAs) + defer cleanup() + + client := httpClientViaProxy(t, proxyAddr, ca.Certificate, upstreamCA) + + resp, err := client.Get(upstream.URL + "/test") + if err != nil { + t.Fatalf("request failed: %v", err) + } + defer resp.Body.Close() + + if resp.StatusCode != http.StatusUnauthorized { + t.Errorf("expected 401 to pass through unmodified, got %d", resp.StatusCode) + } + if got := atomic.LoadInt32(&hits); got != 1 { + t.Errorf("non-refreshable injector should not trigger a retry, expected 1 hit, got %d", got) + } +} + func TestIntegration_ProxyTransport_ResponseHeaderTimeout(t *testing.T) { skipIntegration(t) diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index e415a7a..c90675e 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -1,9 +1,11 @@ package proxy import ( + "bytes" "crypto/tls" "crypto/x509" "fmt" + "io" "log" "net" "net/http" @@ -487,7 +489,19 @@ func New(cfg Config) *http.Server { if cfg.CredStore != nil { switch cfg.CredStore.InjectCredentials(req) { case credentials.InjectOK: - ctx.UserData = credInjectedFlag{} + // Buffer the body so a Refresher-backed route can be retried + // once after an upstream 401; abort if buffering fails. + if resp := prepareRefreshRetry(req, ctx, cfg.CredStore); resp != nil { + return req, resp + } + // Mark the request as credential-injected so the OnResponse + // handler logs upstream errors for it. prepareRefreshRetry may + // have already stashed retry state on ctx.UserData; only set + // the flag when it didn't, so a buffer-failure abort (UserData + // still nil) isn't mislogged as an upstream error. + if ctx.UserData == nil { + ctx.UserData = credInjectedFlag{} + } case credentials.InjectAuthRequired: return req, goproxy.NewResponse(req, goproxy.ContentTypeText, @@ -511,6 +525,20 @@ func New(cfg Config) *http.Server { }, ) + // Retry once with a forced-fresh credential if upstream itself rejects + // the request as unauthorized. This recovers from a credential that + // InjectCredentials considered locally valid but that the real server + // disagreed with (e.g. a cached OAuth token whose local expiry check + // disagreed with Google's after a host suspend/resume cycle). + // + // Registered before the upstream-error logger below so a successful + // recovery replaces the 401 before it would be logged as an upstream error. + proxy.OnResponse(goproxy.StatusCodeIs(http.StatusUnauthorized)).DoFunc( + func(resp *http.Response, ctx *goproxy.ProxyCtx) *http.Response { + return retryWithForcedRefresh(resp, ctx, cfg.CredStore) + }, + ) + // Log upstream error responses for credential-injected requests. // This distinguishes proxy-generated errors from upstream errors. proxy.OnResponse().DoFunc( @@ -540,6 +568,132 @@ func New(cfg Config) *http.Server { } } +// refreshRetryState is stashed on ctx.UserData by prepareRefreshRetry when a +// request was handled by a credentials.Refresher, so the response handler +// can retry once with a forced-fresh credential after an upstream 401. +type refreshRetryState struct { + injector credentials.Refresher + getBody func() (io.ReadCloser, error) + retried bool +} + +// maxRetryBufferBytes caps how much of a request body the proxy buffers in +// memory to enable the upstream-401 retry. Requests up to this size get the +// forced-refresh retry safety net; larger ones (e.g. big multimodal Vertex +// inference payloads) are still forwarded intact but stream through without +// buffering, so they can't be retried. This bounds the memory a hostile agent +// can force the proxy to allocate by POSTing huge bodies to a Refresher-backed +// (*.googleapis.com) route. +const maxRetryBufferBytes = 10 << 20 // 10 MiB + +// prepareRefreshRetry buffers req's body (if any) and stashes retry state on +// ctx.UserData when the matched injector supports ForceRefresh. Buffering +// the body is required because the original io.ReadCloser is drained by the +// first round trip and can't be replayed as-is; it's skipped for injectors +// that can't force-refresh, since those requests will never be retried. +// +// It returns a non-nil response only when buffering fails: at that point the +// body has already been partially drained and can't be forwarded intact, so +// the caller must abort with that response rather than send a truncated +// request upstream. It returns nil in all other cases (proceed normally). +func prepareRefreshRetry(req *http.Request, ctx *goproxy.ProxyCtx, store *credentials.Store) *http.Response { + injector, ok := store.MatchInjector(req).(credentials.Refresher) + if !ok { + return nil + } + + if req.Body == nil || req.Body == http.NoBody { + ctx.UserData = &refreshRetryState{ + injector: injector, + getBody: func() (io.ReadCloser, error) { return http.NoBody, nil }, + } + return nil + } + + // Read up to the cap plus one byte so we can tell whether the body fit. + bodyBytes, err := io.ReadAll(io.LimitReader(req.Body, maxRetryBufferBytes+1)) + if err != nil { + // The body is already partially consumed; forwarding it now would + // send a truncated request upstream. Abort with a 502 instead. + req.Body.Close() + log.Printf("ERROR buffering request body for retry: %v", err) + return goproxy.NewResponse(req, + goproxy.ContentTypeText, + http.StatusBadGateway, + "Proxy failed to buffer request body", + ) + } + + if len(bodyBytes) > maxRetryBufferBytes { + // Too large to hold in memory for a retry. Forward the request intact + // by stitching the already-read prefix back in front of the unread + // remainder and skip retry setup: an oversized body loses the 401 + // retry safety net but is never truncated nor fully buffered. The + // struct fields both capture the original req.Body (evaluated before + // the assignment), so the remainder still streams and Close still + // closes the underlying body. + log.Printf("GCLOUD_TOKEN_REFRESH_RETRY host=%s body exceeds %d-byte cap; forwarding without retry", req.URL.Host, maxRetryBufferBytes) + req.Body = struct { + io.Reader + io.Closer + }{ + Reader: io.MultiReader(bytes.NewReader(bodyBytes), req.Body), + Closer: req.Body, + } + return nil + } + + req.Body.Close() + req.Body = io.NopCloser(bytes.NewReader(bodyBytes)) + getBody := func() (io.ReadCloser, error) { + return io.NopCloser(bytes.NewReader(bodyBytes)), nil + } + req.GetBody = getBody + ctx.UserData = &refreshRetryState{injector: injector, getBody: getBody} + return nil +} + +// retryWithForcedRefresh retries a 401 response exactly once with a +// forced-fresh credential, if the original request was prepared for retry +// by prepareRefreshRetry. Returns the original response unchanged otherwise +// (including when the retry itself fails), so callers always get a response. +func retryWithForcedRefresh(resp *http.Response, ctx *goproxy.ProxyCtx, store *credentials.Store) *http.Response { + state, ok := ctx.UserData.(*refreshRetryState) + if !ok || state.retried || ctx.Req == nil { + return resp + } + state.retried = true // at most one retry, even if this attempt also fails + + if err := state.injector.ForceRefresh(); err != nil { + log.Printf("GCLOUD_TOKEN_REFRESH_RETRY force refresh failed: %v", err) + return resp + } + + body, err := state.getBody() + if err != nil { + log.Printf("GCLOUD_TOKEN_REFRESH_RETRY rebuilding request body failed: %v", err) + return resp + } + ctx.Req.Body = body + + if store.InjectCredentials(ctx.Req) != credentials.InjectOK { + return resp + } + + log.Printf("GCLOUD_TOKEN_REFRESH_RETRY host=%s (upstream 401, retrying with forced-fresh token)", ctx.Req.URL.Host) + newResp, err := ctx.RoundTrip(ctx.Req) + if err != nil { + log.Printf("GCLOUD_TOKEN_REFRESH_RETRY retry round-trip failed: %v", err) + return resp + } + // We're discarding the original 401 in favor of newResp — close its body + // so the underlying connection can be reused instead of leaking. + if resp.Body != nil { + resp.Body.Close() + } + return newResp +} + func stripPort(host string) string { if idx := strings.LastIndex(host, ":"); idx != -1 { return host[:idx] diff --git a/internal/proxy/refresh_retry_test.go b/internal/proxy/refresh_retry_test.go new file mode 100644 index 0000000..3a05d9f --- /dev/null +++ b/internal/proxy/refresh_retry_test.go @@ -0,0 +1,165 @@ +package proxy + +import ( + "io" + "net/http" + "strings" + "testing" + + "github.com/elazarl/goproxy" + + "github.com/bbrowning/paude-proxy/internal/credentials" +) + +// errReader is an io.ReadCloser that always fails on Read, simulating a client +// connection that stalls or resets while the proxy is buffering the body. +type errReader struct{} + +func (errReader) Read([]byte) (int, error) { return 0, io.ErrUnexpectedEOF } +func (errReader) Close() error { return nil } + +// newRefresherStore returns a store with a single example.com route backed by +// a fresh refreshableInjector — the common setup for the retry-buffer tests. +func newRefresherStore() *credentials.Store { + store := credentials.NewStore() + store.AddRoute(credentials.Route{ + ExactDomain: "example.com", + Injector: &refreshableInjector{}, + }) + return store +} + +// When buffering the request body for a possible retry fails, the body has +// already been partially drained and can't be forwarded intact. The request +// must be aborted with a 502 rather than sent upstream with a truncated body. +func TestPrepareRefreshRetry_BodyBufferFailureReturns502(t *testing.T) { + store := newRefresherStore() + + req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/generate", errReader{}) + if err != nil { + t.Fatal(err) + } + ctx := &goproxy.ProxyCtx{} + + resp := prepareRefreshRetry(req, ctx, store) + if resp == nil { + t.Fatal("expected a non-nil abort response when body buffering fails") + } + if resp.StatusCode != http.StatusBadGateway { + t.Errorf("expected 502 Bad Gateway, got %d", resp.StatusCode) + } + if ctx.UserData != nil { + t.Error("retry state should not be stashed when buffering fails") + } +} + +// A successful buffer stashes retry state and leaves the body replayable. +func TestPrepareRefreshRetry_BuffersBodyForRetry(t *testing.T) { + store := newRefresherStore() + + req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/generate", strings.NewReader("payload")) + if err != nil { + t.Fatal(err) + } + ctx := &goproxy.ProxyCtx{} + + if resp := prepareRefreshRetry(req, ctx, store); resp != nil { + t.Fatalf("expected nil response on success, got status %d", resp.StatusCode) + } + state, ok := ctx.UserData.(*refreshRetryState) + if !ok { + t.Fatal("expected refreshRetryState to be stashed on ctx.UserData") + } + + // The buffered body must be readable, and getBody must replay it. + first, err := io.ReadAll(req.Body) + if err != nil { + t.Fatalf("read buffered body: %v", err) + } + if string(first) != "payload" { + t.Errorf("buffered body = %q, want %q", first, "payload") + } + replay, err := state.getBody() + if err != nil { + t.Fatalf("getBody: %v", err) + } + replayed, _ := io.ReadAll(replay) + if string(replayed) != "payload" { + t.Errorf("replayed body = %q, want %q", replayed, "payload") + } +} + +// A body larger than the retry buffer cap must be forwarded intact but +// without retry state — the proxy must never truncate it, and must not hold +// the whole body in memory for a retry it won't perform. +func TestPrepareRefreshRetry_OversizedBodyForwardedWithoutRetry(t *testing.T) { + store := newRefresherStore() + + // One byte over the cap is enough to trip the oversized path. + payload := strings.Repeat("a", maxRetryBufferBytes+1) + req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/generate", strings.NewReader(payload)) + if err != nil { + t.Fatal(err) + } + ctx := &goproxy.ProxyCtx{} + + if resp := prepareRefreshRetry(req, ctx, store); resp != nil { + t.Fatalf("oversized body should be forwarded, not aborted; got status %d", resp.StatusCode) + } + if ctx.UserData != nil { + t.Error("oversized body should not stash retry state") + } + + // The full body must still be forwardable intact (prefix + remainder). + forwarded, err := io.ReadAll(req.Body) + if err != nil { + t.Fatalf("read forwarded body: %v", err) + } + if len(forwarded) != len(payload) { + t.Errorf("forwarded body length = %d, want %d (body was truncated)", len(forwarded), len(payload)) + } + if err := req.Body.Close(); err != nil { + t.Errorf("closing forwarded body: %v", err) + } +} + +// A body exactly at the cap is still buffered for retry (the cap is inclusive). +func TestPrepareRefreshRetry_BodyAtCapBuffered(t *testing.T) { + store := newRefresherStore() + + payload := strings.Repeat("a", maxRetryBufferBytes) + req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/generate", strings.NewReader(payload)) + if err != nil { + t.Fatal(err) + } + ctx := &goproxy.ProxyCtx{} + + if resp := prepareRefreshRetry(req, ctx, store); resp != nil { + t.Fatalf("expected nil response at cap, got status %d", resp.StatusCode) + } + if _, ok := ctx.UserData.(*refreshRetryState); !ok { + t.Fatal("body at cap should be buffered with retry state stashed") + } +} + +// A non-refresher route is never prepared for retry. +func TestPrepareRefreshRetry_NonRefresherSkipped(t *testing.T) { + store := credentials.NewStore() + store.AddRoute(credentials.Route{ + ExactDomain: "example.com", + Injector: &credentials.BearerInjector{Token: "x"}, + }) + + req, err := http.NewRequest(http.MethodPost, "https://example.com/v1/generate", errReader{}) + if err != nil { + t.Fatal(err) + } + ctx := &goproxy.ProxyCtx{} + + if resp := prepareRefreshRetry(req, ctx, store); resp != nil { + t.Errorf("non-refresher route should not abort, got status %d", resp.StatusCode) + } + if ctx.UserData != nil { + t.Error("non-refresher route should not stash retry state") + } +}