feat: add OAuth2 session listing, revocation, family-revocation, VK liveness check, and sweep worker#4509
Conversation
|
|
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (10)
🚧 Files skipped from review as they are similar to previous changes (10)
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds OAuth2 connected-clients session management with refresh-token breach detection and session HTTP endpoints. Extends ChangesOAuth2 Session Management & Breach Detection
Graceful Shutdown Pattern for Sweep Workers
Sequence DiagramssequenceDiagram
participant Client
participant OAuth2SessionsHandler
participant ConfigStore
participant DB
rect rgba(173, 216, 230, 0.5)
note over Client,DB: List Active Sessions
Client->>OAuth2SessionsHandler: GET /api/oauth2/sessions
OAuth2SessionsHandler->>ConfigStore: ListOAuth2Sessions(ctx)
ConfigStore->>DB: SELECT rt JOIN oauth2_clients JOIN virtual_keys
DB-->>ConfigStore: []OAuth2SessionRow
OAuth2SessionsHandler-->>Client: 200 {"sessions": [...]}
end
rect rgba(255, 200, 150, 0.5)
note over Client,DB: Revoke Session
Client->>OAuth2SessionsHandler: DELETE /api/oauth2/sessions/{id}
OAuth2SessionsHandler->>ConfigStore: GetOAuth2SessionByID(ctx, id)
ConfigStore-->>OAuth2SessionsHandler: session (non-revoked)
OAuth2SessionsHandler->>ConfigStore: RevokeOAuth2Session(ctx, id)
ConfigStore->>DB: UPDATE revoked_at=now WHERE id AND revoked_at IS NULL
OAuth2SessionsHandler-->>Client: 204 No Content
end
rect rgba(200, 230, 180, 0.5)
note over Client,DB: Token Refresh with Breach Detection
Client->>handleTokenRefresh: POST /token (refresh_token)
handleTokenRefresh->>ConfigStore: GetOAuth2RefreshTokenByHash(active)
alt not found
handleTokenRefresh->>ConfigStore: GetOAuth2RefreshTokenByHashAny
alt was revoked
handleTokenRefresh->>ConfigStore: RevokeOAuth2RefreshTokensByFamilyID
handleTokenRefresh-->>Client: 400 invalid_grant (breach)
else never issued
handleTokenRefresh-->>Client: 400 invalid_grant
end
else found, check VK if mode=vk
handleTokenRefresh->>ConfigStore: GetVirtualKey(bf_sub)
alt missing or inactive
handleTokenRefresh-->>Client: 400 invalid_grant
else active
handleTokenRefresh-->>Client: 200 new tokens
end
end
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 golangci-lint (2.12.2)level=error msg="[linters_context] typechecking error: pattern ./...: directory prefix . does not contain main module or its selected dependencies" Comment |
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
framework/configstore/rdb.go (1)
6882-6888:⚠️ Potential issue | 🔴 CriticalAdd single-use guard to refresh-token rotation under concurrency.
The rotation update on
idalone allows two concurrent refresh requests to both read the old active token, both update it, and both insert active descendants. This bypasses the revoked-token reuse detection because neither transaction observesErrNotFound.The current error handling (line 374 in oauth2_issuance.go) also maps any rotation failure to
"server_error"without distinguishing a rotation conflict from other errors, preventing proper family revocation when rotation races are detected.Add
AND revoked_at IS NULLto the rotation update's Where clause and checkRowsAffected == 0to returnErrNotFoundwhen the old token is no longer active. The caller must then map this to"invalid_grant"and revoke the family using the already-loadedrt.FamilyID.🔒 Required fix
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 { - if err := tx.Model(&tables.TableOAuth2RefreshToken{}). - Where("id = ?", oldID). - Update("revoked_at", &now).Error; err != nil { - return fmt.Errorf("revoke old refresh token: %w", err) + 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)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@framework/configstore/rdb.go` around lines 6882 - 6888, The Update operation on tables.TableOAuth2RefreshToken only checks the token id in its Where clause, allowing concurrent refresh requests to both read and update the same old token. Add AND revoked_at IS NULL to the Where clause to ensure only active tokens are updated, then check if the tx result RowsAffected equals 0 after the update and return ErrNotFound in that case. This prevents concurrent rotation races from both inserting new active descendants and allows the caller to properly distinguish rotation conflicts from other errors for family revocation.
🧹 Nitpick comments (1)
framework/configstore/rdb.go (1)
6960-6973: 🏗️ Heavy liftBound the connected-clients session query.
ListOAuth2Sessionsscans every active refresh token and performs two joins with no limit/offset. Since active grants have no expiry column inTableOAuth2RefreshToken, this new API path can grow unbounded; add pagination with a stable tiebreaker before exposing it broadly.♻️ Suggested shape
-func (s *RDBConfigStore) ListOAuth2Sessions(ctx context.Context) ([]OAuth2SessionRow, error) { +func (s *RDBConfigStore) ListOAuth2Sessions(ctx context.Context, limit, offset int) ([]OAuth2SessionRow, error) { + if limit <= 0 { + limit = 25 + } else if limit > 100 { + limit = 100 + } + if offset < 0 { + offset = 0 + } rows := []struct { tables.TableOAuth2RefreshToken ClientName string `gorm:"column:client_name"` VKName string `gorm:"column:vk_name"` }{} err := s.ScopedDB(ctx). @@ Joins("LEFT JOIN oauth2_clients c ON c.client_id = rt.client_id"). Joins("LEFT JOIN governance_virtual_keys vk ON vk.id = rt.bf_sub AND rt.bf_mode = 'vk'"). Where("rt.revoked_at IS NULL"). - Order("rt.created_at DESC"). + Order("rt.created_at DESC, rt.id DESC"). + Limit(limit). + Offset(offset). Scan(&rows).ErrorApply the same contract through the interface and handler response. As per coding guidelines,
framework/**changes should use bounded resource usage.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@framework/configstore/rdb.go` around lines 6960 - 6973, The ListOAuth2Sessions function in RDBConfigStore queries all active refresh tokens with no pagination limits, which can result in unbounded resource usage. Add pagination support by accepting limit and offset parameters to the function, apply these to the database query before the Scan call, and add a stable tiebreaker to the Order clause (using rt.id in addition to rt.created_at DESC) for consistent ordering across paginated requests. Update the interface contract and handler response to accept and return pagination parameters consistently with other bounded API endpoints.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@framework/configstore/rdb.go`:
- Around line 6934-6939: The SweepOAuth2RefreshTokens method does not validate
that revokedOlderThan is a positive duration. When revokedOlderThan is zero or
negative, the cutoff time becomes now or in the future, causing the deletion
query to immediately remove nearly all revoked tokens instead of just old ones,
which breaks replay detection. Add a guard at the beginning of the
SweepOAuth2RefreshTokens method to check if revokedOlderThan is positive and
return an error if it is not.
In `@transports/bifrost-http/handlers/oauth2_issuance.go`:
- Around line 327-328: The code at the breach-detection block suppresses
critical store errors instead of handling them explicitly. The call to
GetOAuth2RefreshTokenByHashAny should be handled to address the lookupErr case
beyond just the success path, and the RevokeOAuth2RefreshTokensByFamilyID call
should not discard its error using the blank identifier. Instead, capture and
properly handle the error returned from RevokeOAuth2RefreshTokensByFamilyID to
ensure fail-closed behavior for this security-critical family token revocation
operation. According to coding guidelines, transport persistence and security
changes must have explicit error handling.
- Around line 346-353: In the GetVirtualKey error handling block in
oauth2_issuance.go, reorder the condition checks so that non-ErrNotFound lookup
errors are handled first by returning server_error, before checking if the
virtual key is nil or inactive. The current logic combines the nil check with
ErrNotFound in a single condition, causing transient store failures that return
(nil, err) to incorrectly send invalid_grant instead of server_error. Separate
the vkErr != nil check (excluding ErrNotFound) to execute first with
server_error response, then check vk == nil or !vk.IsActiveValue() for the
invalid_grant response.
In `@transports/bifrost-http/handlers/oauth2_sessions.go`:
- Around line 44-45: The sonic.Marshal call on line 44 ignores the error by
using underscore assignment, which means if JSON serialization fails, the code
will still set an empty or invalid body while returning a success status.
Capture the error returned by sonic.Marshal along with the data variable, check
if the error is not nil, and if so return a 500 status code using ctx methods
before proceeding to ctx.SetBody. Only call ctx.SetBody(data) when the marshal
operation succeeds to ensure the API fails closed on serialization errors.
In `@transports/bifrost-http/server/oauth2.go`:
- Around line 34-53: The oauth2SweepWorker does not cancel in-flight sweep
operations when stop() is called, only closing stopCh. If sweep(ctx) is blocked
on a DB operation, the goroutine remains alive indefinitely. Create a
cancellable context derived from the passed ctx parameter in the run method
using context.WithCancel, store the resulting cancel function as a field in the
worker, and modify the stop() method to call this cancel function in addition to
closing stopCh. This ensures any blocking sweep(ctx) call receives the
cancellation signal through ctx.Done() and can terminate promptly during
shutdown.
In `@transports/bifrost-http/server/server.go`:
- Around line 1764-1767: The OAuth2SweepWorker cleanup block is incorrectly
nested under a condition checking if TempTokenSweepWorker exists, which means
OAuth2 cleanup will be skipped if only OAuth2SweepWorker was created but
TempTokenSweepWorker is nil. Move the OAuth2SweepWorker cleanup code (the check
for s.OAuth2SweepWorker != nil, followed by stop() and nil assignment) to
execute independently without being gated by any TempTokenSweepWorker condition.
Apply this fix at all four locations: the block at lines 1764-1767 and the three
additional locations mentioned (1798-1801, 1836-1839, 1933-1936).
---
Outside diff comments:
In `@framework/configstore/rdb.go`:
- Around line 6882-6888: The Update operation on tables.TableOAuth2RefreshToken
only checks the token id in its Where clause, allowing concurrent refresh
requests to both read and update the same old token. Add AND revoked_at IS NULL
to the Where clause to ensure only active tokens are updated, then check if the
tx result RowsAffected equals 0 after the update and return ErrNotFound in that
case. This prevents concurrent rotation races from both inserting new active
descendants and allows the caller to properly distinguish rotation conflicts
from other errors for family revocation.
---
Nitpick comments:
In `@framework/configstore/rdb.go`:
- Around line 6960-6973: The ListOAuth2Sessions function in RDBConfigStore
queries all active refresh tokens with no pagination limits, which can result in
unbounded resource usage. Add pagination support by accepting limit and offset
parameters to the function, apply these to the database query before the Scan
call, and add a stable tiebreaker to the Order clause (using rt.id in addition
to rt.created_at DESC) for consistent ordering across paginated requests. Update
the interface contract and handler response to accept and return pagination
parameters consistently with other bounded API endpoints.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: ae0ec375-e9c9-4ecc-ab89-f1fe7617bb4d
📒 Files selected for processing (6)
framework/configstore/rdb.goframework/configstore/store.gotransports/bifrost-http/handlers/oauth2_issuance.gotransports/bifrost-http/handlers/oauth2_sessions.gotransports/bifrost-http/server/oauth2.gotransports/bifrost-http/server/server.go
d6a46f9 to
5d1f810
Compare
98169a4 to
be01ec8
Compare
5d1f810 to
bc90e18
Compare
be01ec8 to
9282613
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (1)
framework/configstore/rdb.go (1)
7016-7025: 💤 Low valueConsider adding a guard for non-positive duration for consistency.
SweepOAuth2RefreshTokensguards againstrevokedOlderThan <= 0to prevent unintended mass deletion. This method lacks an equivalent guard—ifregisteredOlderThan <= 0, the cutoff becomes now or later, deleting all orphaned clients immediately rather than respecting a grace window.The risk is lower here (no security detection dependency), but adding a similar guard would maintain consistency across sweep methods.
♻️ Suggested guard for consistency
func (s *RDBConfigStore) SweepOrphanedOAuth2Clients(ctx context.Context, registeredOlderThan time.Duration) (int64, error) { + if registeredOlderThan <= 0 { + return 0, nil + } cutoff := time.Now().Add(-registeredOlderThan)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@framework/configstore/rdb.go` around lines 7016 - 7025, The SweepOrphanedOAuth2Clients method lacks a guard against non-positive duration values for the registeredOlderThan parameter, which could lead to unintended mass deletion of all orphaned clients. Add a validation guard at the beginning of the SweepOrphanedOAuth2Clients function (similar to what exists in SweepOAuth2RefreshTokens) to check if registeredOlderThan is less than or equal to zero and return an error in such cases to maintain consistency across sweep methods and prevent accidental data loss.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@framework/configstore/rdb.go`:
- Around line 7016-7025: The SweepOrphanedOAuth2Clients method lacks a guard
against non-positive duration values for the registeredOlderThan parameter,
which could lead to unintended mass deletion of all orphaned clients. Add a
validation guard at the beginning of the SweepOrphanedOAuth2Clients function
(similar to what exists in SweepOAuth2RefreshTokens) to check if
registeredOlderThan is less than or equal to zero and return an error in such
cases to maintain consistency across sweep methods and prevent accidental data
loss.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro Plus
Run ID: ba94ffc9-a891-4373-b716-137921d834cb
📒 Files selected for processing (10)
framework/configstore/rdb.goframework/configstore/store.goframework/mcp_headers/sweep.goframework/oauth2/sync.goframework/temptoken/sweeper.gotransports/bifrost-http/handlers/oauth2_issuance.gotransports/bifrost-http/handlers/oauth2_sessions.gotransports/bifrost-http/lib/config_test.gotransports/bifrost-http/server/oauth2.gotransports/bifrost-http/server/server.go
🚧 Files skipped from review as they are similar to previous changes (5)
- transports/bifrost-http/server/oauth2.go
- framework/configstore/store.go
- transports/bifrost-http/handlers/oauth2_sessions.go
- transports/bifrost-http/server/server.go
- transports/bifrost-http/handlers/oauth2_issuance.go
bc90e18 to
39907be
Compare
2a42dc9 to
1c19dfb
Compare
1f9653e to
b156095
Compare
1c19dfb to
0ecdc90
Compare
…cks, sessions api
b156095 to
8a84174
Compare
0ecdc90 to
4178622
Compare

Summary
This PR adds OAuth2 session management capabilities including stolen-token detection with automatic family revocation, a Connected Clients API for listing and revoking active downstream grants, VK liveness checks on token refresh, and a background sweep worker to clean up revoked refresh tokens.
Changes
GetOAuth2RefreshTokenByHashAnyto look up refresh tokens including revoked ones, enabling detection of token reuse attacks where a previously rotated token is re-presentedRevokeOAuth2RefreshTokensByFamilyIDandRevokeOAuth2RefreshTokensByModefor bulk revocation of token families and mode-scoped grants respectivelySweepOAuth2RefreshTokensto delete revoked tokens older than a configurable retention window (default 30 days), swept every 10 minutes via a newoauth2SweepWorkerListOAuth2Sessionswhich joins refresh token rows with client names and VK names for human-readable display in the Connected Clients UIGetOAuth2SessionByIDandRevokeOAuth2Sessionfor single-session lookup and revocationOAuth2SessionsHandlerexposingGET /api/oauth2/sessionsandDELETE /api/oauth2/sessions/{id}; user-mode sessions enforce that the caller's identity matchesbf_subbefore allowing revocationinvalid_grantoauth2SweepWorkeris started during server bootstrap and stopped cleanly on all error and shutdown paths alongside the existing temp-token sweep workerType of change
Affected areas
How to test
go test ./...POST /token— expectinvalid_grantand all sibling tokens in the family to be revoked.invalid_grantwith "virtual key is no longer active".GET /api/oauth2/sessions— expect a JSON list of active grants with client names and VK display names populated.DELETE /api/oauth2/sessions/{id}with a user-mode session ID using a mismatched caller identity — expect403 Forbidden.DELETE /api/oauth2/sessions/{id}with the correct caller identity — expect204 No Contentand the session absent from subsequent list calls.Breaking changes
The
ConfigStoreinterface has new required methods. Any custom implementations ofConfigStoremust addGetOAuth2RefreshTokenByHashAny,RevokeOAuth2RefreshTokensByFamilyID,RevokeOAuth2RefreshTokensByMode,SweepOAuth2RefreshTokens,ListOAuth2Sessions,GetOAuth2SessionByID, andRevokeOAuth2Session.Related issues
Security considerations
Checklist
docs/contributing/README.mdand followed the guidelines