From d9728a88f3cdd50a5809730b5d346b74cddb6ad0 Mon Sep 17 00:00:00 2001 From: Thomas Bechtold Date: Tue, 16 Jun 2026 16:42:53 +0200 Subject: [PATCH] auth: bind refreshing token source to a background context, not the request ctx (#988) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem `AuthorizationCodeHandler.exchangeAuthorizationCode` builds the long-lived token source from the **request context** passed to `Authorize()`: ```go clientCtx := context.WithValue(ctx, oauth2.HTTPClient, h.config.Client) ... h.tokenSource = cfg.TokenSource(clientCtx, token) ``` `golang.org/x/oauth2` captures that context and reuses it for **every** future refresh: - `oauth2.go`: `Config.TokenSource(ctx, …)` stores `ctx` in `tokenRefresher{ctx}` - `tokenRefresher.Token()` (no ctx arg) calls `retrieveToken(tf.ctx, …)` - `internal/token.go`: `ContextClient(ctx).Do(req.WithContext(ctx))` `Authorize()` is normally invoked from a request- or connect-scoped context — e.g. the streamable client transport calls it during `setMCPHeaders`/`initialize`, and that context is cancelled once the operation completes. The cached access token works until expiry; the **first refresh after expiry** then fails instantly with `context canceled`, before any HTTP request is sent. The connection can no longer refresh. `setMCPHeaders` only re-invokes `Authorize()` for an `invalid_grant` `oauth2.RetrieveError` (#917). A `context canceled` is not a `RetrieveError`, so it is never recovered — the connection is wedged until the client reconnects. ### Reproduce 1. Connect a `StreamableClientTransport` with an `AuthorizationCodeHandler` to a server whose `initialize` returns 401 (so `Authorize()` runs during connect, under the connect context). 2. Let the connect context be cancelled after connect returns (the common pattern: a `context.WithTimeout(parent, …)` + `defer cancel()` bounding the interactive flow). 3. Wait for the access token to expire and issue a request. 4. Every request fails with `Post "": context canceled` and `dur≈0` — the refresh never hits the network. ## Fix A token source outlives the request that created it, so bind its refreshes to `context.Background()` (still carrying the configured HTTP client via `oauth2.HTTPClient`). The one-shot token exchange keeps using the request context. ```go refreshCtx := context.WithValue(context.Background(), oauth2.HTTPClient, h.config.Client) h.tokenSource = cfg.TokenSource(refreshCtx, token) ``` `auth` package tests pass. Behaviour is otherwise unchanged: the same HTTP client value is carried; only the cancellation parent differs. 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: Guglielmo Colombo --- auth/authorization_code.go | 12 ++- auth/authorization_code_test.go | 95 +++++++++++++++++++ .../oauthtest/fake_authorization_server.go | 48 +++++++++- 3 files changed, 153 insertions(+), 2 deletions(-) diff --git a/auth/authorization_code.go b/auth/authorization_code.go index a3daeec..963884d 100644 --- a/auth/authorization_code.go +++ b/auth/authorization_code.go @@ -605,7 +605,17 @@ func (h *AuthorizationCodeHandler) exchangeAuthorizationCode(ctx context.Context if err != nil { return fmt.Errorf("token exchange failed: %w", err) } - h.tokenSource = cfg.TokenSource(clientCtx, token) + // The token source outlives this authorization request: it is stored on the + // handler and used by the transport for the lifetime of the connection. The + // oauth2 library captures the context passed to TokenSource and reuses it for + // every subsequent token refresh (see golang.org/x/oauth2: tokenRefresher + // retains the context and passes it to each refresh round-trip). Binding it to + // the per-request ctx makes all later refreshes fail with "context canceled" + // once that request (or the connect operation that triggered authorization) + // completes. Use a background context that still carries the configured HTTP + // client so refreshes keep working for the life of the token source. + refreshCtx := context.WithValue(context.Background(), oauth2.HTTPClient, h.config.Client) + h.tokenSource = cfg.TokenSource(refreshCtx, token) return nil } diff --git a/auth/authorization_code_test.go b/auth/authorization_code_test.go index c84b003..15ca69f 100644 --- a/auth/authorization_code_test.go +++ b/auth/authorization_code_test.go @@ -115,6 +115,101 @@ func TestAuthorize(t *testing.T) { } } +// TestAuthorize_RefreshAfterContextCancel verifies that the token source built +// by Authorize keeps refreshing after the context passed to Authorize is +// cancelled. The token source is stored on the handler and used by the +// transport for the whole connection lifetime, whereas Authorize is typically +// called from a request- or connect-scoped context. golang.org/x/oauth2 +// captures the context handed to Config.TokenSource and reuses it for every +// refresh, so binding it to the request context made every refresh after the +// access token expired fail with "context canceled". Regression test for that. +func TestAuthorize_RefreshAfterContextCancel(t *testing.T) { + authServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + // expires_in below oauth2's 10s expiry delta, so the reuse token source + // treats the access token as expired immediately and must refresh. + AccessTokenTTL: 1, + IssueRefreshToken: true, + RegistrationConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "test_client_id": { + Secret: "test_client_secret", + RedirectURIs: []string{"http://localhost:12345/callback"}, + }, + }, + }, + }) + authServer.Start(t) + + resourceMux := http.NewServeMux() + resourceServer := httptest.NewServer(resourceMux) + t.Cleanup(resourceServer.Close) + resourceURL := resourceServer.URL + "/resource" + resourceMux.Handle("/.well-known/oauth-protected-resource/resource", ProtectedResourceMetadataHandler(&oauthex.ProtectedResourceMetadata{ + Resource: resourceURL, + AuthorizationServers: []string{authServer.URL()}, + })) + + handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{ + RedirectURL: "http://localhost:12345/callback", + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "test_client_id", + ClientSecretAuth: &oauthex.ClientSecretAuth{ClientSecret: "test_client_secret"}, + }, + AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + client := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }} + resp, err := client.Get(args.URL) + if err != nil { + return nil, fmt.Errorf("failed to visit auth URL: %v", err) + } + defer resp.Body.Close() + location, err := resp.Location() + if err != nil { + return nil, fmt.Errorf("failed to get location header: %v", err) + } + return &AuthorizationResult{ + Code: location.Query().Get("code"), + State: location.Query().Get("state"), + Iss: location.Query().Get("iss"), + }, nil + }, + }) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler failed: %v", err) + } + + req := httptest.NewRequest(http.MethodGet, resourceURL, nil) + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: make(http.Header), + Body: http.NoBody, + Request: req, + } + resp.Header.Set("WWW-Authenticate", "Bearer resource_metadata="+resourceServer.URL+"/.well-known/oauth-protected-resource/resource") + + // Authorize under a context that we cancel immediately afterwards, mimicking + // the request/connect context the transport passes and that is already done + // by the time a later token refresh runs. + ctx, cancel := context.WithCancel(context.Background()) + if err := handler.Authorize(ctx, req, resp); err != nil { + t.Fatalf("Authorize failed: %v", err) + } + cancel() + + tokenSource, err := handler.TokenSource(context.Background()) + if err != nil { + t.Fatalf("Failed to get token source: %v", err) + } + // The access token is already expired, so this forces a refresh round-trip. + // It must not fail with "context canceled" from the cancelled Authorize ctx. + token, err := tokenSource.Token() + if err != nil { + t.Fatalf("token refresh after Authorize context cancellation failed: %v", err) + } + if token.AccessToken != "test_access_token_refreshed" { + t.Errorf("expected refreshed access token %q, got %q", "test_access_token_refreshed", token.AccessToken) + } +} + func TestAuthorize_ScopeAccumulation(t *testing.T) { authServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ RegistrationConfig: &oauthtest.RegistrationConfig{ diff --git a/internal/oauthtest/fake_authorization_server.go b/internal/oauthtest/fake_authorization_server.go index c180f86..e5134fb 100644 --- a/internal/oauthtest/fake_authorization_server.go +++ b/internal/oauthtest/fake_authorization_server.go @@ -89,6 +89,26 @@ type Config struct { // TokenScopeFunc, if set, is called with the scope from the authorization // request and returns the scope string to include in the token response. TokenScopeFunc func(requestedScope string) string + // AccessTokenTTL, if non-zero, is the expires_in (in seconds) returned by the + // /token endpoint for both the authorization_code and refresh_token grants. + // When zero a default of 3600 is used. Set it small to force a client's reuse + // token source to treat the access token as expired and refresh it. + AccessTokenTTL int + // IssueRefreshToken, if true, includes a refresh_token in token responses and + // enables grant_type=refresh_token at the /token endpoint. + IssueRefreshToken bool +} + +// testRefreshToken is the refresh token issued and accepted by the fake server +// when Config.IssueRefreshToken is set. +const testRefreshToken = "test_refresh_token" + +// accessTokenExpiresIn returns the expires_in value to use in token responses. +func (s *FakeAuthorizationServer) accessTokenExpiresIn() int { + if s.config.AccessTokenTTL != 0 { + return s.config.AccessTokenTTL + } + return 3600 } // FakeAuthorizationServer is a fake OAuth 2.0 Authorization Server for testing. @@ -298,6 +318,8 @@ func (s *FakeAuthorizationServer) handleToken(w http.ResponseWriter, r *http.Req s.handleJWTBearerGrant(w, r) case "client_credentials": s.handleClientCredentialsGrant(w, r) + case "refresh_token": + s.handleRefreshTokenGrant(w, r) default: http.Error(w, fmt.Sprintf("unsupported grant_type: %s", grantType), http.StatusBadRequest) } @@ -329,7 +351,10 @@ func (s *FakeAuthorizationServer) handleAuthorizationCodeGrant(w http.ResponseWr resp := map[string]any{ "access_token": "test_access_token", "token_type": "Bearer", - "expires_in": 3600, + "expires_in": s.accessTokenExpiresIn(), + } + if s.config.IssueRefreshToken { + resp["refresh_token"] = testRefreshToken } if s.config.TokenScopeFunc != nil { if scope := s.config.TokenScopeFunc(codeInfo.Scope); scope != "" { @@ -340,6 +365,27 @@ func (s *FakeAuthorizationServer) handleAuthorizationCodeGrant(w http.ResponseWr json.NewEncoder(w).Encode(resp) } +// handleRefreshTokenGrant implements grant_type=refresh_token (RFC 6749 Section +// 6) when Config.IssueRefreshToken is set, returning a distinct access token so +// callers can observe that a refresh occurred. +func (s *FakeAuthorizationServer) handleRefreshTokenGrant(w http.ResponseWriter, r *http.Request) { + if !s.config.IssueRefreshToken { + http.Error(w, "refresh_token grant not supported", http.StatusBadRequest) + return + } + if r.Form.Get("refresh_token") != testRefreshToken { + http.Error(w, "invalid refresh_token", http.StatusBadRequest) + return + } + w.Header().Set("Content-Type", "application/json") + json.NewEncoder(w).Encode(map[string]any{ + "access_token": "test_access_token_refreshed", + "token_type": "Bearer", + "expires_in": s.accessTokenExpiresIn(), + "refresh_token": testRefreshToken, + }) +} + func (s *FakeAuthorizationServer) handleJWTBearerGrant(w http.ResponseWriter, r *http.Request) { if s.config.JWTBearerConfig == nil { http.Error(w, "JWT bearer grant not supported", http.StatusBadRequest)