258 lines
7.7 KiB
Go
258 lines
7.7 KiB
Go
// 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 oauthtest
|
|
|
|
import (
|
|
"crypto/rand"
|
|
"crypto/sha256"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"maps"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"slices"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/modelcontextprotocol/go-sdk/oauthex"
|
|
)
|
|
|
|
// TokenExchangeConfig configures RFC 8693 token exchange support on a [FakeIdPServer].
|
|
type TokenExchangeConfig struct {
|
|
// IDJAGToken is the ID-JAG value returned from token exchange.
|
|
// Defaults to "test-id-jag-token" if empty.
|
|
IDJAGToken string
|
|
}
|
|
|
|
// IdPConfig holds configuration for [FakeIdPServer].
|
|
type IdPConfig struct {
|
|
// PreregisteredClients maps client IDs to their info.
|
|
PreregisteredClients map[string]ClientInfo
|
|
|
|
// TokenExchangeConfig enables RFC 8693 token exchange at the /token endpoint.
|
|
// If non-nil, the server accepts grant_type=urn:ietf:params:oauth:grant-type:token-exchange.
|
|
TokenExchangeConfig *TokenExchangeConfig
|
|
}
|
|
|
|
// FakeIdPServer is a fake OIDC Identity Provider for testing.
|
|
// It supports:
|
|
// - OIDC discovery (/.well-known/openid-configuration)
|
|
// - Authorization Code Grant with PKCE
|
|
// - ID Token issuance (fake JWTs)
|
|
// - RFC 8693 Token Exchange (ID Token → ID-JAG), if configured
|
|
type FakeIdPServer struct {
|
|
server *httptest.Server
|
|
mux *http.ServeMux
|
|
config IdPConfig
|
|
clients map[string]ClientInfo
|
|
codes map[string]codeInfo
|
|
}
|
|
|
|
// NewFakeIdPServer creates a new FakeIdPServer.
|
|
func NewFakeIdPServer(config IdPConfig) *FakeIdPServer {
|
|
s := &FakeIdPServer{
|
|
mux: http.NewServeMux(),
|
|
config: config,
|
|
clients: make(map[string]ClientInfo),
|
|
codes: make(map[string]codeInfo),
|
|
}
|
|
maps.Copy(s.clients, config.PreregisteredClients)
|
|
|
|
s.mux.HandleFunc("/.well-known/openid-configuration", s.handleMetadata)
|
|
s.mux.HandleFunc("/authorize", s.handleAuthorize)
|
|
s.mux.HandleFunc("/token", s.handleToken)
|
|
s.server = httptest.NewUnstartedServer(s.mux)
|
|
|
|
return s
|
|
}
|
|
|
|
// Start starts the HTTP server and registers a cleanup function on t.
|
|
func (s *FakeIdPServer) Start(t testing.TB) {
|
|
s.server.Start()
|
|
t.Cleanup(s.server.Close)
|
|
}
|
|
|
|
// URL returns the base URL of the server (issuer).
|
|
func (s *FakeIdPServer) URL() string {
|
|
return s.server.URL
|
|
}
|
|
|
|
func (s *FakeIdPServer) handleMetadata(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodGet {
|
|
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
grantTypes := []string{"authorization_code"}
|
|
if s.config.TokenExchangeConfig != nil {
|
|
grantTypes = append(grantTypes, oauthex.GrantTypeTokenExchange)
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]any{
|
|
"issuer": s.URL(),
|
|
"authorization_endpoint": s.URL() + "/authorize",
|
|
"token_endpoint": s.URL() + "/token",
|
|
"jwks_uri": s.URL() + "/.well-known/jwks.json",
|
|
"response_types_supported": []string{"code"},
|
|
"code_challenge_methods_supported": []string{"S256"},
|
|
"grant_types_supported": grantTypes,
|
|
})
|
|
}
|
|
|
|
func (s *FakeIdPServer) handleAuthorize(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodGet {
|
|
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
clientID := r.URL.Query().Get("client_id")
|
|
clientInfo, ok := s.clients[clientID]
|
|
if !ok {
|
|
http.Error(w, "unknown client_id", http.StatusBadRequest)
|
|
return
|
|
}
|
|
redirectURI := r.URL.Query().Get("redirect_uri")
|
|
if redirectURI == "" {
|
|
http.Error(w, "missing redirect_uri", http.StatusBadRequest)
|
|
return
|
|
}
|
|
if !slices.Contains(clientInfo.RedirectURIs, redirectURI) {
|
|
http.Error(w, "invalid redirect_uri", http.StatusBadRequest)
|
|
return
|
|
}
|
|
codeChallenge := r.URL.Query().Get("code_challenge")
|
|
if codeChallenge == "" {
|
|
http.Error(w, "missing code_challenge", http.StatusBadRequest)
|
|
return
|
|
}
|
|
code := rand.Text()
|
|
s.codes[code] = codeInfo{CodeChallenge: codeChallenge}
|
|
state := r.URL.Query().Get("state")
|
|
redirectURL := fmt.Sprintf("%s?code=%s&state=%s", redirectURI, code, state)
|
|
http.Redirect(w, r, redirectURL, http.StatusFound)
|
|
}
|
|
|
|
func (s *FakeIdPServer) handleToken(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodPost {
|
|
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
if err := r.ParseForm(); err != nil {
|
|
http.Error(w, "failed to parse form", http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
if err := s.authenticateClient(r); err != nil {
|
|
http.Error(w, err.Error(), http.StatusUnauthorized)
|
|
return
|
|
}
|
|
|
|
grantType := r.Form.Get("grant_type")
|
|
switch grantType {
|
|
case "authorization_code":
|
|
s.handleAuthorizationCodeGrant(w, r)
|
|
case oauthex.GrantTypeTokenExchange:
|
|
s.handleTokenExchangeGrant(w, r)
|
|
default:
|
|
http.Error(w, fmt.Sprintf("unsupported grant_type: %s", grantType), http.StatusBadRequest)
|
|
}
|
|
}
|
|
|
|
func (s *FakeIdPServer) handleAuthorizationCodeGrant(w http.ResponseWriter, r *http.Request) {
|
|
code := r.Form.Get("code")
|
|
if code == "" {
|
|
http.Error(w, "missing code", http.StatusBadRequest)
|
|
return
|
|
}
|
|
ci, ok := s.codes[code]
|
|
if !ok {
|
|
http.Error(w, "unknown authorization code", http.StatusBadRequest)
|
|
return
|
|
}
|
|
verifier := r.Form.Get("code_verifier")
|
|
if verifier == "" {
|
|
http.Error(w, "missing code_verifier", http.StatusBadRequest)
|
|
return
|
|
}
|
|
sha := sha256.Sum256([]byte(verifier))
|
|
expectedChallenge := base64.RawURLEncoding.EncodeToString(sha[:])
|
|
if expectedChallenge != ci.CodeChallenge {
|
|
http.Error(w, "PKCE verification failed", http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
clientID := r.Form.Get("client_id")
|
|
now := time.Now().Unix()
|
|
idToken := fakeIDToken(s.URL(), clientID, now)
|
|
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]any{
|
|
"access_token": "test_access_token",
|
|
"token_type": "Bearer",
|
|
"expires_in": 3600,
|
|
"refresh_token": "test_refresh_token",
|
|
"id_token": idToken,
|
|
})
|
|
}
|
|
|
|
func (s *FakeIdPServer) handleTokenExchangeGrant(w http.ResponseWriter, r *http.Request) {
|
|
if s.config.TokenExchangeConfig == nil {
|
|
http.Error(w, "token exchange not supported", http.StatusBadRequest)
|
|
return
|
|
}
|
|
if r.Form.Get("requested_token_type") != oauthex.TokenTypeIDJAG {
|
|
http.Error(w, "invalid requested_token_type", http.StatusBadRequest)
|
|
return
|
|
}
|
|
if r.Form.Get("subject_token_type") != oauthex.TokenTypeIDToken {
|
|
http.Error(w, "invalid subject_token_type", http.StatusBadRequest)
|
|
return
|
|
}
|
|
if r.Form.Get("subject_token") == "" {
|
|
http.Error(w, "missing subject_token", http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
idJAG := s.config.TokenExchangeConfig.IDJAGToken
|
|
if idJAG == "" {
|
|
idJAG = "test-id-jag-token"
|
|
}
|
|
|
|
w.Header().Set("Content-Type", "application/json")
|
|
json.NewEncoder(w).Encode(map[string]any{
|
|
"access_token": idJAG,
|
|
"issued_token_type": oauthex.TokenTypeIDJAG,
|
|
"token_type": "N_A",
|
|
"expires_in": 300,
|
|
})
|
|
}
|
|
|
|
func (s *FakeIdPServer) authenticateClient(r *http.Request) error {
|
|
clientID := r.Form.Get("client_id")
|
|
clientSecret := r.Form.Get("client_secret")
|
|
clientInfo, ok := s.clients[clientID]
|
|
if !ok {
|
|
return fmt.Errorf("unknown client")
|
|
}
|
|
if clientInfo.Secret != clientSecret {
|
|
return fmt.Errorf("invalid client credentials")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// fakeIDToken creates a fake JWT ID token for testing.
|
|
// The token has a valid structure but is not cryptographically signed.
|
|
func fakeIDToken(issuer, audience string, now int64) string {
|
|
header := base64.RawURLEncoding.EncodeToString([]byte(`{"alg":"RS256","typ":"JWT"}`))
|
|
claimsJSON, _ := json.Marshal(map[string]any{
|
|
"iss": issuer,
|
|
"sub": "test-user",
|
|
"aud": audience,
|
|
"exp": now + 3600,
|
|
"iat": now,
|
|
"email": "test@example.com",
|
|
})
|
|
claims := base64.RawURLEncoding.EncodeToString(claimsJSON)
|
|
return header + "." + claims + ".mock-signature"
|
|
}
|