consume PKCE verifiers atomically via a GetDel cache op

management/server/cache/memory.go

@@ -2,6 +2,8 @@ package cache
 
 import (
 	"context"
+	"fmt"
+	"sync"
 	"time"
 
 	"github.com/eko/gocache/lib/v4/store"
@@ -12,6 +14,7 @@ import (
 type goCacheStore struct {
 	store.StoreInterface
 	client *gocache.Cache
+	mu     sync.Mutex
 }
 
 func newMemoryStore(maxTimeout, cleanupInterval time.Duration) Store {
@@ -25,7 +28,24 @@ func newMemoryStore(maxTimeout, cleanupInterval time.Duration) Store {
 func (s *goCacheStore) SetNX(_ context.Context, key, value string, ttl time.Duration) (bool, error) {
 	// Add only returns an error when a non-expired entry already exists.
 	if err := s.client.Add(key, value, ttl); err != nil {
-		return false, nil
+		return false, nil //nolint:nilerr
 	}
 	return true, nil
 }
+
+func (s *goCacheStore) GetDel(_ context.Context, key string) (string, bool, error) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	value, found := s.client.Get(key)
+	if !found {
+		return "", false, nil
+	}
+	s.client.Delete(key)
+
+	str, ok := value.(string)
+	if !ok {
+		return "", false, fmt.Errorf("cached value is %T, not a string", value)
+	}
+	return str, true, nil
+}

management/server/cache/memory_test.go

@@ -47,3 +47,44 @@ func TestMemoryStore(t *testing.T) {
 		t.Error("value should not be found")
 	}
 }
+
+func TestMemoryStoreGetDel(t *testing.T) {
+	ctx := context.Background()
+	newStore := func(t *testing.T) cache.Store {
+		t.Helper()
+		memStore, err := cache.NewStore(ctx, time.Minute, time.Minute, 100)
+		if err != nil {
+			t.Fatalf("couldn't create memory store: %s", err)
+		}
+		return memStore
+	}
+
+	const (
+		key   = "consume"
+		value = "verifier"
+	)
+
+	t.Run("exactly one concurrent caller consumes the key", func(t *testing.T) {
+		memStore := newStore(t)
+		if err := memStore.Set(ctx, key, value); err != nil {
+			t.Fatalf("couldn't set testing data: %s", err)
+		}
+
+		assertGetDelConsumedOnce(ctx, t, []cache.Store{memStore}, key, value)
+		assertGetDelMisses(ctx, t, memStore, key)
+	})
+
+	t.Run("missing key is not an error", func(t *testing.T) {
+		assertGetDelMisses(ctx, t, newStore(t), "never-set")
+	})
+
+	t.Run("expired key is not found", func(t *testing.T) {
+		memStore := newStore(t)
+		if _, err := memStore.SetNX(ctx, key, value, 50*time.Millisecond); err != nil {
+			t.Fatalf("couldn't set testing data: %s", err)
+		}
+
+		time.Sleep(100 * time.Millisecond)
+		assertGetDelMisses(ctx, t, memStore, key)
+	})
+}

management/server/cache/redis.go

@@ -2,6 +2,7 @@ package cache
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"math"
 	"time"
@@ -49,3 +50,14 @@ func getRedisStore(ctx context.Context, redisEnvAddr string, maxConn int) (Store
 func (s *redisStore) SetNX(ctx context.Context, key, value string, ttl time.Duration) (bool, error) {
 	return s.client.SetNX(ctx, key, value, ttl).Result()
 }
+
+func (s *redisStore) GetDel(ctx context.Context, key string) (string, bool, error) {
+	value, err := s.client.GetDel(ctx, key).Result()
+	if errors.Is(err, redis.Nil) {
+		return "", false, nil
+	}
+	if err != nil {
+		return "", false, err
+	}
+	return value, true, nil
+}

management/server/cache/redis_test.go

@@ -7,11 +7,41 @@ import (
 
 	"github.com/eko/gocache/lib/v4/store"
 	"github.com/redis/go-redis/v9"
+	"github.com/stretchr/testify/require"
 	testcontainersredis "github.com/testcontainers/testcontainers-go/modules/redis"
 
 	"github.com/netbirdio/netbird/management/server/cache"
 )
 
+func startRedis(t *testing.T) string {
+	t.Helper()
+
+	ctx := context.Background()
+	redisContainer, err := testcontainersredis.Run(ctx, "redis:7")
+	require.NoError(t, err, "couldn't start redis container")
+
+	t.Cleanup(func() {
+		if err := redisContainer.Terminate(ctx); err != nil {
+			t.Logf("failed to terminate container: %s", err)
+		}
+	})
+
+	redisURL, err := redisContainer.ConnectionString(ctx)
+	require.NoError(t, err, "couldn't get connection string")
+
+	t.Setenv(cache.RedisStoreEnvVar, redisURL)
+	return redisURL
+}
+
+func newRedisStore(t *testing.T) cache.Store {
+	t.Helper()
+
+	redisStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100)
+	require.NoError(t, err)
+
+	return redisStore
+}
+
 func TestRedisStoreConnectionFailure(t *testing.T) {
 	t.Setenv(cache.RedisStoreEnvVar, "redis://127.0.0.1:6379")
 	_, err := cache.NewStore(context.Background(), 10*time.Millisecond, 30*time.Millisecond, 100)
@@ -22,28 +52,11 @@ func TestRedisStoreConnectionFailure(t *testing.T) {
 
 func TestRedisStoreConnectionSuccess(t *testing.T) {
 	ctx := context.Background()
-	redisContainer, err := testcontainersredis.Run(ctx, "redis:7")
-	if err != nil {
-		t.Fatalf("couldn't start redis container: %s", err)
-	}
-	defer func() {
-		if err := redisContainer.Terminate(ctx); err != nil {
-			t.Logf("failed to terminate container: %s", err)
-		}
-	}()
-	redisURL, err := redisContainer.ConnectionString(ctx)
-	if err != nil {
-		t.Fatalf("couldn't get connection string: %s", err)
-	}
-
-	t.Setenv(cache.RedisStoreEnvVar, redisURL)
-	redisStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100)
-	if err != nil {
-		t.Fatalf("couldn't create redis store: %s", err)
-	}
+	redisURL := startRedis(t)
+	redisStore := newRedisStore(t)
 
 	key, value := "testing", "tested"
-	err = redisStore.Set(ctx, key, value, store.WithExpiration(100*time.Millisecond))
+	err := redisStore.Set(ctx, key, value, store.WithExpiration(100*time.Millisecond))
 	if err != nil {
 		t.Errorf("couldn't set testing data: %s", err)
 	}
@@ -69,10 +82,24 @@ func TestRedisStoreConnectionSuccess(t *testing.T) {
 		t.Errorf("value returned from redis doesn't match testing data, got %s, expected %s", r, value)
 	}
 
-	secondRedisStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100)
-	if err != nil {
-		t.Fatalf("couldn't create second redis store: %s", err)
+	// test expiration
+	time.Sleep(300 * time.Millisecond)
+	_, err = redisStore.Get(ctx, key)
+	if err == nil {
+		t.Error("value should not be found")
 	}
+}
+
+func TestRedisStoreSetNX(t *testing.T) {
+	ctx := context.Background()
+	redisURL := startRedis(t)
+	redisStore, secondRedisStore := newRedisStore(t), newRedisStore(t)
+
+	const (
+		key   = "atomic"
+		value = "tested"
+	)
+
 	start := make(chan struct{})
 	type setResult struct {
 		created bool
@@ -82,7 +109,7 @@ func TestRedisStoreConnectionSuccess(t *testing.T) {
 	for _, cacheStore := range []cache.Store{redisStore, secondRedisStore} {
 		go func() {
 			<-start
-			created, err := cacheStore.SetNX(ctx, "atomic", value, time.Second)
+			created, err := cacheStore.SetNX(ctx, key, value, time.Minute)
 			results <- setResult{created: created, err: err}
 		}()
 	}
@@ -101,18 +128,51 @@ func TestRedisStoreConnectionSuccess(t *testing.T) {
 	if created != 1 {
 		t.Fatalf("expected exactly one redis client to create the entry, got %d", created)
 	}
-	ttl, err := redisClient.PTTL(ctx, "atomic").Result()
+
+	options, err := redis.ParseURL(redisURL)
+	if err != nil {
+		t.Fatalf("parsing redis cache url: %s", err)
+	}
+	ttl, err := redis.NewClient(options).PTTL(ctx, key).Result()
 	if err != nil {
 		t.Fatalf("couldn't read atomic entry TTL: %s", err)
 	}
 	if ttl <= 0 {
 		t.Fatalf("atomic entry should have a positive TTL, got %s", ttl)
 	}
+}
 
-	// test expiration
-	time.Sleep(300 * time.Millisecond)
-	_, err = redisStore.Get(ctx, key)
-	if err == nil {
-		t.Error("value should not be found")
-	}
+func TestRedisStoreGetDel(t *testing.T) {
+	ctx := context.Background()
+	startRedis(t)
+	redisStore, secondRedisStore := newRedisStore(t), newRedisStore(t)
+
+	const (
+		key   = "consume"
+		value = "verifier"
+	)
+
+	t.Run("exactly one caller across independent clients consumes the key", func(t *testing.T) {
+		// A generous TTL: the key is consumed explicitly, so expiry racing the
+		// concurrent callers would only make the test flaky on a loaded runner.
+		if err := redisStore.Set(ctx, key, value, store.WithExpiration(time.Minute)); err != nil {
+			t.Fatalf("couldn't set value to consume: %s", err)
+		}
+
+		assertGetDelConsumedOnce(ctx, t, []cache.Store{redisStore, secondRedisStore}, key, value)
+		assertGetDelMisses(ctx, t, secondRedisStore, key)
+	})
+
+	t.Run("missing key is not an error", func(t *testing.T) {
+		assertGetDelMisses(ctx, t, redisStore, "never-set")
+	})
+
+	t.Run("expired key is not found", func(t *testing.T) {
+		if err := redisStore.Set(ctx, key, value, store.WithExpiration(50*time.Millisecond)); err != nil {
+			t.Fatalf("couldn't set value to consume: %s", err)
+		}
+
+		time.Sleep(100 * time.Millisecond)
+		assertGetDelMisses(ctx, t, redisStore, key)
+	})
 }

management/server/cache/store.go

@@ -24,11 +24,13 @@ const (
 	DefaultStoreMaxConn = 1000
 )
 
-// Store extends the shared cache interface with atomic insertion support.
+// Store extends the shared cache interface with conditional and consuming operations.
 type Store interface {
 	store.StoreInterface
-	// SetNX atomically stores a value with a TTL only when the key does not exist.
+	// SetNX stores a value with a TTL only when the key does not exist.
 	SetNX(ctx context.Context, key, value string, ttl time.Duration) (bool, error)
+	// GetDel reads a value and removes it, so only one caller can consume a key.
+	GetDel(ctx context.Context, key string) (value string, found bool, err error)
 }
 
 // NewStore creates a new cache store with the given max timeout and cleanup interval. It checks for the environment Variable RedisStoreEnvVar

management/server/cache/store_test.go

@@ -0,0 +1,55 @@
+package cache_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/server/cache"
+)
+
+func assertGetDelConsumedOnce(ctx context.Context, t *testing.T, stores []cache.Store, key, value string) {
+	t.Helper()
+
+	const getDelAttempts = 64
+
+	type getDelResult struct {
+		value string
+		found bool
+		err   error
+	}
+
+	start := make(chan struct{})
+	results := make(chan getDelResult, getDelAttempts)
+	for i := range getDelAttempts {
+		cacheStore := stores[i%len(stores)]
+		go func() {
+			<-start
+			value, found, err := cacheStore.GetDel(ctx, key)
+			results <- getDelResult{value: value, found: found, err: err}
+		}()
+	}
+	close(start)
+
+	consumers := 0
+	for range getDelAttempts {
+		result := <-results
+		require.NoError(t, result.err, "concurrent GetDel failed")
+		if !result.found {
+			continue
+		}
+		consumers++
+		require.Equal(t, value, result.value, "consumed value doesn't match testing data")
+	}
+	require.Equal(t, 1, consumers, "expected exactly one consumer")
+}
+
+func assertGetDelMisses(ctx context.Context, t *testing.T, cacheStore cache.Store, key string) {
+	t.Helper()
+
+	value, found, err := cacheStore.GetDel(ctx, key)
+	require.NoError(t, err, "GetDel on a missing key should not error")
+	require.False(t, found, "GetDel should not find key %q, got value %q", key, value)
+	require.Empty(t, value, "GetDel should return an empty value when not found")
+}