Retry once with a forced credential refresh on upstream 401

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

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)
+	}
+}

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 {