consume PKCE verifiers atomically via a GetDel cache op

management/internals/shared/grpc/pkce_verifier.go

@@ -5,22 +5,23 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/eko/gocache/lib/v4/cache"
 	"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 *cache.Cache[string]
+	cache nbcache.Store
 	ctx   context.Context
 }
 
 // NewPKCEVerifierStore creates a PKCE verifier store using the provided shared cache store.
-func NewPKCEVerifierStore(ctx context.Context, cacheStore store.StoreInterface) *PKCEVerifierStore {
+func NewPKCEVerifierStore(ctx context.Context, cacheStore nbcache.Store) *PKCEVerifierStore {
 	return &PKCEVerifierStore{
-		cache: cache.New[string](cacheStore),
+		cache: cacheStore,
 		ctx:   ctx,
 	}
 }
@@ -40,14 +41,14 @@ func (s *PKCEVerifierStore) Store(state, verifier string, ttl time.Duration) err
 // 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, err := s.cache.Get(s.ctx, state)
+	verifier, found, err := s.cache.GetDel(s.ctx, state)
 	if err != nil {
-		log.Debugf("PKCE verifier not found for state")
+		log.Warnf("Failed to consume PKCE verifier: %v", err)
 		return "", false
 	}
-
-	if err := s.cache.Delete(s.ctx, state); err != nil {
-		log.Warnf("Failed to delete PKCE verifier for state: %v", err)
+	if !found {
+		log.Debug("PKCE verifier not found for state")
+		return "", false
 	}
 
 	return verifier, true

management/internals/shared/grpc/pkce_verifier_test.go

@@ -0,0 +1,85 @@
+package grpc
+
+import (
+	"context"
+	"testing"
+	"time"
+)
+
+func TestPKCEVerifierStoreLoadAndDelete(t *testing.T) {
+	const (
+		state    = "state"
+		verifier = "verifier"
+		attempts = 64
+	)
+
+	t.Run("exactly one concurrent caller consumes the verifier", func(t *testing.T) {
+		store := NewPKCEVerifierStore(context.Background(), testCacheStore(t))
+		if err := store.Store(state, verifier, time.Minute); err != nil {
+			t.Fatalf("couldn't store PKCE verifier: %s", err)
+		}
+
+		start := make(chan struct{})
+		type result struct {
+			verifier string
+			found    bool
+		}
+		results := make(chan result, attempts)
+		for range attempts {
+			go func() {
+				<-start
+				verifier, found := store.LoadAndDelete(state)
+				results <- result{verifier: verifier, found: found}
+			}()
+		}
+		close(start)
+
+		winners := 0
+		for range attempts {
+			result := <-results
+			if result.found {
+				winners++
+				if result.verifier != verifier {
+					t.Fatalf("unexpected verifier: got %q, expected %q", result.verifier, verifier)
+				}
+			}
+		}
+		if winners != 1 {
+			t.Fatalf("expected exactly one PKCE verifier consumer, got %d", winners)
+		}
+	})
+
+	t.Run("replayed state is rejected", func(t *testing.T) {
+		store := NewPKCEVerifierStore(context.Background(), testCacheStore(t))
+		if err := store.Store(state, verifier, time.Minute); err != nil {
+			t.Fatalf("couldn't store PKCE verifier: %s", err)
+		}
+
+		if got, found := store.LoadAndDelete(state); !found || got != verifier {
+			t.Fatalf("first load should return the verifier, got %q, found %t", got, found)
+		}
+		if got, found := store.LoadAndDelete(state); found {
+			t.Fatalf("replayed state should not resolve, got %q", got)
+		}
+	})
+
+	t.Run("unknown state is rejected", func(t *testing.T) {
+		store := NewPKCEVerifierStore(context.Background(), testCacheStore(t))
+
+		if got, found := store.LoadAndDelete("never-stored"); found {
+			t.Fatalf("unknown state should not resolve, got %q", got)
+		}
+	})
+
+	t.Run("expired verifier is rejected", func(t *testing.T) {
+		store := NewPKCEVerifierStore(context.Background(), testCacheStore(t))
+		if err := store.Store(state, verifier, 50*time.Millisecond); err != nil {
+			t.Fatalf("couldn't store PKCE verifier: %s", err)
+		}
+
+		time.Sleep(100 * time.Millisecond)
+		if got, found := store.LoadAndDelete(state); found {
+			t.Fatalf("expired verifier should not resolve, got %q", got)
+		}
+	})
+}

management/internals/shared/grpc/proxy_test.go

@@ -9,7 +9,6 @@ import (
 	"testing"
 	"time"
 
-	cachestore "github.com/eko/gocache/lib/v4/store"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc/codes"
@@ -21,7 +20,7 @@ import (
 	"github.com/netbirdio/netbird/shared/management/proto"
 )
 
-func testCacheStore(t *testing.T) cachestore.StoreInterface {
+func testCacheStore(t *testing.T) nbcache.Store {
 	t.Helper()
 	s, err := nbcache.NewStore(context.Background(), 30*time.Minute, 10*time.Minute, 100)
 	require.NoError(t, err)