extauth: refactor enterprise auth tests (#879)

2026-04-08 13:42:59 +02:00 · 2026-04-08 13:42:59 +02:00 · f36591bf8a
commit f36591bf8a
parent f71c803175
2 changed files with 677 additions and 742 deletions
--- a/internal/oauthtest/fake_authorization_server.go
+++ b/internal/oauthtest/fake_authorization_server.go
@ -49,6 +49,14 @@ type RegistrationConfig struct {
 	DynamicClientRegistrationEnabled bool
 }

+// JWTBearerConfig configures support for the JWT Bearer grant type (RFC 7523)
+// on a [FakeAuthorizationServer].
+type JWTBearerConfig struct {
+	// ValidAssertions is the set of assertion values that are accepted.
+	// If empty, any non-empty assertion is accepted.
+	ValidAssertions []string
+}
+
 // Config holds configuration for FakeAuthorizationServer.
 type Config struct {
 	// The optional path component of the issuer URL.
@ -59,6 +67,9 @@ type Config struct {
 	MetadataEndpointConfig *MetadataEndpointConfig
 	// Configuration for client registration.
 	RegistrationConfig *RegistrationConfig
+	// JWTBearerConfig enables RFC 7523 JWT Bearer grant at the /token endpoint.
+	// If non-nil, the server accepts grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer.
+	JWTBearerConfig *JWTBearerConfig
 }

 // FakeAuthorizationServer is a fake OAuth 2.0 Authorization Server for testing.
@ -253,10 +264,19 @@ func (s *FakeAuthorizationServer) handleToken(w http.ResponseWriter, r *http.Req
 		http.Error(w, err.Error(), http.StatusUnauthorized)
 		return
 	}
-	if r.Form.Get("grant_type") != "authorization_code" {
-		http.Error(w, "invalid grant_type", http.StatusBadRequest)
-		return
+
+	grantType := r.Form.Get("grant_type")
+	switch grantType {
+	case "authorization_code":
+		s.handleAuthorizationCodeGrant(w, r)
+	case "urn:ietf:params:oauth:grant-type:jwt-bearer":
+		s.handleJWTBearerGrant(w, r)
+	default:
+		http.Error(w, fmt.Sprintf("unsupported grant_type: %s", grantType), http.StatusBadRequest)
 	}
+}
+
+func (s *FakeAuthorizationServer) handleAuthorizationCodeGrant(w http.ResponseWriter, r *http.Request) {
 	code := r.Form.Get("code")
 	if code == "" {
 		http.Error(w, "missing code", http.StatusBadRequest)
@ -287,6 +307,31 @@ func (s *FakeAuthorizationServer) handleToken(w http.ResponseWriter, r *http.Req
 	})
 }

+func (s *FakeAuthorizationServer) handleJWTBearerGrant(w http.ResponseWriter, r *http.Request) {
+	if s.config.JWTBearerConfig == nil {
+		http.Error(w, "JWT bearer grant not supported", http.StatusBadRequest)
+		return
+	}
+	assertion := r.Form.Get("assertion")
+	if assertion == "" {
+		http.Error(w, "missing assertion", http.StatusBadRequest)
+		return
+	}
+	if len(s.config.JWTBearerConfig.ValidAssertions) > 0 {
+		if !slices.Contains(s.config.JWTBearerConfig.ValidAssertions, assertion) {
+			http.Error(w, "invalid assertion", http.StatusBadRequest)
+			return
+		}
+	}
+
+	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,
+	})
+}
+
 func (s *FakeAuthorizationServer) authenticateClient(r *http.Request) error {
 	clientID, clientSecret, ok := r.BasicAuth()
 	if !ok {
--- a/internal/oauthtest/fake_idp_server.go
+++ b/internal/oauthtest/fake_idp_server.go
@ -0,0 +1,258 @@
+// 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"
+}