diff --git a/management/internals/modules/reverseproxy/service/manager/manager_test.go b/management/internals/modules/reverseproxy/service/manager/manager_test.go index 29a117921..14de05fd4 100644 --- a/management/internals/modules/reverseproxy/service/manager/manager_test.go +++ b/management/internals/modules/reverseproxy/service/manager/manager_test.go @@ -7,7 +7,6 @@ import ( "testing" "time" - cachestore "github.com/eko/gocache/lib/v4/store" "github.com/golang/mock/gomock" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -31,7 +30,7 @@ import ( "github.com/netbirdio/netbird/shared/management/status" ) -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) @@ -295,6 +294,7 @@ func TestPersistNewService(t *testing.T) { assert.Equal(t, status.AlreadyExists, sErr.Type()) }) } + func TestPreserveExistingAuthSecrets(t *testing.T) { mgr := &Manager{} diff --git a/management/internals/shared/grpc/pkce_verifier.go b/management/internals/shared/grpc/pkce_verifier.go index a1325256c..18155dc1d 100644 --- a/management/internals/shared/grpc/pkce_verifier.go +++ b/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 diff --git a/management/internals/shared/grpc/pkce_verifier_test.go b/management/internals/shared/grpc/pkce_verifier_test.go new file mode 100644 index 000000000..e7175b6c5 --- /dev/null +++ b/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) + } + }) +} diff --git a/management/internals/shared/grpc/proxy_test.go b/management/internals/shared/grpc/proxy_test.go index 0379edc6d..29b7c9523 100644 --- a/management/internals/shared/grpc/proxy_test.go +++ b/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) diff --git a/management/server/cache/memory.go b/management/server/cache/memory.go index ba3b4cc89..62ab62c74 100644 --- a/management/server/cache/memory.go +++ b/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 +} diff --git a/management/server/cache/memory_test.go b/management/server/cache/memory_test.go index 226da8a45..16fb8c27c 100644 --- a/management/server/cache/memory_test.go +++ b/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) + }) +} diff --git a/management/server/cache/redis.go b/management/server/cache/redis.go index 15b936d06..0cd921c92 100644 --- a/management/server/cache/redis.go +++ b/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 +} diff --git a/management/server/cache/redis_test.go b/management/server/cache/redis_test.go index 6e663e790..ae2265057 100644 --- a/management/server/cache/redis_test.go +++ b/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) + }) } diff --git a/management/server/cache/store.go b/management/server/cache/store.go index 63c3e0bb7..a0c093e5d 100644 --- a/management/server/cache/store.go +++ b/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 diff --git a/management/server/cache/store_test.go b/management/server/cache/store_test.go new file mode 100644 index 000000000..a59be8393 --- /dev/null +++ b/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") +}