From 1008d605a338425702ed14ed5cb8b83dd3f1b4de Mon Sep 17 00:00:00 2001 From: Pratham-Mishra04 Date: Wed, 17 Jun 2026 14:14:01 +0530 Subject: [PATCH] feat: adds mcp oauth2 authorization server --- framework/configstore/migrations.go | 55 ++ framework/configstore/rdb.go | 165 +++++ framework/configstore/store.go | 24 + .../configstore/tables/mcpoauth2issuance.go | 129 ++++ .../handlers/mcpoauth2issuance.go | 698 ++++++++++++++++++ .../bifrost-http/handlers/temptokens.go | 16 + transports/bifrost-http/lib/config_test.go | 30 + transports/bifrost-http/server/server.go | 9 + transports/go.mod | 2 +- 9 files changed, 1127 insertions(+), 1 deletion(-) create mode 100644 framework/configstore/tables/mcpoauth2issuance.go create mode 100644 transports/bifrost-http/handlers/mcpoauth2issuance.go diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index bdace0cdc83..119f4ef6431 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -430,6 +430,7 @@ var configstoreMigrationSteps = []migrationStep{ {IDs: []string{"null_legacy_customer_budget_id_refs"}, run: migrationNullLegacyCustomerBudgetID}, {IDs: []string{"add_skills_repo_tables"}, run: migrationAddSkillsRepoTables}, {IDs: []string{"add_oauth2_server_tables"}, run: migrationAddOAuth2ServerTables}, + {IDs: []string{"add_oauth2_issuance_tables"}, run: migrationAddOAuth2IssuanceTables}, {IDs: []string{"add_dump_errors_in_console_logs_column"}, run: migrationAddDumpErrorsInConsoleLogsColumn}, {IDs: []string{"add_bedrock_mantle_key_columns"}, run: migrationAddBedrockMantleKeyColumns}, } @@ -10157,3 +10158,57 @@ func migrationAddOAuth2ServerTables(ctx context.Context, db *gorm.DB, logger sch } return nil } + +func migrationAddOAuth2IssuanceTables(ctx context.Context, db *gorm.DB, logger schemas.Logger) error { + migrationName := "add_oauth2_issuance_tables" + logger.Info("[configstore] starting migration %s", migrationName) + defer logger.Info("[configstore] finished migration %s", migrationName) + m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + ID: migrationName, + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + mg := tx.Migrator() + if !mg.HasTable(&tables.TableOAuth2Client{}) { + if err := mg.CreateTable(&tables.TableOAuth2Client{}); err != nil { + return fmt.Errorf("create oauth2_clients table: %w", err) + } + } + if !mg.HasTable(&tables.TableOAuth2AuthorizeRequest{}) { + if err := mg.CreateTable(&tables.TableOAuth2AuthorizeRequest{}); err != nil { + return fmt.Errorf("create oauth2_authorize_requests table: %w", err) + } + } + if !mg.HasTable(&tables.TableOAuth2RefreshToken{}) { + if err := mg.CreateTable(&tables.TableOAuth2RefreshToken{}); err != nil { + return fmt.Errorf("create oauth2_refresh_tokens table: %w", err) + } + } + return nil + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + mg := tx.Migrator() + // Drop in reverse creation order. + if mg.HasTable(&tables.TableOAuth2RefreshToken{}) { + if err := mg.DropTable(&tables.TableOAuth2RefreshToken{}); err != nil { + return fmt.Errorf("drop oauth2_refresh_tokens table: %w", err) + } + } + if mg.HasTable(&tables.TableOAuth2AuthorizeRequest{}) { + if err := mg.DropTable(&tables.TableOAuth2AuthorizeRequest{}); err != nil { + return fmt.Errorf("drop oauth2_authorize_requests table: %w", err) + } + } + if mg.HasTable(&tables.TableOAuth2Client{}) { + if err := mg.DropTable(&tables.TableOAuth2Client{}); err != nil { + return fmt.Errorf("drop oauth2_clients table: %w", err) + } + } + return nil + }, + }}) + if err := m.Migrate(); err != nil { + return fmt.Errorf("error while running db migration %s: %w", migrationName, err) + } + return nil +} diff --git a/framework/configstore/rdb.go b/framework/configstore/rdb.go index 01d25132285..9a029a2e1a6 100644 --- a/framework/configstore/rdb.go +++ b/framework/configstore/rdb.go @@ -6837,3 +6837,168 @@ func (s *RDBConfigStore) createOAuth2SigningKey(ctx context.Context) (*tables.OA key.PrivateKeyPEM = privPEM return key, nil } + +// --- OAuth2 Clients (DCR) --- + +// CreateOAuth2Client persists a new DCR registration. +func (s *RDBConfigStore) CreateOAuth2Client(ctx context.Context, client *tables.TableOAuth2Client) error { + if err := s.DB().WithContext(ctx).Create(client).Error; err != nil { + return fmt.Errorf("create oauth2 client: %w", err) + } + return nil +} + +// GetOAuth2ClientByClientID returns the client with the given client_id, or nil +// if not found. +func (s *RDBConfigStore) GetOAuth2ClientByClientID(ctx context.Context, clientID string) (*tables.TableOAuth2Client, error) { + var c tables.TableOAuth2Client + err := s.DB().WithContext(ctx).Where("client_id = ?", clientID).First(&c).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, ErrNotFound + } + if err != nil { + return nil, fmt.Errorf("get oauth2 client: %w", err) + } + return &c, nil +} + +// --- OAuth2 Authorize Requests --- + +// CreateOAuth2AuthorizeRequest persists a new pending authorize request. +func (s *RDBConfigStore) CreateOAuth2AuthorizeRequest(ctx context.Context, req *tables.TableOAuth2AuthorizeRequest) error { + if err := s.DB().WithContext(ctx).Create(req).Error; err != nil { + return fmt.Errorf("create oauth2 authorize request: %w", err) + } + return nil +} + +// GetOAuth2AuthorizeRequestByID returns the authorize request with the given ID. +func (s *RDBConfigStore) GetOAuth2AuthorizeRequestByID(ctx context.Context, id string) (*tables.TableOAuth2AuthorizeRequest, error) { + var req tables.TableOAuth2AuthorizeRequest + err := s.DB().WithContext(ctx).Where("id = ?", id).First(&req).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, ErrNotFound + } + if err != nil { + return nil, fmt.Errorf("get oauth2 authorize request: %w", err) + } + return &req, nil +} + +// GetOAuth2AuthorizeRequestByCodeHash finds a consented authorize request by +// the hash of the auth code. Used by the token endpoint. +func (s *RDBConfigStore) GetOAuth2AuthorizeRequestByCodeHash(ctx context.Context, codeHash string) (*tables.TableOAuth2AuthorizeRequest, error) { + var req tables.TableOAuth2AuthorizeRequest + err := s.DB().WithContext(ctx). + Where("code_hash = ? AND status = ?", codeHash, tables.OAuth2AuthorizeRequestStatusConsented). + First(&req).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, ErrNotFound + } + if err != nil { + return nil, fmt.Errorf("get oauth2 authorize request by code hash: %w", err) + } + return &req, nil +} + +// ConsentOAuth2AuthorizeRequest atomically transitions a still-pending authorize +// request to consented, recording the minted code hash and resolved identity in +// a single conditional update. The status guard makes the transition idempotent +// under concurrency: a second consent for the same flow matches zero rows and +// returns ErrNotFound rather than overwriting the code hash the first one minted. +func (s *RDBConfigStore) ConsentOAuth2AuthorizeRequest(ctx context.Context, req *tables.TableOAuth2AuthorizeRequest) error { + result := s.DB().WithContext(ctx).Model(&tables.TableOAuth2AuthorizeRequest{}). + Where("id = ? AND status = ?", req.ID, tables.OAuth2AuthorizeRequestStatusPending). + Updates(map[string]any{ + "status": tables.OAuth2AuthorizeRequestStatusConsented, + "code_hash": req.CodeHash, + "bf_mode": req.BfMode, + "bf_sub": req.BfSub, + "updated_at": req.UpdatedAt, + }) + if result.Error != nil { + return fmt.Errorf("consent authorize request: %w", result.Error) + } + if result.RowsAffected == 0 { + return ErrNotFound + } + return nil +} + +// SweepExpiredOAuth2AuthorizeRequests deletes pending/consented requests past +// their TTL. Safe to call periodically. +func (s *RDBConfigStore) SweepExpiredOAuth2AuthorizeRequests(ctx context.Context) error { + return s.DB().WithContext(ctx). + Where("expires_at < ? AND status != ?", time.Now(), tables.OAuth2AuthorizeRequestStatusCodeIssued). + Delete(&tables.TableOAuth2AuthorizeRequest{}).Error +} + +// --- OAuth2 Refresh Tokens --- + +// GetOAuth2RefreshTokenByHash returns the refresh token row for the given hash. +func (s *RDBConfigStore) GetOAuth2RefreshTokenByHash(ctx context.Context, hash string) (*tables.TableOAuth2RefreshToken, error) { + var rt tables.TableOAuth2RefreshToken + err := s.DB().WithContext(ctx).Where("token_hash = ? AND revoked_at IS NULL", hash).First(&rt).Error + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, ErrNotFound + } + if err != nil { + return nil, fmt.Errorf("get oauth2 refresh token: %w", err) + } + return &rt, nil +} + +// ConsumeOAuth2AuthorizeRequest atomically marks the authorize request as +// code_issued and creates the refresh token in a single transaction. +// If either operation fails the transaction is rolled back — the authorize +// request stays in "consented" state and the client can retry the token exchange. +func (s *RDBConfigStore) ConsumeOAuth2AuthorizeRequest(ctx context.Context, requestID string, rt *tables.TableOAuth2RefreshToken) error { + now := time.Now() + return s.DB().WithContext(ctx).Transaction(func(tx *gorm.DB) error { + // Conditional update guards single-use: only a still-consented, unexpired + // request transitions. A zero-row result means the code was already + // consumed, expired, or never consented — reject before minting a token so + // a racing second exchange can't double-spend one authorization code. + result := tx.Model(&tables.TableOAuth2AuthorizeRequest{}). + Where("id = ? AND status = ? AND expires_at > ?", requestID, tables.OAuth2AuthorizeRequestStatusConsented, now). + Updates(map[string]any{ + "status": tables.OAuth2AuthorizeRequestStatusCodeIssued, + "updated_at": now, + }) + if result.Error != nil { + return fmt.Errorf("consume authorize request: %w", result.Error) + } + if result.RowsAffected == 0 { + return ErrNotFound + } + if err := tx.Create(rt).Error; err != nil { + return fmt.Errorf("create refresh token: %w", err) + } + return nil + }) +} + +// RotateOAuth2RefreshToken atomically revokes the old refresh token and creates +// the new one in a single transaction. If either operation fails the transaction +// is rolled back — the old token stays active and the client can retry the refresh. +func (s *RDBConfigStore) RotateOAuth2RefreshToken(ctx context.Context, oldID string, newRT *tables.TableOAuth2RefreshToken) error { + now := time.Now() + return s.DB().WithContext(ctx).Transaction(func(tx *gorm.DB) error { + // Only an active (not-yet-revoked) token may be rotated. A zero-row result + // means the token was already revoked — either by a concurrent rotation or + // as a replay — so reject before minting a replacement. + result := tx.Model(&tables.TableOAuth2RefreshToken{}). + Where("id = ? AND revoked_at IS NULL", oldID). + Update("revoked_at", &now) + if result.Error != nil { + return fmt.Errorf("revoke old refresh token: %w", result.Error) + } + if result.RowsAffected == 0 { + return ErrNotFound + } + if err := tx.Create(newRT).Error; err != nil { + return fmt.Errorf("create new refresh token: %w", err) + } + return nil + }) +} diff --git a/framework/configstore/store.go b/framework/configstore/store.go index 4979c4ebe14..51776002638 100644 --- a/framework/configstore/store.go +++ b/framework/configstore/store.go @@ -660,6 +660,30 @@ type ConfigStore interface { // on first call. Always returns a usable key — never nil on a nil error. GetOAuth2SigningKey(ctx context.Context) (*tables.OAuth2SigningKey, error) + // OAuth2 clients (DCR) + CreateOAuth2Client(ctx context.Context, client *tables.TableOAuth2Client) error + GetOAuth2ClientByClientID(ctx context.Context, clientID string) (*tables.TableOAuth2Client, error) + + // OAuth2 authorize requests + CreateOAuth2AuthorizeRequest(ctx context.Context, req *tables.TableOAuth2AuthorizeRequest) error + GetOAuth2AuthorizeRequestByID(ctx context.Context, id string) (*tables.TableOAuth2AuthorizeRequest, error) + GetOAuth2AuthorizeRequestByCodeHash(ctx context.Context, codeHash string) (*tables.TableOAuth2AuthorizeRequest, error) + // ConsentOAuth2AuthorizeRequest atomically transitions a still-pending request + // to consented (recording the code hash and resolved identity) — returns + // ErrNotFound when no longer pending, so concurrent double-consent can't + // overwrite an already-minted code. + ConsentOAuth2AuthorizeRequest(ctx context.Context, req *tables.TableOAuth2AuthorizeRequest) error + SweepExpiredOAuth2AuthorizeRequests(ctx context.Context) error + + // OAuth2 refresh tokens + GetOAuth2RefreshTokenByHash(ctx context.Context, hash string) (*tables.TableOAuth2RefreshToken, error) + // ConsumeOAuth2AuthorizeRequest atomically marks the authorize request as + // code_issued and creates the refresh token — if either fails the client can retry. + ConsumeOAuth2AuthorizeRequest(ctx context.Context, requestID string, rt *tables.TableOAuth2RefreshToken) error + // RotateOAuth2RefreshToken atomically revokes the old token and creates the + // new one — if either fails the old token stays active and the client can retry. + RotateOAuth2RefreshToken(ctx context.Context, oldID string, newRT *tables.TableOAuth2RefreshToken) error + // Cleanup Close(ctx context.Context) error } diff --git a/framework/configstore/tables/mcpoauth2issuance.go b/framework/configstore/tables/mcpoauth2issuance.go new file mode 100644 index 00000000000..bc34a6e62dd --- /dev/null +++ b/framework/configstore/tables/mcpoauth2issuance.go @@ -0,0 +1,129 @@ +package tables + +import ( + "encoding/json" + "time" + + "gorm.io/gorm" +) + +// TableOAuth2Client holds a registered OAuth2 client created via Dynamic Client +// Registration (RFC 7591). Bifrost only supports public clients +// (token_endpoint_auth_method=none) — no client secrets. +type TableOAuth2Client struct { + ID string `gorm:"type:varchar(255);primaryKey" json:"id"` + ClientID string `gorm:"type:varchar(255);uniqueIndex;not null" json:"client_id"` + ClientName string `gorm:"type:varchar(255)" json:"client_name"` + RedirectURIsJSON string `gorm:"type:text;not null" json:"-"` // JSON []string + GrantTypesJSON string `gorm:"type:text;not null" json:"-"` // JSON []string + Scope string `gorm:"type:varchar(255)" json:"scope"` + CreatedAt time.Time `gorm:"index;not null" json:"created_at"` + + // Virtual fields + RedirectURIs []string `gorm:"-" json:"redirect_uris"` + GrantTypes []string `gorm:"-" json:"grant_types"` +} + +func (TableOAuth2Client) TableName() string { return "oauth2_clients" } + +func (c *TableOAuth2Client) BeforeSave(tx *gorm.DB) error { + if c.RedirectURIs != nil { + data, err := json.Marshal(c.RedirectURIs) + if err != nil { + return err + } + c.RedirectURIsJSON = string(data) + } + if c.GrantTypes != nil { + data, err := json.Marshal(c.GrantTypes) + if err != nil { + return err + } + c.GrantTypesJSON = string(data) + } + return nil +} + +func (c *TableOAuth2Client) AfterFind(tx *gorm.DB) error { + if c.RedirectURIsJSON != "" { + if err := json.Unmarshal([]byte(c.RedirectURIsJSON), &c.RedirectURIs); err != nil { + return err + } + } + if c.GrantTypesJSON != "" { + if err := json.Unmarshal([]byte(c.GrantTypesJSON), &c.GrantTypes); err != nil { + return err + } + } + return nil +} + +// OAuth2AuthorizeRequestStatus is the status of a downstream authorize request. +type OAuth2AuthorizeRequestStatus string + +const ( + OAuth2AuthorizeRequestStatusPending OAuth2AuthorizeRequestStatus = "pending" // waiting for consent + OAuth2AuthorizeRequestStatusConsented OAuth2AuthorizeRequestStatus = "consented" // identity resolved, code minted + OAuth2AuthorizeRequestStatusCodeIssued OAuth2AuthorizeRequestStatus = "code_issued" // token exchanged, one-time consumed +) + +// TableOAuth2AuthorizeRequest tracks a pending downstream OAuth2 authorization +// request from creation at /oauth2/authorize through consent to token exchange +// at /oauth2/token. +// +// State transitions: +// - pending — request created; browser redirected to consent page +// - consented — user approved; identity resolved; auth code minted (CodeHash set) +// - code_issued — auth code exchanged at /oauth2/token; row is consumed (single-use) +type TableOAuth2AuthorizeRequest struct { + ID string `gorm:"type:varchar(255);primaryKey" json:"id"` + ClientID string `gorm:"type:varchar(255);not null;index" json:"client_id"` + RedirectURI string `gorm:"type:text;not null" json:"-"` + State string `gorm:"type:varchar(512);not null" json:"-"` // CSRF; returned in redirect + Scope string `gorm:"type:varchar(255)" json:"scope"` + Resource string `gorm:"type:text;not null" json:"-"` // RFC 8707 resource indicator + CodeChallenge string `gorm:"type:varchar(512);not null" json:"-"` // PKCE S256 challenge + CodeChallengeMethod string `gorm:"type:varchar(10);not null" json:"-"` // always "S256" + Status OAuth2AuthorizeRequestStatus `gorm:"type:varchar(20);not null;index" json:"status"` + // Set by the consent flow once the user approves: + BfMode string `gorm:"type:varchar(20)" json:"bf_mode,omitempty"` // user|vk|session + BfSub string `gorm:"type:varchar(255)" json:"bf_sub,omitempty"` // resolved identity + // nil while pending; set to SHA256(auth_code) at consent. A pointer so unset + // rows store SQL NULL — NULLs are distinct under the unique index, letting many + // requests stay pending at once while still enforcing uniqueness for real hashes. + CodeHash *string `gorm:"type:varchar(255);uniqueIndex" json:"-"` + // TTL: + ExpiresAt time.Time `gorm:"index;not null" json:"expires_at"` + CreatedAt time.Time `gorm:"not null" json:"created_at"` + UpdatedAt time.Time `gorm:"not null" json:"updated_at"` +} + +func (TableOAuth2AuthorizeRequest) TableName() string { return "oauth2_authorize_requests" } + +// TableOAuth2RefreshToken stores a hashed rotating refresh token. The plaintext +// token is only returned to the client once at issuance; only the SHA256 hash +// is persisted. Invalidation paths: +// - rotation on use: old token revoked atomically when a new one is issued +// - bf_sub liveness: VK deleted or user deactivated → invalid_grant on next refresh +// - explicit revocation via the Connected Clients UI +// +// FamilyID links all tokens descended from the same authorization grant (set to +// the authorize request ID at first issuance, propagated on every rotation). +// When a revoked token is presented — indicating the token was stolen and used +// after the legitimate client already rotated — all tokens sharing the FamilyID +// are revoked immediately, per the OAuth 2.0 Security BCP (RFC 9700 §2.2.2). +type TableOAuth2RefreshToken struct { + ID string `gorm:"type:varchar(255);primaryKey" json:"id"` + TokenHash string `gorm:"type:varchar(255);uniqueIndex;not null" json:"-"` // SHA256 hex + FamilyID string `gorm:"type:varchar(255);not null;index" json:"family_id"` // authorize request ID + ClientID string `gorm:"type:varchar(255);not null;index" json:"client_id"` + BfMode string `gorm:"type:varchar(20);not null" json:"bf_mode"` // user|vk|session + BfSub string `gorm:"type:varchar(255);not null" json:"bf_sub"` // resolved identity + Scope string `gorm:"type:varchar(255)" json:"scope"` + Resource string `gorm:"type:text;not null" json:"-"` // RFC 8707 resource indicator; preserved across rotations for the JWT aud claim + RevokedAt *time.Time `gorm:"index" json:"revoked_at,omitempty"` + LastUsedAt *time.Time `gorm:"index" json:"last_used_at,omitempty"` + CreatedAt time.Time `gorm:"not null" json:"created_at"` +} + +func (TableOAuth2RefreshToken) TableName() string { return "oauth2_refresh_tokens" } diff --git a/transports/bifrost-http/handlers/mcpoauth2issuance.go b/transports/bifrost-http/handlers/mcpoauth2issuance.go new file mode 100644 index 00000000000..1f81dbea0eb --- /dev/null +++ b/transports/bifrost-http/handlers/mcpoauth2issuance.go @@ -0,0 +1,698 @@ +package handlers + +import ( + "crypto/rand" + "crypto/rsa" + "crypto/sha256" + "crypto/subtle" + "crypto/x509" + "encoding/base64" + "encoding/hex" + "encoding/pem" + "errors" + "fmt" + "net/url" + "strings" + "time" + + "github.com/bytedance/sonic" + "github.com/fasthttp/router" + "github.com/golang-jwt/jwt/v5" + "github.com/google/uuid" + "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/configstore" + configtables "github.com/maximhq/bifrost/framework/configstore/tables" + "github.com/maximhq/bifrost/framework/temptoken" + "github.com/maximhq/bifrost/transports/bifrost-http/lib" + "github.com/valyala/fasthttp" +) + +// OAuth2IssuanceHandler implements the three downstream OAuth2 endpoints: +// +// - POST /oauth2/register — RFC 7591 Dynamic Client Registration +// - GET /oauth2/authorize — Authorization endpoint (PKCE-S256, RFC 8707) +// - POST /oauth2/token — Token endpoint (auth-code + refresh grants) +type OAuth2IssuanceHandler struct { + store *lib.Config + tempTokens *temptoken.Service // optional; nil = no consent temp-token minted +} + +// NewOAuth2IssuanceHandler creates a new issuance handler. +func NewOAuth2IssuanceHandler(store *lib.Config, tempTokens *temptoken.Service) *OAuth2IssuanceHandler { + return &OAuth2IssuanceHandler{store: store, tempTokens: tempTokens} +} + +// RegisterRoutes wires the three OAuth2 issuance routes. +func (h *OAuth2IssuanceHandler) RegisterRoutes(r *router.Router, middlewares ...schemas.BifrostHTTPMiddleware) { + // These routes are public — no auth middleware applied. + r.POST("/oauth2/register", h.handleRegister) + r.GET("/oauth2/authorize", h.handleAuthorize) + r.POST("/oauth2/token", h.handleToken) +} + +// --- POST /oauth2/register (RFC 7591 DCR) --- + +type dcrRequest struct { + ClientName string `json:"client_name"` + RedirectURIs []string `json:"redirect_uris"` + GrantTypes []string `json:"grant_types"` + ResponseTypes []string `json:"response_types"` + TokenEndpointAuthMethod string `json:"token_endpoint_auth_method"` + Scope string `json:"scope"` +} + +func (h *OAuth2IssuanceHandler) handleRegister(ctx *fasthttp.RequestCtx) { + if h.store.ConfigStore == nil { + sendOAuthError(ctx, fasthttp.StatusServiceUnavailable, "server_error", "config store unavailable") + return + } + + var req dcrRequest + if err := sonic.Unmarshal(ctx.PostBody(), &req); err != nil { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_request", "malformed request body") + return + } + if len(req.RedirectURIs) == 0 { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_redirect_uri", "redirect_uris is required") + return + } + // Registration is public and unauthenticated, so reject dangerous schemes + // (javascript:, data:, etc.) here at the source. Only https is allowed, with + // http permitted exclusively for loopback addresses (RFC 9700 §4.1.3). + for _, uri := range req.RedirectURIs { + if !isAllowedRedirectScheme(uri) { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_redirect_uri", "redirect_uris must use https (or http for loopback addresses)") + return + } + } + // Only public clients supported. + if req.TokenEndpointAuthMethod != "" && req.TokenEndpointAuthMethod != "none" { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_client_metadata", "only token_endpoint_auth_method=none is supported") + return + } + + grantTypes := req.GrantTypes + if len(grantTypes) == 0 { + grantTypes = []string{"authorization_code"} + } + // The token endpoint only implements the authorization-code and refresh-token + // grants. Reject anything else here so a registration never advertises a flow + // that would later be refused at /oauth2/token. + for _, gt := range grantTypes { + if gt != "authorization_code" && gt != "refresh_token" { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_client_metadata", "unsupported grant_type") + return + } + } + responseTypes := req.ResponseTypes + if len(responseTypes) == 0 { + responseTypes = []string{"code"} + } + // The authorize endpoint only implements response_type=code. + for _, rt := range responseTypes { + if rt != "code" { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_client_metadata", "unsupported response_type") + return + } + } + scope := req.Scope + if scope == "" { + scope = "mcp" + } + + clientID := uuid.New().String() + client := &configtables.TableOAuth2Client{ + ID: uuid.New().String(), + ClientID: clientID, + ClientName: req.ClientName, + RedirectURIs: req.RedirectURIs, + GrantTypes: grantTypes, + Scope: scope, + CreatedAt: time.Now(), + } + if err := h.store.ConfigStore.CreateOAuth2Client(ctx, client); err != nil { + sendOAuthError(ctx, fasthttp.StatusInternalServerError, "server_error", "failed to register client") + return + } + + ctx.SetStatusCode(fasthttp.StatusCreated) + ctx.SetContentType("application/json") + data, err := sonic.Marshal(map[string]any{ + "client_id": clientID, + "client_id_issued_at": client.CreatedAt.Unix(), + "grant_types": grantTypes, + "response_types": responseTypes, + "redirect_uris": req.RedirectURIs, + "token_endpoint_auth_method": "none", + "scope": scope, + }) + if err != nil { + sendOAuthError(ctx, fasthttp.StatusInternalServerError, "server_error", "failed to marshal response") + return + } + ctx.SetBody(data) +} + +// --- GET /oauth2/authorize --- + +func (h *OAuth2IssuanceHandler) handleAuthorize(ctx *fasthttp.RequestCtx) { + if h.store.ConfigStore == nil { + sendOAuthError(ctx, fasthttp.StatusServiceUnavailable, "server_error", "config store unavailable") + return + } + + q := ctx.QueryArgs() + clientID := string(q.Peek("client_id")) + redirectURIRaw := string(q.Peek("redirect_uri")) + state := string(q.Peek("state")) + codeChallenge := string(q.Peek("code_challenge")) + codeChallengeMethod := string(q.Peek("code_challenge_method")) + resource := string(q.Peek("resource")) + scope := string(q.Peek("scope")) + + // Validate client exists before using redirect_uri. + client, err := h.store.ConfigStore.GetOAuth2ClientByClientID(ctx, clientID) + if err != nil || client == nil { + if errors.Is(err, configstore.ErrNotFound) { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_client", "unknown client_id") + return + } + sendOAuthError(ctx, fasthttp.StatusInternalServerError, "server_error", "failed to look up client") + return + } + + // Validate redirect_uri (loopback any-port per RFC 8252 §7.3). + if !matchRedirectURI(redirectURIRaw, client.RedirectURIs) { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_redirect_uri", "redirect_uri not registered for this client") + return + } + + // From here errors redirect to the client. + redirectError := func(errCode, description string) { + redirectWithParams(ctx, redirectURIRaw, map[string]string{ + "error": errCode, + "error_description": description, + "state": state, + }) + } + + // Constrain the requested scope to what the client registered for. scope is + // later copied into the access-token and refresh-token state, so a client must + // not be able to request a broader scope than it registered. An omitted scope + // defaults to the client's registered scope. + registeredScope := client.Scope + if registeredScope == "" { + registeredScope = "mcp" + } + if scope == "" { + scope = registeredScope + } else if !scopeWithinRegistered(scope, registeredScope) { + redirectError("invalid_scope", "requested scope exceeds the scope registered for this client") + return + } + + if string(q.Peek("response_type")) != "code" { + redirectError("unsupported_response_type", "only response_type=code is supported") + return + } + if codeChallengeMethod != "S256" { + redirectError("invalid_request", "code_challenge_method must be S256") + return + } + if codeChallenge == "" { + redirectError("invalid_request", "code_challenge is required") + return + } + // RFC 8707: bind the grant to this server's single protected resource. /mcp + // is the only resource we issue tokens for and token verification pins the + // audience to it. A client that omits resource (e.g. one that doesn't fetch + // the protected-resource metadata) defaults to the canonical /mcp resource + // since there is exactly one; a client that does send it must match. + canonicalResource := oauth2MCPResourceURL(ctx, h.store) + if resource == "" { + resource = canonicalResource + } else if resource != canonicalResource { + redirectError("invalid_target", "resource does not identify this MCP server") + return + } + + cfg := oauth2ServerCfg(h.store) + authCodeTTL := cfg.AuthCodeTTL + if authCodeTTL <= 0 { + authCodeTTL = configtables.DefaultAuthCodeTTL + } + + req := &configtables.TableOAuth2AuthorizeRequest{ + ID: uuid.New().String(), + ClientID: client.ClientID, + RedirectURI: redirectURIRaw, + State: state, + Scope: scope, + Resource: resource, + CodeChallenge: codeChallenge, + CodeChallengeMethod: codeChallengeMethod, + Status: configtables.OAuth2AuthorizeRequestStatusPending, + ExpiresAt: time.Now().Add(time.Duration(authCodeTTL) * time.Second), + CreatedAt: time.Now(), + UpdatedAt: time.Now(), + } + if err := h.store.ConfigStore.CreateOAuth2AuthorizeRequest(ctx, req); err != nil { + // Keep the DB error server-side — it can carry table/constraint names that + // must not leak to the (untrusted) redirect target. + logger.Error("failed to create oauth2 authorize request: %v", err) + redirectError("server_error", "failed to create authorization request") + return + } + + // Mint a temp token scoping the consent page to this request. Without it the + // consent-page API calls (GET/PUT /api/oauth2/consent/flows/{id}) have no auth + // credential and fail with 401, leaving the user on a dead-end page — so a mint + // failure must abort with a well-formed error redirect rather than be swallowed. + tempToken := "" + if h.tempTokens != nil { + tok, err := h.tempTokens.Mint(ctx, temptoken.OAuth2ConsentScopeName, req.ID, time.Duration(authCodeTTL)*time.Second) + if err != nil { + logger.Error("failed to mint oauth2 consent temp token: %v", err) + redirectError("server_error", "failed to prepare consent flow") + return + } + tempToken = tok + } + + base := oauth2IssuerURL(ctx, h.store) + consentURL := fmt.Sprintf("%s/oauth/consent?flow=%s", base, url.QueryEscape(req.ID)) + if tempToken != "" { + consentURL += "#t=" + url.QueryEscape(tempToken) + } + + ctx.Response.Header.Set("Location", consentURL) + ctx.SetStatusCode(fasthttp.StatusFound) +} + +// --- POST /oauth2/token --- + +func (h *OAuth2IssuanceHandler) handleToken(ctx *fasthttp.RequestCtx) { + if h.store.ConfigStore == nil { + sendOAuthError(ctx, fasthttp.StatusServiceUnavailable, "server_error", "config store unavailable") + return + } + + grantType := string(ctx.FormValue("grant_type")) + switch grantType { + case "authorization_code": + h.handleTokenAuthCode(ctx) + case "refresh_token": + h.handleTokenRefresh(ctx) + default: + sendOAuthError(ctx, fasthttp.StatusBadRequest, "unsupported_grant_type", fmt.Sprintf("grant_type %q not supported", grantType)) + } +} + +func (h *OAuth2IssuanceHandler) handleTokenAuthCode(ctx *fasthttp.RequestCtx) { + code := string(ctx.FormValue("code")) + codeVerifier := string(ctx.FormValue("code_verifier")) + redirectURI := string(ctx.FormValue("redirect_uri")) + clientID := clientIDFromRequest(ctx) + resource := string(ctx.FormValue("resource")) + + if code == "" || codeVerifier == "" || clientID == "" { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_request", "code, code_verifier and client_id are required") + return + } + + // Look up the authorize request by hashing the received code. + codeHash := hashSHA256Hex(code) + req, err := h.store.ConfigStore.GetOAuth2AuthorizeRequestByCodeHash(ctx, codeHash) + if err != nil || req == nil { + if errors.Is(err, configstore.ErrNotFound) { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_grant", "authorization code not found or already used") + return + } + sendOAuthError(ctx, fasthttp.StatusInternalServerError, "server_error", "failed to look up authorization code") + return + } + if time.Now().After(req.ExpiresAt) { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_grant", "authorization code expired") + return + } + if req.ClientID != clientID { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_grant", "client_id mismatch") + return + } + // The authorization request always binds a redirect_uri (it is validated + // against the client's registered URIs at /oauth2/authorize), so per RFC 6749 + // §4.1.3 the token request must present it and it must match exactly. Accepting + // a missing redirect_uri would let a code be exchanged outside its bound redirect. + if redirectURI == "" || req.RedirectURI != redirectURI { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_grant", "redirect_uri mismatch") + return + } + if resource != "" && req.Resource != resource { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_grant", "resource mismatch") + return + } + + // Verify PKCE: SHA256(verifier) must equal stored challenge. + if !verifyPKCES256(codeVerifier, req.CodeChallenge) { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_grant", "PKCE verification failed") + return + } + + accessToken, refreshToken, refreshTokenObj, err := h.issueTokenPair(ctx, req.ID, req.ClientID, req.BfMode, req.BfSub, req.Scope, req.Resource) + if err != nil { + return + } + // Atomically mark the code as consumed and create the refresh token. + // If this fails the authorize request stays "consented" and the client can retry. + if err := h.store.ConfigStore.ConsumeOAuth2AuthorizeRequest(ctx, req.ID, refreshTokenObj); err != nil { + if errors.Is(err, configstore.ErrNotFound) { + // The code was concurrently consumed or expired between lookup and consume. + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_grant", "authorization code not found or already used") + return + } + sendOAuthError(ctx, fasthttp.StatusInternalServerError, "server_error", "failed to issue token") + return + } + sendTokenResponse(ctx, accessToken, refreshToken, req.Scope, oauth2ServerCfg(h.store).AccessTokenTTL) +} + +func (h *OAuth2IssuanceHandler) handleTokenRefresh(ctx *fasthttp.RequestCtx) { + refreshToken := string(ctx.FormValue("refresh_token")) + clientID := clientIDFromRequest(ctx) + resource := string(ctx.FormValue("resource")) + + if refreshToken == "" || clientID == "" { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_request", "refresh_token and client_id are required") + return + } + + tokenHash := hashSHA256Hex(refreshToken) + rt, err := h.store.ConfigStore.GetOAuth2RefreshTokenByHash(ctx, tokenHash) + if err != nil || rt == nil { + if errors.Is(err, configstore.ErrNotFound) { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_grant", "refresh token not found or revoked") + return + } + sendOAuthError(ctx, fasthttp.StatusInternalServerError, "server_error", "failed to look up refresh token") + return + } + if rt.ClientID != clientID { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_grant", "client_id mismatch") + return + } + + // RFC 8707: resource (audience URI) is distinct from scope. When the client + // omits it on refresh, carry forward the original resource captured at + // authorization — never substitute the scope string. When the client does + // provide it, it must match the resource bound at authorization: issuing a + // token for a different resource than originally authorized would let a + // refresh-token holder escape the original audience binding. Mirrors the + // auth-code handler's check. + if resource == "" { + resource = rt.Resource + } else if resource != rt.Resource { + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_grant", "resource mismatch") + return + } + + accessToken, newRefreshToken, newRefreshTokenObj, err := h.issueTokenPair(ctx, rt.FamilyID, rt.ClientID, rt.BfMode, rt.BfSub, rt.Scope, resource) + if err != nil { + return + } + // Carry the original grant's creation time forward across rotations so the + // grant's "Created" timestamp stays anchored to when it was first authorized, + // and stamp last_used_at to mark this refresh as the grant's latest activity. + usedAt := time.Now() + newRefreshTokenObj.CreatedAt = rt.CreatedAt + newRefreshTokenObj.LastUsedAt = &usedAt + // Atomically revoke the old token and create the new one. + // If this fails the old token stays active and the client can retry the refresh. + if err := h.store.ConfigStore.RotateOAuth2RefreshToken(ctx, rt.ID, newRefreshTokenObj); err != nil { + if errors.Is(err, configstore.ErrNotFound) { + // The token was concurrently rotated/revoked between lookup and rotate. + sendOAuthError(ctx, fasthttp.StatusBadRequest, "invalid_grant", "refresh token not found or revoked") + return + } + sendOAuthError(ctx, fasthttp.StatusInternalServerError, "server_error", "failed to rotate token") + return + } + sendTokenResponse(ctx, accessToken, newRefreshToken, rt.Scope, oauth2ServerCfg(h.store).AccessTokenTTL) +} + +// issueTokenPair mints a signed JWT access token and builds a refresh token row. +// It is a pure function — no DB writes. The caller is responsible for atomically +// persisting the refresh token row alongside any grant-specific side-effects +// (e.g. marking the auth code consumed, or revoking the previous refresh token). +// +// familyID traces the token back to its original authorization grant; all +// rotated descendants share the same ID for stolen-token detection (RFC 9700 §2.2.2). +// +// On error, issueTokenPair writes an OAuth error response to ctx and returns a +// non-nil error so the caller can return immediately without writing again. +func (h *OAuth2IssuanceHandler) issueTokenPair( + ctx *fasthttp.RequestCtx, + familyID, clientID, bfMode, bfSub, scope, resource string, +) (accessToken, refreshTokenPlain string, rt *configtables.TableOAuth2RefreshToken, err error) { + cfg := oauth2ServerCfg(h.store) + accessTokenTTL := cfg.AccessTokenTTL + if accessTokenTTL <= 0 { + accessTokenTTL = configtables.DefaultAccessTokenTTL + } + + signingKey, err := h.store.ConfigStore.GetOAuth2SigningKey(ctx) + if err != nil { + sendOAuthError(ctx, fasthttp.StatusInternalServerError, "server_error", "signing key unavailable") + return + } + privKey, err := parseRSAPrivateKeyPEM(signingKey.PrivateKeyPEM) + if err != nil { + sendOAuthError(ctx, fasthttp.StatusInternalServerError, "server_error", "invalid signing key") + return + } + + issuer := oauth2IssuerURL(ctx, h.store) + now := time.Now() + claims := jwt.MapClaims{ + "iss": issuer, + "aud": jwt.ClaimStrings{resource}, + "sub": bfSub, + "bf_mode": bfMode, + "scope": scope, + "iat": now.Unix(), + "nbf": now.Unix(), + "exp": now.Add(time.Duration(accessTokenTTL) * time.Second).Unix(), + } + tok := jwt.NewWithClaims(jwt.SigningMethodRS256, claims) + tok.Header["kid"] = signingKey.KID + accessToken, err = tok.SignedString(privKey) + if err != nil { + sendOAuthError(ctx, fasthttp.StatusInternalServerError, "server_error", "failed to sign access token") + return + } + + refreshTokenPlain, err = generateSecureToken(32) + if err != nil { + sendOAuthError(ctx, fasthttp.StatusInternalServerError, "server_error", "failed to generate refresh token") + return + } + rt = &configtables.TableOAuth2RefreshToken{ + ID: uuid.New().String(), + TokenHash: hashSHA256Hex(refreshTokenPlain), + FamilyID: familyID, + ClientID: clientID, + BfMode: bfMode, + BfSub: bfSub, + Scope: scope, + Resource: resource, + CreatedAt: now, + } + return +} + +// sendTokenResponse writes the RFC 6749 token response to ctx. +func sendTokenResponse(ctx *fasthttp.RequestCtx, accessToken, refreshToken, scope string, accessTokenTTL int) { + if accessTokenTTL <= 0 { + accessTokenTTL = configtables.DefaultAccessTokenTTL + } + ctx.SetStatusCode(fasthttp.StatusOK) + ctx.SetContentType("application/json") + ctx.Response.Header.Set("Cache-Control", "no-store") + ctx.Response.Header.Set("Pragma", "no-cache") + data, err := sonic.Marshal(map[string]any{ + "access_token": accessToken, + "token_type": "Bearer", + "expires_in": accessTokenTTL, + "refresh_token": refreshToken, + "scope": scope, + }) + if err != nil { + sendOAuthError(ctx, fasthttp.StatusInternalServerError, "server_error", "failed to marshal token response") + return + } + ctx.SetBody(data) +} + +// clientIDFromRequest resolves the OAuth client_id for a token request. Public +// clients (token_endpoint_auth_method=none) may send it either as a request +// parameter or as the username of an HTTP Basic Authorization header +// (RFC 6749 §2.3.1); some clients default to the header form. Both are accepted; +// the Basic password is ignored since only public clients are supported. +func clientIDFromRequest(ctx *fasthttp.RequestCtx) string { + if v := string(ctx.FormValue("client_id")); v != "" { + return v + } + auth := string(ctx.Request.Header.Peek("Authorization")) + const prefix = "Basic " + if len(auth) > len(prefix) && strings.EqualFold(auth[:len(prefix)], prefix) { + if decoded, err := base64.StdEncoding.DecodeString(strings.TrimSpace(auth[len(prefix):])); err == nil { + // Basic credentials are "client_id:client_secret", each + // application/x-www-form-urlencoded; take and decode the username. + username, _, _ := strings.Cut(string(decoded), ":") + if id, uerr := url.QueryUnescape(username); uerr == nil { + return id + } + return username + } + } + return "" +} + +// --- Helpers --- + +// matchRedirectURI validates redirect_uri against registered URIs. +// For loopback addresses (localhost / 127.0.0.1), port is ignored per RFC 8252 §7.3. +// isAllowedRedirectScheme reports whether a redirect URI uses a safe scheme: +// https for any host, or http only for loopback addresses (localhost/127.0.0.1). +// This rejects javascript:, data:, and other schemes that could be abused when a +// Location header is built from the URI. +func isAllowedRedirectScheme(candidate string) bool { + parsed, err := url.Parse(candidate) + if err != nil { + return false + } + switch parsed.Scheme { + case "https": + return true + case "http": + host := parsed.Hostname() + return host == "localhost" || host == "127.0.0.1" + default: + return false + } +} + +func matchRedirectURI(candidate string, registered []string) bool { + parsed, err := url.Parse(candidate) + if err != nil { + return false + } + isLoopback := parsed.Hostname() == "localhost" || parsed.Hostname() == "127.0.0.1" + + for _, r := range registered { + rParsed, err := url.Parse(r) + if err != nil { + continue + } + if isLoopback && (rParsed.Hostname() == "localhost" || rParsed.Hostname() == "127.0.0.1") { + // Loopback: match scheme + host (without port) + path. + if parsed.Scheme == rParsed.Scheme && parsed.Path == rParsed.Path { + return true + } + } else { + // Non-loopback: exact match. + if candidate == r { + return true + } + } + } + return false +} + +// redirectWithParams builds a redirect URL with the given query params and redirects. +func redirectWithParams(ctx *fasthttp.RequestCtx, base string, params map[string]string) { + u, err := url.Parse(base) + if err != nil { + ctx.SetStatusCode(fasthttp.StatusInternalServerError) + return + } + q := u.Query() + for k, v := range params { + if v != "" { + q.Set(k, v) + } + } + u.RawQuery = q.Encode() + ctx.Response.Header.Set("Location", u.String()) + ctx.SetStatusCode(fasthttp.StatusFound) +} + +// scopeWithinRegistered reports whether every space-delimited token in requested +// is present in the registered scope set. Callers default an empty requested scope +// to the registered scope before calling, so an empty requested scope is vacuously +// within bounds. +func scopeWithinRegistered(requested, registered string) bool { + allowed := make(map[string]struct{}) + for s := range strings.FieldsSeq(registered) { + allowed[s] = struct{}{} + } + for s := range strings.FieldsSeq(requested) { + if _, ok := allowed[s]; !ok { + return false + } + } + return true +} + +// verifyPKCES256 verifies a PKCE S256 code_verifier against a stored challenge. +func verifyPKCES256(verifier, challenge string) bool { + h := sha256.Sum256([]byte(verifier)) + computed := base64.RawURLEncoding.EncodeToString(h[:]) + return subtle.ConstantTimeCompare([]byte(computed), []byte(challenge)) == 1 +} + +// hashSHA256Hex returns the hex-encoded SHA-256 hash of the input. +func hashSHA256Hex(input string) string { + h := sha256.Sum256([]byte(input)) + return hex.EncodeToString(h[:]) +} + +// generateSecureToken returns a cryptographically secure URL-safe random token. +func generateSecureToken(length int) (string, error) { + b := make([]byte, length) + if _, err := rand.Read(b); err != nil { + return "", err + } + return base64.RawURLEncoding.EncodeToString(b), nil +} + +// parseRSAPrivateKeyPEM decodes and parses a PKCS8 RSA private key PEM. +func parseRSAPrivateKeyPEM(pemStr string) (*rsa.PrivateKey, error) { + block, rest := pem.Decode([]byte(pemStr)) + if block == nil || len(rest) > 0 { + return nil, fmt.Errorf("malformed private key PEM") + } + key, err := x509.ParsePKCS8PrivateKey(block.Bytes) + if err != nil { + return nil, fmt.Errorf("parse private key: %w", err) + } + rsaKey, ok := key.(*rsa.PrivateKey) + if !ok { + return nil, fmt.Errorf("expected RSA private key, got %T", key) + } + return rsaKey, nil +} + +// sendOAuthError writes an RFC 6749 §5.2 error response. +func sendOAuthError(ctx *fasthttp.RequestCtx, statusCode int, errCode, description string) { + ctx.SetStatusCode(statusCode) + ctx.SetContentType("application/json") + data, err := sonic.Marshal(map[string]string{ + "error": errCode, + "error_description": description, + }) + if err != nil { + ctx.Error("failed to marshal error response", fasthttp.StatusInternalServerError) + return + } + ctx.SetBody(data) +} diff --git a/transports/bifrost-http/handlers/temptokens.go b/transports/bifrost-http/handlers/temptokens.go index 5a5cb631321..ea7d6635873 100644 --- a/transports/bifrost-http/handlers/temptokens.go +++ b/transports/bifrost-http/handlers/temptokens.go @@ -51,6 +51,19 @@ var mcpHeadersAuthScope = temptoken.Scope{ MaxTTL: 15 * time.Minute, } +// oauth2ConsentScope authorizes the public OAuth2 consent page to call the +// consent flow APIs. The flow request ID is substituted into {id} at +// validation time, binding each token to exactly one authorize request. +var oauth2ConsentScope = temptoken.Scope{ + Name: temptoken.OAuth2ConsentScopeName, + AllowedRoutes: []temptoken.RoutePattern{ + {Method: "GET", Path: "/api/oauth2/consent/flows/{id}"}, + {Method: "PUT", Path: "/api/oauth2/consent/flows/{id}"}, + }, + ResourceIDInPath: "{id}", + MaxTTL: 15 * time.Minute, +} + // RegisterTempTokenScopes registers every scope owned by this handlers // package on the given service. Called at server startup once the service // has been constructed. Returns an error if any scope is invalid or has @@ -65,5 +78,8 @@ func RegisterTempTokenScopes(svc *temptoken.Service) error { if err := svc.Registry().Register(mcpHeadersAuthScope); err != nil { return fmt.Errorf("temp_token_scopes: register mcp_headers_auth: %w", err) } + if err := svc.Registry().Register(oauth2ConsentScope); err != nil { + return fmt.Errorf("temp_token_scopes: register oauth2_consent: %w", err) + } return nil } diff --git a/transports/bifrost-http/lib/config_test.go b/transports/bifrost-http/lib/config_test.go index 950a1d3b9cc..f78d7f10e6d 100644 --- a/transports/bifrost-http/lib/config_test.go +++ b/transports/bifrost-http/lib/config_test.go @@ -428,6 +428,36 @@ func (m *MockConfigStore) RefreshConnectionPool(ctx context.Context) error { func (m *MockConfigStore) GetOAuth2SigningKey(ctx context.Context) (*tables.OAuth2SigningKey, error) { return &tables.OAuth2SigningKey{}, nil } +func (m *MockConfigStore) CreateOAuth2Client(ctx context.Context, client *tables.TableOAuth2Client) error { + return nil +} +func (m *MockConfigStore) GetOAuth2ClientByClientID(ctx context.Context, clientID string) (*tables.TableOAuth2Client, error) { + return nil, configstore.ErrNotFound +} +func (m *MockConfigStore) CreateOAuth2AuthorizeRequest(ctx context.Context, req *tables.TableOAuth2AuthorizeRequest) error { + return nil +} +func (m *MockConfigStore) GetOAuth2AuthorizeRequestByID(ctx context.Context, id string) (*tables.TableOAuth2AuthorizeRequest, error) { + return nil, configstore.ErrNotFound +} +func (m *MockConfigStore) GetOAuth2AuthorizeRequestByCodeHash(ctx context.Context, codeHash string) (*tables.TableOAuth2AuthorizeRequest, error) { + return nil, configstore.ErrNotFound +} +func (m *MockConfigStore) ConsentOAuth2AuthorizeRequest(ctx context.Context, req *tables.TableOAuth2AuthorizeRequest) error { + return nil +} +func (m *MockConfigStore) SweepExpiredOAuth2AuthorizeRequests(ctx context.Context) error { + return nil +} +func (m *MockConfigStore) GetOAuth2RefreshTokenByHash(ctx context.Context, hash string) (*tables.TableOAuth2RefreshToken, error) { + return nil, configstore.ErrNotFound +} +func (m *MockConfigStore) ConsumeOAuth2AuthorizeRequest(ctx context.Context, requestID string, rt *tables.TableOAuth2RefreshToken) error { + return nil +} +func (m *MockConfigStore) RotateOAuth2RefreshToken(ctx context.Context, oldID string, newRT *tables.TableOAuth2RefreshToken) error { + return nil +} func (m *MockConfigStore) Ping(ctx context.Context) error { return nil } func (m *MockConfigStore) EncryptPlaintextRows(ctx context.Context) error { return nil } func (m *MockConfigStore) Close(ctx context.Context) error { return nil } diff --git a/transports/bifrost-http/server/server.go b/transports/bifrost-http/server/server.go index 2740682d29f..2c8d5fbd305 100644 --- a/transports/bifrost-http/server/server.go +++ b/transports/bifrost-http/server/server.go @@ -1402,6 +1402,7 @@ func (s *BifrostHTTPServer) RegisterAPIRoutes(ctx context.Context, callbacks Ser featureFlagsHandler := handlers.NewFeatureFlagsHandler(s.Config.FeatureFlags, s.Config.ConfigStore) // Going ahead with API handlers handlers.NewOAuth2DiscoveryHandler(s.Config).RegisterRoutes(s.Router, middlewares...) + handlers.NewOAuth2IssuanceHandler(s.Config, s.TempTokens).RegisterRoutes(s.Router) healthHandler.RegisterRoutes(s.Router, middlewares...) providerHandler.RegisterRoutes(s.Router, middlewares...) mcpHandler.RegisterRoutes(s.Router, middlewares...) @@ -1771,6 +1772,14 @@ func (s *BifrostHTTPServer) Bootstrap(ctx context.Context) error { if err == nil && semanticCachePlugin != nil { semanticCachePlugin.SetEmbeddingRequestExecutor(s.Client.EmbeddingRequest) } + // Bootstrap OAuth2 signing key when discovery is enabled — ensures JWKS + // and JWT signing are ready before the first request arrives. + if s.Config.ConfigStore != nil && s.Config.ClientConfig.IsMCPOAuthDiscoveryEnabled() { + if _, keyErr := s.Config.ConfigStore.GetOAuth2SigningKey(s.Ctx); keyErr != nil { + logger.Warn("oauth2: failed to bootstrap signing key: %v", keyErr) + } + } + // Register routes err = s.RegisterAPIRoutes(s.Ctx, s, apiMiddlewares...) if err != nil { diff --git a/transports/go.mod b/transports/go.mod index c1dfbe9bf77..6f5e6c889a0 100644 --- a/transports/go.mod +++ b/transports/go.mod @@ -10,6 +10,7 @@ require ( github.com/fasthttp/websocket v1.5.12 github.com/go-git/go-billy/v5 v5.9.0 github.com/go-git/go-git/v5 v5.19.1 + github.com/golang-jwt/jwt/v5 v5.3.1 github.com/google/pprof v0.0.0-20251213031049-b05bdaca462f github.com/google/uuid v1.6.0 github.com/klauspost/compress v1.18.6 @@ -123,7 +124,6 @@ require ( github.com/go-openapi/swag/yamlutils v0.25.4 // indirect github.com/go-openapi/validate v0.25.1 // indirect github.com/go-viper/mapstructure/v2 v2.5.0 // indirect - github.com/golang-jwt/jwt/v5 v5.3.1 // indirect github.com/golang/groupcache v0.0.0-20241129210726-2c02b8208cf8 // indirect github.com/google/cel-go v0.28.1 // indirect github.com/google/s2a-go v0.1.9 // indirect