address review feedback on cache and session store

Split the memory and Redis cache implementations into separate files and
provide backend-native atomic set-if-absent support.

client/firewall/firewalld/firewalld.go

@@ -2,8 +2,8 @@
 // its wg interface into firewalld's "trusted" zone. This is required because
 // firewalld's nftables chains are created with NFT_CHAIN_OWNER on recent
 // versions, which returns EPERM to any other process that tries to insert
-// rules into them. Trusting the interface makes firewalld itself add the
-// accept rules to its own chains instead.
+// rules into them. The workaround mirrors what Tailscale does: let firewalld
+// itself add the accept rules to its own chains by trusting the interface.
 package firewalld
 
 // TrustedZone is the firewalld zone name used for interfaces whose traffic

go.mod

@@ -239,6 +239,7 @@ require (
 	github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect
 	github.com/jackc/puddle/v2 v2.2.1 // indirect
 	github.com/jackpal/go-nat-pmp v1.0.2 // indirect
+	github.com/jchv/go-winloader v0.0.0-20250406163304-c1995be93bd1 // indirect
 	github.com/jinzhu/inflection v1.0.0 // indirect
 	github.com/jinzhu/now v1.1.5 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
@@ -339,4 +340,4 @@ replace github.com/dexidp/dex/api/v2 => github.com/netbirdio/dex/api/v2 v2.0.0-2
 
 replace github.com/mailru/easyjson => github.com/netbirdio/easyjson v0.9.0
 
-replace github.com/wailsapp/wails/v3 => github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20260807062023-6c9243052a40
+replace github.com/wailsapp/wails/v3 => github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20260803205919-ad21e92381f4

go.sum

@@ -349,6 +349,8 @@ github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk
 github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
 github.com/jackpal/go-nat-pmp v1.0.2 h1:KzKSgb7qkJvOUTqYl9/Hg/me3pWgBmERKrTGD7BdWus=
 github.com/jackpal/go-nat-pmp v1.0.2/go.mod h1:QPH045xvCAeXUZOxsnwmrtiCoxIr9eob+4orBN1SBKc=
+github.com/jchv/go-winloader v0.0.0-20250406163304-c1995be93bd1 h1:njuLRcjAuMKr7kI3D85AXWkw6/+v9PwtV6M6o11sWHQ=
+github.com/jchv/go-winloader v0.0.0-20250406163304-c1995be93bd1/go.mod h1:alcuEEnZsY1WQsagKhZDsoPCRoOijYqhZvPwLG0kzVs=
 github.com/jcmturner/aescts/v2 v2.0.0 h1:9YKLH6ey7H4eDBXW8khjYslgyqG2xZikXP0EQFKrle8=
 github.com/jcmturner/aescts/v2 v2.0.0/go.mod h1:AiaICIRyfYg35RUkr8yESTqvSy7csK90qZ5xfvvsoNs=
 github.com/jcmturner/dnsutils/v2 v2.0.0 h1:lltnkeZGL0wILNvrNiVCR6Ro5PGU/SeBvVO/8c/iPbo=
@@ -488,8 +490,8 @@ github.com/netbirdio/service v0.0.0-20240911161631-f62744f42502 h1:3tHlFmhTdX9ax
 github.com/netbirdio/service v0.0.0-20240911161631-f62744f42502/go.mod h1:CIMRFEJVL+0DS1a3Nx06NaMn4Dz63Ng6O7dl0qH0zVM=
 github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20250805121659-6b4ac470ca45 h1:ujgviVYmx243Ksy7NdSwrdGPSRNE3pb8kEDSpH0QuAQ=
 github.com/netbirdio/signal-dispatcher/dispatcher v0.0.0-20250805121659-6b4ac470ca45/go.mod h1:5/sjFmLb8O96B5737VCqhHyGRzNFIaN/Bu7ZodXc3qQ=
-github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20260807062023-6c9243052a40 h1:Oz7axni2nYnoXA0P88r76pphshjSS6kgTjwuk++Tf0M=
-github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20260807062023-6c9243052a40/go.mod h1:bsdahLwBQxXjlmdPPeQyrTcDJfcqAr/ymFj0RXhwtWI=
+github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20260803205919-ad21e92381f4 h1:UKztc3QjWvzU5DZk+uYaOWN0x62NSe/pkxuPvzqZIy4=
+github.com/netbirdio/wails/v3 v3.0.0-beta.3.0.20260803205919-ad21e92381f4/go.mod h1:BzATbK71VFikMMMCo434wAi0QcaI03P+xeaWgDvQvjw=
 github.com/netbirdio/wireguard-go v0.0.0-20260628102922-2834bebf6c1a h1:3CWK+yTvRKOcC0Q8VCTGy4l60TEb27CQVS7LkMxwjmw=
 github.com/netbirdio/wireguard-go v0.0.0-20260628102922-2834bebf6c1a/go.mod h1:rpwXGsirqLqN2L0JDJQlwOboGHmptD5ZD6T2VmcqhTw=
 github.com/nxadm/tail v1.4.4/go.mod h1:kenIhsEOeOJmVchQTgglprH7qJGnHDVpk1VPCcaMI8A=
@@ -791,6 +793,7 @@ golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20191005200804-aed5e4c7ecf9/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20191120155948-bd437916bb0e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200323222414-85ca7c5b95cd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.0.0-20200810151505-1b9f1253b3ed/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201015000850-e3ed0017c211/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

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{}
 

management/internals/server/boot.go

@@ -20,8 +20,6 @@ import (
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/keepalive"
 
-	cachestore "github.com/eko/gocache/lib/v4/store"
-
 	"github.com/netbirdio/netbird/encryption"
 	"github.com/netbirdio/netbird/formatter/hook"
 	"github.com/netbirdio/netbird/management/internals/modules/agentnetwork"
@@ -70,8 +68,8 @@ func (s *BaseServer) Metrics() telemetry.AppMetrics {
 
 // CacheStore returns a shared cache store backed by Redis or in-memory depending on the environment.
 // All consumers should reuse this store to avoid creating multiple Redis connections.
-func (s *BaseServer) CacheStore() cachestore.StoreInterface {
-	return Create(s, func() cachestore.StoreInterface {
+func (s *BaseServer) CacheStore() nbcache.Store {
+	return Create(s, func() nbcache.Store {
 		cs, err := nbcache.NewStore(context.Background(), nbcache.DefaultStoreMaxTimeout, nbcache.DefaultStoreCleanupInterval, nbcache.DefaultStoreMaxConn)
 		if err != nil {
 			log.Fatalf("failed to create shared cache store: %v", err)

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)

management/server/auth/session.go

@@ -7,9 +7,6 @@ import (
 	"errors"
 	"fmt"
 	"time"
-
-	"github.com/eko/gocache/lib/v4/cache"
-	"github.com/eko/gocache/lib/v4/store"
 )
 
 const (
@@ -22,12 +19,17 @@ var (
 	ErrTokenExpired     = errors.New("JWT expired")
 )
 
-type SessionStore struct {
-	cache *cache.Cache[string]
+// TokenCache atomically records used JWTs until their expiration.
+type TokenCache interface {
+	SetNX(ctx context.Context, key, value string, ttl time.Duration) (bool, error)
 }
 
-func NewSessionStore(cacheStore store.StoreInterface) *SessionStore {
-	return &SessionStore{cache: cache.New[string](cacheStore)}
+type SessionStore struct {
+	cache TokenCache
+}
+
+func NewSessionStore(cacheStore TokenCache) *SessionStore {
+	return &SessionStore{cache: cacheStore}
 }
 
 // RegisterToken records a JWT until its exp time and rejects reuse.
@@ -38,20 +40,14 @@ func (s *SessionStore) RegisterToken(ctx context.Context, token string, expiresA
 	}
 
 	key := usedTokenKeyPrefix + hashToken(token)
-	_, err := s.cache.Get(ctx, key)
-	if err == nil {
+	created, err := s.cache.SetNX(ctx, key, usedTokenMarker, ttl)
+	if err != nil {
+		return fmt.Errorf("store used token entry: %w", err)
+	}
+	if !created {
 		return ErrTokenAlreadyUsed
 	}
 
-	var notFound *store.NotFound
-	if !errors.As(err, &notFound) {
-		return fmt.Errorf("failed to lookup used token entry: %w", err)
-	}
-
-	if err := s.cache.Set(ctx, key, usedTokenMarker, store.WithExpiration(ttl)); err != nil {
-		return fmt.Errorf("failed to store used token entry: %w", err)
-	}
-
 	return nil
 }
 

management/server/auth/session_test.go

@@ -2,6 +2,7 @@ package auth
 
 import (
 	"context"
+	"errors"
 	"testing"
 	"time"
 
@@ -38,6 +39,39 @@ func TestSessionStore_RegisterSameTokenTwiceIsRejected(t *testing.T) {
 	assert.ErrorIs(t, err, ErrTokenAlreadyUsed)
 }
 
+func TestSessionStore_ConcurrentRegistrationAllowsOneCaller(t *testing.T) {
+	s := newTestSessionStore(t)
+	ctx := context.Background()
+	const attempts = 100
+
+	start := make(chan struct{})
+	results := make(chan error, attempts)
+	for range attempts {
+		go func() {
+			<-start
+			results <- s.RegisterToken(ctx, "token", time.Now().Add(time.Hour))
+		}()
+	}
+	close(start)
+
+	succeeded := 0
+	alreadyUsed := 0
+	for range attempts {
+		err := <-results
+		switch {
+		case err == nil:
+			succeeded++
+		case errors.Is(err, ErrTokenAlreadyUsed):
+			alreadyUsed++
+		default:
+			require.NoError(t, err, "concurrent registration returned an unexpected error")
+		}
+	}
+
+	assert.Equal(t, 1, succeeded, "exactly one concurrent caller should register the token")
+	assert.Equal(t, attempts-1, alreadyUsed, "every other caller should be rejected as already used")
+}
+
 func TestSessionStore_RegisterDifferentTokensAreIndependent(t *testing.T) {
 	s := newTestSessionStore(t)
 	ctx := context.Background()
@@ -72,6 +106,23 @@ func TestSessionStore_EntryEvictsAtTTLAndAllowsReRegistration(t *testing.T) {
 	require.NoError(t, s.RegisterToken(ctx, token, time.Now().Add(time.Hour)))
 }
 
+type failingTokenCache struct {
+	err error
+}
+
+func (f failingTokenCache) SetNX(context.Context, string, string, time.Duration) (bool, error) {
+	return false, f.err
+}
+
+func TestSessionStore_CacheErrorIsReturned(t *testing.T) {
+	cacheErr := errors.New("cache unavailable")
+	s := NewSessionStore(failingTokenCache{err: cacheErr})
+
+	err := s.RegisterToken(context.Background(), "token", time.Now().Add(time.Hour))
+	require.Error(t, err, "cache failure should be surfaced to the caller")
+	assert.ErrorIs(t, err, cacheErr, "cache error should be wrapped, not replaced")
+}
+
 func TestHashToken_StableAndDoesNotLeak(t *testing.T) {
 	a := hashToken("tokenA")
 	b := hashToken("tokenB")

management/server/cache/memory.go

@@ -0,0 +1,57 @@
+package cache
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/eko/gocache/lib/v4/store"
+	gocachestore "github.com/eko/gocache/store/go_cache/v4"
+	gocache "github.com/patrickmn/go-cache"
+)
+
+type goCacheStore struct {
+	store.StoreInterface
+	client *gocache.Cache
+	mu     sync.Mutex
+}
+
+func newMemoryStore(maxTimeout, cleanupInterval time.Duration) Store {
+	client := gocache.New(maxTimeout, cleanupInterval)
+	return &goCacheStore{
+		StoreInterface: gocachestore.NewGoCache(client),
+		client:         client,
+	}
+}
+
+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 //nolint:nilerr
+	}
+	return true, nil
+}
+
+// GetDel reads the value under key and removes it. go-cache has no native read-and-delete
+// and releases its own lock between the two calls, so mu holds the pair together and no
+// value is consumed twice.
+//
+// Writes do not take mu: a Set landing mid-pair is lost, since GetDel returns the prior
+// value and deletes the new one. Callers must write a consumed key only once.
+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

@@ -0,0 +1,76 @@
+package cache_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/netbirdio/netbird/management/server/cache"
+)
+
+func TestMemoryStore(t *testing.T) {
+	memStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100)
+	require.NoError(t, err, "couldn't create memory store")
+
+	ctx := context.Background()
+	key, value := "testing", "tested"
+	err = memStore.Set(ctx, key, value)
+	assert.NoError(t, err, "couldn't set testing data")
+
+	result, err := memStore.Get(ctx, key)
+	assert.NoError(t, err, "couldn't get testing data")
+	assert.Equal(t, value, result, "value returned doesn't match testing data")
+
+	created, err := memStore.SetNX(ctx, "conditional", value, 100*time.Millisecond)
+	require.NoError(t, err, "couldn't conditionally set testing data")
+	require.True(t, created, "first conditional set should create the entry")
+
+	created, err = memStore.SetNX(ctx, "conditional", value, 100*time.Millisecond)
+	require.NoError(t, err, "couldn't conditionally check testing data")
+	require.False(t, created, "second conditional set should not replace the entry")
+
+	// test expiration
+	time.Sleep(300 * time.Millisecond)
+	_, err = memStore.Get(ctx, key)
+	assert.Error(t, err, "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)
+		require.NoError(t, err, "couldn't create memory store")
+
+		return memStore
+	}
+
+	const (
+		key   = "consume"
+		value = "verifier"
+	)
+
+	t.Run("exactly one concurrent caller consumes the key", func(t *testing.T) {
+		memStore := newStore(t)
+		require.NoError(t, memStore.Set(ctx, key, value), "couldn't set testing data")
+
+		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)
+		_, err := memStore.SetNX(ctx, key, value, 50*time.Millisecond)
+		require.NoError(t, err, "couldn't set testing data")
+
+		time.Sleep(100 * time.Millisecond)
+		assertGetDelMisses(ctx, t, memStore, key)
+	})
+}

management/server/cache/redis.go

@@ -0,0 +1,63 @@
+package cache
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"math"
+	"time"
+
+	"github.com/eko/gocache/lib/v4/store"
+	redisstore "github.com/eko/gocache/store/redis/v4"
+	"github.com/redis/go-redis/v9"
+	log "github.com/sirupsen/logrus"
+)
+
+type redisStore struct {
+	store.StoreInterface
+	client *redis.Client
+}
+
+func getRedisStore(ctx context.Context, redisEnvAddr string, maxConn int) (Store, error) {
+	options, err := redis.ParseURL(redisEnvAddr)
+	if err != nil {
+		return nil, fmt.Errorf("parsing redis cache url: %s", err)
+	}
+
+	options.MaxIdleConns = int(math.Ceil(float64(maxConn) * 0.5)) // 50% of max conns
+	options.MinIdleConns = int(math.Ceil(float64(maxConn) * 0.1)) // 10% of max conns
+	options.MaxActiveConns = maxConn
+	options.ConnMaxIdleTime = 30 * time.Minute
+	options.ConnMaxLifetime = 0
+	options.PoolTimeout = 10 * time.Second
+	redisClient := redis.NewClient(options)
+	subCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
+	defer cancel()
+
+	_, err = redisClient.Ping(subCtx).Result()
+	if err != nil {
+		return nil, err
+	}
+
+	log.WithContext(subCtx).Infof("using redis cache at %s", redisEnvAddr)
+
+	return &redisStore{
+		StoreInterface: redisstore.NewRedis(redisClient),
+		client:         redisClient,
+	}, nil
+}
+
+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

@@ -0,0 +1,153 @@
+package cache_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/eko/gocache/lib/v4/store"
+	"github.com/redis/go-redis/v9"
+	"github.com/stretchr/testify/assert"
+	"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)
+	require.Error(t, err, "getting redis cache store should return error")
+}
+
+func TestRedisStoreConnectionSuccess(t *testing.T) {
+	ctx := context.Background()
+	redisURL := startRedis(t)
+	redisStore := newRedisStore(t)
+
+	key, value := "testing", "tested"
+	err := redisStore.Set(ctx, key, value, store.WithExpiration(100*time.Millisecond))
+	assert.NoError(t, err, "couldn't set testing data")
+
+	result, err := redisStore.Get(ctx, key)
+	assert.NoError(t, err, "couldn't get testing data")
+	assert.Equal(t, value, result, "value returned doesn't match testing data")
+
+	options, err := redis.ParseURL(redisURL)
+	require.NoError(t, err, "parsing redis cache url")
+
+	redisClient := redis.NewClient(options)
+	r, err := redisClient.Get(ctx, key).Result()
+	assert.NoError(t, err, "couldn't get testing data from redis")
+	assert.Equal(t, value, r, "value returned from redis doesn't match testing data")
+
+	// test expiration
+	time.Sleep(300 * time.Millisecond)
+	_, err = redisStore.Get(ctx, key)
+	assert.Error(t, err, "value should not be found")
+}
+
+func TestRedisStoreSetNX(t *testing.T) {
+	ctx := context.Background()
+	redisURL := startRedis(t)
+	redisStore, secondRedisStore := newRedisStore(t), newRedisStore(t)
+
+	const (
+		key   = "conditional"
+		value = "tested"
+	)
+
+	start := make(chan struct{})
+	type setResult struct {
+		created bool
+		err     error
+	}
+	results := make(chan setResult, 2)
+	for _, cacheStore := range []cache.Store{redisStore, secondRedisStore} {
+		go func() {
+			<-start
+			created, err := cacheStore.SetNX(ctx, key, value, time.Minute)
+			results <- setResult{created: created, err: err}
+		}()
+	}
+	close(start)
+
+	created := 0
+	for range 2 {
+		result := <-results
+		require.NoError(t, result.err, "conditional redis set failed")
+		if result.created {
+			created++
+		}
+	}
+	require.Equal(t, 1, created, "expected exactly one redis client to create the entry")
+
+	options, err := redis.ParseURL(redisURL)
+	require.NoError(t, err, "parsing redis cache url")
+
+	ttl, err := redis.NewClient(options).PTTL(ctx, key).Result()
+	require.NoError(t, err, "couldn't read entry TTL")
+	require.Positive(t, ttl, "created entry should have a positive TTL")
+}
+
+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.
+		err := redisStore.Set(ctx, key, value, store.WithExpiration(time.Minute))
+		require.NoError(t, err, "couldn't set value to consume")
+
+		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) {
+		err := redisStore.Set(ctx, key, value, store.WithExpiration(50*time.Millisecond))
+		require.NoError(t, err, "couldn't set value to consume")
+
+		time.Sleep(100 * time.Millisecond)
+		assertGetDelMisses(ctx, t, redisStore, key)
+	})
+}

management/server/cache/store.go

@@ -2,17 +2,10 @@ package cache
 
 import (
 	"context"
-	"fmt"
-	"math"
 	"os"
 	"time"
 
 	"github.com/eko/gocache/lib/v4/store"
-	gocache_store "github.com/eko/gocache/store/go_cache/v4"
-	redis_store "github.com/eko/gocache/store/redis/v4"
-	gocache "github.com/patrickmn/go-cache"
-	"github.com/redis/go-redis/v9"
-	log "github.com/sirupsen/logrus"
 )
 
 // RedisStoreEnvVar is the environment variable that determines if a redis store should be used.
@@ -31,15 +24,23 @@ const (
 	DefaultStoreMaxConn = 1000
 )
 
+// Store extends the shared cache interface with conditional and consuming operations.
+type Store interface {
+	store.StoreInterface
+	// 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
 // to determine if a redis store should be used. If the environment variable is set, it will attempt to connect to the redis store.
-func NewStore(ctx context.Context, maxTimeout, cleanupInterval time.Duration, maxConn int) (store.StoreInterface, error) {
+func NewStore(ctx context.Context, maxTimeout, cleanupInterval time.Duration, maxConn int) (Store, error) {
 	redisAddr := GetAddrFromEnv()
 	if redisAddr != "" {
 		return getRedisStore(ctx, redisAddr, maxConn)
 	}
-	goc := gocache.New(maxTimeout, cleanupInterval)
-	return gocache_store.NewGoCache(goc), nil
+	return newMemoryStore(maxTimeout, cleanupInterval), nil
 }
 
 // GetAddrFromEnv returns the redis address from the environment variable RedisStoreEnvVar or its legacy counterpart.
@@ -50,29 +51,3 @@ func GetAddrFromEnv() string {
 	}
 	return addr
 }
-
-func getRedisStore(ctx context.Context, redisEnvAddr string, maxConn int) (store.StoreInterface, error) {
-	options, err := redis.ParseURL(redisEnvAddr)
-	if err != nil {
-		return nil, fmt.Errorf("parsing redis cache url: %s", err)
-	}
-
-	options.MaxIdleConns = int(math.Ceil(float64(maxConn) * 0.5)) // 50% of max conns
-	options.MinIdleConns = int(math.Ceil(float64(maxConn) * 0.1)) // 10% of max conns
-	options.MaxActiveConns = maxConn
-	options.ConnMaxIdleTime = 30 * time.Minute
-	options.ConnMaxLifetime = 0
-	options.PoolTimeout = 10 * time.Second
-	redisClient := redis.NewClient(options)
-	subCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
-	defer cancel()
-
-	_, err = redisClient.Ping(subCtx).Result()
-	if err != nil {
-		return nil, err
-	}
-
-	log.WithContext(subCtx).Infof("using redis cache at %s", redisEnvAddr)
-
-	return redis_store.NewRedis(redisClient), nil
-}

management/server/cache/store_test.go

@@ -3,101 +3,53 @@ package cache_test
 import (
 	"context"
 	"testing"
-	"time"
 
-	"github.com/eko/gocache/lib/v4/store"
-	"github.com/redis/go-redis/v9"
-	testcontainersredis "github.com/testcontainers/testcontainers-go/modules/redis"
+	"github.com/stretchr/testify/require"
 
 	"github.com/netbirdio/netbird/management/server/cache"
 )
 
-func TestMemoryStore(t *testing.T) {
-	memStore, err := cache.NewStore(context.Background(), 100*time.Millisecond, 300*time.Millisecond, 100)
-	if err != nil {
-		t.Fatalf("couldn't create memory store: %s", err)
-	}
-	ctx := context.Background()
-	key, value := "testing", "tested"
-	err = memStore.Set(ctx, key, value)
-	if err != nil {
-		t.Errorf("couldn't set testing data: %s", err)
-	}
-	result, err := memStore.Get(ctx, key)
-	if err != nil {
-		t.Errorf("couldn't get testing data: %s", err)
-	}
-	if value != result.(string) {
-		t.Errorf("value returned doesn't match testing data, got %s, expected %s", result, value)
-	}
-	// test expiration
-	time.Sleep(300 * time.Millisecond)
-	_, err = memStore.Get(ctx, key)
-	if err == nil {
-		t.Error("value should not be found")
-	}
-}
+func assertGetDelConsumedOnce(ctx context.Context, t *testing.T, stores []cache.Store, key, value string) {
+	t.Helper()
 
-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)
-	if err == nil {
-		t.Fatal("getting redis cache store should return error")
-	}
-}
+	const getDelAttempts = 64
 
-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)
+	type getDelResult struct {
+		value string
+		found bool
+		err   error
 	}
-	defer func() {
-		if err := redisContainer.Terminate(ctx); err != nil {
-			t.Logf("failed to terminate container: %s", err)
+
+	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
 		}
-	}()
-	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)
-	}
-
-	key, value := "testing", "tested"
-	err = redisStore.Set(ctx, key, value, store.WithExpiration(100*time.Millisecond))
-	if err != nil {
-		t.Errorf("couldn't set testing data: %s", err)
-	}
-	result, err := redisStore.Get(ctx, key)
-	if err != nil {
-		t.Errorf("couldn't get testing data: %s", err)
-	}
-	if value != result.(string) {
-		t.Errorf("value returned doesn't match testing data, got %s, expected %s", result, value)
-	}
-
-	options, err := redis.ParseURL(redisURL)
-	if err != nil {
-		t.Errorf("parsing redis cache url: %s", err)
-	}
-
-	redisClient := redis.NewClient(options)
-	r, e := redisClient.Get(ctx, key).Result()
-	if e != nil {
-		t.Errorf("couldn't get testing data from redis: %s", e)
-	}
-	if value != r {
-		t.Errorf("value returned from redis doesn't match testing data, got %s, expected %s", r, value)
-	}
-	// test expiration
-	time.Sleep(300 * time.Millisecond)
-	_, err = redisStore.Get(ctx, key)
-	if err == nil {
-		t.Error("value should not be found")
+		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")
 }

release_files/darwin_pkg/preinstall

@@ -4,67 +4,9 @@ set -x
 
 LOG_FILE=/var/log/netbird/client_pre_install.log
 AGENT=/usr/local/bin/netbird
-UI_PROCESS=netbird-ui
 
 mkdir -p /var/log/netbird/
 
-# wait_for_ui_exit polls for up to $1 seconds, returning 0 as soon as no UI
-# process is left and 1 if one is still running when the time is up.
-wait_for_ui_exit() {
-    waited=0
-    while [ "$waited" -lt "$1" ]; do
-        pgrep -x "$UI_PROCESS" > /dev/null 2>&1 || return 0
-        sleep 1
-        waited=$((waited + 1))
-    done
-    return 1
-}
-
-# request_ui_quit asks the UI to quit from inside the console user's session and
-# reports whether the request could be sent at all. The installer runs as root
-# outside that session, so a quit Apple event sent straight from here always
-# fails with -600.
-request_ui_quit() {
-    console_user=$(stat -f%Su /dev/console 2>/dev/null)
-    case "$console_user" in
-        ""|root|loginwindow|_mbsetupuser)
-            echo "No active GUI user session (console user: '${console_user:-none}'); skipping the quit request."
-            return 1
-            ;;
-    esac
-
-    uid=$(id -u "$console_user" 2>/dev/null)
-    if [ -z "$uid" ]; then
-        echo "Could not resolve uid for console user '$console_user'; skipping the quit request."
-        return 1
-    fi
-
-    echo "Asking the NetBird UI to quit as console user $console_user (uid $uid)."
-    launchctl asuser "$uid" sudo -u "$console_user" -H osascript -e 'quit app "NetBird"' || true
-}
-
-# quit_ui stops a running UI so the app bundle can be replaced underneath it. A
-# UI process that survives the install keeps serving the old binary until it is
-# quit by hand, so anything still running once the quit request is out of the
-# way is signalled. Waiting for a graceful exit only makes sense when a quit
-# request was actually sent.
-quit_ui() {
-    if request_ui_quit && wait_for_ui_exit 10; then
-        return 0
-    fi
-
-    pgrep -x "$UI_PROCESS" > /dev/null 2>&1 || return 0
-
-    echo "NetBird UI still running; terminating it."
-    pkill -x "$UI_PROCESS" || true
-    if wait_for_ui_exit 3; then
-        return 0
-    fi
-
-    echo "NetBird UI ignored SIGTERM; killing it."
-    pkill -KILL -x "$UI_PROCESS" || true
-}
-
 {
     # check if it was installed with brew
     brew list --formula | grep netbird
@@ -73,9 +15,10 @@ quit_ui() {
         echo "NetBird has been installed with Brew. Please use Brew to update the package."
         exit 1
     fi
-    quit_ui
+    osascript -e 'quit app "Netbird"' || true
     $AGENT service stop || true
 
     echo "Preinstall complete"
     exit 0 # all good
 } &> $LOG_FILE
+