package idp
import (
"context"
"crypto/rsa"
"encoding/base64"
"encoding/json"
"errors"
"fmt"
"html"
"math/big"
"net/http"
"net/url"
"os"
"strings"
"time"
"github.com/gin-contrib/sessions"
"github.com/gin-gonic/gin"
"github.com/ory/fosite"
"github.com/ory/fosite/compose"
"github.com/ory/fosite/token/jwt"
"github.com/sigbit/mcp-auth-proxy/v2/pkg/auth"
"github.com/sigbit/mcp-auth-proxy/v2/pkg/repository"
"github.com/sigbit/mcp-auth-proxy/v2/pkg/utils"
"go.uber.org/zap"
"golang.org/x/crypto/bcrypt"
)
type IDPRouter struct {
repo repository.Repository
privKey *rsa.PrivateKey
logger *zap.Logger
externalURL string
hasher fosite.Hasher
provider fosite.OAuth2Provider
signer *jwt.DefaultSigner
authRouter *auth.AuthRouter
redirectPrefixes string
redirectRestricted bool
}
func NewIDPRouter(
repo repository.Repository,
privKey *rsa.PrivateKey,
logger *zap.Logger,
externalURL string,
secret []byte,
authRouter *auth.AuthRouter,
) (*IDPRouter, error) {
hasher := &fosite.BCrypt{
Config: &fosite.Config{
HashCost: bcrypt.DefaultCost,
},
}
config := &fosite.Config{
GlobalSecret: secret,
AccessTokenLifespan: 24 * time.Hour,
RefreshTokenLifespan: 30 * 24 * time.Hour,
RefreshTokenScopes: []string{},
AccessTokenIssuer: externalURL,
EnforcePKCE: true,
EnforcePKCEForPublicClients: true,
EnablePKCEPlainChallengeMethod: false,
ScopeStrategy: fosite.HierarchicScopeStrategy,
MinParameterEntropy: fosite.MinParameterEntropy,
ClientSecretsHasher: hasher,
}
provider, signer := customCompose(config, repo, privKey)
prefixes, restricted := os.LookupEnv("MCP_REDIRECT_PREFIXES")
return &IDPRouter{
repo: repo,
privKey: privKey,
logger: logger,
externalURL: externalURL,
hasher: hasher,
provider: provider,
signer: signer,
authRouter: authRouter,
redirectPrefixes: prefixes,
redirectRestricted: restricted,
}, nil
}
func customCompose(config *fosite.Config, storage any, key any) (fosite.OAuth2Provider, *jwt.DefaultSigner) {
keyGetter := func(context.Context) (any, error) { return key, nil }
signer := &jwt.DefaultSigner{GetPrivateKey: keyGetter}
provider := compose.Compose(
config,
storage,
&compose.CommonStrategy{
CoreStrategy: compose.NewOAuth2JWTStrategy(keyGetter, compose.NewOAuth2HMACStrategy(config), config),
OpenIDConnectTokenStrategy: compose.NewOpenIDConnectStrategy(keyGetter, config),
Signer: signer,
},
compose.OAuth2AuthorizeExplicitFactory,
compose.OAuth2RefreshTokenGrantFactory,
compose.OAuth2TokenIntrospectionFactory,
compose.OAuth2PKCEFactory,
)
return provider, signer
}
const (
AuthorizationEndpoint = "/.idp/auth"
AuthorizationReturnEndpoint = "/.idp/auth/:ar_id"
TokenEndpoint = "/.idp/token"
IntrospectionEndpoint = "/.idp/introspect"
RegistrationEndpoint = "/.idp/register"
OauthAuthorizationServerEndpoint = "/.well-known/oauth-authorization-server"
JWKSEndpoint = "/.well-known/jwks.json"
sessionKeyAuthorizeRequestIDs = "idp_authorize_request_ids"
)
func (a *IDPRouter) SetupRoutes(router gin.IRouter) {
router.GET(AuthorizationEndpoint, a.handleAuth)
router.GET(AuthorizationReturnEndpoint, a.authRouter.RequireAuth(), a.handleAuthorizationReturnForm)
router.POST(AuthorizationReturnEndpoint, a.authRouter.RequireAuth(), a.handleAuthorizationReturn)
router.POST(TokenEndpoint, a.handleToken)
router.POST(IntrospectionEndpoint, a.handleIntrospect)
router.POST(RegistrationEndpoint, a.handleRegister)
router.GET(OauthAuthorizationServerEndpoint, a.handleOauthAuthorizationServer)
router.GET(JWKSEndpoint, a.handleJWKS)
}
func (a *IDPRouter) handleAuth(c *gin.Context) {
ctx := c.Request.Context()
// RFC 6749 makes state RECOMMENDED, not REQUIRED, but fosite enforces
// minimum entropy (8 chars). Generate a server-side state for clients
// that omit it (e.g., MCP Inspector, Cursor CLI) so they can complete
// the OAuth flow. The generated state is echoed back in the redirect;
// clients that didn't send state will simply ignore it.
if c.Request.URL.Query().Get("state") == "" {
state, err := utils.GenerateState()
if err != nil {
a.provider.WriteAuthorizeError(ctx, c.Writer, nil, fosite.ErrServerError.WithWrap(err))
return
}
q := c.Request.URL.Query()
q.Set("state", state)
c.Request.URL.RawQuery = q.Encode()
}
ar, err := a.provider.NewAuthorizeRequest(ctx, c.Request)
if err != nil {
a.provider.WriteAuthorizeError(ctx, c.Writer, ar, err)
return
}
if err := a.repo.CreateAuthorizeRequest(ctx, ar); err != nil {
a.logger.Error("Failed to create authorize requester", zap.Error(err))
a.provider.WriteAuthorizeError(ctx, c.Writer, ar, fosite.ErrServerError.WithWrap(err))
return
}
session := sessions.Default(c)
addAuthorizeRequestID(session, ar.GetID())
if err := session.Save(); err != nil {
a.logger.Error("Failed to save authorize request in session", zap.Error(err))
_ = a.repo.DeleteAuthorizeRequest(ctx, ar.GetID())
a.provider.WriteAuthorizeError(ctx, c.Writer, ar, fosite.ErrServerError.WithWrap(err))
return
}
c.Redirect(302, strings.ReplaceAll(AuthorizationReturnEndpoint, ":ar_id", ar.GetID()))
}
func (a *IDPRouter) handleAuthorizationReturnForm(c *gin.Context) {
arID := c.Param("ar_id")
if !hasAuthorizeRequestID(sessions.Default(c), arID) {
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "invalid authorization session"})
return
}
ar, err := a.repo.GetAuthorizeRequest(c.Request.Context(), arID)
if err != nil {
c.AbortWithStatusJSON(http.StatusNotFound, gin.H{"error": "invalid authorization request"})
return
}
page := fmt.Sprintf(`<!doctype html><html><body><p>Client: %s</p><p>Redirect URI: %s</p><form method="post"><button type="submit">Authorize</button></form></body></html>`, html.EscapeString(ar.GetClient().GetID()), html.EscapeString(ar.GetRedirectURI().String()))
c.Data(http.StatusOK, "text/html; charset=utf-8", []byte(page))
}
func (a *IDPRouter) handleAuthorizationReturn(c *gin.Context) {
ctx := c.Request.Context()
arID := c.Param("ar_id")
session := sessions.Default(c)
if !hasAuthorizeRequestID(session, arID) {
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "invalid authorization session"})
return
}
ar, err := a.repo.GetAuthorizeRequest(ctx, arID)
if err != nil {
a.logger.Error("Failed to get authorize requester", zap.Error(err))
c.AbortWithStatusJSON(500, gin.H{"error": "Internal Server Error"})
return
}
defer func() {
if err := a.repo.DeleteAuthorizeRequest(ctx, arID); err != nil {
a.logger.Error("Failed to delete authorize requester", zap.Error(err))
}
}()
for _, scope := range ar.GetRequestedScopes() {
ar.GrantScope(scope)
}
ar.GrantAudience(utils.MCPResource(a.externalURL))
subject := "user"
if userID, ok := session.Get(auth.SessionKeyUserID).(string); ok && userID != "" {
subject = userID
}
var userInfo map[string]any
if userInfoJSON, ok := session.Get(auth.SessionKeyUserInfo).(string); ok && userInfoJSON != "" {
json.Unmarshal([]byte(userInfoJSON), &userInfo)
}
jwtSession, err := NewJWTSessionWithKey(a.externalURL, subject, a.privKey, userInfo)
if err != nil {
a.logger.With(utils.Err(err)...).Error("Failed to create JWT session", zap.Error(err))
a.provider.WriteAuthorizeError(ctx, c.Writer, ar, err)
return
}
response, err := a.provider.NewAuthorizeResponse(ctx, ar, jwtSession)
if err != nil {
a.logger.With(utils.Err(err)...).Error("Failed to generate authorization response", zap.Error(err))
a.provider.WriteAuthorizeError(ctx, c.Writer, ar, err)
return
}
removeAuthorizeRequestID(session, arID)
if err := session.Save(); err != nil {
a.logger.Error("Failed to remove authorize request from session", zap.Error(err))
a.provider.WriteAuthorizeError(ctx, c.Writer, ar, fosite.ErrServerError.WithWrap(err))
return
}
a.provider.WriteAuthorizeResponse(ctx, c.Writer, ar, response)
}
func authorizeRequestIDs(session sessions.Session) []string {
value, ok := session.Get(sessionKeyAuthorizeRequestIDs).(string)
if !ok || value == "" {
return nil
}
var ids []string
if err := json.Unmarshal([]byte(value), &ids); err != nil {
return nil
}
return ids
}
func addAuthorizeRequestID(session sessions.Session, arID string) {
ids := authorizeRequestIDs(session)
if hasAuthorizeRequestID(session, arID) {
return
}
ids = append(ids, arID)
data, _ := json.Marshal(ids)
session.Set(sessionKeyAuthorizeRequestIDs, string(data))
}
func hasAuthorizeRequestID(session sessions.Session, arID string) bool {
for _, id := range authorizeRequestIDs(session) {
if id == arID {
return true
}
}
return false
}
func removeAuthorizeRequestID(session sessions.Session, arID string) {
ids := authorizeRequestIDs(session)
remaining := ids[:0]
for _, id := range ids {
if id != arID {
remaining = append(remaining, id)
}
}
if len(remaining) == 0 {
session.Delete(sessionKeyAuthorizeRequestIDs)
return
}
data, _ := json.Marshal(remaining)
session.Set(sessionKeyAuthorizeRequestIDs, string(data))
}
func (a *IDPRouter) handleToken(c *gin.Context) {
ctx := c.Request.Context()
session, err := NewJWTSessionWithKey("", "", a.privKey, nil)
if err != nil {
a.logger.With(utils.Err(err)...).Error("Failed to create JWT session for token", zap.Error(err))
a.provider.WriteAccessError(ctx, c.Writer, nil, fosite.ErrServerError.WithWrap(err))
return
}
accessRequest, err := a.provider.NewAccessRequest(ctx, c.Request, session)
if err != nil {
a.logger.With(utils.Err(err)...).Error("Failed to create access request", zap.String("grant_type", c.PostForm("grant_type")))
a.provider.WriteAccessError(ctx, c.Writer, accessRequest, err)
return
}
response, err := a.provider.NewAccessResponse(ctx, accessRequest)
if err != nil {
a.logger.With(utils.Err(err)...).Error("Failed to create access response", zap.String("grant_type", c.PostForm("grant_type")), zap.Error(err))
a.provider.WriteAccessError(ctx, c.Writer, accessRequest, err)
return
}
a.provider.WriteAccessResponse(ctx, c.Writer, accessRequest, response)
}
func (a *IDPRouter) handleIntrospect(c *gin.Context) {
ctx := c.Request.Context()
session, err := NewJWTSessionWithKey("", "", a.privKey, nil)
if err != nil {
a.provider.WriteIntrospectionError(ctx, c.Writer, fosite.ErrServerError.WithWrap(err))
return
}
ir, err := a.provider.NewIntrospectionRequest(ctx, c.Request, session)
if err != nil {
a.provider.WriteIntrospectionError(ctx, c.Writer, err)
return
}
a.provider.WriteIntrospectionResponse(ctx, c.Writer, ir)
}
type registrationRequest struct {
ClientName string `json:"client_name"`
GrantTypes []string `json:"grant_types"`
ResponseTypes []string `json:"response_types"`
TokenEndpointAuthMethod string `json:"token_endpoint_auth_method"`
Scope string `json:"scope"`
RedirectURIs []string `json:"redirect_uris"`
}
type registrationResponse struct {
ClientID string `json:"client_id"`
ClientSecret string `json:"client_secret,omitempty"`
RedirectURIs []string `json:"redirect_uris"`
ClientName string `json:"client_name"`
GrantTypes []string `json:"grant_types"`
ResponseTypes []string `json:"response_types"`
TokenEndpointAuthMethod string `json:"token_endpoint_auth_method"`
RegistrationClientURI string `json:"registration_client_uri"`
ClientIDIssuedAt int64 `json:"client_id_issued_at"`
}
func (a *IDPRouter) handleRegister(c *gin.Context) {
ctx := c.Request.Context()
c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 32*1024)
var req registrationRequest
if err := c.ShouldBindJSON(&req); err != nil {
a.logRegistrationRejection("invalid_json", req, err)
c.JSON(400, gin.H{"error": "invalid_request", "error_description": err.Error()})
return
}
if (a.redirectRestricted && a.redirectPrefixes == "") || !validRedirectURIs(req.RedirectURIs, a.redirectPrefixes) {
a.logRegistrationRejection("redirect_policy", req, nil)
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid_redirect_uri"})
return
}
clientID, err := utils.GenerateClientID()
if err != nil {
a.logger.Error("Failed to generate client ID", zap.Error(err))
c.JSON(500, gin.H{"error": "server_error", "error_description": err.Error()})
return
}
var clientSecret string
var hashedSecret []byte
isPublic := req.TokenEndpointAuthMethod == "none"
if !isPublic {
// Generate client secret for confidential clients
clientSecret, err = utils.GenerateClientSecret()
if err != nil {
a.logger.Error("Failed to generate client secret", zap.Error(err))
c.JSON(500, gin.H{"error": "server_error", "error_description": err.Error()})
return
}
hashedSecret, err = a.hasher.Hash(ctx, []byte(clientSecret))
if err != nil {
a.logger.Error("Failed to hash client secret", zap.Error(err))
c.JSON(500, gin.H{"error": "server_error", "error_description": err.Error()})
return
}
}
client := &fosite.DefaultClient{
ID: clientID,
Secret: hashedSecret,
RedirectURIs: req.RedirectURIs,
GrantTypes: req.GrantTypes,
ResponseTypes: req.ResponseTypes,
Scopes: strings.Fields(req.Scope),
Audience: []string{utils.MCPResource(a.externalURL)},
Public: isPublic,
}
if err := a.repo.RegisterClient(ctx, client); err != nil {
a.logger.Error("Failed to register client", zap.String("client_id", clientID), zap.Error(err))
c.JSON(500, gin.H{"error": "server_error", "error_description": err.Error()})
return
}
registrationClientURI, err := url.JoinPath(RegistrationEndpoint, clientID)
if err != nil {
a.logger.Error("Failed to create registration client URI", zap.String("client_id", clientID), zap.Error(err))
c.JSON(500, gin.H{"error": "server_error", "error_description": err.Error()})
return
}
response := registrationResponse{
ClientID: clientID,
RedirectURIs: req.RedirectURIs,
ClientName: req.ClientName,
GrantTypes: req.GrantTypes,
ResponseTypes: req.ResponseTypes,
TokenEndpointAuthMethod: req.TokenEndpointAuthMethod,
RegistrationClientURI: registrationClientURI,
ClientIDIssuedAt: time.Now().Unix(),
}
if !isPublic {
response.ClientSecret = clientSecret
}
c.JSON(201, response)
}
func (a *IDPRouter) logRegistrationRejection(reason string, req registrationRequest, err error) {
limit := min(len(req.RedirectURIs), 8)
origins := make([]string, 0, limit)
issues := make([]string, 0, limit)
for _, raw := range req.RedirectURIs[:limit] {
issues = append(issues, redirectURIRejection([]string{raw}, a.redirectPrefixes))
u, parseErr := url.Parse(raw)
if parseErr != nil || u.Host == "" {
origins = append(origins, "invalid")
continue
}
origins = append(origins, (&url.URL{Scheme: u.Scheme, Host: u.Host}).String())
}
fields := []zap.Field{
zap.String("reason", reason),
zap.Int("redirect_count", len(req.RedirectURIs)),
zap.Strings("redirect_origins", origins),
zap.Strings("redirect_issues", issues),
}
var typeErr *json.UnmarshalTypeError
if errors.As(err, &typeErr) {
receivedType, _, _ := strings.Cut(typeErr.Value, " ")
fields = append(fields,
zap.String("field", typeErr.Field),
zap.String("expected_type", typeErr.Type.String()),
zap.String("received_type", receivedType),
)
}
a.logger.Warn("OAuth registration rejected", fields...)
}
type authorizationServerResponse struct {
Issuer string `json:"issuer"`
AuthorizationEndpoint string `json:"authorization_endpoint"`
TokenEndpoint string `json:"token_endpoint"`
RegistrationEndpoint string `json:"registration_endpoint"`
ScopesSupported []string `json:"scopes_supported"`
ResponseTypesSupported []string `json:"response_types_supported"`
ResponseModesSupported []string `json:"response_modes_supported"`
GrantTypesSupported []string `json:"grant_types_supported"`
TokenEndpointAuthMethodsSupported []string `json:"token_endpoint_auth_methods_supported"`
CodeChallengeMethodsSupported []string `json:"code_challenge_methods_supported"`
}
func (a *IDPRouter) handleOauthAuthorizationServer(c *gin.Context) {
authorizationEndpoint, err := url.JoinPath(a.externalURL, AuthorizationEndpoint)
if err != nil {
a.logger.Error("Failed to create authorization endpoint URL", zap.Error(err))
c.JSON(500, gin.H{"error": "server_error", "error_description": err.Error()})
return
}
tokenEndpoint, err := url.JoinPath(a.externalURL, TokenEndpoint)
if err != nil {
a.logger.Error("Failed to create token endpoint URL", zap.Error(err))
c.JSON(500, gin.H{"error": "server_error", "error_description": err.Error()})
return
}
registrationEndpoint, err := url.JoinPath(a.externalURL, RegistrationEndpoint)
if err != nil {
a.logger.Error("Failed to create registration endpoint URL", zap.Error(err))
c.JSON(500, gin.H{"error": "server_error", "error_description": err.Error()})
return
}
res := &authorizationServerResponse{
Issuer: a.externalURL,
AuthorizationEndpoint: authorizationEndpoint,
TokenEndpoint: tokenEndpoint,
RegistrationEndpoint: registrationEndpoint,
ScopesSupported: []string{},
ResponseTypesSupported: []string{"code"},
ResponseModesSupported: []string{"query"},
GrantTypesSupported: []string{"authorization_code", "refresh_token"},
TokenEndpointAuthMethodsSupported: []string{"client_secret_basic", "client_secret_post", "none"},
CodeChallengeMethodsSupported: []string{"S256"},
}
c.JSON(200, res)
}
type jwk struct {
Kty string `json:"kty"`
Use string `json:"use"`
Kid string `json:"kid"`
Alg string `json:"alg"`
N string `json:"n"`
E string `json:"e"`
}
type jwks struct {
Keys []jwk `json:"keys"`
}
func (a *IDPRouter) handleJWKS(c *gin.Context) {
publicKey := &a.privKey.PublicKey
// Convert RSA public key components to base64url
nBytes := publicKey.N.Bytes()
eBytes := big.NewInt(int64(publicKey.E)).Bytes()
n := base64.RawURLEncoding.EncodeToString(nBytes)
e := base64.RawURLEncoding.EncodeToString(eBytes)
keyID, err := utils.GenerateKeyID(&a.privKey.PublicKey)
if err != nil {
a.logger.Error("Failed to generate key ID for JWKS", zap.Error(err))
c.JSON(500, gin.H{"error": "failed to generate key ID"})
return
}
k := jwk{
Kty: "RSA",
Use: "sig",
Kid: keyID,
Alg: "RS256",
N: n,
E: e,
}
ks := jwks{Keys: []jwk{k}}
c.JSON(200, ks)
}
func NewJWTSessionWithKey(iss string, subject string, privateKey *rsa.PrivateKey, userInfo map[string]any) (*Session, error) {
keyID, err := utils.GenerateKeyID(&privateKey.PublicKey)
if err != nil {
return nil, err
}
var extra map[string]any
if userInfo != nil {
extra = map[string]any{"userinfo": userInfo}
}
return &Session{
DefaultSession: &fosite.DefaultSession{
Username: subject,
Subject: subject,
},
JWTClaims: &jwt.JWTClaims{
Issuer: iss,
Subject: subject,
Audience: []string{},
ExpiresAt: time.Now().Add(time.Hour),
IssuedAt: time.Now(),
NotBefore: time.Now(),
Extra: extra,
},
JWTHeader: &jwt.Headers{
Extra: map[string]any{
"kid": keyID,
},
},
}, nil
}
type Session struct {
*fosite.DefaultSession
JWTClaims *jwt.JWTClaims
JWTHeader *jwt.Headers
}
func (s *Session) GetJWTClaims() jwt.JWTClaimsContainer {
return s.JWTClaims
}
func (s *Session) GetJWTHeader() *jwt.Headers {
return s.JWTHeader
}
func (s *Session) Clone() fosite.Session {
if s == nil {
return nil
}
clone := &Session{
DefaultSession: &fosite.DefaultSession{
Username: s.DefaultSession.Username,
Subject: s.DefaultSession.Subject,
ExpiresAt: s.DefaultSession.ExpiresAt,
},
JWTClaims: &jwt.JWTClaims{
Issuer: s.JWTClaims.Issuer,
Subject: s.JWTClaims.Subject,
Audience: s.JWTClaims.Audience,
ExpiresAt: s.JWTClaims.ExpiresAt,
IssuedAt: s.JWTClaims.IssuedAt,
NotBefore: s.JWTClaims.NotBefore,
Extra: s.JWTClaims.Extra,
},
JWTHeader: &jwt.Headers{
Extra: make(map[string]any),
},
}
for k, v := range s.JWTHeader.Extra {
clone.JWTHeader.Extra[k] = v
}
return clone
}