[management] use the cache for the pkce state#5516
Merged
pascal-fischer merged 3 commits intomainfrom Mar 9, 2026
Merged
Conversation
Contributor
📝 WalkthroughWalkthroughThis PR introduces a PKCEVerifierStore (in-memory/Redis-backed), wires it into server bootstrap and ProxyServiceServer, removes the per-server pkce sync.Map and cleanup goroutine, and updates all constructor call sites and tests to pass and use the new store. Changes
Sequence Diagram(s)sequenceDiagram
participant Client as OAuth Client
participant Proxy as ProxyServiceServer
participant Store as PKCEVerifierStore
participant Cache as Cache Backend (Redis / In-Memory)
participant IDP as Identity Provider
Client->>Proxy: GetOIDCURL request
Proxy->>Proxy: generate code_verifier & state
Proxy->>Store: Store(state, verifier, ttl)
Store->>Cache: Set(state -> verifier)
Cache-->>Store: OK
Store-->>Proxy: Store OK
Proxy-->>Client: Redirect URL with state
Client->>IDP: Authorize with code_challenge
IDP-->>Client: Authorization code + state
Client->>Proxy: Callback (state, code)
Proxy->>Store: LoadAndDelete(state)
Store->>Cache: Get & Delete state
Cache-->>Store: verifier
Store-->>Proxy: verifier (atomically deleted)
Proxy->>Proxy: verify code_challenge, exchange code
Proxy-->>Client: Session established
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
management/server/account_test.go (1)
3136-3136: AvoidnilPKCE store in shared test setupUsing
nilhere makescreateManagerbrittle—any future test path touching OIDC/PKCE can panic. Prefer constructing a real PKCE store in this helper, like other test setups in this PR.♻️ Suggested change
- proxyGrpcServer := nbgrpc.NewProxyServiceServer(nil, nil, nil, nbgrpc.ProxyOIDCConfig{}, peersManager, nil, proxyManager) + pkceStore, err := nbgrpc.NewPKCEVerifierStore(ctx, 10*time.Minute, 10*time.Minute, 100) + if err != nil { + return nil, nil, err + } + proxyGrpcServer := nbgrpc.NewProxyServiceServer(nil, nil, pkceStore, nbgrpc.ProxyOIDCConfig{}, peersManager, nil, proxyManager)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@management/server/account_test.go` at line 3136, The shared test setup passes a nil PKCE store to nbgrpc.NewProxyServiceServer which makes createManager brittle and can cause panics when OIDC/PKCE paths are exercised; instead instantiate the real PKCE store used elsewhere in the PR and pass it into nbgrpc.NewProxyServiceServer (replace the nil PKCE store argument) so tests use a concrete PKCE store; update the helper createManager (and its call sites if needed) to construct and supply that store alongside ProxyOIDCConfig, peersManager and proxyManager.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@management/internals/shared/grpc/pkce_verifier.go`:
- Around line 49-60: The current PKCEVerifierStore.LoadAndDelete uses separate
cache.Get and cache.Delete calls (cache.Get, cache.Delete on s.ctx) so
concurrent callers can both read the verifier before deletion and delete errors
are logged but still return success; change LoadAndDelete to perform an atomic
get-and-delete (use a cache.GetAndDelete / GetDelete method if the cache
supports it) or, if not supported, serialize access per state (e.g., a per-key
mutex map inside PKCEVerifierStore) so only one goroutine can call Get then
Delete for a given state, and treat delete failures as fatal by returning ("",
false) instead of true (also update log messages to include error context).
---
Nitpick comments:
In `@management/server/account_test.go`:
- Line 3136: The shared test setup passes a nil PKCE store to
nbgrpc.NewProxyServiceServer which makes createManager brittle and can cause
panics when OIDC/PKCE paths are exercised; instead instantiate the real PKCE
store used elsewhere in the PR and pass it into nbgrpc.NewProxyServiceServer
(replace the nil PKCE store argument) so tests use a concrete PKCE store; update
the helper createManager (and its call sites if needed) to construct and supply
that store alongside ProxyOIDCConfig, peersManager and proxyManager.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: d002198f-e5b1-4dda-a597-bdbc3ecfb340
📒 Files selected for processing (9)
management/internals/modules/reverseproxy/service/manager/manager_test.gomanagement/internals/server/boot.gomanagement/internals/shared/grpc/pkce_verifier.gomanagement/internals/shared/grpc/proxy.gomanagement/internals/shared/grpc/validate_session_test.gomanagement/server/account_test.gomanagement/server/http/handlers/proxy/auth_callback_integration_test.gomanagement/server/http/testing/testing_tools/channel/channel.goproxy/management_integration_test.go
|
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
management/internals/shared/grpc/proxy_test.go (1)
96-101: Use cancellable contexts in tests for store lifecycle cleanup.These tests pass
context.Background()into store constructors. If those stores run cleanup goroutines, they cannot be stopped per-test, which can accumulate background work across the suite.♻️ Suggested change pattern
- ctx := context.Background() + ctx, cancel := context.WithCancel(context.Background()) + t.Cleanup(cancel) tokenStore, err := NewOneTimeTokenStore(ctx, time.Hour, 10*time.Minute, 100) require.NoError(t, err)Apply the same pattern in tests that only create
pkceStore.Also applies to: 158-163, 197-202, 258-260, 287-289, 308-310
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@management/internals/shared/grpc/proxy_test.go` around lines 96 - 101, Replace uses of context.Background() when constructing stores (e.g., NewOneTimeTokenStore and NewPKCEVerifierStore) with a per-test cancellable context created via ctx, cancel := context.WithCancel(context.Background()), pass ctx into the store constructor, and defer cancel() at the end of the test so any cleanup goroutines are stopped; apply the same pattern for all tests that only create pkceStore or similar stores so each test controls store lifecycle and avoids accumulating background work.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@management/internals/shared/grpc/proxy_test.go`:
- Around line 96-101: Replace uses of context.Background() when constructing
stores (e.g., NewOneTimeTokenStore and NewPKCEVerifierStore) with a per-test
cancellable context created via ctx, cancel :=
context.WithCancel(context.Background()), pass ctx into the store constructor,
and defer cancel() at the end of the test so any cleanup goroutines are stopped;
apply the same pattern for all tests that only create pkceStore or similar
stores so each test controls store lifecycle and avoids accumulating background
work.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 03572854-b203-450c-be38-1f07ae064832
📒 Files selected for processing (1)
management/internals/shared/grpc/proxy_test.go
mlsmaycon
approved these changes
Mar 9, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Describe your changes
Issue ticket number and link
Stack
Checklist
Documentation
Select exactly one:
Docs PR URL (required if "docs added" is checked)
Paste the PR link from https://github.com/netbirdio/docs here:
https://github.com/netbirdio/docs/pull/__
Summary by CodeRabbit
New Features
Refactor