Merge branch 'main' into agent-network-validate-proxy-cluster

management/internals/shared/grpc/peer_update_handler.go

@@ -0,0 +1,135 @@
+package grpc
+
+import (
+	"context"
+	"time"
+
+	"github.com/netbirdio/netbird/encryption"
+	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	"github.com/netbirdio/netbird/management/server/telemetry"
+	"github.com/netbirdio/netbird/shared/management/proto"
+	log "github.com/sirupsen/logrus"
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+func PeerUpdateHandlerFactory(
+	peerKey wgtypes.Key,
+	updates chan *network_map.UpdateMessage,
+	secretsManager SecretsManager,
+	srv proto.ManagementService_SyncServer,
+	cleanupfunc func()) *PeerUpdateHandler {
+	return &PeerUpdateHandler{
+		peerKey:        peerKey,
+		updates:        updates,
+		secretsManager: secretsManager,
+		srv:            srv,
+		encrypter:      encryption.DefaultEncrypter{},
+		debouncer:      NewUpdateDebouncer(1000 * time.Millisecond),
+		cleanupFunc:    cleanupfunc,
+	}
+}
+
+// PeerUpdateHandler sends updates to the connected peer until the updates channel is closed.
+// It implements a backpressure mechanism that sends the first update immediately,
+// then debounces subsequent rapid updates, ensuring only the latest update is sent
+// after a quiet period.
+type PeerUpdateHandler struct {
+	peerKey        wgtypes.Key
+	updates        chan *network_map.UpdateMessage
+	appMetrics     telemetry.AppMetrics
+	secretsManager SecretsManager
+	srv            syncSender
+	encrypter      encryption.Encrypter
+	debouncer      Debouncer
+	cleanupFunc    func()
+}
+
+func (pu *PeerUpdateHandler) WithMetrics(appMetrics telemetry.AppMetrics) *PeerUpdateHandler {
+	pu.appMetrics = appMetrics
+	return pu
+}
+
+//go:generate go tool mockgen -source=./peer_update_handler.go -destination=./sync_sender_mock.go -package=grpc
+type syncSender interface {
+	Send(*proto.EncryptedMessage) error
+	Context() context.Context
+}
+
+func (pu *PeerUpdateHandler) HandleUpdates(ctx context.Context) error {
+	log.WithContext(ctx).Tracef("starting to handle updates for peer %s", pu.peerKey.String())
+
+	defer pu.debouncer.Stop()
+
+	for {
+		select {
+		// condition when there are some updates
+		// todo set the updates channel size to 1
+		case update, open := <-pu.updates:
+			if pu.appMetrics != nil {
+				pu.appMetrics.GRPCMetrics().UpdateChannelQueueLength(len(pu.updates) + 1)
+			}
+
+			if !open {
+				log.WithContext(ctx).Debugf("updates channel for peer %s was closed", pu.peerKey.String())
+				pu.cleanupFunc()
+				return nil
+			}
+
+			log.WithContext(ctx).Tracef("received an update for peer %s", pu.peerKey.String())
+			if pu.debouncer.ProcessUpdate(update) {
+				// Send immediately (first update or after quiet period)
+				if err := pu.SendUpdate(ctx, update); err != nil {
+					log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", pu.peerKey.String(), err)
+					return err
+				}
+			}
+
+		// Timer expired - quiet period reached, send pending updates if any
+		case <-pu.debouncer.TimerChannel():
+			pendingUpdates := pu.debouncer.GetPendingUpdates()
+			if len(pendingUpdates) == 0 {
+				continue
+			}
+			log.WithContext(ctx).Debugf("sending %d debounced update(s) for peer %s", len(pendingUpdates), pu.peerKey.String())
+			for _, pendingUpdate := range pendingUpdates {
+				if err := pu.SendUpdate(ctx, pendingUpdate); err != nil {
+					log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", pu.peerKey.String(), err)
+					return err
+				}
+			}
+
+		// condition when client <-> server connection has been terminated
+		case <-pu.srv.Context().Done():
+			// happens when connection drops, e.g. client disconnects
+			log.WithContext(ctx).Debugf("stream of peer %s has been closed", pu.peerKey.String())
+			pu.cleanupFunc()
+			return pu.srv.Context().Err()
+		}
+	}
+}
+
+func (pu *PeerUpdateHandler) SendUpdate(ctx context.Context, update *network_map.UpdateMessage) error {
+	key, err := pu.secretsManager.GetWGKey()
+	if err != nil {
+		pu.cleanupFunc()
+		return status.Errorf(codes.Internal, "failed processing update message")
+	}
+
+	encryptedResp, err := pu.encrypter.EncryptMessage(pu.peerKey, key, update.Update)
+	if err != nil {
+		pu.cleanupFunc()
+		return status.Errorf(codes.Internal, "failed processing update message")
+	}
+	err = pu.srv.Send(&proto.EncryptedMessage{
+		WgPubKey: key.PublicKey().String(),
+		Body:     encryptedResp,
+	})
+	if err != nil {
+		pu.cleanupFunc()
+		return status.Errorf(codes.Internal, "failed sending update message")
+	}
+	log.WithContext(ctx).Tracef("sent an update to peer %s", pu.peerKey.String())
+	return nil
+}

management/internals/shared/grpc/peer_update_handler_test.go

@@ -0,0 +1,155 @@
+package grpc
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	pb "github.com/golang/protobuf/proto" //nolint
+	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	"github.com/netbirdio/netbird/shared/management/proto"
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/mock/gomock"
+	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+)
+
+func TestSendPeerUpdates_FirstUpdate(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	secretsManager := NewMockSecretsManager(ctrl)
+	updateDebouncer := NewMockDebouncer(ctrl)
+	syncSender := NewMocksyncSender(ctrl)
+
+	pu := PeerUpdateHandler{
+		peerKey:        mustGenerateKey(t),
+		updates:        make(chan *network_map.UpdateMessage),
+		secretsManager: secretsManager,
+		encrypter:      testEncrypter{},
+		debouncer:      updateDebouncer,
+		srv:            syncSender,
+		cleanupFunc:    func() {},
+	}
+
+	msg := network_map.UpdateMessage{
+		Update: &proto.SyncResponse{Version: 1},
+	}
+
+	timeCh := make(chan time.Time)
+	srvCtx := context.TODO()
+	srvKey := mustGenerateKey(t)
+	// mock a first update, should send it right away
+	updateDebouncer.EXPECT().ProcessUpdate(gomock.Eq(&msg)).Return(true)
+	updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh)
+	syncSender.EXPECT().Context().AnyTimes().Return(srvCtx)
+	secretsManager.EXPECT().GetWGKey().Return(srvKey, nil)
+	syncSender.EXPECT().Send(pbMatcher{x: &proto.EncryptedMessage{WgPubKey: srvKey.PublicKey().String(), Body: mustMarshal(t, &msg)}})
+	updateDebouncer.EXPECT().Stop()
+
+	var wg sync.WaitGroup
+	wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck
+	pu.updates <- &msg
+	close(pu.updates)
+	wg.Wait()
+}
+
+func TestSendPeerUpdates_TimerUpdate(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	secretsManager := NewMockSecretsManager(ctrl)
+	updateDebouncer := NewMockDebouncer(ctrl)
+	syncSender := NewMocksyncSender(ctrl)
+
+	pu := PeerUpdateHandler{
+		peerKey:        mustGenerateKey(t),
+		updates:        make(chan *network_map.UpdateMessage),
+		secretsManager: secretsManager,
+		encrypter:      testEncrypter{},
+		debouncer:      updateDebouncer,
+		srv:            syncSender,
+		cleanupFunc:    func() {},
+	}
+
+	msg := network_map.UpdateMessage{
+		Update: &proto.SyncResponse{Version: 1},
+	}
+
+	timeCh := make(chan time.Time)
+	srvCtx := context.TODO()
+	srvKey := mustGenerateKey(t)
+	updateDebouncer.EXPECT().GetPendingUpdates().Return([]*network_map.UpdateMessage{&msg})
+	updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh)
+	syncSender.EXPECT().Context().AnyTimes().Return(srvCtx)
+	secretsManager.EXPECT().GetWGKey().Return(srvKey, nil)
+	syncSender.EXPECT().Send(pbMatcher{x: &proto.EncryptedMessage{WgPubKey: srvKey.PublicKey().String(), Body: mustMarshal(t, &msg)}})
+	updateDebouncer.EXPECT().Stop()
+
+	var wg sync.WaitGroup
+	wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck
+	timeCh <- time.Now()
+	close(pu.updates)
+	wg.Wait()
+}
+
+func TestSendPeerUpdates_ServerContextDone(t *testing.T) {
+	ctrl := gomock.NewController(t)
+	secretsManager := NewMockSecretsManager(ctrl)
+	updateDebouncer := NewMockDebouncer(ctrl)
+	syncSender := NewMocksyncSender(ctrl)
+
+	pu := PeerUpdateHandler{
+		peerKey:        mustGenerateKey(t),
+		updates:        make(chan *network_map.UpdateMessage),
+		secretsManager: secretsManager,
+		encrypter:      testEncrypter{},
+		debouncer:      updateDebouncer,
+		srv:            syncSender,
+		cleanupFunc:    func() {},
+	}
+
+	timeCh := make(chan time.Time)
+	srvCtx, cancel := context.WithCancel(context.TODO())
+	updateDebouncer.EXPECT().TimerChannel().AnyTimes().Return(timeCh)
+	syncSender.EXPECT().Context().AnyTimes().Return(srvCtx)
+	updateDebouncer.EXPECT().Stop()
+
+	var wg sync.WaitGroup
+	wg.Go(func() { pu.HandleUpdates(context.TODO()) }) //nolint:errcheck
+	cancel()
+	wg.Wait()
+}
+
+func mustGenerateKey(t *testing.T) wgtypes.Key {
+	t.Helper()
+	k, err := wgtypes.GenerateKey()
+	assert.NoError(t, err)
+	return k
+}
+
+func mustMarshal(t *testing.T, msg *network_map.UpdateMessage) []byte {
+	t.Helper()
+	r, err := pb.Marshal(msg.Update)
+	assert.NoError(t, err)
+	return r
+}
+
+type testEncrypter struct{}
+
+func (testEncrypter) EncryptMessage(remotePubKey wgtypes.Key, ourPrivateKey wgtypes.Key, message pb.Message) ([]byte, error) {
+	return pb.Marshal(message)
+}
+
+type pbMatcher struct {
+	x pb.Message
+}
+
+func (pbm pbMatcher) Matches(x any) bool {
+	msg, ok := x.(pb.Message)
+	if !ok {
+		return false
+	}
+	return pb.Equal(pbm.x, msg)
+}
+
+func (pbm pbMatcher) String() string {
+	return fmt.Sprintf("is equal to %s (%T)", pbm.x, pbm.x)
+}

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.go

@@ -1681,6 +1681,10 @@ var (
 	// ErrUserBlocked reports a blocked user, who may not hold a proxy session.
 	ErrUserBlocked = errors.New("user blocked")
 
+	// ErrUserNotInGroup reports a user outside the service's distribution
+	// groups, who may not hold a proxy session for it.
+	ErrUserNotInGroup = errors.New("user not in allowed groups")
+
 	errUserUnresolved = errors.New("user could not be resolved")
 )
 
@@ -1719,8 +1723,10 @@ func sameAccount(userAccountID, serviceAccountID string) bool {
 // GenerateSessionToken creates a signed session JWT for the given domain and
 // user. The user's group memberships are embedded in the token so policy-aware
 // middlewares on the proxy can authorise without an extra management round-trip.
-// A user the store cannot resolve, or whose account is pending approval or
-// blocked, gets no token at all, so the browser never receives a session cookie.
+// A user the store cannot resolve, whose account is pending approval or blocked,
+// or who is outside the service's distribution groups, gets no token at all: the
+// token is a bearer credential for the service, so authorisation has to run
+// before it is signed rather than only when the proxy presents it back.
 func (s *ProxyServiceServer) GenerateSessionToken(ctx context.Context, domain, userID string, method proxyauth.Method) (string, error) {
 	service, err := s.getServiceByDomain(ctx, domain)
 	if err != nil {
@@ -1756,6 +1762,14 @@ func (s *ProxyServiceServer) GenerateSessionToken(ctx context.Context, domain, u
 		return "", fmt.Errorf("session token for user %s: %w", userID, err)
 	}
 
+	if err := s.checkGroupAccess(service, user); err != nil {
+		log.WithContext(ctx).WithFields(log.Fields{
+			"domain":  domain,
+			"user_id": userID,
+		}).Debug("GenerateSessionToken: user not in the service's distribution groups")
+		return "", fmt.Errorf("session token for user %s: %w", userID, ErrUserNotInGroup)
+	}
+
 	groupIDs, groupNames := pairGroupIDsAndNames(userGroups)
 
 	token, err := sessionkey.SignToken(

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/internals/shared/grpc/server.go

@@ -247,17 +247,8 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
 	sRealIP := realIP.String()
 	peerMeta := extractPeerMeta(ctx, syncReq.GetMeta())
 
-	userID, err := s.accountManager.GetUserIDByPeerKey(ctx, peerKey.String())
-	if err != nil {
-		s.syncSem.Add(-1)
-		if errStatus, ok := internalStatus.FromError(err); ok && errStatus.Type() == internalStatus.NotFound {
-			return status.Errorf(codes.PermissionDenied, "peer is not registered")
-		}
-		return mapError(ctx, err)
-	}
-
 	metahashed := metaHash(peerMeta)
-	if userID == "" && !s.loginFilter.allowLogin(peerKey.String(), metahashed) {
+	if !s.loginFilter.allowLogin(peerKey.String(), metahashed) {
 		if s.appMetrics != nil {
 			s.appMetrics.GRPCMetrics().CountSyncRequestBlocked()
 		}
@@ -346,7 +337,8 @@ func (s *Server) Sync(req *proto.EncryptedMessage, srv proto.ManagementService_S
 
 	s.syncSem.Add(-1)
 
-	return s.handleUpdates(ctx, accountID, peerKey, peer, updates, srv, syncStart)
+	return PeerUpdateHandlerFactory(peerKey, updates, s.secretsManager, srv, func() { s.cancelPeerRoutines(ctx, accountID, peer, syncStart) }).
+		WithMetrics(s.appMetrics).HandleUpdates(ctx)
 }
 
 func (s *Server) handleHandshake(ctx context.Context, srv proto.ManagementService_JobServer) (wgtypes.Key, error) {
@@ -413,91 +405,6 @@ func (s *Server) sendJobsLoop(ctx context.Context, accountID string, peerKey wgt
 	}
 }
 
-// handleUpdates sends updates to the connected peer until the updates channel is closed.
-// It implements a backpressure mechanism that sends the first update immediately,
-// then debounces subsequent rapid updates, ensuring only the latest update is sent
-// after a quiet period.
-func (s *Server) handleUpdates(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, updates chan *network_map.UpdateMessage, srv proto.ManagementService_SyncServer, streamStartTime time.Time) error {
-	log.WithContext(ctx).Tracef("starting to handle updates for peer %s", peerKey.String())
-
-	// Create a debouncer for this peer connection
-	debouncer := NewUpdateDebouncer(1000 * time.Millisecond)
-	defer debouncer.Stop()
-
-	for {
-		select {
-		// condition when there are some updates
-		// todo set the updates channel size to 1
-		case update, open := <-updates:
-			if s.appMetrics != nil {
-				s.appMetrics.GRPCMetrics().UpdateChannelQueueLength(len(updates) + 1)
-			}
-
-			if !open {
-				log.WithContext(ctx).Debugf("updates channel for peer %s was closed", peerKey.String())
-				s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-				return nil
-			}
-
-			log.WithContext(ctx).Tracef("received an update for peer %s", peerKey.String())
-			if debouncer.ProcessUpdate(update) {
-				// Send immediately (first update or after quiet period)
-				if err := s.sendUpdate(ctx, accountID, peerKey, peer, update, srv, streamStartTime); err != nil {
-					log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", peerKey.String(), err)
-					return err
-				}
-			}
-
-		// Timer expired - quiet period reached, send pending updates if any
-		case <-debouncer.TimerChannel():
-			pendingUpdates := debouncer.GetPendingUpdates()
-			if len(pendingUpdates) == 0 {
-				continue
-			}
-			log.WithContext(ctx).Debugf("sending %d debounced update(s) for peer %s", len(pendingUpdates), peerKey.String())
-			for _, pendingUpdate := range pendingUpdates {
-				if err := s.sendUpdate(ctx, accountID, peerKey, peer, pendingUpdate, srv, streamStartTime); err != nil {
-					log.WithContext(ctx).Debugf("error while sending an update to peer %s: %v", peerKey.String(), err)
-					return err
-				}
-			}
-
-		// condition when client <-> server connection has been terminated
-		case <-srv.Context().Done():
-			// happens when connection drops, e.g. client disconnects
-			log.WithContext(ctx).Debugf("stream of peer %s has been closed", peerKey.String())
-			s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-			return srv.Context().Err()
-		}
-	}
-}
-
-// sendUpdate encrypts the update message using the peer key and the server's wireguard key,
-// then sends the encrypted message to the connected peer via the sync server.
-func (s *Server) sendUpdate(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, update *network_map.UpdateMessage, srv proto.ManagementService_SyncServer, streamStartTime time.Time) error {
-	key, err := s.secretsManager.GetWGKey()
-	if err != nil {
-		s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-		return status.Errorf(codes.Internal, "failed processing update message")
-	}
-
-	encryptedResp, err := encryption.EncryptMessage(peerKey, key, update.Update)
-	if err != nil {
-		s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-		return status.Errorf(codes.Internal, "failed processing update message")
-	}
-	err = srv.Send(&proto.EncryptedMessage{
-		WgPubKey: key.PublicKey().String(),
-		Body:     encryptedResp,
-	})
-	if err != nil {
-		s.cancelPeerRoutines(ctx, accountID, peer, streamStartTime)
-		return status.Errorf(codes.Internal, "failed sending update message")
-	}
-	log.WithContext(ctx).Tracef("sent an update to peer %s", peerKey.String())
-	return nil
-}
-
 // sendJob encrypts the update message using the peer key and the server's wireguard key,
 // then sends the encrypted message to the connected peer via the sync server.
 func (s *Server) sendJob(ctx context.Context, peerKey wgtypes.Key, job *job.Event, srv proto.ManagementService_JobServer) error {
@@ -676,6 +583,7 @@ func extractPeerMeta(ctx context.Context, meta *proto.PeerSystemMeta) nbpeer.Pee
 			RosenpassEnabled:      meta.GetFlags().GetRosenpassEnabled(),
 			RosenpassPermissive:   meta.GetFlags().GetRosenpassPermissive(),
 			ServerSSHAllowed:      meta.GetFlags().GetServerSSHAllowed(),
+			RemoteJobsAllowed:     meta.GetFlags().GetRemoteJobsAllowed(),
 			DisableClientRoutes:   meta.GetFlags().GetDisableClientRoutes(),
 			DisableServerRoutes:   meta.GetFlags().GetDisableServerRoutes(),
 			DisableDNS:            meta.GetFlags().GetDisableDNS(),

management/internals/shared/grpc/sync_sender_mock.go

@@ -0,0 +1,70 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: ./peer_update_handler.go
+//
+// Generated by this command:
+//
+//	mockgen -source=./peer_update_handler.go -destination=./sync_sender_mock.go -package=grpc
+//
+
+// Package grpc is a generated GoMock package.
+package grpc
+
+import (
+	context "context"
+	reflect "reflect"
+
+	proto "github.com/netbirdio/netbird/shared/management/proto"
+	gomock "go.uber.org/mock/gomock"
+)
+
+// MocksyncSender is a mock of syncSender interface.
+type MocksyncSender struct {
+	ctrl     *gomock.Controller
+	recorder *MocksyncSenderMockRecorder
+	isgomock struct{}
+}
+
+// MocksyncSenderMockRecorder is the mock recorder for MocksyncSender.
+type MocksyncSenderMockRecorder struct {
+	mock *MocksyncSender
+}
+
+// NewMocksyncSender creates a new mock instance.
+func NewMocksyncSender(ctrl *gomock.Controller) *MocksyncSender {
+	mock := &MocksyncSender{ctrl: ctrl}
+	mock.recorder = &MocksyncSenderMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MocksyncSender) EXPECT() *MocksyncSenderMockRecorder {
+	return m.recorder
+}
+
+// Context mocks base method.
+func (m *MocksyncSender) Context() context.Context {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Context")
+	ret0, _ := ret[0].(context.Context)
+	return ret0
+}
+
+// Context indicates an expected call of Context.
+func (mr *MocksyncSenderMockRecorder) Context() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Context", reflect.TypeOf((*MocksyncSender)(nil).Context))
+}
+
+// Send mocks base method.
+func (m *MocksyncSender) Send(arg0 *proto.EncryptedMessage) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Send", arg0)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Send indicates an expected call of Send.
+func (mr *MocksyncSenderMockRecorder) Send(arg0 any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Send", reflect.TypeOf((*MocksyncSender)(nil).Send), arg0)
+}

management/internals/shared/grpc/token_mgr.go

@@ -25,6 +25,8 @@ import (
 const defaultDuration = 12 * time.Hour
 
 // SecretsManager used to manage TURN and relay secrets
+//
+//go:generate go tool mockgen -source=./token_mgr.go -destination=./token_mgr_mock.go -package=grpc
 type SecretsManager interface {
 	GenerateTurnToken() (*Token, error)
 	GenerateRelayToken() (*Token, error)

management/internals/shared/grpc/token_mgr_mock.go

@@ -0,0 +1,111 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: ./token_mgr.go
+//
+// Generated by this command:
+//
+//	mockgen -source=./token_mgr.go -destination=./token_mgr_mock.go -package=grpc
+//
+
+// Package grpc is a generated GoMock package.
+package grpc
+
+import (
+	context "context"
+	reflect "reflect"
+
+	gomock "go.uber.org/mock/gomock"
+	wgtypes "golang.zx2c4.com/wireguard/wgctrl/wgtypes"
+)
+
+// MockSecretsManager is a mock of SecretsManager interface.
+type MockSecretsManager struct {
+	ctrl     *gomock.Controller
+	recorder *MockSecretsManagerMockRecorder
+	isgomock struct{}
+}
+
+// MockSecretsManagerMockRecorder is the mock recorder for MockSecretsManager.
+type MockSecretsManagerMockRecorder struct {
+	mock *MockSecretsManager
+}
+
+// NewMockSecretsManager creates a new mock instance.
+func NewMockSecretsManager(ctrl *gomock.Controller) *MockSecretsManager {
+	mock := &MockSecretsManager{ctrl: ctrl}
+	mock.recorder = &MockSecretsManagerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockSecretsManager) EXPECT() *MockSecretsManagerMockRecorder {
+	return m.recorder
+}
+
+// CancelRefresh mocks base method.
+func (m *MockSecretsManager) CancelRefresh(peerKey string) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "CancelRefresh", peerKey)
+}
+
+// CancelRefresh indicates an expected call of CancelRefresh.
+func (mr *MockSecretsManagerMockRecorder) CancelRefresh(peerKey any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "CancelRefresh", reflect.TypeOf((*MockSecretsManager)(nil).CancelRefresh), peerKey)
+}
+
+// GenerateRelayToken mocks base method.
+func (m *MockSecretsManager) GenerateRelayToken() (*Token, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GenerateRelayToken")
+	ret0, _ := ret[0].(*Token)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GenerateRelayToken indicates an expected call of GenerateRelayToken.
+func (mr *MockSecretsManagerMockRecorder) GenerateRelayToken() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GenerateRelayToken", reflect.TypeOf((*MockSecretsManager)(nil).GenerateRelayToken))
+}
+
+// GenerateTurnToken mocks base method.
+func (m *MockSecretsManager) GenerateTurnToken() (*Token, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GenerateTurnToken")
+	ret0, _ := ret[0].(*Token)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GenerateTurnToken indicates an expected call of GenerateTurnToken.
+func (mr *MockSecretsManagerMockRecorder) GenerateTurnToken() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GenerateTurnToken", reflect.TypeOf((*MockSecretsManager)(nil).GenerateTurnToken))
+}
+
+// GetWGKey mocks base method.
+func (m *MockSecretsManager) GetWGKey() (wgtypes.Key, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetWGKey")
+	ret0, _ := ret[0].(wgtypes.Key)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// GetWGKey indicates an expected call of GetWGKey.
+func (mr *MockSecretsManagerMockRecorder) GetWGKey() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetWGKey", reflect.TypeOf((*MockSecretsManager)(nil).GetWGKey))
+}
+
+// SetupRefresh mocks base method.
+func (m *MockSecretsManager) SetupRefresh(ctx context.Context, accountID, peerKey string) {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "SetupRefresh", ctx, accountID, peerKey)
+}
+
+// SetupRefresh indicates an expected call of SetupRefresh.
+func (mr *MockSecretsManagerMockRecorder) SetupRefresh(ctx, accountID, peerKey any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetupRefresh", reflect.TypeOf((*MockSecretsManager)(nil).SetupRefresh), ctx, accountID, peerKey)
+}

management/internals/shared/grpc/update_debouncer.go

@@ -6,6 +6,14 @@ import (
 	"github.com/netbirdio/netbird/management/internals/controllers/network_map"
 )
 
+//go:generate go tool mockgen -source=./update_debouncer.go -destination=./update_debouncer_mock.go -package=grpc
+type Debouncer interface {
+	Stop()
+	TimerChannel() <-chan time.Time
+	ProcessUpdate(update *network_map.UpdateMessage) bool
+	GetPendingUpdates() []*network_map.UpdateMessage
+}
+
 // UpdateDebouncer implements a backpressure mechanism that:
 // - Sends the first update immediately
 // - Coalesces rapid subsequent network map updates (only latest matters)

management/internals/shared/grpc/update_debouncer_mock.go

@@ -0,0 +1,96 @@
+// Code generated by MockGen. DO NOT EDIT.
+// Source: ./update_debouncer.go
+//
+// Generated by this command:
+//
+//	mockgen -source=./update_debouncer.go -destination=./update_debouncer_mock.go -package=grpc
+//
+
+// Package grpc is a generated GoMock package.
+package grpc
+
+import (
+	reflect "reflect"
+	time "time"
+
+	network_map "github.com/netbirdio/netbird/management/internals/controllers/network_map"
+	gomock "go.uber.org/mock/gomock"
+)
+
+// MockDebouncer is a mock of Debouncer interface.
+type MockDebouncer struct {
+	ctrl     *gomock.Controller
+	recorder *MockDebouncerMockRecorder
+	isgomock struct{}
+}
+
+// MockDebouncerMockRecorder is the mock recorder for MockDebouncer.
+type MockDebouncerMockRecorder struct {
+	mock *MockDebouncer
+}
+
+// NewMockDebouncer creates a new mock instance.
+func NewMockDebouncer(ctrl *gomock.Controller) *MockDebouncer {
+	mock := &MockDebouncer{ctrl: ctrl}
+	mock.recorder = &MockDebouncerMockRecorder{mock}
+	return mock
+}
+
+// EXPECT returns an object that allows the caller to indicate expected use.
+func (m *MockDebouncer) EXPECT() *MockDebouncerMockRecorder {
+	return m.recorder
+}
+
+// GetPendingUpdates mocks base method.
+func (m *MockDebouncer) GetPendingUpdates() []*network_map.UpdateMessage {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "GetPendingUpdates")
+	ret0, _ := ret[0].([]*network_map.UpdateMessage)
+	return ret0
+}
+
+// GetPendingUpdates indicates an expected call of GetPendingUpdates.
+func (mr *MockDebouncerMockRecorder) GetPendingUpdates() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetPendingUpdates", reflect.TypeOf((*MockDebouncer)(nil).GetPendingUpdates))
+}
+
+// ProcessUpdate mocks base method.
+func (m *MockDebouncer) ProcessUpdate(update *network_map.UpdateMessage) bool {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "ProcessUpdate", update)
+	ret0, _ := ret[0].(bool)
+	return ret0
+}
+
+// ProcessUpdate indicates an expected call of ProcessUpdate.
+func (mr *MockDebouncerMockRecorder) ProcessUpdate(update any) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "ProcessUpdate", reflect.TypeOf((*MockDebouncer)(nil).ProcessUpdate), update)
+}
+
+// Stop mocks base method.
+func (m *MockDebouncer) Stop() {
+	m.ctrl.T.Helper()
+	m.ctrl.Call(m, "Stop")
+}
+
+// Stop indicates an expected call of Stop.
+func (mr *MockDebouncerMockRecorder) Stop() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Stop", reflect.TypeOf((*MockDebouncer)(nil).Stop))
+}
+
+// TimerChannel mocks base method.
+func (m *MockDebouncer) TimerChannel() <-chan time.Time {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "TimerChannel")
+	ret0, _ := ret[0].(<-chan time.Time)
+	return ret0
+}
+
+// TimerChannel indicates an expected call of TimerChannel.
+func (mr *MockDebouncerMockRecorder) TimerChannel() *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "TimerChannel", reflect.TypeOf((*MockDebouncer)(nil).TimerChannel))
+}

management/internals/shared/grpc/validate_session_test.go

@@ -431,6 +431,57 @@ func TestValidateSession_MissingToken(t *testing.T) {
 	assert.Contains(t, resp.DeniedReason, "missing")
 }
 
+// TestGenerateSessionToken_UserNotInAllowedGroupGetsNoToken is the regression
+// guard for the group-authorisation bypass: the callback used to hand a signed
+// token to a user the service denies, and the proxy honoured that token as soon
+// as the user moved it into the nb_session cookie themselves. Authorisation has
+// to run before the token is signed.
+func TestGenerateSessionToken_UserNotInAllowedGroupGetsNoToken(t *testing.T) {
+	setup := setupValidateSessionTest(t)
+	defer setup.cleanup()
+
+	token, err := setup.proxyService.GenerateSessionToken(context.Background(), "restricted-proxy.example.com", "nonGroupUserId", auth.MethodOIDC)
+
+	require.Error(t, err, "a user outside the distribution groups must not receive a token")
+	assert.ErrorIs(t, err, ErrUserNotInGroup, "the callback maps this sentinel onto the access denied page")
+	assert.Empty(t, token, "no token may reach the browser")
+}
+
+func TestGenerateSessionToken_UserInAllowedGroupGetsTokenWithGroups(t *testing.T) {
+	setup := setupValidateSessionTest(t)
+	defer setup.cleanup()
+
+	ctx := context.Background()
+	svc, err := setup.store.GetServiceByID(ctx, store.LockingStrengthNone, "testAccountId", "restrictedProxyId")
+	require.NoError(t, err)
+
+	token, err := setup.proxyService.GenerateSessionToken(ctx, "restricted-proxy.example.com", "allowedUserId", auth.MethodOIDC)
+	require.NoError(t, err)
+	require.NotEmpty(t, token)
+
+	pubKey, err := base64.StdEncoding.DecodeString(svc.SessionPublicKey)
+	require.NoError(t, err)
+
+	userID, _, method, groups, _, err := auth.ValidateSessionJWT(token, "restricted-proxy.example.com", pubKey)
+	require.NoError(t, err)
+	assert.Equal(t, "allowedUserId", userID)
+	assert.Equal(t, auth.MethodOIDC.String(), method)
+	assert.Equal(t, []string{"allowedGroupId"}, groups, "the proxy gates the cookie on this claim, so it must carry the matched group")
+}
+
+// TestGenerateSessionToken_UnrestrictedServiceAllowsAnyAccountUser keeps the new
+// gate scoped: a service without distribution groups is open to every user of
+// its account, as before.
+func TestGenerateSessionToken_UnrestrictedServiceAllowsAnyAccountUser(t *testing.T) {
+	setup := setupValidateSessionTest(t)
+	defer setup.cleanup()
+
+	token, err := setup.proxyService.GenerateSessionToken(context.Background(), "test-proxy.example.com", "nonGroupUserId", auth.MethodOIDC)
+
+	require.NoError(t, err, "an unrestricted service must keep working for any user of the account")
+	assert.NotEmpty(t, token)
+}
+
 type testValidateSessionServiceManager struct {
 	store store.Store
 }