sendnrw/netbird
2
0
Fork 0
mirror of https://github.com/netbirdio/netbird.git synced 2026-10-25 05:49:06 +01:00
Code Issues Packages Projects Releases Wiki Activity

[management] Refactor PKCE verifier store into a reusable single-use store (#7634)

Browse Source 
* Generalize PKCE verifier store into SingleUseStore

* Use the single-use store in integration tests

* Remove unrelated session code test stub
This commit is contained in:
Bethuel Mmbaga
2026-09-29 13:25:56 +03:00
committed by GitHub
parent a2919e26dd
commit 082aca1556
12 changed files with 96 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
management/internals/shared/grpc/pkce_verifier.go
-55
View File
@@ -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
}
management/internals/shared/grpc/proxy.go
+5 -5
View File
@@ -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")
}
management/internals/shared/grpc/proxy_test.go
+17 -17
View File
@@ -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")
management/internals/shared/grpc/single_use.go
+47
View File
@@ -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
}
management/internals/shared/grpc/pkce_verifier_test.go → management/internals/shared/grpc/single_use_test.go
+5 -5
View File
@@ -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)
}
management/internals/shared/grpc/validate_session_test.go
+2 -2
View File
@@ -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)