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
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in a new issue