auth: reuse the dynamically registered client across authorizations

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,14 +172,18 @@ 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.
 	tokenSource oauth2.TokenSource
 
 	// 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.

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"