Repository navigation
Expand file tree
/
Copy pathenterprise_handler.go
More file actions
269 lines (234 loc) · 9.3 KB
/
Copy pathenterprise_handler.go
File metadata and controls
269 lines (234 loc) · 9.3 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
// Copyright 2026 The Go MCP SDK Authors. All rights reserved.
// Use of this source code is governed by an MIT-style
// license that can be found in the LICENSE file.
// Package extauth provides OAuth handler implementations for MCP authorization extensions.
// This package implements Enterprise Managed Authorization as defined in SEP-990.
package extauth
import (
"context"
"errors"
"fmt"
"io"
"net/http"
"sync"
"github.com/modelcontextprotocol/go-sdk/auth"
"github.com/modelcontextprotocol/go-sdk/internal/util"
"github.com/modelcontextprotocol/go-sdk/oauthex"
"golang.org/x/oauth2"
)
// grantTypeJWTBearer is the grant type for RFC 7523 JWT Bearer authorization grant.
const grantTypeJWTBearer = "urn:ietf:params:oauth:grant-type:jwt-bearer"
// IDTokenFetcher is called to obtain an ID Token from the enterprise IdP.
// This is typically done via OIDC login flow where the user authenticates
// with their enterprise identity provider.
//
// Returns an oauth2.Token where Extra("id_token") contains the OpenID Connect ID Token (JWT).
type IDTokenFetcher func(ctx context.Context) (*oauth2.Token, error)
// EnterpriseHandlerConfig is the configuration for [EnterpriseHandler].
type EnterpriseHandlerConfig struct {
// IdP configuration (where the user authenticates)
// IdPIssuerURL is the enterprise IdP's issuer URL (e.g., "https://acme.okta.com").
// Used for OIDC discovery to find the token endpoint.
// REQUIRED.
IdPIssuerURL string
// IdPCredentials contains the MCP Client's credentials registered at the IdP.
// REQUIRED. These credentials are used for token exchange at the IdP.
// The ClientID is always required. ClientSecretAuth is optional and only needed
// if the IdP requires client authentication (confidential clients).
IdPCredentials *oauthex.ClientCredentials
// MCP Server configuration (the resource being accessed)
// MCPAuthServerURL is the MCP Server's authorization server issuer URL.
// Used as the audience for token exchange and for metadata discovery.
// REQUIRED.
MCPAuthServerURL string
// MCPResourceURI is the MCP Server's resource identifier (RFC 9728).
// Used as the resource parameter in token exchange.
// REQUIRED.
MCPResourceURI string
// MCPCredentials contains the MCP Client's credentials registered at the MCP Server.
// REQUIRED. These credentials are used for JWT Bearer grant at the MCP Server.
// The ClientID is always required. ClientSecretAuth is optional and only needed
// if the MCP Server requires client authentication.
MCPCredentials *oauthex.ClientCredentials
// MCPScopes is the list of scopes to request at the MCP Server.
// OPTIONAL.
MCPScopes []string
// IDTokenFetcher is called to obtain an ID Token when authorization is needed.
// The implementation should handle the OIDC login flow (e.g., browser redirect,
// callback handling) and return the ID token.
// REQUIRED.
IDTokenFetcher IDTokenFetcher
// HTTPClient is an optional HTTP client for customization.
// If nil, http.DefaultClient is used.
// Unless HTTPClient sets CheckRedirect, token requests follow a redirect
// only within the origin of the request and only if it keeps the method.
// OPTIONAL.
HTTPClient *http.Client
}
// EnterpriseHandler is an implementation of [auth.OAuthHandler] that uses
// Enterprise Managed Authorization (SEP-990) to obtain access tokens.
//
// The flow consists of:
// 1. OIDC Login: User authenticates with enterprise IdP → ID Token
// 2. Token Exchange (RFC 8693): ID Token → ID-JAG at IdP
// 3. JWT Bearer Grant (RFC 7523): ID-JAG → Access Token at MCP Server
type EnterpriseHandler struct {
config *EnterpriseHandlerConfig
// mu protects tokenSource. Authorize can run for one request while the
// transport reads the token source for another.
mu sync.RWMutex
// tokenSource is the token source obtained after authorization.
tokenSource oauth2.TokenSource
}
// Compile-time check that EnterpriseHandler implements auth.OAuthHandler.
var _ auth.OAuthHandler = (*EnterpriseHandler)(nil)
// NewEnterpriseHandler creates a new EnterpriseHandler.
// It performs validation of the configuration and returns an error if invalid.
func NewEnterpriseHandler(config *EnterpriseHandlerConfig) (*EnterpriseHandler, error) {
if config == nil {
return nil, errors.New("config must be provided")
}
if config.IdPIssuerURL == "" {
return nil, errors.New("IdPIssuerURL is required")
}
if config.IdPCredentials == nil {
return nil, errors.New("IdPCredentials is required")
}
if err := config.IdPCredentials.Validate(); err != nil {
return nil, fmt.Errorf("invalid IdPCredentials: %w", err)
}
if config.MCPAuthServerURL == "" {
return nil, errors.New("MCPAuthServerURL is required")
}
if config.MCPResourceURI == "" {
return nil, errors.New("MCPResourceURI is required")
}
if config.MCPCredentials == nil {
return nil, errors.New("MCPCredentials is required")
}
if err := config.MCPCredentials.Validate(); err != nil {
return nil, fmt.Errorf("invalid MCPCredentials: %w", err)
}
if config.IDTokenFetcher == nil {
return nil, errors.New("IDTokenFetcher is required")
}
return &EnterpriseHandler{config: config}, nil
}
// TokenSource returns the token source for outgoing requests.
// Returns nil if authorization has not been performed yet.
func (h *EnterpriseHandler) TokenSource(ctx context.Context) (oauth2.TokenSource, error) {
h.mu.RLock()
defer h.mu.RUnlock()
return h.tokenSource, nil
}
// Authorize performs the Enterprise Managed Authorization flow.
// It is called when a request fails with 401 or 403.
func (h *EnterpriseHandler) Authorize(ctx context.Context, req *http.Request, resp *http.Response) error {
defer resp.Body.Close()
defer io.Copy(io.Discard, resp.Body)
httpClient := h.config.HTTPClient
if httpClient == nil {
httpClient = http.DefaultClient
}
// Step 1: Get ID Token via the configured fetcher (e.g., OIDC login)
oidcToken, err := h.config.IDTokenFetcher(ctx)
if err != nil {
return fmt.Errorf("failed to obtain ID token: %w", err)
}
// Extract ID token from the oauth2.Token
idToken, ok := oidcToken.Extra("id_token").(string)
if !ok || idToken == "" {
return fmt.Errorf("id_token not found in OIDC token response")
}
// Step 2: Discover IdP token endpoint via OIDC discovery
idpMeta, err := auth.GetAuthServerMetadata(ctx, h.config.IdPIssuerURL, httpClient)
if err != nil {
return fmt.Errorf("failed to discover IdP metadata: %w", err)
}
if idpMeta == nil {
return fmt.Errorf("no authorization server metadata found for IdP: %s", h.config.IdPIssuerURL)
}
// Step 3: Token Exchange (ID Token → ID-JAG)
tokenExchangeReq := &oauthex.TokenExchangeRequest{
RequestedTokenType: oauthex.TokenTypeIDJAG,
Audience: h.config.MCPAuthServerURL,
Resource: h.config.MCPResourceURI,
Scope: h.config.MCPScopes,
SubjectToken: idToken,
SubjectTokenType: oauthex.TokenTypeIDToken,
}
idJAGToken, err := oauthex.ExchangeToken(
ctx,
idpMeta.TokenEndpoint,
tokenExchangeReq,
h.config.IdPCredentials,
httpClient,
)
if err != nil {
return fmt.Errorf("token exchange failed: %w", err)
}
// The token is sent on to the MCP Server's authorization server, a different
// party than the IdP. Only forward it if the IdP says it is an ID-JAG, so
// that any other token the IdP issued, such as an access token, does not
// leave it.
if issued, _ := idJAGToken.Extra("issued_token_type").(string); issued != oauthex.TokenTypeIDJAG {
return fmt.Errorf("token exchange issued a token of type %q, want %q", issued, oauthex.TokenTypeIDJAG)
}
// Step 4: Discover MCP Server token endpoint
mcpMeta, err := auth.GetAuthServerMetadata(ctx, h.config.MCPAuthServerURL, httpClient)
if err != nil {
return fmt.Errorf("failed to discover MCP auth server metadata: %w", err)
}
if mcpMeta == nil {
return fmt.Errorf("no authorization server metadata found for MCP server: %s", h.config.MCPAuthServerURL)
}
// Step 5: JWT Bearer Grant (ID-JAG → Access Token)
// The ID-JAG is in the AccessToken field of the token (despite the name)
accessToken, err := exchangeJWTBearer(
ctx,
mcpMeta.TokenEndpoint,
idJAGToken.AccessToken,
h.config.MCPCredentials,
httpClient,
)
if err != nil {
return fmt.Errorf("JWT bearer grant failed: %w", err)
}
// Store the token source for subsequent requests
h.mu.Lock()
h.tokenSource = oauth2.StaticTokenSource(accessToken)
h.mu.Unlock()
return nil
}
// exchangeJWTBearer exchanges an Identity Assertion JWT Authorization Grant (ID-JAG)
// for an access token using JWT Bearer Grant per RFC 7523.
func exchangeJWTBearer(
ctx context.Context,
tokenEndpoint string,
assertion string,
clientCreds *oauthex.ClientCredentials,
httpClient *http.Client,
) (*oauth2.Token, error) {
cfg := &oauth2.Config{
ClientID: clientCreds.ClientID,
Endpoint: oauth2.Endpoint{
TokenURL: tokenEndpoint,
AuthStyle: oauth2.AuthStyleInParams,
},
}
// Set ClientSecret if ClientSecretAuth is configured
if clientCreds.ClientSecretAuth != nil {
cfg.ClientSecret = clientCreds.ClientSecretAuth.ClientSecret
}
ctxWithClient := context.WithValue(ctx, oauth2.HTTPClient, util.RestrictRedirects(httpClient))
token, err := cfg.Exchange(
ctxWithClient,
"",
oauth2.SetAuthURLParam("grant_type", grantTypeJWTBearer),
oauth2.SetAuthURLParam("assertion", assertion),
)
if err != nil {
return nil, fmt.Errorf("JWT bearer grant request failed: %w", err)
}
return token, nil
}