diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index 06299e3d5b..7be0acd810 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -329,6 +329,8 @@ const ( IsLocalAdminContextKey BifrostContextKey = "is_local_admin" // bool (set by auth middleware when password-based auth succeeds - local admin user bypasses RBAC) BifrostContextKeyPassthroughOverridesPresent BifrostContextKey = "passthrough_overrides_present" // bool (set by HTTP transport) - passthrough raw request requested BifrostContextKeyConnectionClosed BifrostContextKey = "connection_closed" + BifrostContextKeyTempTokenScope BifrostContextKey = "bifrost-temp-token-scope" // string (set by auth middleware when a temp token authorized the request - names the scope from the temptoken registry) + BifrostContextKeyTempTokenResourceID BifrostContextKey = "bifrost-temp-token-resource-id" // string (set by auth middleware alongside the scope - the resource_id the token is bound to, e.g. an OAuth flow ID for mcp_auth) ) const ( diff --git a/framework/configstore/encryption.go b/framework/configstore/encryption.go index b2de668abe..82475e3ce0 100644 --- a/framework/configstore/encryption.go +++ b/framework/configstore/encryption.go @@ -46,6 +46,13 @@ func (s *RDBConfigStore) EncryptPlaintextRows(ctx context.Context) error { } totalEncrypted += count + // temp_tokens + count, err = s.encryptPlaintextTempTokens(ctx) + if err != nil { + return fmt.Errorf("failed to encrypt temp_tokens: %w", err) + } + totalEncrypted += count + // oauth_tokens count, err = s.encryptPlaintextOAuthTokens(ctx) if err != nil { @@ -185,6 +192,36 @@ func (s *RDBConfigStore) encryptPlaintextSessions(ctx context.Context) (int, err return count, nil } +// encryptPlaintextTempTokens finds all temp_tokens rows with plaintext encryption status +// and re-saves them in batches. The TempToken.BeforeSave hook handles encryption. +func (s *RDBConfigStore) encryptPlaintextTempTokens(ctx context.Context) (int, error) { + var count int + for { + var tokens []tables.TempToken + if err := s.DB().WithContext(ctx). + Where("(encryption_status = ? OR encryption_status IS NULL OR encryption_status = '') AND token != ''", encryptionStatusPlainText). + Limit(encryptionBatchSize). + Find(&tokens).Error; err != nil { + return count, err + } + if len(tokens) == 0 { + break + } + if err := s.DB().WithContext(ctx).Transaction(func(tx *gorm.DB) error { + for i := range tokens { + if err := tx.Save(&tokens[i]).Error; err != nil { + return err + } + } + return nil + }); err != nil { + return count, err + } + count += len(tokens) + } + return count, nil +} + // encryptPlaintextOAuthTokens finds all oauth_tokens rows with plaintext encryption status // and re-saves them in batches. The TableOauthToken.BeforeSave hook handles encryption. func (s *RDBConfigStore) encryptPlaintextOAuthTokens(ctx context.Context) (int, error) { diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index fdc63eeeb8..0757d36322 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -779,6 +779,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error { if err := migrationAddClientConfigMetadataColumn(ctx, db); err != nil { return err } + if err := migrationAddTempTokensTable(ctx, db); err != nil { + return err + } // Runs LAST in this batch — the refresh does tx.Find(&clientConfigs) // which GORM projects to every column in TableClientConfig, so it has // to come after every config_client column-add above (otherwise the @@ -8483,3 +8486,35 @@ func migrationAddFrameworkConfigHashColumn(ctx context.Context, db *gorm.DB) err } return nil } + +// migrationAddTempTokensTable creates the temp_tokens table that backs the +// temptoken service. +func migrationAddTempTokensTable(ctx context.Context, db *gorm.DB) error { + m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + ID: "add_temp_tokens_table", + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + mig := tx.Migrator() + if !mig.HasTable(&tables.TempToken{}) { + if err := mig.CreateTable(&tables.TempToken{}); err != nil { + return fmt.Errorf("failed to create temp_tokens table: %w", err) + } + } + return nil + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + mig := tx.Migrator() + if mig.HasTable(&tables.TempToken{}) { + if err := mig.DropTable(&tables.TempToken{}); err != nil { + return err + } + } + return nil + }, + }}) + if err := m.Migrate(); err != nil { + return fmt.Errorf("error running add_temp_tokens_table migration: %s", err.Error()) + } + return nil +} diff --git a/framework/configstore/rdb.go b/framework/configstore/rdb.go index 3779826cb1..dc95f5f39b 100644 --- a/framework/configstore/rdb.go +++ b/framework/configstore/rdb.go @@ -4443,6 +4443,64 @@ func (s *RDBConfigStore) FlushSessions(ctx context.Context) error { return s.DB().WithContext(ctx).Session(&gorm.Session{AllowGlobalUpdate: true}).Delete(&tables.SessionsTable{}).Error } +// CreateTempToken inserts a new temp_tokens row. The plaintext token must be +// set on the struct; the BeforeSave hook populates token_hash and (when +// encryption is enabled) encrypts the plaintext in place. The optional tx +// lets callers fold this write into an existing transaction (mirrors the +// pattern used by other mutating configstore methods). +func (s *RDBConfigStore) CreateTempToken(ctx context.Context, token *tables.TempToken, tx ...*gorm.DB) error { + db := s.DB() + if len(tx) > 0 { + db = tx[0] + } + return db.WithContext(ctx).Create(token).Error +} + +// GetTempTokenByHash retrieves a temp_tokens row by the SHA-256 hash of its +// plaintext. Returns (nil, nil) when no row matches — callers should treat that +// as "no such token" rather than an error. +func (s *RDBConfigStore) GetTempTokenByHash(ctx context.Context, tokenHash string) (*tables.TempToken, error) { + var token tables.TempToken + err := s.DB().WithContext(ctx).First(&token, "token_hash = ?", tokenHash).Error + if err != nil { + if errors.Is(err, gorm.ErrRecordNotFound) { + return nil, nil + } + return nil, err + } + return &token, nil +} + +// DeleteTempTokensByResourceID hard-deletes every row matching (scope, +// resource_id). Used by lifecycle owners (e.g. OAuth provider on flow +// termination) to invalidate the link as soon as the work it authorized +// completes. The (scope, resource_id) pair — not resource_id alone — keeps +// future scopes that happen to reuse the same opaque ID untouched. The +// optional tx lets callers fold the delete into an existing transaction. +func (s *RDBConfigStore) DeleteTempTokensByResourceID(ctx context.Context, scope, resourceID string, tx ...*gorm.DB) (int64, error) { + db := s.DB() + if len(tx) > 0 { + db = tx[0] + } + res := db.WithContext(ctx). + Where("scope = ? AND resource_id = ?", scope, resourceID). + Delete(&tables.TempToken{}) + if res.Error != nil { + return 0, res.Error + } + return res.RowsAffected, nil +} + +// DeleteExpiredTempTokens hard-deletes rows whose expires_at is at or before +// the given cutoff. Returns the number of rows deleted. +func (s *RDBConfigStore) DeleteExpiredTempTokens(ctx context.Context, before time.Time) (int64, error) { + res := s.DB().WithContext(ctx).Where("expires_at <= ?", before).Delete(&tables.TempToken{}) + if res.Error != nil { + return 0, res.Error + } + return res.RowsAffected, nil +} + // ExecuteTransaction executes a transaction. func (s *RDBConfigStore) ExecuteTransaction(ctx context.Context, fn func(tx *gorm.DB) error) error { return s.DB().WithContext(ctx).Transaction(fn) @@ -5098,4 +5156,3 @@ func (s *RDBConfigStore) DeleteOrphanedOauthUserTokens(ctx context.Context, olde } return result.RowsAffected, nil } - diff --git a/framework/configstore/store.go b/framework/configstore/store.go index 021b3eb882..451ba86c58 100644 --- a/framework/configstore/store.go +++ b/framework/configstore/store.go @@ -258,6 +258,15 @@ type ConfigStore interface { DeleteSession(ctx context.Context, token string) error FlushSessions(ctx context.Context) error + // Temp token CRUD + CreateTempToken(ctx context.Context, token *tables.TempToken, tx ...*gorm.DB) error + GetTempTokenByHash(ctx context.Context, tokenHash string) (*tables.TempToken, error) + // DeleteTempTokensByResourceID removes every row matching (scope, resource_id). + // Used by lifecycle owners (e.g. OAuth provider on flow termination) to burn + // the link as soon as the work it authorized is finished. + DeleteTempTokensByResourceID(ctx context.Context, scope, resourceID string, tx ...*gorm.DB) (int64, error) + DeleteExpiredTempTokens(ctx context.Context, before time.Time) (int64, error) + // Model pricing CRUD GetModelPrices(ctx context.Context) ([]tables.TableModelPricing, error) UpsertModelPrices(ctx context.Context, pricing *tables.TableModelPricing, tx ...*gorm.DB) error diff --git a/framework/configstore/tables/temp_token.go b/framework/configstore/tables/temp_token.go new file mode 100644 index 0000000000..8af30886b7 --- /dev/null +++ b/framework/configstore/tables/temp_token.go @@ -0,0 +1,57 @@ +package tables + +import ( + "fmt" + "time" + + "github.com/maximhq/bifrost/framework/encrypt" + "gorm.io/gorm" +) + +// TempToken is a short-lived, narrow-scope credential that authorizes access +// to a specific set of endpoints without requiring dashboard login. +// +// Each row is bound to a (scope, resource_id) pair: the scope names a set of +// allowed routes (registered in framework/temptoken), and the resource_id ties +// the token to the specific resource those routes act on (e.g. the OAuth flow +// ID for the mcp_auth scope). The plaintext token is hashed for lookup and +// encrypted at rest, matching the SessionsTable pattern. +type TempToken struct { + ID string `gorm:"type:varchar(255);primaryKey" json:"id"` // UUID + Token string `gorm:"type:text;not null" json:"-"` // encrypted at rest when encryption is enabled + TokenHash string `gorm:"type:varchar(64);uniqueIndex:idx_temp_token_hash" json:"-"` // SHA-256 of plaintext for lookup + Scope string `gorm:"type:varchar(64);index;not null" json:"scope"` // e.g. "mcp_auth" — keys into the scope registry + ResourceID string `gorm:"type:text;index" json:"resource_id,omitempty"` // resource the scope binds to (semantics per scope); indexed for lifecycle-driven deletes + ExpiresAt time.Time `gorm:"index;not null" json:"expires_at"` + CreatedAt time.Time `gorm:"index;not null" json:"created_at"` + UpdatedAt time.Time `gorm:"index;not null" json:"updated_at"` + EncryptionStatus string `gorm:"type:varchar(20);default:'plain_text'" json:"-"` +} + +// TableName sets the table name for the model. +func (TempToken) TableName() string { return "temp_tokens" } + +// BeforeSave hashes the plaintext for lookup and encrypts it for storage. +// Hash must be computed before encryption so it always covers the plaintext. +func (t *TempToken) BeforeSave(tx *gorm.DB) error { + if t.Token != "" { + t.TokenHash = encrypt.HashSHA256(t.Token) + } + if encrypt.IsEnabled() && t.Token != "" { + if err := encryptString(&t.Token); err != nil { + return fmt.Errorf("failed to encrypt temp token: %w", err) + } + t.EncryptionStatus = EncryptionStatusEncrypted + } + return nil +} + +// AfterFind decrypts the stored plaintext when encryption is in effect. +func (t *TempToken) AfterFind(tx *gorm.DB) error { + if t.EncryptionStatus == EncryptionStatusEncrypted { + if err := decryptString(&t.Token); err != nil { + return fmt.Errorf("failed to decrypt temp token: %w", err) + } + } + return nil +} diff --git a/framework/oauth2/main.go b/framework/oauth2/main.go index de545b1bd1..45ca86c6cd 100644 --- a/framework/oauth2/main.go +++ b/framework/oauth2/main.go @@ -20,6 +20,7 @@ import ( "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/framework/configstore" "github.com/maximhq/bifrost/framework/configstore/tables" + "github.com/maximhq/bifrost/framework/temptoken" ) const ( @@ -33,6 +34,13 @@ type OAuth2Provider struct { configStore configstore.ConfigStore mu sync.RWMutex retryBaseDelay time.Duration // base delay for token endpoint retry backoff; doubles each attempt (1×, 2×, 4×) + + // tempTokens, when non-nil, is used by InitiateUserOAuthFlow to mint a + // short-lived mcp_auth temp token and embed it in the returned auth-page + // URL as a fragment. Optional — when nil, the URL is returned without a + // fragment and the page works only for callers already authenticated to + // the dashboard. + tempTokens *temptoken.Service } // NewOAuth2Provider creates a new OAuth provider instance @@ -47,6 +55,38 @@ func NewOAuth2Provider(configStore configstore.ConfigStore, logger schemas.Logge } } +// SetTempTokenService installs the temp-token service used by +// InitiateUserOAuthFlow to mint the mcp_auth token embedded in the +// auth-page URL fragment. Called by server startup once both services +// have been constructed (the provider is built first by lib/config.go, +// the service later by the HTTP transport). +func (p *OAuth2Provider) SetTempTokenService(svc *temptoken.Service) { + p.mu.Lock() + defer p.mu.Unlock() + p.tempTokens = svc +} + +// cleanupFlow deletes the flow row and any temp tokens minted for it. Called +// on every terminal transition (success or any failure) so the auth-page +// link stops working as soon as the work it authorized ends. +// +// Detached from the caller's context via WithoutCancel so a client cancellation +// (e.g. the browser closing the tab after the upstream OAuth bounce) can't +// short-circuit the deletes and leave the rows alive until the sweep. Mirrors +// the pattern used by markExpiredIfPermanent in this file. +func (p *OAuth2Provider) cleanupFlow(ctx context.Context, sessionID string) { + cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 5*time.Second) + defer cancel() + if err := p.configStore.DeleteOauthUserSession(cleanupCtx, sessionID); err != nil { + logger.Warn("per-user OAuth flow row cleanup failed: session_id=%s err=%v", sessionID, err) + } + if p.tempTokens != nil { + if _, err := p.tempTokens.DeleteByResourceID(cleanupCtx, temptoken.MCPAuthScopeName, sessionID); err != nil { + logger.Warn("per-user OAuth temp-token cleanup failed: session_id=%s err=%v", sessionID, err) + } + } +} + // GetAccessToken retrieves the access token for a given oauth_config_id func (p *OAuth2Provider) GetAccessToken(ctx context.Context, oauthConfigID string) (string, error) { // Load oauth_config by ID @@ -984,6 +1024,23 @@ func (p *OAuth2Provider) InitiateUserOAuthFlow(ctx context.Context, oauthConfigI // elsewhere in the dashboard UI. frontendURL := strings.TrimSuffix(redirectURI, "/api/oauth/callback") + "/workspace/mcp-sessions/auth?flow=" + sessionID + // Mint a mcp_auth temp token bound to this flow's row ID and embed it in + // the URL as a fragment so a browser hitting the auth page without a + // dashboard session can still call the per-user flow endpoints. The + // fragment never leaves the browser (not in server logs, not in the + // upstream-OAuth Referer), unlike a query param. + if p.tempTokens != nil { + ttl := time.Until(expiresAt) + if ttl > 0 { + plaintext, mintErr := p.tempTokens.Mint(ctx, temptoken.MCPAuthScopeName, sessionID, ttl) + if mintErr != nil { + logger.Warn("Failed to mint mcp_auth temp token for flow %s: %v (link still usable for dashboard-authenticated callers)", sessionID, mintErr) + } else { + frontendURL = frontendURL + "#t=" + plaintext + } + } + } + logger.Debug("Per-user OAuth flow initiated: session_id=%s, mcp_client_id=%s", sessionID, mcpClientID) return &schemas.OAuth2FlowInitiation{ @@ -1012,7 +1069,7 @@ func (p *OAuth2Provider) CompleteUserOAuthFlow(ctx context.Context, state string // dead rows. The UI sees 404 on flow-detail and renders "expired // or already completed" — no audit trail to preserve here. if time.Now().After(session.ExpiresAt) { - _ = p.configStore.DeleteOauthUserSession(ctx, session.ID) + p.cleanupFlow(ctx, session.ID) return "", fmt.Errorf("per-user oauth flow expired") } @@ -1022,7 +1079,7 @@ func (p *OAuth2Provider) CompleteUserOAuthFlow(ctx context.Context, state string // (config row missing) is hidden. templateConfig, err := p.configStore.GetOauthConfigByID(ctx, session.OauthConfigID) if err != nil { - _ = p.configStore.DeleteOauthUserSession(ctx, session.ID) + p.cleanupFlow(ctx, session.ID) return "", fmt.Errorf("failed to load template oauth config: %w", err) } if templateConfig == nil { @@ -1046,7 +1103,7 @@ func (p *OAuth2Provider) CompleteUserOAuthFlow(ctx context.Context, state string session.CodeVerifier, ) if err != nil { - _ = p.configStore.DeleteOauthUserSession(ctx, session.ID) + p.cleanupFlow(ctx, session.ID) return "", fmt.Errorf("per-user token exchange failed: %w", err) } @@ -1082,7 +1139,7 @@ func (p *OAuth2Provider) CompleteUserOAuthFlow(ctx context.Context, state string } } if tokenUserID == nil || *tokenUserID == "" { - _ = p.configStore.DeleteOauthUserSession(ctx, session.ID) + p.cleanupFlow(ctx, session.ID) return "", fmt.Errorf("user-mode oauth flow has no user_id at completion (neither flow nor completer context)") } // Stamp the resolved user_id back on the flow for audit. @@ -1090,7 +1147,7 @@ func (p *OAuth2Provider) CompleteUserOAuthFlow(ctx context.Context, state string case schemas.MCPAuthModeVK: tokenVKID = session.VirtualKeyID if tokenVKID == nil || *tokenVKID == "" { - _ = p.configStore.DeleteOauthUserSession(ctx, session.ID) + p.cleanupFlow(ctx, session.ID) return "", fmt.Errorf("vk-mode oauth flow has no virtual_key_id at completion") } case schemas.MCPAuthModeSession: @@ -1132,16 +1189,13 @@ func (p *OAuth2Provider) CompleteUserOAuthFlow(ctx context.Context, state string return "", fmt.Errorf("failed to create per-user oauth token: %w", err) } - // Token row is written; the flow row's purpose ends here. Delete it - // instead of marking 'authorized' — the token row is the durable - // record of the binding; the flow row is just the transient PKCE/ - // state carrier. The UI shows "expired or already completed" on the - // resulting 404, which is the truthful read. - if err := p.configStore.DeleteOauthUserSession(ctx, session.ID); err != nil { - // Non-fatal: log but proceed. The token row is already created - // and usable; a lingering flow row is hygienic, not correctness. - logger.Warn("per-user OAuth flow row cleanup failed (token already created): session_id=%s err=%v", session.ID, err) - } + // Token row is written; the flow row's purpose ends here. cleanupFlow + // deletes both the flow row (transient PKCE/state carrier — the token + // row is the durable record now) and the mcp_auth temp token bound to + // it, so the auth-page link stops working immediately. The UI shows + // "expired or already completed" on the resulting 404, which is the + // truthful read. + p.cleanupFlow(ctx, session.ID) logger.Debug("Per-user OAuth flow completed: session_id=%s, mcp_client_id=%s", session.ID, session.MCPClientID) diff --git a/framework/temptoken/scope.go b/framework/temptoken/scope.go new file mode 100644 index 0000000000..490fedbc78 --- /dev/null +++ b/framework/temptoken/scope.go @@ -0,0 +1,115 @@ +package temptoken + +import ( + "fmt" + "strings" + "sync" + "time" +) + +// Well-known scope names. The framework reserves the canonical strings here +// so that issuers (e.g. framework/oauth2) and registrants (the handlers +// package that defines the routes) reference a single source of truth. The +// actual Scope definition — routes, TTL, single-use — is owned by the +// transports/handlers layer that registers the scope with a Service. +const ( + // MCPAuthScopeName names the scope that authorizes the MCP per-user OAuth + // auth page to call the per-user flow endpoints. Bound resource_id is the + // OAuth flow ID. See transports/bifrost-http/handlers/temp_token_scopes.go. + MCPAuthScopeName = "mcp_auth" +) + +// RoutePattern is one (method, path) pair a Scope grants access to. The path +// may include the placeholder declared in the owning Scope's ResourceIDInPath +// field; at validation time the placeholder is substituted with the row's +// resource_id and the result is exact-matched against the request path. +type RoutePattern struct { + Method string // "GET", "POST", ... — case-insensitive on the wire, normalized to upper-case here. + Path string // e.g. "/api/oauth/per-user/flows/{id}" +} + +// Scope describes one class of temp token. The Name is the contract key stored +// on each row's scope column; AllowedRoutes + ResourceIDInPath define which +// requests the token authorizes; MaxTTL caps how long Mint will set +// expires_at. Invalidation is by row deletion — callers either let the +// expiry pass, the sweep worker collect, or explicitly Delete / +// DeleteByResourceID when the work the token authorized is complete. +type Scope struct { + Name string + AllowedRoutes []RoutePattern + ResourceIDInPath string // e.g. "{id}" — empty means routes have no resource_id substitution + MaxTTL time.Duration +} + +// matchesRequest reports whether the (method, path) pair satisfies any of the +// scope's AllowedRoutes after substituting ResourceIDInPath with resourceID. +// Method comparison is case-insensitive; path comparison is exact after +// substitution. An empty ResourceIDInPath means the patterns are matched as-is. +func (s Scope) matchesRequest(method, path, resourceID string) bool { + method = strings.ToUpper(method) + for _, r := range s.AllowedRoutes { + if strings.ToUpper(r.Method) != method { + continue + } + expected := r.Path + if s.ResourceIDInPath != "" { + expected = strings.ReplaceAll(expected, s.ResourceIDInPath, resourceID) + } + if expected == path { + return true + } + } + return false +} + +// Registry holds the process-global set of declared scopes. Scopes register at +// startup; the middleware looks them up on every validation. Registration is +// keyed by Scope.Name and double-registration is an error so misconfiguration +// fails loudly at boot rather than silently overwriting. +type Registry struct { + mu sync.RWMutex + scopes map[string]Scope +} + +// NewRegistry constructs an empty registry. +func NewRegistry() *Registry { + return &Registry{scopes: make(map[string]Scope)} +} + +// Register adds a Scope. Returns an error if a scope with the same Name is +// already registered, or if the Scope is invalid (missing Name, missing +// routes, missing placeholder when ResourceIDInPath references one). +func (r *Registry) Register(s Scope) error { + if s.Name == "" { + return fmt.Errorf("temptoken: scope name is required") + } + if len(s.AllowedRoutes) == 0 { + return fmt.Errorf("temptoken: scope %q must declare at least one allowed route", s.Name) + } + if s.MaxTTL <= 0 { + return fmt.Errorf("temptoken: scope %q must declare a positive MaxTTL", s.Name) + } + if s.ResourceIDInPath != "" { + for _, r := range s.AllowedRoutes { + if !strings.Contains(r.Path, s.ResourceIDInPath) { + return fmt.Errorf("temptoken: scope %q declares ResourceIDInPath %q but route %q %q does not contain it", + s.Name, s.ResourceIDInPath, r.Method, r.Path) + } + } + } + r.mu.Lock() + defer r.mu.Unlock() + if _, exists := r.scopes[s.Name]; exists { + return fmt.Errorf("temptoken: scope %q already registered", s.Name) + } + r.scopes[s.Name] = s + return nil +} + +// Lookup returns the Scope registered under the given name, or false if none. +func (r *Registry) Lookup(name string) (Scope, bool) { + r.mu.RLock() + defer r.mu.RUnlock() + s, ok := r.scopes[name] + return s, ok +} diff --git a/framework/temptoken/service.go b/framework/temptoken/service.go new file mode 100644 index 0000000000..e3b7ccec1a --- /dev/null +++ b/framework/temptoken/service.go @@ -0,0 +1,164 @@ +package temptoken + +import ( + "context" + "crypto/rand" + "encoding/base64" + "errors" + "fmt" + "time" + + "github.com/google/uuid" + "github.com/maximhq/bifrost/framework/configstore" + "github.com/maximhq/bifrost/framework/configstore/tables" + "github.com/maximhq/bifrost/framework/encrypt" +) + +// Errors returned by the service. Callers (notably the auth middleware) should +// treat these as opaque "not authorized" signals — they're typed so tests can +// assert on them without coupling to error message text. +var ( + // ErrTokenNotFound is returned when the presented token does not match any row. + ErrTokenNotFound = errors.New("temptoken: token not found") + // ErrTokenExpired is returned when the matched row's expires_at is in the past. + ErrTokenExpired = errors.New("temptoken: token expired") + // ErrScopeUnknown is returned when the row's scope column does not match any + // registered scope. Indicates either a stale row (scope was deregistered) or + // a corrupted row. + ErrScopeUnknown = errors.New("temptoken: token scope is not registered") + // ErrRouteNotAllowed is returned when the (method, path) of the request does + // not satisfy any of the scope's AllowedRoutes (after resource_id substitution). + ErrRouteNotAllowed = errors.New("temptoken: request method and path are not allowed by token scope") + // ErrTTLExceedsMax is returned by Mint when the caller-requested TTL is + // larger than the scope's MaxTTL. + ErrTTLExceedsMax = errors.New("temptoken: requested TTL exceeds scope MaxTTL") +) + +// ValidatedToken is the result of a successful Validate call. Callers attach +// these values to the request context so handlers can apply defense-in-depth +// checks. +type ValidatedToken struct { + ID string + Scope string + ResourceID string +} + +// Service mints and validates temp tokens. It composes the scope registry with +// the configstore-backed persistence layer; nothing else needs the row format +// directly. +type Service struct { + store configstore.ConfigStore + registry *Registry + now func() time.Time // injectable for tests +} + +// NewService constructs a Service backed by the given store and registry. The +// registry can be empty at construction time — scopes can be Register()'d +// before the first Validate call (in practice, at server startup). +func NewService(store configstore.ConfigStore, registry *Registry) *Service { + return &Service{store: store, registry: registry, now: time.Now} +} + +// Registry exposes the underlying scope registry so callers can register +// scopes without holding a separate reference. Useful for transports that +// receive only the Service from server startup. +func (s *Service) Registry() *Registry { return s.registry } + +// Mint creates a new temp token under the given scope, bound to resourceID, +// with the requested TTL. The TTL must be > 0 and <= the scope's MaxTTL. The +// returned plaintext is the value the caller embeds in URLs or hands back to +// the user; it is never persisted in plaintext when encryption is enabled. +func (s *Service) Mint(ctx context.Context, scopeName, resourceID string, ttl time.Duration) (string, error) { + scope, ok := s.registry.Lookup(scopeName) + if !ok { + return "", fmt.Errorf("%w: %s", ErrScopeUnknown, scopeName) + } + if ttl <= 0 || ttl > scope.MaxTTL { + return "", fmt.Errorf("%w: requested=%s max=%s", ErrTTLExceedsMax, ttl, scope.MaxTTL) + } + plaintext, err := generatePlaintext() + if err != nil { + return "", fmt.Errorf("temptoken: failed to generate token: %w", err) + } + row := &tables.TempToken{ + ID: uuid.New().String(), + Token: plaintext, + Scope: scope.Name, + ResourceID: resourceID, + ExpiresAt: s.now().Add(ttl), + } + if err := s.store.CreateTempToken(ctx, row); err != nil { + return "", fmt.Errorf("temptoken: failed to persist token: %w", err) + } + return plaintext, nil +} + +// Validate authenticates the presented plaintext for the given request +// (method, path). +func (s *Service) Validate(ctx context.Context, plaintext, method, path string) (*ValidatedToken, error) { + if plaintext == "" { + return nil, ErrTokenNotFound + } + hash := encrypt.HashSHA256(plaintext) + row, err := s.store.GetTempTokenByHash(ctx, hash) + if err != nil { + return nil, fmt.Errorf("temptoken: lookup failed: %w", err) + } + if row == nil { + return nil, ErrTokenNotFound + } + if !row.ExpiresAt.After(s.now()) { + return nil, ErrTokenExpired + } + scope, ok := s.registry.Lookup(row.Scope) + if !ok { + return nil, fmt.Errorf("%w: %s", ErrScopeUnknown, row.Scope) + } + if !scope.matchesRequest(method, path, row.ResourceID) { + return nil, ErrRouteNotAllowed + } + return &ValidatedToken{ + ID: row.ID, + Scope: row.Scope, + ResourceID: row.ResourceID, + }, nil +} + +// DeleteExpired removes every token row whose expires_at is at or before +// `before`. Called by [SweepWorker] on its tick; callers passing time.Now() +// reap everything currently past its TTL. Returns the number of rows removed. +func (s *Service) DeleteExpired(ctx context.Context, before time.Time) (int64, error) { + n, err := s.store.DeleteExpiredTempTokens(ctx, before) + if err != nil { + return 0, fmt.Errorf("temptoken: delete expired failed: %w", err) + } + return n, nil +} + +// DeleteByResourceID removes every token row matching (scope, resourceID). +// Lifecycle owners call this when the underlying resource the token authorized +// is finished — e.g. the OAuth provider after a per-user flow terminates +// (success or failure) so the link stops working immediately instead of +// waiting for TTL. Returns the number of rows removed; both 0 and N are +// considered successful outcomes — callers should not treat 0 as an error. +func (s *Service) DeleteByResourceID(ctx context.Context, scope, resourceID string) (int64, error) { + if scope == "" || resourceID == "" { + return 0, nil + } + n, err := s.store.DeleteTempTokensByResourceID(ctx, scope, resourceID) + if err != nil { + return 0, fmt.Errorf("temptoken: delete by resource_id failed: %w", err) + } + return n, nil +} + +// generatePlaintext returns a cryptographically random URL-safe string. 32 +// bytes of entropy yields ~43 base64url characters — plenty against any +// realistic brute-force budget given the 15-minute TTL ceiling. +func generatePlaintext() (string, error) { + buf := make([]byte, 32) + if _, err := rand.Read(buf); err != nil { + return "", err + } + return base64.RawURLEncoding.EncodeToString(buf), nil +} diff --git a/framework/temptoken/service_test.go b/framework/temptoken/service_test.go new file mode 100644 index 0000000000..ec15271631 --- /dev/null +++ b/framework/temptoken/service_test.go @@ -0,0 +1,258 @@ +package temptoken + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/maximhq/bifrost/framework/configstore" + "github.com/maximhq/bifrost/framework/configstore/tables" + "github.com/maximhq/bifrost/framework/encrypt" + "gorm.io/gorm" +) + +// fakeStore is a minimal in-memory configstore that only implements the +// temp-token CRUD surface. Other methods panic via the embedded interface so +// accidental dependencies fail loudly. +type fakeStore struct { + configstore.ConfigStore + + mu sync.Mutex + byID map[string]*tables.TempToken + byHash map[string]*tables.TempToken +} + +func newFakeStore() *fakeStore { + return &fakeStore{ + byID: make(map[string]*tables.TempToken), + byHash: make(map[string]*tables.TempToken), + } +} + +func (f *fakeStore) CreateTempToken(_ context.Context, tok *tables.TempToken, _ ...*gorm.DB) error { + f.mu.Lock() + defer f.mu.Unlock() + // Mimic BeforeSave: compute hash from plaintext for lookup. We don't run + // the actual GORM hook here, so do the same computation the hook does. + hash := encrypt.HashSHA256(tok.Token) + tok.TokenHash = hash + f.byID[tok.ID] = tok + f.byHash[hash] = tok + return nil +} + +func (f *fakeStore) GetTempTokenByHash(_ context.Context, hash string) (*tables.TempToken, error) { + f.mu.Lock() + defer f.mu.Unlock() + if t, ok := f.byHash[hash]; ok { + // Return a copy so caller mutations don't bleed into the store. + cp := *t + return &cp, nil + } + return nil, nil +} + +func (f *fakeStore) DeleteTempTokensByResourceID(_ context.Context, scope, resourceID string, _ ...*gorm.DB) (int64, error) { + f.mu.Lock() + defer f.mu.Unlock() + var deleted int64 + for hash, tok := range f.byHash { + if tok.Scope == scope && tok.ResourceID == resourceID { + delete(f.byHash, hash) + delete(f.byID, tok.ID) + deleted++ + } + } + return deleted, nil +} + +func (f *fakeStore) DeleteExpiredTempTokens(_ context.Context, before time.Time) (int64, error) { + f.mu.Lock() + defer f.mu.Unlock() + var deleted int64 + for hash, tok := range f.byHash { + if !tok.ExpiresAt.After(before) { + delete(f.byHash, hash) + delete(f.byID, tok.ID) + deleted++ + } + } + return deleted, nil +} + +// reusable scope used across most tests +func mcpAuthScope() Scope { + return Scope{ + Name: "mcp_auth", + AllowedRoutes: []RoutePattern{ + {Method: "GET", Path: "/api/oauth/per-user/flows/{id}"}, + {Method: "GET", Path: "/api/oauth/per-user/flows/{id}/start"}, + }, + ResourceIDInPath: "{id}", + MaxTTL: 15 * time.Minute, + } +} + +func newServiceWithMcpAuth(t *testing.T) (*Service, *fakeStore) { + t.Helper() + reg := NewRegistry() + if err := reg.Register(mcpAuthScope()); err != nil { + t.Fatalf("register scope: %v", err) + } + store := newFakeStore() + return NewService(store, reg), store +} + +func TestMintRejectsUnknownScope(t *testing.T) { + svc, _ := newServiceWithMcpAuth(t) + _, err := svc.Mint(context.Background(), "no_such_scope", "flow-1", time.Minute) + if !errors.Is(err, ErrScopeUnknown) { + t.Fatalf("expected ErrScopeUnknown, got %v", err) + } +} + +func TestMintRejectsTTLOverMax(t *testing.T) { + svc, _ := newServiceWithMcpAuth(t) + _, err := svc.Mint(context.Background(), "mcp_auth", "flow-1", time.Hour) + if !errors.Is(err, ErrTTLExceedsMax) { + t.Fatalf("expected ErrTTLExceedsMax, got %v", err) + } +} + +func TestMintRejectsNonPositiveTTL(t *testing.T) { + svc, _ := newServiceWithMcpAuth(t) + if _, err := svc.Mint(context.Background(), "mcp_auth", "flow-1", 0); !errors.Is(err, ErrTTLExceedsMax) { + t.Fatalf("expected ErrTTLExceedsMax for zero TTL, got %v", err) + } +} + +func TestValidateHappyPath(t *testing.T) { + svc, _ := newServiceWithMcpAuth(t) + tok, err := svc.Mint(context.Background(), "mcp_auth", "flow-abc", 5*time.Minute) + if err != nil { + t.Fatalf("mint: %v", err) + } + got, err := svc.Validate(context.Background(), tok, "GET", "/api/oauth/per-user/flows/flow-abc") + if err != nil { + t.Fatalf("validate detail: %v", err) + } + if got.Scope != "mcp_auth" || got.ResourceID != "flow-abc" { + t.Fatalf("got %+v", got) + } + if _, err := svc.Validate(context.Background(), tok, "GET", "/api/oauth/per-user/flows/flow-abc/start"); err != nil { + t.Fatalf("validate start: %v", err) + } +} + +func TestValidateRejectsWrongResourceID(t *testing.T) { + svc, _ := newServiceWithMcpAuth(t) + tok, err := svc.Mint(context.Background(), "mcp_auth", "flow-abc", 5*time.Minute) + if err != nil { + t.Fatalf("mint: %v", err) + } + _, err = svc.Validate(context.Background(), tok, "GET", "/api/oauth/per-user/flows/flow-xyz") + if !errors.Is(err, ErrRouteNotAllowed) { + t.Fatalf("expected ErrRouteNotAllowed, got %v", err) + } +} + +func TestValidateRejectsWrongMethod(t *testing.T) { + svc, _ := newServiceWithMcpAuth(t) + tok, _ := svc.Mint(context.Background(), "mcp_auth", "flow-abc", 5*time.Minute) + _, err := svc.Validate(context.Background(), tok, "POST", "/api/oauth/per-user/flows/flow-abc") + if !errors.Is(err, ErrRouteNotAllowed) { + t.Fatalf("expected ErrRouteNotAllowed, got %v", err) + } +} + +func TestValidateRejectsUnrelatedPath(t *testing.T) { + svc, _ := newServiceWithMcpAuth(t) + tok, _ := svc.Mint(context.Background(), "mcp_auth", "flow-abc", 5*time.Minute) + _, err := svc.Validate(context.Background(), tok, "GET", "/api/config/core") + if !errors.Is(err, ErrRouteNotAllowed) { + t.Fatalf("expected ErrRouteNotAllowed, got %v", err) + } +} + +func TestValidateRejectsExpired(t *testing.T) { + svc, _ := newServiceWithMcpAuth(t) + // Freeze time inside the service so we can advance it past the TTL. + now := time.Now() + svc.now = func() time.Time { return now } + tok, err := svc.Mint(context.Background(), "mcp_auth", "flow-abc", 1*time.Minute) + if err != nil { + t.Fatalf("mint: %v", err) + } + svc.now = func() time.Time { return now.Add(2 * time.Minute) } + _, err = svc.Validate(context.Background(), tok, "GET", "/api/oauth/per-user/flows/flow-abc") + if !errors.Is(err, ErrTokenExpired) { + t.Fatalf("expected ErrTokenExpired, got %v", err) + } +} + +func TestValidateRejectsUnknownToken(t *testing.T) { + svc, _ := newServiceWithMcpAuth(t) + _, err := svc.Validate(context.Background(), "definitely-not-a-real-token", "GET", "/api/oauth/per-user/flows/flow-abc") + if !errors.Is(err, ErrTokenNotFound) { + t.Fatalf("expected ErrTokenNotFound, got %v", err) + } +} + +func TestDeleteExpiredReapsPastTTL(t *testing.T) { + svc, store := newServiceWithMcpAuth(t) + now := time.Now() + svc.now = func() time.Time { return now } + // Two live + one already-expired. + if _, err := svc.Mint(context.Background(), "mcp_auth", "flow-a", 5*time.Minute); err != nil { + t.Fatalf("mint a: %v", err) + } + if _, err := svc.Mint(context.Background(), "mcp_auth", "flow-b", 5*time.Minute); err != nil { + t.Fatalf("mint b: %v", err) + } + svc.now = func() time.Time { return now.Add(-10 * time.Minute) } + if _, err := svc.Mint(context.Background(), "mcp_auth", "flow-c", 1*time.Minute); err != nil { + t.Fatalf("mint c: %v", err) + } + svc.now = func() time.Time { return now } + + n, err := svc.DeleteExpired(context.Background(), now) + if err != nil { + t.Fatalf("DeleteExpired: %v", err) + } + if n != 1 { + t.Fatalf("expected 1 row reaped, got %d", n) + } + if len(store.byID) != 2 { + t.Fatalf("expected 2 rows remaining, got %d", len(store.byID)) + } +} + +func TestRegistryRejectsDuplicate(t *testing.T) { + reg := NewRegistry() + if err := reg.Register(mcpAuthScope()); err != nil { + t.Fatalf("first register: %v", err) + } + if err := reg.Register(mcpAuthScope()); err == nil { + t.Fatalf("expected error on duplicate Register") + } +} + +func TestRegistryRejectsInvalidScope(t *testing.T) { + reg := NewRegistry() + if err := reg.Register(Scope{Name: ""}); err == nil { + t.Fatalf("expected error on empty Name") + } + if err := reg.Register(Scope{Name: "x", MaxTTL: time.Minute}); err == nil { + t.Fatalf("expected error on missing routes") + } + if err := reg.Register(Scope{ + Name: "x", + AllowedRoutes: []RoutePattern{{Method: "GET", Path: "/static"}}, + ResourceIDInPath: "{id}", + MaxTTL: time.Minute, + }); err == nil { + t.Fatalf("expected error when ResourceIDInPath is declared but routes don't contain the placeholder") + } +} diff --git a/framework/temptoken/sweeper.go b/framework/temptoken/sweeper.go new file mode 100644 index 0000000000..57f48896fe --- /dev/null +++ b/framework/temptoken/sweeper.go @@ -0,0 +1,100 @@ +package temptoken + +import ( + "context" + "sync" + "time" + + "github.com/maximhq/bifrost/core/schemas" +) + +// SweepWorker periodically deletes temp_tokens rows whose expires_at is in the +// past. It's the centralized expiry janitor for the temp-token table, mirroring +// the pattern used by PerUserOAuthSweepWorker for oauth_user_sessions. +// +// Tokens are also deleted eagerly by lifecycle owners (see Service.DeleteByResourceID, +// called from OAuth flow terminal transitions). The sweeper exists to catch rows +// whose owning resource timed out before any terminal transition fired, or whose +// scope no longer participates in lifecycle-driven cleanup at all. +type SweepWorker struct { + service *Service + sweepInterval time.Duration + stopCh chan struct{} + stopOnce sync.Once + logger schemas.Logger +} + +// NewSweepWorker constructs a worker bound to the given service. Returns nil +// when service is nil so callers can wire it unconditionally and check the +// result before starting. +func NewSweepWorker(service *Service, logger schemas.Logger) *SweepWorker { + if service == nil { + if logger != nil { + logger.Warn("temp-token sweep worker not started: service is nil") + } + return nil + } + return &SweepWorker{ + service: service, + sweepInterval: 5 * time.Minute, + stopCh: make(chan struct{}), + logger: logger, + } +} + +// Start begins the sweep loop in a background goroutine. +func (w *SweepWorker) Start(ctx context.Context) { + go w.run(ctx) + if w.logger != nil { + w.logger.Info("temp-token sweep worker started (interval=%s)", w.sweepInterval) + } +} + +// Stop gracefully stops the sweep worker. sync.Once guards against double-close +// panics from redundant shutdown paths. +func (w *SweepWorker) Stop() { + w.stopOnce.Do(func() { + close(w.stopCh) + if w.logger != nil { + w.logger.Info("temp-token sweep worker stopped") + } + }) +} + +func (w *SweepWorker) run(ctx context.Context) { + ticker := time.NewTicker(w.sweepInterval) + defer ticker.Stop() + + // Run once on start so a deploy doesn't have to wait a full interval to + // reap rows that expired while the process was down. + w.sweepExpired(ctx) + + for { + select { + case <-ticker.C: + w.sweepExpired(ctx) + case <-w.stopCh: + return + case <-ctx.Done(): + return + } + } +} + +func (w *SweepWorker) sweepExpired(ctx context.Context) { + n, err := w.service.DeleteExpired(ctx, time.Now()) + if err != nil { + if w.logger != nil { + w.logger.Error("temp-token sweep failed: %v", err) + } + return + } + if n > 0 && w.logger != nil { + w.logger.Debug("temp-token sweep removed %d expired rows", n) + } +} + +// SetSweepInterval updates the sweep cadence (for testing). +func (w *SweepWorker) SetSweepInterval(d time.Duration) { + w.sweepInterval = d +} diff --git a/transports/bifrost-http/handlers/middlewares.go b/transports/bifrost-http/handlers/middlewares.go index 58c2b86777..46b4d22540 100644 --- a/transports/bifrost-http/handlers/middlewares.go +++ b/transports/bifrost-http/handlers/middlewares.go @@ -17,6 +17,7 @@ import ( "github.com/maximhq/bifrost/core/schemas" "github.com/maximhq/bifrost/framework/configstore" "github.com/maximhq/bifrost/framework/encrypt" + "github.com/maximhq/bifrost/framework/temptoken" "github.com/maximhq/bifrost/framework/tracing" "github.com/maximhq/bifrost/transports/bifrost-http/integrations" "github.com/maximhq/bifrost/transports/bifrost-http/lib" @@ -698,10 +699,13 @@ type AuthMiddleware struct { whitelistedRoutes atomic.Pointer[[]string] authConfig atomic.Pointer[configstore.AuthConfig] wsTicketStore *WSTicketStore + tempTokensService *temptoken.Service // optional; when nil, temp-token fallback is disabled } -// InitAuthMiddleware initializes the auth middleware. -func InitAuthMiddleware(store configstore.ConfigStore, wsTicketStore *WSTicketStore) (*AuthMiddleware, error) { +// InitAuthMiddleware initializes the auth middleware. The tempTokens service +// is optional — when nil, the temp-token fallback path is disabled and the +// middleware behaves exactly as before. +func InitAuthMiddleware(store configstore.ConfigStore, wsTicketStore *WSTicketStore, tempTokensService *temptoken.Service) (*AuthMiddleware, error) { if store == nil { return nil, fmt.Errorf("store is not present") } @@ -710,9 +714,10 @@ func InitAuthMiddleware(store configstore.ConfigStore, wsTicketStore *WSTicketSt return nil, fmt.Errorf("failed to get auth config from store: %v", err) } am := &AuthMiddleware{ - store: store, - authConfig: atomic.Pointer[configstore.AuthConfig]{}, - wsTicketStore: wsTicketStore, + store: store, + authConfig: atomic.Pointer[configstore.AuthConfig]{}, + wsTicketStore: wsTicketStore, + tempTokensService: tempTokensService, } am.authConfig.Store(authConfig) @@ -738,6 +743,34 @@ func (m *AuthMiddleware) UpdateWhitelistedRoutes(routes []string) { m.whitelistedRoutes.Store(&routes) } +// tryTempTokenOrUnauthorized is the last-resort auth path: a request that +// failed every conventional credential check (no Authorization header, no +// valid cookie) is given one more chance to present an X-Bifrost-Temp-Token +// header that authorizes the specific (method, path) being requested. On +// success the validated scope and resource_id are attached to ctx for +// handler-side defense-in-depth checks, and the next handler runs. On +// failure (no header, expired, route-mismatch, etc.) a 401 is written. +// +// Temp-token validation is intentionally *not* attempted when an +// Authorization header or session cookie is present — those paths have +// their own success/failure semantics and silently rescuing a bad password +// with a temp token would be surprising. +func (m *AuthMiddleware) tryTempTokenOrUnauthorized(ctx *fasthttp.RequestCtx, next fasthttp.RequestHandler) { + if m.tempTokensService != nil { + token := string(ctx.Request.Header.Peek("X-Bifrost-Temp-Token")) + if token != "" { + validated, err := m.tempTokensService.Validate(ctx, token, string(ctx.Method()), string(ctx.Path())) + if err == nil && validated != nil { + ctx.SetUserValue(schemas.BifrostContextKeyTempTokenScope, validated.Scope) + ctx.SetUserValue(schemas.BifrostContextKeyTempTokenResourceID, validated.ResourceID) + next(ctx) + return + } + } + } + SendError(ctx, fasthttp.StatusUnauthorized, "Unauthorized") +} + // InferenceMiddleware is for inference requests (including MCP routes) if authConfig is set, it will skip authentication if disableAuthOnInference is true. func (m *AuthMiddleware) InferenceMiddleware() schemas.BifrostHTTPMiddleware { return m.middleware(func(authConfig *configstore.AuthConfig, url string) bool { @@ -770,8 +803,13 @@ func (m *AuthMiddleware) APIMiddleware() schemas.BifrostHTTPMiddleware { "/api/version", } whitelistedPrefixes := []string{ - "/api/oauth/callback", - "/api/oauth", + // "/api/oauth/callback" is also in systemWhitelistedRoutes above as an + // exact match — that's the only OAuth route that must be public (it's + // hit by the browser after the upstream provider redirects back, with + // no cookie context). DO NOT add a broad "/api/oauth" prefix here: + // it would whitelist /api/oauth/per-user/* (auth-via-temp-token) and + // /api/oauth/config/* (admin-only) and bypass the temp-token fallback + // in tryTempTokenOrUnauthorized. "/api/dev", } return m.middleware(func(authConfig *configstore.AuthConfig, url string) bool { @@ -882,7 +920,10 @@ func (m *AuthMiddleware) middleware(shouldSkip func(*configstore.AuthConfig, str next(ctx) return } - SendError(ctx, fasthttp.StatusUnauthorized, "Unauthorized") + // Last-resort: a scoped temp token (e.g. for the MCP per-user + // OAuth auth page accessed by a non-admin browser) can rescue + // this request when it targets a route the token authorizes. + m.tryTempTokenOrUnauthorized(ctx, next) return } // Split the authorization header into the scheme and the token diff --git a/transports/bifrost-http/handlers/oauth2.go b/transports/bifrost-http/handlers/oauth2.go index 40c03e2b29..714dd58acf 100644 --- a/transports/bifrost-http/handlers/oauth2.go +++ b/transports/bifrost-http/handlers/oauth2.go @@ -11,6 +11,7 @@ import ( "github.com/fasthttp/router" bifrost "github.com/maximhq/bifrost/core" "github.com/maximhq/bifrost/core/schemas" + "github.com/maximhq/bifrost/framework/configstore" "github.com/maximhq/bifrost/framework/oauth2" "github.com/maximhq/bifrost/transports/bifrost-http/lib" "github.com/valyala/fasthttp" @@ -79,11 +80,12 @@ func (h *OAuthHandler) handleOAuthCallback(ctx *fasthttp.RequestCtx) { // The OAuth state is the CSRF token — never log it raw; it could // be replayed by anyone with log access while the flow is alive. logger.Error("[oauth] per-user callback completion failed: err=%v", perUserErr) - ctx.Redirect("/workspace/mcp-sessions?error="+url.QueryEscape("OAuth authentication failed. Please try again."), fasthttp.StatusFound) + const userMsg = "OAuth authentication failed. Please try again." + ctx.Redirect(perUserCallbackRedirect(ctx, h.store.ConfigStore, userMsg, false), fasthttp.StatusFound) return } if perUserErr == nil { - ctx.Redirect("/workspace/mcp-sessions?completed=1", fasthttp.StatusFound) + ctx.Redirect(perUserCallbackRedirect(ctx, h.store.ConfigStore, "", true), fasthttp.StatusFound) return } @@ -134,7 +136,29 @@ func (h *OAuthHandler) handleCallbackError(ctx *fasthttp.RequestCtx, state, erro ctx.Redirect("/workspace/mcp-registry/oauth-callback?status=failed&error="+url.QueryEscape(userMsg), fasthttp.StatusFound) return } - ctx.Redirect("/workspace/mcp-sessions?error="+url.QueryEscape(userMsg), fasthttp.StatusFound) + ctx.Redirect(perUserCallbackRedirect(ctx, h.store.ConfigStore, userMsg, false), fasthttp.StatusFound) +} + +// perUserCallbackRedirect picks the post-callback destination for a per-user +// flow based on whether the visitor has a valid dashboard session. Admins land +// back on the sessions list (full chrome) with either ?completed=1 (success) +// or ?error=... (failure). Anonymous temp-token visitors land on the public +// MinimalShell pages (/auth-success or /auth-failed) which don't require a +// cookie. This mirrors the model used by the temp-token-aware UI: keep admins +// in the dashboard, route end users to chrome-less landings. +func perUserCallbackRedirect(ctx *fasthttp.RequestCtx, store configstore.ConfigStore, userMsg string, success bool) string { + cookieToken := string(ctx.Request.Header.Cookie("token")) + authenticated := cookieToken != "" && validateSession(ctx, store, cookieToken) + if success { + if authenticated { + return "/workspace/mcp-sessions?completed=1" + } + return "/workspace/mcp-sessions/auth-success" + } + if authenticated { + return "/workspace/mcp-sessions?error=" + url.QueryEscape(userMsg) + } + return "/workspace/mcp-sessions/auth-failed?error=" + url.QueryEscape(userMsg) } // getOAuthConfigStatus returns the current status of an OAuth config diff --git a/transports/bifrost-http/handlers/temp_token_scopes.go b/transports/bifrost-http/handlers/temp_token_scopes.go new file mode 100644 index 0000000000..35b387810b --- /dev/null +++ b/transports/bifrost-http/handlers/temp_token_scopes.go @@ -0,0 +1,47 @@ +package handlers + +import ( + "fmt" + "time" + + "github.com/maximhq/bifrost/framework/temptoken" +) + +// mcpAuthScope declares the routes the mcp_auth scope grants access to. The +// flow ID is substituted into {id} at validation time, binding each token to +// exactly one flow. +// +// The canonical scope name lives in framework/temptoken so issuers (e.g. the +// OAuth provider that mints these tokens) and registrants reference a single +// source of truth. +// +// The page also calls /api/version and /api/session/is-auth-enabled — those +// are unconditionally whitelisted in APIMiddleware so they do not need to +// appear here. +// The page makes flowDetail then flowStart, so the token must remain valid for +// multiple requests within its TTL. Invalidation isn't single-use — it happens +// at OAuth completion when CompleteUserOAuthFlow deletes the token by +// resource_id. +var mcpAuthScope = temptoken.Scope{ + Name: temptoken.MCPAuthScopeName, + AllowedRoutes: []temptoken.RoutePattern{ + {Method: "GET", Path: "/api/oauth/per-user/flows/{id}"}, + {Method: "GET", Path: "/api/oauth/per-user/flows/{id}/start"}, + }, + 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 +// already been registered. +func RegisterTempTokenScopes(svc *temptoken.Service) error { + if svc == nil { + return fmt.Errorf("temp_token_scopes: service is nil") + } + if err := svc.Registry().Register(mcpAuthScope); err != nil { + return fmt.Errorf("temp_token_scopes: register mcp_auth: %w", err) + } + return nil +} diff --git a/transports/bifrost-http/lib/config_test.go b/transports/bifrost-http/lib/config_test.go index bc08c574ec..c52951e049 100644 --- a/transports/bifrost-http/lib/config_test.go +++ b/transports/bifrost-http/lib/config_test.go @@ -986,6 +986,23 @@ func (m *MockConfigStore) DeleteSession(ctx context.Context, token string) error return nil } +// Temp token +func (m *MockConfigStore) CreateTempToken(ctx context.Context, token *tables.TempToken, tx ...*gorm.DB) error { + return nil +} + +func (m *MockConfigStore) GetTempTokenByHash(ctx context.Context, tokenHash string) (*tables.TempToken, error) { + return nil, nil +} + +func (m *MockConfigStore) DeleteTempTokensByResourceID(ctx context.Context, scope, resourceID string, tx ...*gorm.DB) (int64, error) { + return 0, nil +} + +func (m *MockConfigStore) DeleteExpiredTempTokens(ctx context.Context, before time.Time) (int64, error) { + return 0, nil +} + // Model pricing func (m *MockConfigStore) GetModelPrices(ctx context.Context) ([]tables.TableModelPricing, error) { return nil, nil diff --git a/transports/bifrost-http/server/server.go b/transports/bifrost-http/server/server.go index bfefeeb47e..fe81ea9869 100644 --- a/transports/bifrost-http/server/server.go +++ b/transports/bifrost-http/server/server.go @@ -23,6 +23,7 @@ import ( "github.com/maximhq/bifrost/framework/encrypt" "github.com/maximhq/bifrost/framework/logstore" dynamicPlugins "github.com/maximhq/bifrost/framework/plugins" + "github.com/maximhq/bifrost/framework/temptoken" "github.com/maximhq/bifrost/framework/tracing" "github.com/maximhq/bifrost/plugins/governance" "github.com/maximhq/bifrost/plugins/logging" @@ -136,7 +137,9 @@ type BifrostHTTPServer struct { AuthMiddleware *handlers.AuthMiddleware TracingMiddleware *handlers.TracingMiddleware - WSTicketStore *handlers.WSTicketStore + WSTicketStore *handlers.WSTicketStore + TempTokens *temptoken.Service + TempTokenSweepWorker *temptoken.SweepWorker wsPool *bfws.Pool } @@ -1458,10 +1461,36 @@ func (s *BifrostHTTPServer) Bootstrap(ctx context.Context) error { // so tickets are verifiable across nodes; otherwise fall back to in-memory. // NewSignedWSTicketStore handles empty key by degrading to in-memory mode. s.WSTicketStore = handlers.NewSignedWSTicketStore(encrypt.Key()) - s.AuthMiddleware, err = handlers.InitAuthMiddleware(s.Config.ConfigStore, s.WSTicketStore) + // Initialize the temp-token service and register all scopes owned by the + // handlers package. The service is the seam every "scoped, anonymous, + // browser-only" workflow plugs into (currently just the MCP per-user OAuth + // auth page + s.TempTokens = temptoken.NewService(s.Config.ConfigStore, temptoken.NewRegistry()) + if regErr := handlers.RegisterTempTokenScopes(s.TempTokens); regErr != nil { + s.WSTicketStore.Stop() + s.WSTicketStore = nil + return fmt.Errorf("failed to register temp token scopes: %v", regErr) + } + // Centralized janitor that reaps expired temp_tokens rows. Independent + // of the per-user OAuth sweep so any future scope (not just mcp_auth) + // benefits from the same cleanup loop without piggybacking on OAuth. + s.TempTokenSweepWorker = temptoken.NewSweepWorker(s.TempTokens, logger) + if s.TempTokenSweepWorker != nil { + s.TempTokenSweepWorker.Start(s.Ctx) + } + // Hand the service to the OAuth provider so InitiateUserOAuthFlow mints + // a mcp_auth token and embeds it as a URL fragment on the auth-page link. + if s.Config.OAuthProvider != nil { + s.Config.OAuthProvider.SetTempTokenService(s.TempTokens) + } + s.AuthMiddleware, err = handlers.InitAuthMiddleware(s.Config.ConfigStore, s.WSTicketStore, s.TempTokens) if err != nil { s.WSTicketStore.Stop() s.WSTicketStore = nil + if s.TempTokenSweepWorker != nil { + s.TempTokenSweepWorker.Stop() + s.TempTokenSweepWorker = nil + } return fmt.Errorf("failed to initialize auth middleware: %v", err) } if ctx.Value(schemas.BifrostContextKeyIsEnterprise) == nil { @@ -1480,6 +1509,10 @@ func (s *BifrostHTTPServer) Bootstrap(ctx context.Context) error { s.WSTicketStore.Stop() s.WSTicketStore = nil } + if s.TempTokenSweepWorker != nil { + s.TempTokenSweepWorker.Stop() + s.TempTokenSweepWorker = nil + } return fmt.Errorf("failed to initialize routes: %v", err) } // Registering inference routes @@ -1510,6 +1543,10 @@ func (s *BifrostHTTPServer) Bootstrap(ctx context.Context) error { s.WSTicketStore.Stop() s.WSTicketStore = nil } + if s.TempTokenSweepWorker != nil { + s.TempTokenSweepWorker.Stop() + s.TempTokenSweepWorker = nil + } return fmt.Errorf("failed to initialize inference routes: %v", err) } // Register UI handler @@ -1589,6 +1626,10 @@ func (s *BifrostHTTPServer) Start() error { logger.Info("stopping ws ticket store...") s.WSTicketStore.Stop() } + if s.TempTokenSweepWorker != nil { + logger.Info("stopping temp-token sweep worker...") + s.TempTokenSweepWorker.Stop() + } if s.devPprofHandler != nil { logger.Info("stopping dev pprof handler...") s.devPprofHandler.Cleanup() diff --git a/ui/app/clientLayout.tsx b/ui/app/clientLayout.tsx index 49d095f501..cb312f7868 100644 --- a/ui/app/clientLayout.tsx +++ b/ui/app/clientLayout.tsx @@ -1,6 +1,5 @@ import FullPageLoader from "@/components/fullPageLoader"; import NotAvailableBanner from "@/components/notAvailableBanner"; -import OnboardingWidget from "@/components/onboardingWidget"; import ProgressProvider from "@/components/progressBar"; import Sidebar from "@/components/sidebar"; import { ThemeProvider } from "@/components/themeProvider"; @@ -8,81 +7,183 @@ import TrialExpiryBanner from "@/components/trialExpiryBanner"; import { SidebarProvider } from "@/components/ui/sidebar"; import { useStoreSync } from "@/hooks/useStoreSync"; import { WebSocketProvider } from "@/hooks/useWebSocket"; -import { getErrorMessage, ReduxProvider, useGetCoreConfigQuery } from "@/lib/store"; +import { + getErrorMessage, + ReduxProvider, + useGetCoreConfigQuery, + useIsAuthEnabledQuery, +} from "@/lib/store"; import { BifrostConfig } from "@/lib/types/config"; import { RbacProvider } from "@enterprise/lib/contexts/rbacContext"; -import { useLocation } from "@tanstack/react-router"; +import { useLocation, useMatches } from "@tanstack/react-router"; import { NuqsAdapter } from "nuqs/adapters/tanstack-router"; -import { lazy, Suspense, useEffect } from "react"; +import { lazy, Suspense, useEffect, useState } from "react"; import { CookiesProvider } from "react-cookie"; import { toast, Toaster } from "sonner"; // Lazy import — only loaded in development, completely excluded from prod bundle -const DevProfilerLazy = lazy(() => import("@/components/devProfiler").then((mod) => ({ default: mod.DevProfiler }))); +const DevProfilerLazy = lazy(() => + import("@/components/devProfiler").then((mod) => ({ + default: mod.DevProfiler, + })), +); const DevProfiler = () => ( - - - + + + ); function StoreSyncInitializer() { - useStoreSync(); - return null; + useStoreSync(); + return null; } function AppContent({ children }: { children: React.ReactNode }) { - const { data: bifrostConfig, error, isLoading } = useGetCoreConfigQuery({}); + // Routes can declare `staticData: { tempTokenScoped: true }` to advertise that + // they're reachable via a server-emitted, temp-token-bearing URL by visitors + // without a dashboard session. The actual layout choice is made per-visitor: + // an authenticated admin still sees the full dashboard chrome, while an + // anonymous visitor arriving with `#t=` gets a stripped MinimalShell. + // The auth-via-temp-token half lives in . + const matches = useMatches(); + const tempTokenScoped = matches.some( + (m) => + (m.staticData as { tempTokenScoped?: boolean } | undefined) + ?.tempTokenScoped === true, + ); + // publicShell: route declares it's a static, auth-free page that should + // always render MinimalShell — no chrome, no auth probe, no API calls. + // Used by the post-OAuth "authentication successful" landing, which has + // neither a fragment nor a cookie to drive the tempTokenScoped per-visitor + // logic. + const publicShell = matches.some( + (m) => + (m.staticData as { publicShell?: boolean } | undefined)?.publicShell === + true, + ); + + // Probe dashboard auth state on opted-in routes. is-auth-enabled is whitelisted + // (no 401 risk) and returns whether the current cookie is a valid session. + const { data: authState, isLoading: authLoading } = useIsAuthEnabledQuery( + undefined, + { skip: !tempTokenScoped }, + ); + + // Snapshot fragment presence at mount: TempTokenScope strips the fragment + // shortly after, so re-reading window.location.hash would flip false on + // re-render. Only fragment-bearing arrivals are MinimalShell candidates. + const [hadFragmentTempToken] = useState(() => { + if (typeof window === "undefined") return false; + const fragment = window.location.hash; + if (!fragment || fragment.length < 2) return false; + return !!new URLSearchParams(fragment.slice(1)).get("t"); + }); + + const useMinimalShell = + tempTokenScoped && + !!authState?.is_auth_enabled && + !authState?.has_valid_token && + hadFragmentTempToken; + + const { + data: bifrostConfig, + error, + isLoading, + } = useGetCoreConfigQuery( + {}, + { + skip: publicShell || useMinimalShell || (tempTokenScoped && authLoading), + }, + ); - useEffect(() => { - if (error) { - toast.error(getErrorMessage(error)); - } - }, [error]); + useEffect(() => { + if (error) { + toast.error(getErrorMessage(error)); + } + }, [error]); - return ( - - - - - -
- -
- {isLoading ? : {children}} -
- {bifrostConfig?.is_db_connected && } -
-
-
-
- ); + if (publicShell) { + return {children}; + } + if (tempTokenScoped && authLoading) { + return ; + } + if (useMinimalShell) { + return {children}; + } + + return ( + + + + + +
+ +
+ {isLoading ? ( + + ) : ( + {children} + )} +
+
+
+
+
+ ); +} + +// MinimalShell renders a centered container without sidebar, websocket, +// store-sync, or any dashboard-config fetches. Used for routes that opt +// in via `staticData.tempTokenScoped` — typically public, scoped pages +// like the MCP per-user OAuth auth page. +function MinimalShell({ children }: { children: React.ReactNode }) { + return ( +
+
+ {children} +
+
+ ); } -function FullPage({ config, children }: { config: BifrostConfig | undefined; children: React.ReactNode }) { - const pathname = useLocation({ select: (l) => l.pathname }); - if (config && config.is_db_connected) { - return children; - } - if (config && config.is_logs_connected && pathname.startsWith("/workspace/logs")) { - return children; - } - return ; +function FullPage({ + config, + children, +}: { + config: BifrostConfig | undefined; + children: React.ReactNode; +}) { + const pathname = useLocation({ select: (l) => l.pathname }); + if (config && config.is_db_connected) { + return children; + } + if ( + config && + config.is_logs_connected && + pathname.startsWith("/workspace/logs") + ) { + return children; + } + return ; } export function ClientLayout({ children }: { children: React.ReactNode }) { - return ( - - - - - - - {children} - {process.env.NODE_ENV === "development" && !process.env.BIFROST_DISABLE_PROFILER && } - - - - - - ); -} \ No newline at end of file + return ( + + + + + + + {children} + {process.env.NODE_ENV === "development" && + !process.env.BIFROST_DISABLE_PROFILER && } + + + + + + ); +} diff --git a/ui/app/workspace/mcp-sessions/auth-failed/layout.tsx b/ui/app/workspace/mcp-sessions/auth-failed/layout.tsx new file mode 100644 index 0000000000..7eb1149425 --- /dev/null +++ b/ui/app/workspace/mcp-sessions/auth-failed/layout.tsx @@ -0,0 +1,11 @@ +import { createFileRoute } from "@tanstack/react-router"; +import MCPSessionsAuthFailedPage from "./page"; + +// Public landing for per-user OAuth callback failures. Symmetric to auth-success: +// the anonymous (temp-token) branch of the callback handler redirects here when +// upstream denied the request or the token exchange failed. publicShell makes +// it MinimalShell-only with no API calls — works without a dashboard cookie. +export const Route = createFileRoute("/workspace/mcp-sessions/auth-failed")({ + staticData: { publicShell: true }, + component: MCPSessionsAuthFailedPage, +}); diff --git a/ui/app/workspace/mcp-sessions/auth-failed/page.tsx b/ui/app/workspace/mcp-sessions/auth-failed/page.tsx new file mode 100644 index 0000000000..1b45b64f6a --- /dev/null +++ b/ui/app/workspace/mcp-sessions/auth-failed/page.tsx @@ -0,0 +1,25 @@ +import { AlertCircle } from "lucide-react"; +import { useQueryState } from "nuqs"; + +export default function MCPSessionsAuthFailedPage() { + const [error] = useQueryState("error"); + return ( +
+
+
+ +
+

+ Authentication failed +

+

+ {error ?? "We couldn't complete the authentication flow."} +

+

+ You can close this tab and retry the original request from your MCP + client to generate a fresh authentication link. +

+
+
+ ); +} diff --git a/ui/app/workspace/mcp-sessions/auth-success/layout.tsx b/ui/app/workspace/mcp-sessions/auth-success/layout.tsx new file mode 100644 index 0000000000..a21fe91d2b --- /dev/null +++ b/ui/app/workspace/mcp-sessions/auth-success/layout.tsx @@ -0,0 +1,13 @@ +import { createFileRoute } from "@tanstack/react-router"; +import MCPSessionsAuthSuccessPage from "./page"; + +// Landing page shown after a per-user OAuth callback completes successfully. +// `publicShell` tells ClientLayout to render the MinimalShell unconditionally +// — the post-OAuth redirect arrives with no fragment and no dashboard cookie +// (temp-token visitors authenticated externally, not against Bifrost), so the +// normal tempTokenScoped logic wouldn't fire. This flag short-circuits all of +// that: no chrome, no auth probe, no API calls, just a static "done" view. +export const Route = createFileRoute("/workspace/mcp-sessions/auth-success")({ + staticData: { publicShell: true }, + component: MCPSessionsAuthSuccessPage, +}); diff --git a/ui/app/workspace/mcp-sessions/auth-success/page.tsx b/ui/app/workspace/mcp-sessions/auth-success/page.tsx new file mode 100644 index 0000000000..24381e3d5d --- /dev/null +++ b/ui/app/workspace/mcp-sessions/auth-success/page.tsx @@ -0,0 +1,21 @@ +import { CheckCircle2 } from "lucide-react"; + +export default function MCPSessionsAuthSuccessPage() { + return ( +
+
+
+ +
+

+ Authentication successful +

+

+ Your credential has been stored. You can close this tab and return to + your MCP client — future requests will use this credential + automatically. +

+
+
+ ); +} diff --git a/ui/app/workspace/mcp-sessions/auth/layout.tsx b/ui/app/workspace/mcp-sessions/auth/layout.tsx index e22a07d441..d05e5c40b0 100644 --- a/ui/app/workspace/mcp-sessions/auth/layout.tsx +++ b/ui/app/workspace/mcp-sessions/auth/layout.tsx @@ -1,14 +1,25 @@ +import TempTokenScope from "@/components/tempTokenScope"; import { createFileRoute } from "@tanstack/react-router"; import MCPSessionsAuthPage from "./page"; +// staticData.tempTokenScoped opts this route out of the dashboard chrome — +// ClientLayout renders a minimal shell and skips the protected +// useGetCoreConfigQuery fetch when this flag is set, so an unauthenticated +// browser can land on this page without bouncing to /login. +// +// TempTokenScope handles the auth half: it reads the `#t=…` fragment the +// server appended to the URL, attaches it as `X-Bifrost-Temp-Token` on +// outbound API calls, and suppresses the global 401-redirect so a stale +// link renders an inline error instead. function RouteComponent() { - // Public-by-policy in OSS: the backend enforces identity match on the flow - // row itself. We route any incoming caller to the page; if their identity - // doesn't match the flow's, the API returns 403 and the page renders an - // appropriate message. - return ; + return ( + + + + ); } export const Route = createFileRoute("/workspace/mcp-sessions/auth")({ - component: RouteComponent, + staticData: { tempTokenScoped: true }, + component: RouteComponent, }); diff --git a/ui/app/workspace/mcp-sessions/auth/page.tsx b/ui/app/workspace/mcp-sessions/auth/page.tsx index 5f1a324f6e..12ef269928 100644 --- a/ui/app/workspace/mcp-sessions/auth/page.tsx +++ b/ui/app/workspace/mcp-sessions/auth/page.tsx @@ -12,270 +12,346 @@ import FullPageLoader from "@/components/fullPageLoader"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { useToast } from "@/hooks/use-toast"; -import { getErrorMessage, useGetMCPFlowDetailQuery, useStartMCPFlowMutation } from "@/lib/store"; +import { + getErrorMessage, + useGetMCPFlowDetailQuery, + useIsAuthEnabledQuery, + useStartMCPFlowMutation, +} from "@/lib/store"; import { MCPFlowDetail } from "@/lib/types/mcpSessions"; import { Link } from "@tanstack/react-router"; -import { ExternalLink, Fingerprint, KeyRound, Loader2, ShieldCheck, UserRound } from "lucide-react"; +import { + ExternalLink, + Fingerprint, + KeyRound, + Loader2, + ShieldCheck, + UserRound, +} from "lucide-react"; import { useQueryState } from "nuqs"; export default function MCPSessionsAuthPage() { - const { toast } = useToast(); - const [flowId] = useQueryState("flow"); - const skip = !flowId; - const { data: flow, isLoading, isError, error, refetch } = useGetMCPFlowDetailQuery(flowId ?? "", { skip }); - const [startFlow, { isLoading: starting }] = useStartMCPFlowMutation(); + const { toast } = useToast(); + const [flowId] = useQueryState("flow"); + const skip = !flowId; + const { + data: flow, + isLoading, + isError, + error, + } = useGetMCPFlowDetailQuery(flowId ?? "", { skip }); + const [startFlow, { isLoading: starting }] = useStartMCPFlowMutation(); - if (!flowId) { - return ( - -

Missing flow identifier

-

- This URL is missing the flow query parameter. Open the link from your - inference response or the sessions tab. -

-
- -
-
- ); - } + if (!flowId) { + return ( + +

Missing flow identifier

+

+ This URL is missing the{" "} + flow query + parameter. Open the link from your inference response or the sessions + tab. +

+
+ +
+
+ ); + } - if (isLoading) { - return ; - } + if (isLoading) { + return ; + } - if (isError || !flow) { - const status = (error as { status?: number } | undefined)?.status; - if (status === 401) { - return ; - } - if (status === 403) { - return ( - -

This authentication flow isn't yours

-

- The pending flow belongs to a different identity. Ask the teammate whose VK or user identity triggered the original request - to complete it, or trigger a new request yourself. -

-
- -
-
- ); - } - if (status === 404) { - return ( - -

This authentication flow has expired or been completed

-

- Pending flows expire after a short window. If you still need to authenticate, trigger the original action again so a fresh - flow is created. -

-
- -
-
- ); - } - return ( - -

Could not load this authentication flow

-

{getErrorMessage(error)}

-
- ); - } + if (isError || !flow) { + const status = (error as { status?: number } | undefined)?.status; + if (status === 401) { + return ; + } + if (status === 403) { + return ( + +

+ This authentication flow isn't yours +

+

+ The pending flow belongs to a different identity. Ask the teammate + whose VK or user identity triggered the original request to complete + it, or trigger a new request yourself. +

+
+ +
+
+ ); + } + if (status === 404) { + return ( + +

+ This authentication flow has expired or been completed +

+

+ Pending flows expire after a short window. If you still need to + authenticate, trigger the original action again so a fresh flow is + created. +

+
+ +
+
+ ); + } + return ( + +

+ Could not load this authentication flow +

+

+ {getErrorMessage(error)} +

+
+ ); + } - // Flow row exists but isn't pending: it's already been completed, failed, - // or expired. Don't show the "Authenticate" button since startFlow would - // reject (BuildUpstreamAuthorizeURL rejects non-pending flows). - if (flow.status !== "pending") { - return ; - } + // Flow row exists but isn't pending: it's already been completed, failed, + // or expired. Don't show the "Authenticate" button since startFlow would + // reject (BuildUpstreamAuthorizeURL rejects non-pending flows). + if (flow.status !== "pending") { + return ; + } - const handleAuthenticate = async () => { - try { - const res = await startFlow(flowId).unwrap(); - window.location.href = res.authorize_url; - } catch (err) { - // The server-side flow could have been completed or expired in - // another tab between the initial query and this click — startFlow - // then rejects. Refetch so the page flips to the completed/expired - // view instead of leaving a stale "Authenticate" CTA visible for a - // retry that can't succeed. - refetch(); - toast({ title: "Failed to start authentication", description: getErrorMessage(err), variant: "destructive" }); - } - }; + const handleAuthenticate = async () => { + try { + const res = await startFlow(flowId).unwrap(); + window.location.href = res.authorize_url; + } catch (err) { + toast({ + title: "Failed to start authentication", + description: getErrorMessage(err), + variant: "destructive", + }); + } + }; - const mcpClientName = flow.mcp_client?.name || flow.mcp_client?.client_id || "MCP server"; - const isReauth = flow.has_active_token === true; + const mcpClientName = + flow.mcp_client?.name || flow.mcp_client?.client_id || "MCP server"; + const isReauth = flow.has_active_token === true; - return ( - -
- -
-

- {isReauth ? "Re-authenticate with" : "Authenticate with"} {mcpClientName} -

-

- {isReauth ? ( - <> - An active credential already exists for the binding below. Completing this flow will replace it with a fresh - credential. You can also close this tab to keep using the existing one. - - ) : ( - <> - You'll be redirected to the provider to sign in and grant access. Bifrost stores the resulting credential against the binding - below so this request and future ones can proceed automatically. - - )} -

+ return ( + +
+ +
+

+ {isReauth ? "Re-authenticate with" : "Authenticate with"}{" "} + {mcpClientName} +

+

+ {isReauth ? ( + <> + An active credential already exists for the binding below. + Completing this flow will replace it with a fresh + credential. You can also close this tab to keep using the existing + one. + + ) : ( + <> + You'll be redirected to the provider to sign in and grant access. + Bifrost stores the resulting credential against the binding below so + this request and future ones can proceed automatically. + + )} +

-
- - } /> - -
+
+ + } /> + +
-
- - -
-
- ); +
+ + +
+
+ ); } function CompletedFlowView({ flow }: { flow: MCPFlowDetail }) { - const mcpClientName = flow.mcp_client?.name || flow.mcp_client?.client_id || "this MCP server"; - // has_active_token wins over the flow's row status: a pending flow with an - // existing active token means OAuth was re-initiated unnecessarily. - const effectivelyAuthorized = flow.status === "authorized" || flow.has_active_token; - const title = effectivelyAuthorized - ? "Already authenticated" - : flow.status === "expired" - ? "This authentication flow has expired" - : "This authentication flow can no longer be completed"; - const body = effectivelyAuthorized - ? `The OAuth credential for ${mcpClientName} is already stored. You can close this tab.` - : "Trigger the original action again so a fresh flow is created."; - return ( - -

{title}

-

{body}

-
- - } /> -
-
- -
-
- ); + const mcpClientName = + flow.mcp_client?.name || flow.mcp_client?.client_id || "this MCP server"; + // has_active_token wins over the flow's row status: a pending flow with an + // existing active token means OAuth was re-initiated unnecessarily. + const effectivelyAuthorized = + flow.status === "authorized" || flow.has_active_token; + const title = effectivelyAuthorized + ? "Already authenticated" + : flow.status === "expired" + ? "This authentication flow has expired" + : "This authentication flow can no longer be completed"; + const body = effectivelyAuthorized + ? `The OAuth credential for ${mcpClientName} is already stored. You can close this tab.` + : "Trigger the original action again so a fresh flow is created."; + return ( + +

{title}

+

{body}

+
+ + } /> +
+
+ +
+
+ ); } -function DetailRow({ label, value, mono = false }: { label: string; value: React.ReactNode; mono?: boolean }) { - return ( -
-
{label}
-
{value}
-
- ); +function DetailRow({ + label, + value, + mono = false, +}: { + label: string; + value: React.ReactNode; + mono?: boolean; +}) { + return ( +
+
+ {label} +
+
+ {value} +
+
+ ); } function BindingValue({ flow }: { flow: MCPFlowDetail }) { - if (flow.flow_mode === "user") { - const userID = flow.user_id; - if (!userID) { - return ( - - - First signed-in user - - ); - } - const displayName = flow.user?.name || flow.user?.email; - return ( - - - {displayName ? {displayName} : {userID}} - - ); - } - if (flow.flow_mode === "vk" && flow.virtual_key) { - return ( - - - {flow.virtual_key.name || flow.virtual_key.id} - - ); - } - if (flow.flow_mode === "session" && flow.session_id) { - return ( - - - {flow.session_id} - - ); - } - return Unknown; + if (flow.flow_mode === "user") { + const userID = flow.user_id; + if (!userID) { + return ( + + + First signed-in user + + ); + } + const displayName = flow.user?.name || flow.user?.email; + return ( + + + {displayName ? ( + {displayName} + ) : ( + {userID} + )} + + ); + } + if (flow.flow_mode === "vk" && flow.virtual_key) { + return ( + + + {flow.virtual_key.name || flow.virtual_key.id} + + ); + } + if (flow.flow_mode === "session" && flow.session_id) { + return ( + + + {flow.session_id} + + ); + } + return Unknown; } function formatExpiry(iso: string): string { - try { - const t = new Date(iso).getTime(); - if (Number.isNaN(t)) return iso; - const diffMs = t - Date.now(); - if (diffMs < 0) return "Expired"; - const mins = Math.floor(diffMs / 60_000); - if (mins < 1) return "in less than a minute"; - if (mins === 1) return "in 1 minute"; - return `in ${mins} minutes`; - } catch { - return iso; - } + try { + const t = new Date(iso).getTime(); + if (Number.isNaN(t)) return iso; + const diffMs = t - Date.now(); + if (diffMs < 0) return "Expired"; + const mins = Math.floor(diffMs / 60_000); + if (mins < 1) return "in less than a minute"; + if (mins === 1) return "in 1 minute"; + return `in ${mins} minutes`; + } catch { + return iso; + } } function CenteredCard({ children }: { children: React.ReactNode }) { - return ( -
-
{children}
-
- ); + return ( +
+
+ {children} +
+
+ ); } -function SessionsTabLink({ variant = "outline" }: { variant?: "outline" | "ghost" }) { - return ( - - ); +function SessionsTabLink({ + variant = "outline", +}: { + variant?: "outline" | "ghost"; +}) { + // Hide the link only when the visitor has no dashboard session — for them, + // /workspace/mcp-sessions would 401 and bounce to /login. Admins (cookie + // present) still see it. ClientLayout already cached this query for the + // route, so this is a free hook call. + const { data: authState } = useIsAuthEnabledQuery(); + if (authState?.is_auth_enabled && !authState.has_valid_token) { + return null; + } + return ( + + ); } -// UnauthenticatedView is the 401 fallback: caller is not logged into the -// dashboard / has no identity in context. Frontend redirects to the dashboard -// login route with a return param so the user lands back here after signing in. -// -// The dashboard-auth-on-but-non-admin-needs-temp-token branch ships in OSS-3 -// alongside the temp-token mint endpoint; this OSS-2 cut just sends the user -// to /login and lets the existing login flow handle it. -function UnauthenticatedView({ flowId }: { flowId: string }) { - const goto = `/workspace/mcp-sessions/auth?flow=${encodeURIComponent(flowId)}`; - const loginURL = `/login?goto=${encodeURIComponent(goto)}`; - return ( - -

Sign in to complete authentication

-

- Bifrost needs to know who you are before linking this OAuth credential. You'll be sent back here after signing in. -

-
- -
-
- ); +// InvalidLinkView renders when the per-user-flow API returns 401, which now +// means the caller arrived without either a valid dashboard session or a +// valid mcp_auth temp token. Most often this is an expired or hand-edited +// link — the temp token embedded in the URL fragment has aged out or the +// fragment was dropped along the way. Trigger the original action again to +// get a fresh URL. +function InvalidLinkView() { + return ( + +

+ This authentication link is no longer valid +

+

+ The link may have expired, been used already, invalid, or had its + short-lived token stripped. Trigger the original action again so a fresh + authentication link is created. +

+
+ ); } diff --git a/ui/components/tempTokenScope.tsx b/ui/components/tempTokenScope.tsx new file mode 100644 index 0000000000..051abeaf1c --- /dev/null +++ b/ui/components/tempTokenScope.tsx @@ -0,0 +1,95 @@ +// TempTokenScope wraps a page that authenticates via a short-lived temp token +// embedded in the URL fragment (`#t=`). It does three things: +// +// 1. On mount, reads the token from `window.location.hash` and installs it +// in the baseApi module state so all RTK Query calls attach a +// `X-Bifrost-Temp-Token` header. +// 2. Strips the fragment from the URL via `history.replaceState` so the +// token does not leak into Referer headers if the user later navigates +// away. +// 3. Sets the suppression flag so a 401 from a wrapped API call does NOT +// trigger the global redirect-to-/login in baseQueryWithErrorHandling. +// The wrapped page renders its own invalid/expired-link UI. +// +// The wrapper is scope-agnostic — the `name` prop only identifies the scope in +// log lines (and is wired into future error UI). Routes that opt in still need +// to declare `staticData: { tempTokenScoped: true }` on their `createFileRoute` +// so ClientLayout skips the protected dashboard fetches; that piece is +// orthogonal to this wrapper. + +import { + setActiveTempToken, + setSuppressGlobal401, +} from "@/lib/store/apis/tempToken"; +import { useEffect, useState } from "react"; + +interface TempTokenScopeProps { + name: string; + children: React.ReactNode; +} + +export default function TempTokenScope({ + name: _name, + children, +}: TempTokenScopeProps) { + // Install the module state synchronously during render — NOT in useEffect. + // React fires child effects before parent effects, so a child API call + // triggered from its own useEffect would race ahead of a parent useEffect + // and go out without the X-Bifrost-Temp-Token header (and without the + // global-401 suppression flag set, so the 401 would force a /login + // redirect). useState's initializer runs once during the parent's render, + // strictly before any descendant render or effect — so by the time the + // child's query effect fires, the module state is already in place. + // + // Both setters are idempotent, which makes this safe under React strict + // mode's double-invocation. + useState(() => { + if (typeof window === "undefined") { + return null; + } + const token = parseTokenFromFragment(window.location.hash); + if (token) { + // Token present: install both. The page authenticates via temp + // token and handles its own 401 display. + setActiveTempToken(token); + setSuppressGlobal401(true); + } + // No token: leave both unset so a 401 (e.g. a dashboard user whose + // session expired mid-page) still triggers the normal /login redirect. + // This preserves the existing reauth-from-sessions-tab flow. + return token; + }); + + useEffect(() => { + // Strip the fragment so the token doesn't end up in Referer headers on + // outbound navigation (e.g. the redirect to the upstream OAuth provider + // when the user clicks Authenticate). Pure URL cosmetics — safe to defer + // to an effect, doesn't affect auth correctness. + if (typeof window !== "undefined" && window.location.hash) { + window.history.replaceState( + null, + "", + window.location.pathname + window.location.search, + ); + } + return () => { + setActiveTempToken(null); + setSuppressGlobal401(false); + }; + }, []); + + return <>{children}; +} + +// parseTokenFromFragment extracts the `t` parameter from a URL fragment like +// `#t=abc123` or `#foo=bar&t=abc123`. Returns null if absent. +function parseTokenFromFragment(fragment: string): string | null { + if (!fragment || fragment.length < 2) { + return null; + } + // URLSearchParams handles `?` and `&` separators; the fragment shape used + // by the server (`#t=...`) parses cleanly after stripping the leading `#`. + const params = new URLSearchParams(fragment.slice(1)); + const token = params.get("t"); + return token && token.length > 0 ? token : null; +} diff --git a/ui/lib/store/apis/baseApi.ts b/ui/lib/store/apis/baseApi.ts index f4a2da2ad8..c040b517ab 100644 --- a/ui/lib/store/apis/baseApi.ts +++ b/ui/lib/store/apis/baseApi.ts @@ -4,6 +4,7 @@ import { getApiBaseUrl } from "@/lib/utils/port"; import { createBaseQueryWithRefresh } from "@enterprise/lib/store/utils/baseQueryWithRefresh"; import { clearOAuthStorage } from "@enterprise/lib/store/utils/tokenManager"; import { createApi, fetchBaseQuery } from "@reduxjs/toolkit/query/react"; +import { getActiveTempToken, getSuppressGlobal401 } from "./tempToken"; // Auth tokens are now stored in HTTP-only cookies (set by server) // No client-side token needed — handled by credentials: "include" @@ -39,23 +40,27 @@ export const clearAuthStorage = () => { // Define the base query with authentication headers const baseQuery = fetchBaseQuery({ - baseUrl: getApiBaseUrl(), - credentials: "include", - prepareHeaders: async (headers) => { - // Default JSON only when an endpoint hasn't already set a content type. - // Forcing application/json unconditionally would clobber multipart - // requests (e.g. FormData uploads) — the browser-generated boundary - // would be lost and the upload would fail. + baseUrl: getApiBaseUrl(), + credentials: "include", + prepareHeaders: async (headers) => { if (!headers.has("Content-Type")) { headers.set("Content-Type", "application/json"); } - // Automatically include token from localStorage in Authorization header - const token = await getTokenFromStorage(); - if (token) { - headers.set("Authorization", `Bearer ${token}`); - } - return headers; - }, + // Automatically include token from localStorage in Authorization header + const token = await getTokenFromStorage(); + if (token) { + headers.set("Authorization", `Bearer ${token}`); + } + // Attach a temp token when a TempTokenScope wrapper is mounted. The + // dashboard cookie (if present) still takes precedence on the server + // side; the temp token is the fallback that rescues unauthenticated + // browsers visiting a scoped page. + const tempToken = getActiveTempToken(); + if (tempToken) { + headers.set("X-Bifrost-Temp-Token", tempToken); + } + return headers; + }, }); // Wrap base query with enterprise refresh logic (or passthrough for non-enterprise) @@ -74,17 +79,20 @@ const baseQueryWithErrorHandling: typeof baseQueryWithRefresh = async ( if (result.error) { const error = result.error as any; - // Handle 401 for non-enterprise (no refresh available) - if (error?.status === 401 && !IS_ENTERPRISE) { - clearAuthStorage(); - if ( - typeof window !== "undefined" && - !window.location.pathname.includes("/login") - ) { - window.location.href = "/login"; - } - return result; - } + // Handle 401 for non-enterprise (no refresh available) + if (error?.status === 401 && !IS_ENTERPRISE) { + // When a TempTokenScope wrapper is active, the wrapped page handles + // its own 401 display (an "invalid/expired link" view). Skip the + // global redirect so the user stays on the page they opened. + if (getSuppressGlobal401()) { + return result; + } + clearAuthStorage(); + if (typeof window !== "undefined" && !window.location.pathname.includes("/login")) { + window.location.href = "/login"; + } + return result; + } // Handle specific error types if (error?.status === "FETCH_ERROR") { diff --git a/ui/lib/store/apis/tempToken.ts b/ui/lib/store/apis/tempToken.ts new file mode 100644 index 0000000000..625aec755e --- /dev/null +++ b/ui/lib/store/apis/tempToken.ts @@ -0,0 +1,33 @@ +// Module-level state for the temp-token scope wrapper. +// +// The wrapper component (components/tempTokenScope.tsx) sets these on mount +// and clears them on unmount. baseApi reads them on every request: +// - prepareHeaders attaches `X-Bifrost-Temp-Token: ` when a token +// is active, so APIs called from inside the scope can authenticate via +// temp token instead of the dashboard session cookie. +// - baseQueryWithErrorHandling consults the suppression flag before +// force-redirecting to /login on a 401, so a scoped page can render its +// own "invalid/expired link" view instead of yanking the user away. +// +// A module-level singleton is fine because we never expect two TempTokenScope +// wrappers to be mounted concurrently in the same tab. The wrapper guards +// against nested mounts via the same-token check on set. + +let activeTempToken: string | null = null; +let suppressGlobal401 = false; + +export function setActiveTempToken(token: string | null): void { + activeTempToken = token; +} + +export function getActiveTempToken(): string | null { + return activeTempToken; +} + +export function setSuppressGlobal401(value: boolean): void { + suppressGlobal401 = value; +} + +export function getSuppressGlobal401(): boolean { + return suppressGlobal401; +}