From 082aca155615f17f0c0ac84e958a34b9681ff55f Mon Sep 17 00:00:00 2001 From: Bethuel Mmbaga Date: Tue, 29 Sep 2026 13:25:56 +0300 Subject: [PATCH] [management] Refactor PKCE verifier store into a reusable single-use store (#7634) * Generalize PKCE verifier store into SingleUseStore * Use the single-use store in integration tests * Remove unrelated session code test stub --- .../service/manager/manager_test.go | 12 ++-- management/internals/server/boot.go | 8 +-- .../internals/shared/grpc/pkce_verifier.go | 55 ------------------- management/internals/shared/grpc/proxy.go | 10 ++-- .../internals/shared/grpc/proxy_test.go | 34 ++++++------ .../internals/shared/grpc/single_use.go | 47 ++++++++++++++++ ...ce_verifier_test.go => single_use_test.go} | 10 ++-- .../shared/grpc/validate_session_test.go | 4 +- .../proxy/auth_callback_integration_test.go | 4 +- .../testing/testing_tools/channel/channel.go | 8 +-- proxy/management_byop_integration_test.go | 4 +- proxy/management_integration_test.go | 4 +- 12 files changed, 96 insertions(+), 104 deletions(-) delete mode 100644 management/internals/shared/grpc/pkce_verifier.go create mode 100644 management/internals/shared/grpc/single_use.go rename management/internals/shared/grpc/{pkce_verifier_test.go => single_use_test.go} (85%) diff --git a/management/internals/modules/reverseproxy/service/manager/manager_test.go b/management/internals/modules/reverseproxy/service/manager/manager_test.go index 9704796be..2ec1af1a2 100644 --- a/management/internals/modules/reverseproxy/service/manager/manager_test.go +++ b/management/internals/modules/reverseproxy/service/manager/manager_test.go @@ -433,8 +433,8 @@ func TestDeletePeerService_SourcePeerValidation(t *testing.T) { newProxyServer := func(t *testing.T) *nbgrpc.ProxyServiceServer { t.Helper() tokenStore := nbgrpc.NewOneTimeTokenStore(context.Background(), testCacheStore(t)) - pkceStore := nbgrpc.NewPKCEVerifierStore(context.Background(), testCacheStore(t)) - srv := nbgrpc.NewProxyServiceServer(nil, tokenStore, pkceStore, nbgrpc.ProxyOIDCConfig{}, nil, nil, nil, nil, nil) + singleUseStore := nbgrpc.NewSingleUseStore(context.Background(), testCacheStore(t)) + srv := nbgrpc.NewProxyServiceServer(nil, tokenStore, singleUseStore, nbgrpc.ProxyOIDCConfig{}, nil, nil, nil, nil, nil) return srv } @@ -726,8 +726,8 @@ func setupIntegrationTest(t *testing.T) (*Manager, store.Store) { } tokenStore := nbgrpc.NewOneTimeTokenStore(ctx, testCacheStore(t)) - pkceStore := nbgrpc.NewPKCEVerifierStore(ctx, testCacheStore(t)) - proxySrv := nbgrpc.NewProxyServiceServer(nil, tokenStore, pkceStore, nbgrpc.ProxyOIDCConfig{}, nil, nil, nil, nil, nil) + singleUseStore := nbgrpc.NewSingleUseStore(ctx, testCacheStore(t)) + proxySrv := nbgrpc.NewProxyServiceServer(nil, tokenStore, singleUseStore, nbgrpc.ProxyOIDCConfig{}, nil, nil, nil, nil, nil) proxyController, err := proxymanager.NewGRPCController(proxySrv, noop.NewMeterProvider().Meter("")) require.NoError(t, err) @@ -1150,8 +1150,8 @@ func TestDeleteService_DeletesTargets(t *testing.T) { mockAcct := account.NewMockManager(ctrl) tokenStore := nbgrpc.NewOneTimeTokenStore(ctx, testCacheStore(t)) - pkceStore := nbgrpc.NewPKCEVerifierStore(ctx, testCacheStore(t)) - proxySrv := nbgrpc.NewProxyServiceServer(nil, tokenStore, pkceStore, nbgrpc.ProxyOIDCConfig{}, nil, nil, nil, nil, nil) + singleUseStore := nbgrpc.NewSingleUseStore(ctx, testCacheStore(t)) + proxySrv := nbgrpc.NewProxyServiceServer(nil, tokenStore, singleUseStore, nbgrpc.ProxyOIDCConfig{}, nil, nil, nil, nil, nil) proxyController, err := proxymanager.NewGRPCController(proxySrv, noop.NewMeterProvider().Meter("")) require.NoError(t, err) diff --git a/management/internals/server/boot.go b/management/internals/server/boot.go index 0f136ce82..34c8363f1 100644 --- a/management/internals/server/boot.go +++ b/management/internals/server/boot.go @@ -248,7 +248,7 @@ func (s *BaseServer) GRPCServer() *grpc.Server { func (s *BaseServer) ReverseProxyGRPCServer() *nbgrpc.ProxyServiceServer { return Create(s, func() *nbgrpc.ProxyServiceServer { - proxyService := nbgrpc.NewProxyServiceServer(s.AccessLogsManager(), s.ProxyTokenStore(), s.PKCEVerifierStore(), s.proxyOIDCConfig(), s.PeersManager(), s.UsersManager(), s.IdpManager(), s.ProxyManager(), s.Store()) + proxyService := nbgrpc.NewProxyServiceServer(s.AccessLogsManager(), s.ProxyTokenStore(), s.SingleUseStore(), s.proxyOIDCConfig(), s.PeersManager(), s.UsersManager(), s.IdpManager(), s.ProxyManager(), s.Store()) s.AfterInit(func(s *BaseServer) { proxyService.SetServiceManager(s.ServiceManager()) proxyService.SetActivityManager(s.ProxyActivityManager()) @@ -305,9 +305,9 @@ func (s *BaseServer) ProxyTokenStore() *nbgrpc.OneTimeTokenStore { }) } -func (s *BaseServer) PKCEVerifierStore() *nbgrpc.PKCEVerifierStore { - return Create(s, func() *nbgrpc.PKCEVerifierStore { - return nbgrpc.NewPKCEVerifierStore(context.Background(), s.CacheStore()) +func (s *BaseServer) SingleUseStore() *nbgrpc.SingleUseStore { + return Create(s, func() *nbgrpc.SingleUseStore { + return nbgrpc.NewSingleUseStore(context.Background(), s.CacheStore()) }) } diff --git a/management/internals/shared/grpc/pkce_verifier.go b/management/internals/shared/grpc/pkce_verifier.go deleted file mode 100644 index 18155dc1d..000000000 --- a/management/internals/shared/grpc/pkce_verifier.go +++ /dev/null @@ -1,55 +0,0 @@ -package grpc - -import ( - "context" - "fmt" - "time" - - "github.com/eko/gocache/lib/v4/store" - log "github.com/sirupsen/logrus" - - nbcache "github.com/netbirdio/netbird/management/server/cache" -) - -// PKCEVerifierStore manages PKCE verifiers for OAuth flows. -// Supports both in-memory and Redis storage via NB_IDP_CACHE_REDIS_ADDRESS env var. -type PKCEVerifierStore struct { - cache nbcache.Store - ctx context.Context -} - -// NewPKCEVerifierStore creates a PKCE verifier store using the provided shared cache store. -func NewPKCEVerifierStore(ctx context.Context, cacheStore nbcache.Store) *PKCEVerifierStore { - return &PKCEVerifierStore{ - cache: cacheStore, - ctx: ctx, - } -} - -// Store saves a PKCE verifier associated with an OAuth state parameter. -// The verifier is stored with the specified TTL and will be automatically deleted after expiration. -func (s *PKCEVerifierStore) Store(state, verifier string, ttl time.Duration) error { - if err := s.cache.Set(s.ctx, state, verifier, store.WithExpiration(ttl)); err != nil { - return fmt.Errorf("failed to store PKCE verifier: %w", err) - } - - log.Debugf("Stored PKCE verifier for state (expires in %s)", ttl) - return nil -} - -// LoadAndDelete retrieves and removes a PKCE verifier for the given state. -// Returns the verifier and true if found, or empty string and false if not found. -// This enforces single-use semantics for PKCE verifiers. -func (s *PKCEVerifierStore) LoadAndDelete(state string) (string, bool) { - verifier, found, err := s.cache.GetDel(s.ctx, state) - if err != nil { - log.Warnf("Failed to consume PKCE verifier: %v", err) - return "", false - } - if !found { - log.Debug("PKCE verifier not found for state") - return "", false - } - - return verifier, true -} diff --git a/management/internals/shared/grpc/proxy.go b/management/internals/shared/grpc/proxy.go index 28df7ed6f..e5d89ab3d 100644 --- a/management/internals/shared/grpc/proxy.go +++ b/management/internals/shared/grpc/proxy.go @@ -143,7 +143,7 @@ type ProxyServiceServer struct { oidcConfig ProxyOIDCConfig // Store for PKCE verifiers - pkceVerifierStore *PKCEVerifierStore + singleUseStore *SingleUseStore // tokenTTL is the lifetime of one-time tokens generated for proxy // authentication. Defaults to defaultProxyTokenTTL when zero. @@ -208,13 +208,13 @@ func enforceAccountScope(ctx context.Context, requestAccountID string) error { } // NewProxyServiceServer creates a new proxy service server. -func NewProxyServiceServer(accessLogMgr accesslogs.Manager, tokenStore *OneTimeTokenStore, pkceStore *PKCEVerifierStore, oidcConfig ProxyOIDCConfig, peersManager peers.Manager, usersManager users.Manager, idpManager idp.Manager, proxyMgr proxy.Manager, tokenChecker ProxyTokenChecker) *ProxyServiceServer { +func NewProxyServiceServer(accessLogMgr accesslogs.Manager, tokenStore *OneTimeTokenStore, singleUseStore *SingleUseStore, oidcConfig ProxyOIDCConfig, peersManager peers.Manager, usersManager users.Manager, idpManager idp.Manager, proxyMgr proxy.Manager, tokenChecker ProxyTokenChecker) *ProxyServiceServer { ctx, cancel := context.WithCancel(context.Background()) s := &ProxyServiceServer{ accessLogManager: accessLogMgr, oidcConfig: oidcConfig, tokenStore: tokenStore, - pkceVerifierStore: pkceStore, + singleUseStore: singleUseStore, peersManager: peersManager, usersManager: usersManager, idpManager: idpManager, @@ -1575,7 +1575,7 @@ func (s *ProxyServiceServer) GetOIDCURL(ctx context.Context, req *proto.GetOIDCU state := fmt.Sprintf("%s|%s|%s", base64.URLEncoding.EncodeToString([]byte(redirectURL.String())), nonceB64, hmacSum) codeVerifier := oauth2.GenerateVerifier() - if err := s.pkceVerifierStore.Store(state, codeVerifier, pkceVerifierTTL); err != nil { + if err := s.singleUseStore.Store(state, codeVerifier, pkceVerifierTTL); err != nil { log.WithContext(ctx).Errorf("failed to store PKCE verifier: %v", err) return nil, status.Errorf(codes.Internal, "store PKCE verifier: %v", err) } @@ -1641,7 +1641,7 @@ func (s *ProxyServiceServer) ValidateState(state string) (verifier, redirectURL } // Consume the PKCE verifier only after HMAC validation passes. - verifier, ok := s.pkceVerifierStore.LoadAndDelete(state) + verifier, ok := s.singleUseStore.LoadAndDelete(state) if !ok { return "", "", errors.New("no verifier for state") } diff --git a/management/internals/shared/grpc/proxy_test.go b/management/internals/shared/grpc/proxy_test.go index 29b7c9523..060689da1 100644 --- a/management/internals/shared/grpc/proxy_test.go +++ b/management/internals/shared/grpc/proxy_test.go @@ -129,11 +129,11 @@ func drainEmpty(ch chan *proto.GetMappingUpdateResponse) bool { func TestSendServiceUpdateToCluster_UniqueTokensPerProxy(t *testing.T) { ctx := context.Background() tokenStore := NewOneTimeTokenStore(ctx, testCacheStore(t)) - pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t)) + singleUseStore := NewSingleUseStore(ctx, testCacheStore(t)) s := &ProxyServiceServer{ - tokenStore: tokenStore, - pkceVerifierStore: pkceStore, + tokenStore: tokenStore, + singleUseStore: singleUseStore, } s.SetProxyController(newTestProxyController()) @@ -186,11 +186,11 @@ func TestSendServiceUpdateToCluster_UniqueTokensPerProxy(t *testing.T) { func TestSendServiceUpdateToCluster_DeleteNoToken(t *testing.T) { ctx := context.Background() tokenStore := NewOneTimeTokenStore(ctx, testCacheStore(t)) - pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t)) + singleUseStore := NewSingleUseStore(ctx, testCacheStore(t)) s := &ProxyServiceServer{ - tokenStore: tokenStore, - pkceVerifierStore: pkceStore, + tokenStore: tokenStore, + singleUseStore: singleUseStore, } s.SetProxyController(newTestProxyController()) @@ -220,11 +220,11 @@ func TestSendServiceUpdateToCluster_DeleteNoToken(t *testing.T) { func TestSendServiceUpdate_UniqueTokensPerProxy(t *testing.T) { ctx := context.Background() tokenStore := NewOneTimeTokenStore(ctx, testCacheStore(t)) - pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t)) + singleUseStore := NewSingleUseStore(ctx, testCacheStore(t)) s := &ProxyServiceServer{ - tokenStore: tokenStore, - pkceVerifierStore: pkceStore, + tokenStore: tokenStore, + singleUseStore: singleUseStore, } s.SetProxyController(newTestProxyController()) @@ -272,13 +272,13 @@ func generateState(s *ProxyServiceServer, redirectURL string) string { func TestOAuthState_NeverTheSame(t *testing.T) { ctx := context.Background() - pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t)) + singleUseStore := NewSingleUseStore(ctx, testCacheStore(t)) s := &ProxyServiceServer{ oidcConfig: ProxyOIDCConfig{ HMACKey: []byte("test-hmac-key"), }, - pkceVerifierStore: pkceStore, + singleUseStore: singleUseStore, } redirectURL := "https://app.example.com/callback" @@ -300,17 +300,17 @@ func TestOAuthState_NeverTheSame(t *testing.T) { func TestValidateState_RejectsOldTwoPartFormat(t *testing.T) { ctx := context.Background() - pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t)) + singleUseStore := NewSingleUseStore(ctx, testCacheStore(t)) s := &ProxyServiceServer{ oidcConfig: ProxyOIDCConfig{ HMACKey: []byte("test-hmac-key"), }, - pkceVerifierStore: pkceStore, + singleUseStore: singleUseStore, } // Old format had only 2 parts: base64(url)|hmac - err := s.pkceVerifierStore.Store("base64url|hmac", "test", 10*time.Minute) + err := s.singleUseStore.Store("base64url|hmac", "test", 10*time.Minute) require.NoError(t, err) _, _, err = s.ValidateState("base64url|hmac") @@ -372,17 +372,17 @@ func TestEnforceAccountScope_AllowsNoTokenInContext(t *testing.T) { func TestValidateState_RejectsInvalidHMAC(t *testing.T) { ctx := context.Background() - pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t)) + singleUseStore := NewSingleUseStore(ctx, testCacheStore(t)) s := &ProxyServiceServer{ oidcConfig: ProxyOIDCConfig{ HMACKey: []byte("test-hmac-key"), }, - pkceVerifierStore: pkceStore, + singleUseStore: singleUseStore, } // Store with tampered HMAC - err := s.pkceVerifierStore.Store("dGVzdA==|nonce|wrong-hmac", "test", 10*time.Minute) + err := s.singleUseStore.Store("dGVzdA==|nonce|wrong-hmac", "test", 10*time.Minute) require.NoError(t, err) _, _, err = s.ValidateState("dGVzdA==|nonce|wrong-hmac") diff --git a/management/internals/shared/grpc/single_use.go b/management/internals/shared/grpc/single_use.go new file mode 100644 index 000000000..1698707d7 --- /dev/null +++ b/management/internals/shared/grpc/single_use.go @@ -0,0 +1,47 @@ +package grpc + +import ( + "context" + "fmt" + "time" + + "github.com/eko/gocache/lib/v4/store" + log "github.com/sirupsen/logrus" + + nbcache "github.com/netbirdio/netbird/management/server/cache" +) + +// SingleUseStore stores short-lived values that can be retrieved only once. +type SingleUseStore struct { + cache nbcache.Store + ctx context.Context +} + +// NewSingleUseStore creates a single-use value store over the shared cache. +func NewSingleUseStore(ctx context.Context, cacheStore nbcache.Store) *SingleUseStore { + return &SingleUseStore{ + cache: cacheStore, + ctx: ctx, + } +} + +// Store saves value under key with the given TTL, after which it is evicted. +func (s *SingleUseStore) Store(key, value string, ttl time.Duration) error { + if err := s.cache.Set(s.ctx, key, value, store.WithExpiration(ttl)); err != nil { + return fmt.Errorf("store single-use value: %w", err) + } + return nil +} + +// LoadAndDelete retrieves and removes the value for a key. +func (s *SingleUseStore) LoadAndDelete(key string) (string, bool) { + value, found, err := s.cache.GetDel(s.ctx, key) + if err != nil { + log.Warnf("failed to consume single-use value: %v", err) + return "", false + } + if !found { + return "", false + } + return value, true +} diff --git a/management/internals/shared/grpc/pkce_verifier_test.go b/management/internals/shared/grpc/single_use_test.go similarity index 85% rename from management/internals/shared/grpc/pkce_verifier_test.go rename to management/internals/shared/grpc/single_use_test.go index e7175b6c5..357046720 100644 --- a/management/internals/shared/grpc/pkce_verifier_test.go +++ b/management/internals/shared/grpc/single_use_test.go @@ -6,7 +6,7 @@ import ( "time" ) -func TestPKCEVerifierStoreLoadAndDelete(t *testing.T) { +func TestSingleUseStoreLoadAndDelete(t *testing.T) { const ( state = "state" verifier = "verifier" @@ -14,7 +14,7 @@ func TestPKCEVerifierStoreLoadAndDelete(t *testing.T) { ) t.Run("exactly one concurrent caller consumes the verifier", func(t *testing.T) { - store := NewPKCEVerifierStore(context.Background(), testCacheStore(t)) + store := NewSingleUseStore(context.Background(), testCacheStore(t)) if err := store.Store(state, verifier, time.Minute); err != nil { t.Fatalf("couldn't store PKCE verifier: %s", err) } @@ -50,7 +50,7 @@ func TestPKCEVerifierStoreLoadAndDelete(t *testing.T) { }) t.Run("replayed state is rejected", func(t *testing.T) { - store := NewPKCEVerifierStore(context.Background(), testCacheStore(t)) + store := NewSingleUseStore(context.Background(), testCacheStore(t)) if err := store.Store(state, verifier, time.Minute); err != nil { t.Fatalf("couldn't store PKCE verifier: %s", err) } @@ -64,7 +64,7 @@ func TestPKCEVerifierStoreLoadAndDelete(t *testing.T) { }) t.Run("unknown state is rejected", func(t *testing.T) { - store := NewPKCEVerifierStore(context.Background(), testCacheStore(t)) + store := NewSingleUseStore(context.Background(), testCacheStore(t)) if got, found := store.LoadAndDelete("never-stored"); found { t.Fatalf("unknown state should not resolve, got %q", got) @@ -72,7 +72,7 @@ func TestPKCEVerifierStoreLoadAndDelete(t *testing.T) { }) t.Run("expired verifier is rejected", func(t *testing.T) { - store := NewPKCEVerifierStore(context.Background(), testCacheStore(t)) + store := NewSingleUseStore(context.Background(), testCacheStore(t)) if err := store.Store(state, verifier, 50*time.Millisecond); err != nil { t.Fatalf("couldn't store PKCE verifier: %s", err) } diff --git a/management/internals/shared/grpc/validate_session_test.go b/management/internals/shared/grpc/validate_session_test.go index b300e8c1d..2af250a85 100644 --- a/management/internals/shared/grpc/validate_session_test.go +++ b/management/internals/shared/grpc/validate_session_test.go @@ -40,9 +40,9 @@ func setupValidateSessionTest(t *testing.T) *validateSessionTestSetup { proxyManager := &testValidateSessionProxyManager{} tokenStore := NewOneTimeTokenStore(ctx, testCacheStore(t)) - pkceStore := NewPKCEVerifierStore(ctx, testCacheStore(t)) + singleUseStore := NewSingleUseStore(ctx, testCacheStore(t)) - proxyService := NewProxyServiceServer(nil, tokenStore, pkceStore, ProxyOIDCConfig{}, nil, usersManager, nil, proxyManager, nil) + proxyService := NewProxyServiceServer(nil, tokenStore, singleUseStore, ProxyOIDCConfig{}, nil, usersManager, nil, proxyManager, nil) proxyService.SetServiceManager(serviceManager) createTestProxies(t, ctx, testStore) diff --git a/management/server/http/handlers/proxy/auth_callback_integration_test.go b/management/server/http/handlers/proxy/auth_callback_integration_test.go index 1dbfca4cd..41ecb5f23 100644 --- a/management/server/http/handlers/proxy/auth_callback_integration_test.go +++ b/management/server/http/handlers/proxy/auth_callback_integration_test.go @@ -197,7 +197,7 @@ func setupAuthCallbackTest(t *testing.T) *testSetup { require.NoError(t, err) tokenStore := nbgrpc.NewOneTimeTokenStore(ctx, cacheStore) - pkceStore := nbgrpc.NewPKCEVerifierStore(ctx, cacheStore) + singleUseStore := nbgrpc.NewSingleUseStore(ctx, cacheStore) usersManager := users.NewManager(testStore) @@ -212,7 +212,7 @@ func setupAuthCallbackTest(t *testing.T) *testSetup { proxyService := nbgrpc.NewProxyServiceServer( &testAccessLogManager{}, tokenStore, - pkceStore, + singleUseStore, oidcConfig, nil, usersManager, diff --git a/management/server/http/testing/testing_tools/channel/channel.go b/management/server/http/testing/testing_tools/channel/channel.go index 598adafcc..c3f6a06e0 100644 --- a/management/server/http/testing/testing_tools/channel/channel.go +++ b/management/server/http/testing/testing_tools/channel/channel.go @@ -110,13 +110,13 @@ func BuildApiBlackBoxWithDBState(t testing_tools.TB, sqlFile string, expectedPee accessLogsManager := accesslogsmanager.NewManager(accesslogsmanager.NewRepository(store.(*nbstore.SqlStore).Conn()), store, permissionsManager, nil) proxyTokenStore := nbgrpc.NewOneTimeTokenStore(ctx, cacheStore) - pkceverifierStore := nbgrpc.NewPKCEVerifierStore(ctx, cacheStore) + singleUseStore := nbgrpc.NewSingleUseStore(ctx, cacheStore) noopMeter := noop.NewMeterProvider().Meter("") proxyMgr, err := proxymanager.NewManager(store, noopMeter) if err != nil { t.Fatalf("Failed to create proxy manager: %v", err) } - proxyServiceServer := nbgrpc.NewProxyServiceServer(accessLogsManager, proxyTokenStore, pkceverifierStore, nbgrpc.ProxyOIDCConfig{}, peersManager, userManager, nil, proxyMgr, nil) + proxyServiceServer := nbgrpc.NewProxyServiceServer(accessLogsManager, proxyTokenStore, singleUseStore, nbgrpc.ProxyOIDCConfig{}, peersManager, userManager, nil, proxyMgr, nil) // NewProxyServiceServer starts cleanupStaleProxies on a context it derives // from context.Background(), independent of the cancellable ctx above; // Close() cancels it so the goroutine does not outlive the test. @@ -250,13 +250,13 @@ func BuildApiBlackBoxWithDBStateAndPeerChannel(t testing_tools.TB, sqlFile strin accessLogsManager := accesslogsmanager.NewManager(accesslogsmanager.NewRepository(store.(*nbstore.SqlStore).Conn()), store, permissionsManager, nil) proxyTokenStore := nbgrpc.NewOneTimeTokenStore(ctx, cacheStore) - pkceverifierStore := nbgrpc.NewPKCEVerifierStore(ctx, cacheStore) + singleUseStore := nbgrpc.NewSingleUseStore(ctx, cacheStore) noopMeter := noop.NewMeterProvider().Meter("") proxyMgr, err := proxymanager.NewManager(store, noopMeter) if err != nil { t.Fatalf("Failed to create proxy manager: %v", err) } - proxyServiceServer := nbgrpc.NewProxyServiceServer(accessLogsManager, proxyTokenStore, pkceverifierStore, nbgrpc.ProxyOIDCConfig{}, peersManager, userManager, nil, proxyMgr, nil) + proxyServiceServer := nbgrpc.NewProxyServiceServer(accessLogsManager, proxyTokenStore, singleUseStore, nbgrpc.ProxyOIDCConfig{}, peersManager, userManager, nil, proxyMgr, nil) // NewProxyServiceServer starts cleanupStaleProxies on a context it derives // from context.Background(), independent of the cancellable ctx above; // Close() cancels it so the goroutine does not outlive the test. diff --git a/proxy/management_byop_integration_test.go b/proxy/management_byop_integration_test.go index d075e47ec..42301254d 100644 --- a/proxy/management_byop_integration_test.go +++ b/proxy/management_byop_integration_test.go @@ -104,7 +104,7 @@ func setupBYOPIntegrationTest(t *testing.T) *byopTestSetup { require.NoError(t, err) tokenStore := nbgrpc.NewOneTimeTokenStore(ctx, cacheStore) - pkceStore := nbgrpc.NewPKCEVerifierStore(ctx, cacheStore) + singleUseStore := nbgrpc.NewSingleUseStore(ctx, cacheStore) meter := noop.NewMeterProvider().Meter("test") realProxyManager, err := proxymanager.NewManager(testStore, meter) @@ -121,7 +121,7 @@ func setupBYOPIntegrationTest(t *testing.T) *byopTestSetup { proxyService := nbgrpc.NewProxyServiceServer( &testAccessLogManager{}, tokenStore, - pkceStore, + singleUseStore, oidcConfig, nil, usersManager, diff --git a/proxy/management_integration_test.go b/proxy/management_integration_test.go index 0e148f858..befd1f1ee 100644 --- a/proxy/management_integration_test.go +++ b/proxy/management_integration_test.go @@ -119,7 +119,7 @@ func setupIntegrationTest(t *testing.T) *integrationTestSetup { require.NoError(t, err) tokenStore := nbgrpc.NewOneTimeTokenStore(ctx, cacheStore) - pkceStore := nbgrpc.NewPKCEVerifierStore(ctx, cacheStore) + singleUseStore := nbgrpc.NewSingleUseStore(ctx, cacheStore) // Create real users manager usersManager := users.NewManager(testStore) @@ -136,7 +136,7 @@ func setupIntegrationTest(t *testing.T) *integrationTestSetup { proxyService := nbgrpc.NewProxyServiceServer( &testAccessLogManager{}, tokenStore, - pkceStore, + singleUseStore, oidcConfig, nil, usersManager,