diff --git a/auth/auth_example_test.go b/auth/auth_example_test.go new file mode 100644 index 00000000..bf97eb87 --- /dev/null +++ b/auth/auth_example_test.go @@ -0,0 +1,222 @@ +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. +// Use of this source code is governed by the license +// that can be found in the LICENSE file. + +package auth_test + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "net/url" + "sync" + + "github.com/modelcontextprotocol/go-sdk/auth" + "github.com/modelcontextprotocol/go-sdk/mcp" + "github.com/modelcontextprotocol/go-sdk/oauthex" + "golang.org/x/oauth2" +) + +// savingTokenSource is an oauth2.TokenSource that passes the oauth2 config and +// token to the given saver function each time the access token value changes. +type savingTokenSource struct { + mu sync.Mutex + src oauth2.TokenSource + saver func(*oauth2.Config, *oauth2.Token) error + config *oauth2.Config + accessToken string +} + +func (s *savingTokenSource) Token() (*oauth2.Token, error) { + s.mu.Lock() + defer s.mu.Unlock() + tok, err := s.src.Token() + if err != nil { + return nil, err + } + if s.accessToken != tok.AccessToken { + s.accessToken = tok.AccessToken + // This saver implementation always returns nil. + _ = s.saver(s.config, tok) + } + return tok, nil +} + +// NewSavingTokenSource persists OAuth 2.0 sessions by intercepting token +// changes from the wrapped oauth2.TokenSource. When this wrapper detects an +// access token change, it calls the provided session saver with the +// oauth2.Config and the new oauth2.Token. +// +// initial is an optional access token the caller may already hold (such as a +// token loaded from storage). It initialises the wrapper's state so that if +// the first wrapped.Token() call returns the same token, does not trigger a +// redundant call to saver(). Pass nil when there is no existing token, in +// which case the first token produced by wrapped.Token() is saved. +func NewSavingTokenSource(wrapped oauth2.TokenSource, config *oauth2.Config, initial *oauth2.Token, saver func(*oauth2.Config, *oauth2.Token) error) oauth2.TokenSource { + if wrapped == nil { + return nil + } + if saver == nil { + return wrapped + } + var accessToken string + if initial != nil { + accessToken = initial.AccessToken + } + return &savingTokenSource{ + src: wrapped, + saver: saver, + config: config, + accessToken: accessToken, + } +} + +// sessionStore is an in-memory OAuth 2.0 session store used to demonstrate +// persistence. A production implementation would persist the session to disk or +// a secret store. +type sessionStore struct { + config *oauth2.Config + token *oauth2.Token +} + +// save persists the given OAuth 2.0 config and token. +func (s *sessionStore) save(config *oauth2.Config, token *oauth2.Token) error { + fmt.Printf("Saving token: %s\n", token.AccessToken) + s.config = config + s.token = token + return nil +} + +// restore loads a previously persisted OAuth 2.0 session, if one exists. +func (s *sessionStore) restore() (*oauth2.Config, *oauth2.Token, error) { + if s.config != nil && s.token != nil { + fmt.Println("Restoring session.") + } else { + fmt.Println("No session found to restore.") + } + return s.config, s.token, nil +} + +// newMockAuthServer returns an httptest.Server that simulates both an MCP +// resource server requiring authorization and its OAuth authorization server. +func newMockAuthServer() *httptest.Server { + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case "/.well-known/oauth-authorization-server": + w.Header().Set("Content-Type", "application/json") + _, _ = fmt.Fprintf(w, `{"issuer": "http://%s", "authorization_endpoint": "http://%s/auth", "token_endpoint": "http://%s/token", "code_challenge_methods_supported": ["S256"]}`, r.Host, r.Host, r.Host) + case "/token": + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"access_token": "mock-token", "token_type": "bearer"}`)) + default: + // The mock MCP endpoint returns 401 until the client presents a valid + // bearer token. + if r.Header.Get("Authorization") == "Bearer mock-token" { + // A real server would return a valid MCP message here. The empty + // response causes Connect to return an error after authorization, + // which is ignored for this example which demonstrates token + // persistence only. + w.WriteHeader(http.StatusOK) + return + } + w.Header().Set("WWW-Authenticate", "Bearer") + w.WriteHeader(http.StatusUnauthorized) + } + })) +} + +// This example shows how OAuth2 session persistence might be implemented. It +// connects twice using a shared in-memory session store: the first connection +// authorizes and saves the session, and the second restores and reuses it. +func Example_persistence() { + // Simulate an MCP server that requires authorization and an OAuth server. + mockServer := newMockAuthServer() + defer mockServer.Close() + + // store persists the OAuth2 session across both connections. + store := &sessionStore{} + + // connect performs a single client connection, restoring any saved session + // beforehand and saving any new session acquired during authorization. + connect := func() { + // Load the OAuth2 session if available, and use it. + var initialTS oauth2.TokenSource + cfg, tok, err := store.restore() + if err == nil && cfg != nil && tok != nil { + initialTS = NewSavingTokenSource( + cfg.TokenSource(context.Background(), tok), + cfg, tok, store.save, + ) + } + + // Configure and initialize the AuthorizationCodeHandler with a + // NewTokenSource that saves the session when the access token changes, + // and the InitialTokenSource set to the session loaded via restore(). + config := &auth.AuthorizationCodeHandlerConfig{ + RedirectURL: "http://localhost/callback", + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "example", + }, + Client: mockServer.Client(), + AuthorizationCodeFetcher: func(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { + fmt.Println("No token source found. Transport is calling Authorize()...") + // Extract the generated state from the authorization URL + u, _ := url.Parse(args.URL) + state := u.Query().Get("state") + return &auth.AuthorizationResult{Code: "mock-code", State: state}, nil + }, + NewTokenSource: func(ctx context.Context, cfg *oauth2.Config, token *oauth2.Token) (oauth2.TokenSource, error) { + // This save implementation always returns nil. + _ = store.save(cfg, token) + return NewSavingTokenSource( + cfg.TokenSource(ctx, token), cfg, token, store.save, + ), nil + }, + InitialTokenSource: initialTS, + } + handler, err := auth.NewAuthorizationCodeHandler(config) + if err != nil { + fmt.Printf("Error creating handler: %v\n", err) + return + } + + // Set the constructed handler on a transport. + transport := &mcp.StreamableClientTransport{ + Endpoint: mockServer.URL + "/sse", + OAuthHandler: handler, + } + + // Create a client and attempt to connect using the configured + // transport. The transport will automatically: + // 1. Call TokenSource() to check for an existing session. This will + // return InitialTokenSource, if set. + // 2. Try the MCP endpoint, and encounter a 401 response from the mock + // server. + // 3. Call Authorize() to perform the OAuth flow, which calls + // NewTokenSource. In this example NewTokenSource saves the newly + // acquired session. + // 4. Retry the MCP endpoint with a valid bearer token and get a 200. + client := mcp.NewClient( + &mcp.Implementation{Name: "example", Version: "1.0.0"}, nil, + ) + // Response ignored: this example asserts authorization only. + _, _ = client.Connect(context.Background(), transport, nil) + } + + // The first connection has no saved session, so it authorizes and saves. + fmt.Println("--- First connect ---") + connect() + // The second connection restores the saved session and reuses it, so no + // further authorization or save occurs. + fmt.Println("--- Second connect ---") + connect() + + // Output: + // --- First connect --- + // No session found to restore. + // No token source found. Transport is calling Authorize()... + // Saving token: mock-token + // --- Second connect --- + // Restoring session. +} diff --git a/auth/authorization_code.go b/auth/authorization_code.go index 963884d0..d897d848 100644 --- a/auth/authorization_code.go +++ b/auth/authorization_code.go @@ -14,6 +14,7 @@ import ( "net/url" "slices" "strings" + "sync" "github.com/modelcontextprotocol/go-sdk/internal/authutil" "github.com/modelcontextprotocol/go-sdk/internal/util" @@ -126,6 +127,24 @@ type AuthorizationCodeHandlerConfig struct { // https://modelcontextprotocol.io/docs/tutorials/security/security_best_practices#server-side-request-forgery-ssrf // If not provided, http.DefaultClient will be used. Client *http.Client + + // NewTokenSource is an optional function that can be set to construct the + // token source that will be used by the [AuthorizationCodeHandler]. If + // non-nil, it is called after the authorization code is successfully + // exchanged for a token in [AuthorizationCodeHandler.Authorize] + // to obtain the [oauth2.TokenSource] returned by + // [AuthorizationCodeHandler.TokenSource]. Implementations must use the + // provided context, which is properly configured for constructing a + // TokenSource. The default is to call [oauth2.Config.TokenSource]. + NewTokenSource func(context.Context, *oauth2.Config, *oauth2.Token) (oauth2.TokenSource, error) + + // InitialTokenSource is an optional field that can be set to inject the + // token source that will be used by the [AuthorizationCodeHandler]. If + // non-nil, it is set as the token source that will be returned by + // [AuthorizationCodeHandler.TokenSource] during handler initialization. + // The default is nil, which means no token source has been set initially, + // and will trigger a call to [AuthorizationCodeHandler.Authorize]. + InitialTokenSource oauth2.TokenSource } // AuthorizationCodeHandler is an implementation of [OAuthHandler] that uses @@ -133,6 +152,9 @@ type AuthorizationCodeHandlerConfig struct { type AuthorizationCodeHandler struct { config *AuthorizationCodeHandlerConfig + // mu protects concurrent access to tokenSource and grantedScopes. + mu sync.RWMutex + // tokenSource is the token source to use for authorization. tokenSource oauth2.TokenSource @@ -143,6 +165,8 @@ type AuthorizationCodeHandler struct { var _ OAuthHandler = (*AuthorizationCodeHandler)(nil) func (h *AuthorizationCodeHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) { + h.mu.RLock() + defer h.mu.RUnlock() return h.tokenSource, nil } @@ -199,6 +223,7 @@ func NewAuthorizationCodeHandler(config *AuthorizationCodeHandlerConfig) (*Autho } return &AuthorizationCodeHandler{ config: config, + tokenSource: config.InitialTokenSource, grantedScopes: make(map[string][]string), }, nil } @@ -304,7 +329,10 @@ func (h *AuthorizationCodeHandler) Authorize(ctx context.Context, req *http.Requ // Accumulate scopes: union previously granted scopes with the newly // challenged scopes so that step-up authorization does not lose // permissions granted in earlier rounds (SEP-2350). - requestedScopes = authutil.UnionScopes(h.grantedScopes[asm.Issuer], requestedScopes) + h.mu.RLock() + granted := h.grantedScopes[asm.Issuer] + h.mu.RUnlock() + requestedScopes = authutil.UnionScopes(granted, requestedScopes) cfg := &oauth2.Config{ ClientID: resolvedClientConfig.clientID, @@ -615,23 +643,42 @@ func (h *AuthorizationCodeHandler) exchangeAuthorizationCode(ctx context.Context // completes. Use a background context that still carries the configured HTTP // client so refreshes keep working for the life of the token source. refreshCtx := context.WithValue(context.Background(), oauth2.HTTPClient, h.config.Client) - h.tokenSource = cfg.TokenSource(refreshCtx, token) + var ts oauth2.TokenSource + if h.config.NewTokenSource == nil { + ts = cfg.TokenSource(refreshCtx, token) + } else { + var err error + ts, err = h.config.NewTokenSource(refreshCtx, cfg, token) + if err != nil { + return fmt.Errorf("constructing token source failed: %w", err) + } + } + h.mu.Lock() + h.tokenSource = ts + h.mu.Unlock() return nil } // updateGrantedScopes updates the granted scopes based on the token source and requested scopes. func (h *AuthorizationCodeHandler) updateGrantedScopes(issuer string, requestedScopes []string) error { - if h.tokenSource == nil { + h.mu.RLock() + ts := h.tokenSource + h.mu.RUnlock() + + if ts == nil { return nil } - tok, err := h.tokenSource.Token() + tok, err := ts.Token() if err != nil { return err } + + h.mu.Lock() if tokenScopes := authutil.ScopesFromToken(tok); tokenScopes == nil { h.grantedScopes[issuer] = requestedScopes } else { h.grantedScopes[issuer] = tokenScopes } + h.mu.Unlock() return nil } diff --git a/auth/authorization_code_test.go b/auth/authorization_code_test.go index 15ca69f0..9833597f 100644 --- a/auth/authorization_code_test.go +++ b/auth/authorization_code_test.go @@ -1233,3 +1233,92 @@ func validConfig() *AuthorizationCodeHandlerConfig { }, } } + +func TestNewTokenSource(t *testing.T) { + // mock the /token endpoint to successfully return an access token on code exchange + mockTS := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/token" { + w.Header().Set("Content-Type", "application/json") + w.Write([]byte(`{"access_token": "test_token", "token_type": "bearer"}`)) + return + } + w.WriteHeader(http.StatusNotFound) + })) + defer mockTS.Close() + + // configure the handler and set NewTokenSource + var called bool + handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{ + RedirectURL: "http://localhost/callback", + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "test_client", + }, + AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + u, _ := url.Parse(args.URL) + return &AuthorizationResult{ + Code: "test_code", + State: u.Query().Get("state"), + }, nil + }, + NewTokenSource: func(ctx context.Context, cfg *oauth2.Config, token *oauth2.Token) (oauth2.TokenSource, error) { + called = true + return oauth2.StaticTokenSource(token), nil + }, + Client: mockTS.Client(), + }) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + // Simulate a 401 response from a resource server. + // The WWW-Authenticate: Bearer header triggers the authorization logic. + req := httptest.NewRequest(http.MethodGet, mockTS.URL, nil) + resp := &http.Response{ + StatusCode: http.StatusUnauthorized, + Header: make(http.Header), + Body: http.NoBody, + Request: req, + } + resp.Header.Set("WWW-Authenticate", "Bearer") + + // Authorize and confirm NewTokenSource was called. + err = handler.Authorize(t.Context(), req, resp) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if !called { + t.Error("expected NewTokenSource to be called") + } +} + +func TestInitialTokenSource(t *testing.T) { + handler, err := NewAuthorizationCodeHandler(&AuthorizationCodeHandlerConfig{ + RedirectURL: "http://localhost:12345/callback", + PreregisteredClient: &oauthex.ClientCredentials{ + ClientID: "test_client_id", + }, + AuthorizationCodeFetcher: func(ctx context.Context, args *AuthorizationArgs) (*AuthorizationResult, error) { + return nil, nil + }, + InitialTokenSource: oauth2.StaticTokenSource(&oauth2.Token{AccessToken: "set_token"}), + }) + if err != nil { + t.Fatalf("NewAuthorizationCodeHandler failed: %v", err) + } + + ts, err := handler.TokenSource(t.Context()) + if err != nil { + t.Fatalf("failed to get token source: %v", err) + } + if ts == nil { + t.Fatal("expected token source to be non-nil") + } + + tok, err := ts.Token() + if err != nil { + t.Fatalf("failed to get Token: %v", err) + } + if tok.AccessToken != "set_token" { + t.Errorf("expected access token 'set_token', got '%s'", tok.AccessToken) + } +}