diff --git a/auth/authorization_code.go b/auth/authorization_code.go index a65931ce7..5b60438bc 100644 --- a/auth/authorization_code.go +++ b/auth/authorization_code.go @@ -15,6 +15,7 @@ import ( "slices" "strings" "sync" + "time" "github.com/modelcontextprotocol/go-sdk/internal/authutil" "github.com/modelcontextprotocol/go-sdk/internal/util" @@ -34,6 +35,11 @@ type ClientIDMetadataDocumentConfig struct { // DynamicClientRegistrationConfig is used to configure dynamic client registration per // https://modelcontextprotocol.io/specification/2025-11-25/basic/authorization#dynamic-client-registration. +// +// The client registered with an authorization server is reused by later +// authorizations with that server, for the lifetime of the +// [AuthorizationCodeHandler]. It is registered again if its client secret has +// expired, or if an authorization that used it fails. type DynamicClientRegistrationConfig struct { // Metadata to be used in dynamic client registration request as per // https://datatracker.ietf.org/doc/html/rfc7591#section-2. @@ -166,7 +172,7 @@ type AuthorizationCodeHandlerConfig struct { type AuthorizationCodeHandler struct { config *AuthorizationCodeHandlerConfig - // mu protects concurrent access to tokenSource and grantedScopes. + // mu protects concurrent access to tokenSource, grantedScopes and registrations. mu sync.RWMutex // tokenSource is the token source to use for authorization. @@ -174,6 +180,10 @@ type AuthorizationCodeHandler struct { // grantedScopes maps authorization server issuer to the list of scopes granted by that issuer. grantedScopes map[string][]string + + // registrations maps authorization server issuer to the client registered + // with it through dynamic client registration. + registrations map[string]*resolvedClientConfig } var _ OAuthHandler = (*AuthorizationCodeHandler)(nil) @@ -241,6 +251,7 @@ func NewAuthorizationCodeHandler(config *AuthorizationCodeHandlerConfig) (*Autho config: config, tokenSource: config.InitialTokenSource, grantedScopes: make(map[string][]string), + registrations: make(map[string]*resolvedClientConfig), }, nil } @@ -300,7 +311,7 @@ func inferApplicationType(redirectURIs []string) string { // Authorize performs the authorization flow. // It is designed to perform the whole Authorization Code Grant flow. // On success, [AuthorizationCodeHandler.TokenSource] will return a token source with the fetched token. -func (h *AuthorizationCodeHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error { +func (h *AuthorizationCodeHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) (err error) { defer resp.Body.Close() defer io.Copy(io.Discard, resp.Body) @@ -343,6 +354,16 @@ func (h *AuthorizationCodeHandler) Authorize(ctx context.Context, req *http.Requ if err != nil { return err } + if resolvedClientConfig.registrationType == registrationTypeDynamic { + // The registered client may be why this attempt failed, for example + // because the authorization server has deleted it since. Forget it, so + // that the next attempt registers a new one. + defer func() { + if err != nil { + h.forgetDynamicRegistration(asm.Issuer, resolvedClientConfig) + } + }() + } requestedScopes := scopesFromChallenges(wwwChallenges) if len(requestedScopes) == 0 && len(prm.ScopesSupported) > 0 { @@ -530,6 +551,9 @@ type resolvedClientConfig struct { clientID string clientSecret string authStyle oauth2.AuthStyle + // secretExpiresAt is when clientSecret expires, or zero if it does not. + // It is only set for dynamic client registration. + secretExpiresAt time.Time } func selectTokenAuthMethod(supported []string) oauth2.AuthStyle { @@ -599,6 +623,9 @@ func (h *AuthorizationCodeHandler) handleRegistration(ctx context.Context, asm * // 3. Attempt to use dynamic client registration. dcrCfg := h.config.DynamicClientRegistrationConfig if dcrCfg != nil && asm.RegistrationEndpoint != "" { + if cfg := h.dynamicRegistration(asm.Issuer); cfg != nil { + return cfg, nil + } regResp, err := oauthex.RegisterClient(ctx, asm.RegistrationEndpoint, dcrCfg.Metadata, h.config.Client) if err != nil { return nil, fmt.Errorf("failed to register client: %w", err) @@ -608,12 +635,38 @@ func (h *AuthorizationCodeHandler) handleRegistration(ctx context.Context, asm * clientID: regResp.ClientID, clientSecret: regResp.ClientSecret, authStyle: authMethodToStyle(regResp.TokenEndpointAuthMethod), + secretExpiresAt: regResp.ClientSecretExpiresAt, } + h.mu.Lock() + h.registrations[asm.Issuer] = cfg + h.mu.Unlock() return cfg, nil } return nil, fmt.Errorf("no configured client registration methods are supported by the authorization server") } +// dynamicRegistration returns the client registered with issuer by an earlier +// authorization, or nil if there is none or its client secret has expired. +func (h *AuthorizationCodeHandler) dynamicRegistration(issuer string) *resolvedClientConfig { + h.mu.RLock() + defer h.mu.RUnlock() + cfg := h.registrations[issuer] + if cfg == nil || (!cfg.secretExpiresAt.IsZero() && !time.Now().Before(cfg.secretExpiresAt)) { + return nil + } + return cfg +} + +// forgetDynamicRegistration drops the client registered with issuer, unless +// another authorization has replaced it since cfg was returned. +func (h *AuthorizationCodeHandler) forgetDynamicRegistration(issuer string, cfg *resolvedClientConfig) { + h.mu.Lock() + defer h.mu.Unlock() + if h.registrations[issuer] == cfg { + delete(h.registrations, issuer) + } +} + type authResult struct { *AuthorizationResult // usedCodeVerifier is the PKCE code verifier used to obtain the authorization code. diff --git a/auth/authorization_code_test.go b/auth/authorization_code_test.go index 61e6856dc..97002aa7e 100644 --- a/auth/authorization_code_test.go +++ b/auth/authorization_code_test.go @@ -14,6 +14,7 @@ import ( "slices" "strings" "testing" + "time" "github.com/google/go-cmp/cmp" "github.com/google/go-cmp/cmp/cmpopts" @@ -1044,6 +1045,161 @@ func TestDynamicRegistration(t *testing.T) { } } +// dcrTestEnv is a handler that registers its client dynamically with a fake +// authorization server, for a resource whose metadata points to that server. +type dcrTestEnv struct { + handler *AuthorizationCodeHandler + issuer string + resourceURL string + metadataURL string + // clientIDs holds the client_id of each authorization request, in order. + clientIDs []string +} + +func newDCRTestEnv(t *testing.T) *dcrTestEnv { + t.Helper() + authServer := oauthtest.NewFakeAuthorizationServer(oauthtest.Config{ + RegistrationConfig: &oauthtest.RegistrationConfig{DynamicClientRegistrationEnabled: true}, + }) + authServer.Start(t) + + resourceMux := http.NewServeMux() + resourceServer := httptest.NewServer(resourceMux) + t.Cleanup(resourceServer.Close) + e := &dcrTestEnv{ + issuer: authServer.URL(), + resourceURL: resourceServer.URL + "/resource", + metadataURL: resourceServer.URL + "/.well-known/oauth-protected-resource/resource", + } + resourceMux.Handle("/.well-known/oauth-protected-resource/resource", ProtectedResourceMetadataHandler(&oauthex.ProtectedResourceMetadata{ + Resource: e.resourceURL, + AuthorizationServers: []string{authServer.URL()}, + })) + + noRedirectClient := &http.Client{CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }} + handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{ + DynamicClientRegistrationConfig: &DynamicClientRegistrationConfig{ + Metadata: &oauthex.ClientRegistrationMetadata{ + RedirectURIs: []string{"http://localhost:12345/callback"}, + }, + }, + AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + u, err := url.Parse(args.URL) + if err != nil { + return nil, err + } + e.clientIDs = append(e.clientIDs, u.Query().Get("client_id")) + resp, err := noRedirectClient.Get(args.URL) + if err != nil { + return nil, err + } + defer resp.Body.Close() + loc, err := resp.Location() + if err != nil { + return nil, err + } + return &AuthorizationResult{ + Code: loc.Query().Get("code"), + State: loc.Query().Get("state"), + Iss: loc.Query().Get("iss"), + }, nil + }, + }) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler failed: %v", err) + } + e.handler = handler + return e +} + +// authorize calls Authorize for a response with the given status and Bearer +// challenge parameters. +func (e *dcrTestEnv) authorize(status int, params string) error { + req := httptest.NewRequest(http.MethodGet, e.resourceURL, nil) + resp := &http.Response{ + StatusCode: status, + Header: make(http.Header), + Body: http.NoBody, + Request: req, + } + resp.Header.Set("WWW-Authenticate", fmt.Sprintf(`Bearer %s, resource_metadata="%s"`, params, e.metadataURL)) + return e.handler.Authorize(context.Background(), req, resp) +} + +func TestAuthorize_ReusesDynamicRegistration(t *testing.T) { + e := newDCRTestEnv(t) + if err := e.authorize(http.StatusUnauthorized, `scope="read"`); err != nil { + t.Fatalf("first Authorize failed: %v", err) + } + // Step-up authorization with the same authorization server. + if err := e.authorize(http.StatusForbidden, `error="insufficient_scope", scope="admin"`); err != nil { + t.Fatalf("second Authorize failed: %v", err) + } + if len(e.clientIDs) != 2 { + t.Fatalf("got %d authorization requests, want 2", len(e.clientIDs)) + } + if e.clientIDs[0] != e.clientIDs[1] { + t.Errorf("step-up authorization used client_id %q, want the registered %q", e.clientIDs[1], e.clientIDs[0]) + } +} + +func TestAuthorize_ForgetsDynamicRegistrationAfterFailure(t *testing.T) { + e := newDCRTestEnv(t) + // A client that the authorization server no longer knows, for example + // because it was deleted after an earlier authorization registered it. + e.handler.registrations[e.issuer] = &resolvedClientConfig{ + registrationType: registrationTypeDynamic, + clientID: "deleted_client_id", + clientSecret: "deleted_client_secret", + } + if err := e.authorize(http.StatusUnauthorized, `scope="read"`); err == nil { + t.Fatal("Authorize with a deleted client succeeded, want an error") + } + if err := e.authorize(http.StatusUnauthorized, `scope="read"`); err != nil { + t.Fatalf("Authorize after a failed attempt failed: %v", err) + } + if len(e.clientIDs) != 2 { + t.Fatalf("got %d authorization requests, want 2", len(e.clientIDs)) + } + if e.clientIDs[1] == "deleted_client_id" { + t.Error("Authorize reused the client of the failed attempt, want a new registration") + } +} + +func TestHandleRegistration_StoredDynamicRegistration(t *testing.T) { + for _, tt := range []struct { + name string + secretExpiresAt time.Time + wantReuse bool + }{ + {"secret never expires", time.Time{}, true}, + {"secret not expired", time.Now().Add(time.Hour), true}, + {"secret expired", time.Now().Add(-time.Hour), false}, + } { + t.Run(tt.name, func(t *testing.T) { + e := newDCRTestEnv(t) + stored := &resolvedClientConfig{ + registrationType: registrationTypeDynamic, + clientID: "stored_client_id", + clientSecret: "stored_client_secret", + secretExpiresAt: tt.secretExpiresAt, + } + e.handler.registrations[e.issuer] = stored + asm, err := GetAuthServerMetadata(t.Context(), e.issuer, http.DefaultClient) + if err != nil { + t.Fatalf("GetAuthServerMetadata() unexpected error = %v", err) + } + got, err := e.handler.handleRegistration(t.Context(), asm) + if err != nil { + t.Fatalf("handleRegistration() error = %v, want nil", err) + } + if reused := got.clientID == stored.clientID; reused != tt.wantReuse { + t.Errorf("handleRegistration() clientID = %q, reused = %v, want reused = %v", got.clientID, reused, tt.wantReuse) + } + }) + } +} + func TestValidateIssuerResponse(t *testing.T) { const expectedIssuer = "https://auth.example.com"