diff --git a/auth/authorization_code.go b/auth/authorization_code.go index a49c100e2..a65931ce7 100644 --- a/auth/authorization_code.go +++ b/auth/authorization_code.go @@ -700,7 +700,15 @@ func (h *AuthorizationCodeHandler) exchangeAuthorizationCode(ctx context.Context refreshCtx := context.WithValue(context.Background(), oauth2.HTTPClient, h.config.Client) var ts oauth2.TokenSource if h.config.NewTokenSource == nil { - ts = cfg.TokenSource(refreshCtx, token) + if token.RefreshToken == "" { + // cfg.TokenSource would fail once the access token expires, and the + // transport would then give up instead of re-authorizing. Keep + // returning the token so that the server's 401 starts a new + // authorization. + ts = oauth2.StaticTokenSource(token) + } else { + ts = cfg.TokenSource(refreshCtx, token) + } } else { var err error ts, err = h.config.NewTokenSource(refreshCtx, cfg, token) diff --git a/auth/authorization_code_test.go b/auth/authorization_code_test.go index ee995e211..61e6856dc 100644 --- a/auth/authorization_code_test.go +++ b/auth/authorization_code_test.go @@ -210,6 +210,89 @@ func TestAuthorize_RefreshAfterContextCancel(t *testing.T) { } } +// TestAuthorize_ExpiredTokenWithoutRefreshToken verifies that the token source +// built by Authorize keeps returning the access token after it expires when the +// authorization server issued no refresh token. Failing instead left the +// streamable transport unable to send any request, so the server's 401 never +// reached Authorize. +func TestAuthorize_ExpiredTokenWithoutRefreshToken(t *testing.T) { + authServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + // expires_in below oauth2's 10s expiry delta, so the access token is + // treated as expired immediately. + AccessTokenTTL: 1, + 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") + if err := handler.Authorize(context.Background(), req, resp); err != nil { + t.Fatalf("Authorize failed: %v", err) + } + + tokenSource, err := handler.TokenSource(context.Background()) + if err != nil { + t.Fatalf("Failed to get token source: %v", err) + } + token, err := tokenSource.Token() + if err != nil { + t.Fatalf("Token() on an expired token without a refresh token failed: %v", err) + } + if token.AccessToken != "test_access_token" { + t.Errorf("expected access token %q, got %q", "test_access_token", token.AccessToken) + } +} + func TestAuthorize_ScopeAccumulation(t *testing.T) { authServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ RegistrationConfig: &oauthtest.RegistrationConfig{ diff --git a/mcp/streamable_client_test.go b/mcp/streamable_client_test.go index 16bd6774d..928b65edc 100644 --- a/mcp/streamable_client_test.go +++ b/mcp/streamable_client_test.go @@ -21,7 +21,9 @@ import ( "github.com/google/go-cmp/cmp" "github.com/modelcontextprotocol/go-sdk/auth" "github.com/modelcontextprotocol/go-sdk/internal/jsonrpc2" + "github.com/modelcontextprotocol/go-sdk/internal/oauthtest" "github.com/modelcontextprotocol/go-sdk/jsonrpc" + "github.com/modelcontextprotocol/go-sdk/oauthex" "golang.org/x/oauth2" ) @@ -1304,6 +1306,98 @@ func TestStreamableClientOAuth_RetrieveError(t *testing.T) { } } +// TestStreamableClientOAuth_ExpiredTokenWithoutRefreshToken verifies that a +// session authorized by an [auth.AuthorizationCodeHandler] keeps working after +// its access token expires when the authorization server issued no refresh +// token: the server's 401 must lead to a new authorization, as it does after a +// rejected refresh. +func TestStreamableClientOAuth_ExpiredTokenWithoutRefreshToken(t *testing.T) { + ctx := context.Background() + + authServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + // expires_in below oauth2's 10s expiry delta, so the access token is + // treated as expired as soon as it is issued. + AccessTokenTTL: 1, + RegistrationConfig: &oauthtest.RegistrationConfig{ + PreregisteredClients: map[string]oauthtest.ClientInfo{ + "test_client_id": { + Secret: "test_client_secret", + RedirectURIs: []string{"http://localhost:12345/callback"}, + }, + }, + }, + }) + authServer.Start(t) + + // rejectToken makes the server reject the current access token, as it + // would once the token has expired. A new authorization clears it. + var rejectToken atomic.Bool + verifier := func(ctx context.Context, token string, req *http.Request) (*auth.TokenInfo, error) { + if token != "test_access_token" || rejectToken.Load() { + return nil, auth.ErrInvalidToken + } + return &auth.TokenInfo{Expiration: time.Now().Add(time.Hour)}, nil + } + mux := http.NewServeMux() + httpServer := httptest.NewServer(mux) + t.Cleanup(httpServer.Close) + mcpURL := httpServer.URL + "/mcp" + metadataURL := httpServer.URL + "/.well-known/oauth-protected-resource/mcp" + mux.Handle("/.well-known/oauth-protected-resource/mcp", auth.ProtectedResourceMetadataHandler(&oauthex.ProtectedResourceMetadata{ + Resource: mcpURL, + AuthorizationServers: []string{authServer.URL()}, + })) + server := NewServer(testImpl, nil) + streamHandler := NewStreamableHTTPHandler(func(*http.Request) *Server { return server }, nil) + mux.Handle("/mcp", auth.RequireBearerToken(verifier, &auth.RequireBearerTokenOptions{ResourceMetadataURL: metadataURL})(streamHandler)) + + var authorizations atomic.Int32 + oauthHandler, err := auth.NewAuthorizationCodeHandler(&auth.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 *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + authorizations.Add(1) + rejectToken.Store(false) + client := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }} + resp, err := client.Get(args.URL) + if err != nil { + return nil, err + } + defer resp.Body.Close() + location, err := resp.Location() + if err != nil { + return nil, err + } + return &auth.AuthorizationResult{ + Code: location.Query().Get("code"), + State: location.Query().Get("state"), + Iss: location.Query().Get("iss"), + }, nil + }, + }) + if err != nil { + t.Fatal(err) + } + + client := NewClient(testImpl, nil) + session, err := client.Connect(ctx, &StreamableClientTransport{Endpoint: mcpURL, OAuthHandler: oauthHandler}, nil) + if err != nil { + t.Fatalf("client.Connect() failed: %v", err) + } + defer session.Close() + + rejectToken.Store(true) + if _, err := session.ListTools(ctx, nil); err != nil { + t.Fatalf("ListTools() after the access token expired failed: %v", err) + } + if got := authorizations.Load(); got < 2 { + t.Errorf("got %d authorizations, want a new one after the access token expired", got) + } +} + // discoverResult is the canned successful DiscoverResult returned by // fakeStreamableServer setups in the tests below. var discoverResult = &DiscoverResult{