|
| 1 | +// Copyright 2026 The Go MCP SDK Authors. All rights reserved. |
| 2 | +// Use of this source code is governed by the license |
| 3 | +// that can be found in the LICENSE file. |
| 4 | + |
| 5 | +package auth_test |
| 6 | + |
| 7 | +import ( |
| 8 | + "context" |
| 9 | + "fmt" |
| 10 | + "net/http" |
| 11 | + "net/http/httptest" |
| 12 | + "net/url" |
| 13 | + "sync" |
| 14 | + |
| 15 | + "github.com/modelcontextprotocol/go-sdk/auth" |
| 16 | + "github.com/modelcontextprotocol/go-sdk/mcp" |
| 17 | + "github.com/modelcontextprotocol/go-sdk/oauthex" |
| 18 | + "golang.org/x/oauth2" |
| 19 | +) |
| 20 | + |
| 21 | +// savingTokenSource is an oauth2.TokenSource that passes the oauth2 config and |
| 22 | +// token to the given saver function each time the access token value changes. |
| 23 | +type savingTokenSource struct { |
| 24 | + mu sync.Mutex |
| 25 | + src oauth2.TokenSource |
| 26 | + saver func(*oauth2.Config, *oauth2.Token) error |
| 27 | + config *oauth2.Config |
| 28 | + accessToken string |
| 29 | +} |
| 30 | + |
| 31 | +func (s *savingTokenSource) Token() (*oauth2.Token, error) { |
| 32 | + s.mu.Lock() |
| 33 | + defer s.mu.Unlock() |
| 34 | + tok, err := s.src.Token() |
| 35 | + if err != nil { |
| 36 | + return nil, err |
| 37 | + } |
| 38 | + if s.accessToken != tok.AccessToken { |
| 39 | + s.accessToken = tok.AccessToken |
| 40 | + // This saver implementation always returns nil. |
| 41 | + _ = s.saver(s.config, tok) |
| 42 | + } |
| 43 | + return tok, nil |
| 44 | +} |
| 45 | + |
| 46 | +// NewSavingTokenSource persists OAuth 2.0 sessions by intercepting token |
| 47 | +// changes from the wrapped oauth2.TokenSource. When this wrapper detects an |
| 48 | +// access token change, it calls the provided session saver with the |
| 49 | +// oauth2.Config and the new oauth2.Token. |
| 50 | +// |
| 51 | +// initial is an optional access token the caller may already hold (such as a |
| 52 | +// token loaded from storage). It initialises the wrapper's state so that if |
| 53 | +// the first wrapped.Token() call returns the same token, does not trigger a |
| 54 | +// redundant call to saver(). Pass nil when there is no existing token, in |
| 55 | +// which case the first token produced by wrapped.Token() is saved. |
| 56 | +func NewSavingTokenSource(wrapped oauth2.TokenSource, config *oauth2.Config, initial *oauth2.Token, saver func(*oauth2.Config, *oauth2.Token) error) oauth2.TokenSource { |
| 57 | + if wrapped == nil { |
| 58 | + return nil |
| 59 | + } |
| 60 | + if saver == nil { |
| 61 | + return wrapped |
| 62 | + } |
| 63 | + var accessToken string |
| 64 | + if initial != nil { |
| 65 | + accessToken = initial.AccessToken |
| 66 | + } |
| 67 | + return &savingTokenSource{ |
| 68 | + src: wrapped, |
| 69 | + saver: saver, |
| 70 | + config: config, |
| 71 | + accessToken: accessToken, |
| 72 | + } |
| 73 | +} |
| 74 | + |
| 75 | +// sessionStore is an in-memory OAuth 2.0 session store used to demonstrate |
| 76 | +// persistence. A production implementation would persist the session to disk or |
| 77 | +// a secret store. |
| 78 | +type sessionStore struct { |
| 79 | + config *oauth2.Config |
| 80 | + token *oauth2.Token |
| 81 | +} |
| 82 | + |
| 83 | +// save persists the given OAuth 2.0 config and token. |
| 84 | +func (s *sessionStore) save(config *oauth2.Config, token *oauth2.Token) error { |
| 85 | + fmt.Printf("Saving token: %s\n", token.AccessToken) |
| 86 | + s.config = config |
| 87 | + s.token = token |
| 88 | + return nil |
| 89 | +} |
| 90 | + |
| 91 | +// restore loads a previously persisted OAuth 2.0 session, if one exists. |
| 92 | +func (s *sessionStore) restore() (*oauth2.Config, *oauth2.Token, error) { |
| 93 | + if s.config != nil && s.token != nil { |
| 94 | + fmt.Println("Restoring session.") |
| 95 | + } else { |
| 96 | + fmt.Println("No session found to restore.") |
| 97 | + } |
| 98 | + return s.config, s.token, nil |
| 99 | +} |
| 100 | + |
| 101 | +// newMockAuthServer returns an httptest.Server that simulates both an MCP |
| 102 | +// resource server requiring authorization and its OAuth authorization server. |
| 103 | +func newMockAuthServer() *httptest.Server { |
| 104 | + return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 105 | + switch r.URL.Path { |
| 106 | + case "/.well-known/oauth-authorization-server": |
| 107 | + w.Header().Set("Content-Type", "application/json") |
| 108 | + _, _ = 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) |
| 109 | + case "/token": |
| 110 | + w.Header().Set("Content-Type", "application/json") |
| 111 | + _, _ = w.Write([]byte(`{"access_token": "mock-token", "token_type": "bearer"}`)) |
| 112 | + default: |
| 113 | + // The mock MCP endpoint returns 401 until the client presents a valid |
| 114 | + // bearer token. |
| 115 | + if r.Header.Get("Authorization") == "Bearer mock-token" { |
| 116 | + // A real server would return a valid MCP message here. The empty |
| 117 | + // response causes Connect to return an error after authorization, |
| 118 | + // which is ignored for this example which demonstrates token |
| 119 | + // persistence only. |
| 120 | + w.WriteHeader(http.StatusOK) |
| 121 | + return |
| 122 | + } |
| 123 | + w.Header().Set("WWW-Authenticate", "Bearer") |
| 124 | + w.WriteHeader(http.StatusUnauthorized) |
| 125 | + } |
| 126 | + })) |
| 127 | +} |
| 128 | + |
| 129 | +// This example shows how OAuth2 session persistence might be implemented. It |
| 130 | +// connects twice using a shared in-memory session store: the first connection |
| 131 | +// authorizes and saves the session, and the second restores and reuses it. |
| 132 | +func Example_persistence() { |
| 133 | + // Simulate an MCP server that requires authorization and an OAuth server. |
| 134 | + mockServer := newMockAuthServer() |
| 135 | + defer mockServer.Close() |
| 136 | + |
| 137 | + // store persists the OAuth2 session across both connections. |
| 138 | + store := &sessionStore{} |
| 139 | + |
| 140 | + // connect performs a single client connection, restoring any saved session |
| 141 | + // beforehand and saving any new session acquired during authorization. |
| 142 | + connect := func() { |
| 143 | + // Load the OAuth2 session if available, and use it. |
| 144 | + var initialTS oauth2.TokenSource |
| 145 | + cfg, tok, err := store.restore() |
| 146 | + if err == nil && cfg != nil && tok != nil { |
| 147 | + initialTS = NewSavingTokenSource( |
| 148 | + cfg.TokenSource(context.Background(), tok), |
| 149 | + cfg, tok, store.save, |
| 150 | + ) |
| 151 | + } |
| 152 | + |
| 153 | + // Configure and initialize the AuthorizationCodeHandler with a |
| 154 | + // NewTokenSource that saves the session when the access token changes, |
| 155 | + // and the InitialTokenSource set to the session loaded via restore(). |
| 156 | + config := &auth.AuthorizationCodeHandlerConfig{ |
| 157 | + RedirectURL: "http://localhost/callback", |
| 158 | + PreregisteredClient: &oauthex.ClientCredentials{ |
| 159 | + ClientID: "example", |
| 160 | + }, |
| 161 | + Client: mockServer.Client(), |
| 162 | + AuthorizationCodeFetcher: func(ctx context.Context, args *auth.AuthorizationArgs) (*auth.AuthorizationResult, error) { |
| 163 | + fmt.Println("No token source found. Transport is calling Authorize()...") |
| 164 | + // Extract the generated state from the authorization URL |
| 165 | + u, _ := url.Parse(args.URL) |
| 166 | + state := u.Query().Get("state") |
| 167 | + return &auth.AuthorizationResult{Code: "mock-code", State: state}, nil |
| 168 | + }, |
| 169 | + NewTokenSource: func(ctx context.Context, cfg *oauth2.Config, token *oauth2.Token) (oauth2.TokenSource, error) { |
| 170 | + // This save implementation always returns nil. |
| 171 | + _ = store.save(cfg, token) |
| 172 | + return NewSavingTokenSource( |
| 173 | + cfg.TokenSource(ctx, token), cfg, token, store.save, |
| 174 | + ), nil |
| 175 | + }, |
| 176 | + InitialTokenSource: initialTS, |
| 177 | + } |
| 178 | + handler, err := auth.NewAuthorizationCodeHandler(config) |
| 179 | + if err != nil { |
| 180 | + fmt.Printf("Error creating handler: %v\n", err) |
| 181 | + return |
| 182 | + } |
| 183 | + |
| 184 | + // Set the constructed handler on a transport. |
| 185 | + transport := &mcp.StreamableClientTransport{ |
| 186 | + Endpoint: mockServer.URL + "/sse", |
| 187 | + OAuthHandler: handler, |
| 188 | + } |
| 189 | + |
| 190 | + // Create a client and attempt to connect using the configured |
| 191 | + // transport. The transport will automatically: |
| 192 | + // 1. Call TokenSource() to check for an existing session. This will |
| 193 | + // return InitialTokenSource, if set. |
| 194 | + // 2. Try the MCP endpoint, and encounter a 401 response from the mock |
| 195 | + // server. |
| 196 | + // 3. Call Authorize() to perform the OAuth flow, which calls |
| 197 | + // NewTokenSource. In this example NewTokenSource saves the newly |
| 198 | + // acquired session. |
| 199 | + // 4. Retry the MCP endpoint with a valid bearer token and get a 200. |
| 200 | + client := mcp.NewClient( |
| 201 | + &mcp.Implementation{Name: "example", Version: "1.0.0"}, nil, |
| 202 | + ) |
| 203 | + // Response ignored: this example asserts authorization only. |
| 204 | + _, _ = client.Connect(context.Background(), transport, nil) |
| 205 | + } |
| 206 | + |
| 207 | + // The first connection has no saved session, so it authorizes and saves. |
| 208 | + fmt.Println("--- First connect ---") |
| 209 | + connect() |
| 210 | + // The second connection restores the saved session and reuses it, so no |
| 211 | + // further authorization or save occurs. |
| 212 | + fmt.Println("--- Second connect ---") |
| 213 | + connect() |
| 214 | + |
| 215 | + // Output: |
| 216 | + // --- First connect --- |
| 217 | + // No session found to restore. |
| 218 | + // No token source found. Transport is calling Authorize()... |
| 219 | + // Saving token: mock-token |
| 220 | + // --- Second connect --- |
| 221 | + // Restoring session. |
| 222 | +} |
0 commit comments