auth: bind refreshing token source to a background context, not the request ctx (#988)
## 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 "<token endpoint>": 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) <noreply@anthropic.com>
Co-authored-by: Guglielmo Colombo <guglielmoc@google.com>
This commit is contained in:
parent
96074c9918
commit
d9728a88f3
1 changed files with 153 additions and 2 deletions
|
|
@ -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{
|
||||
|
|
|
|||
Loading…
Reference in a new issue