Merge branch 'feature/optimize-network-map-updates' into feature/validate-group-association

# Conflicts:
#	management/server/dns_test.go
#	management/server/group.go
#	management/server/nameserver.go
#	management/server/peer.go
#	management/server/peer_test.go
#	management/server/user.go

management/client/client_test.go

@@ -9,7 +9,10 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
+
 	"github.com/netbirdio/netbird/management/server/activity"
+	"github.com/netbirdio/netbird/management/server/telemetry"
 
 	"github.com/netbirdio/netbird/client/system"
 
@@ -71,7 +74,11 @@ func startManagement(t *testing.T) (*grpc.Server, net.Listener) {
 	peersUpdateManager := mgmt.NewPeersUpdateManager(nil)
 	eventStore := &activity.InMemoryEventStore{}
 	ia, _ := integrations.NewIntegratedValidator(context.Background(), eventStore)
-	accountManager, err := mgmt.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, ia)
+
+	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
+	require.NoError(t, err)
+
+	accountManager, err := mgmt.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, ia, metrics)
 	if err != nil {
 		t.Fatal(err)
 	}

management/client/grpc.go

@@ -2,7 +2,6 @@ package client
 
 import (
 	"context"
-	"crypto/tls"
 	"fmt"
 	"io"
 	"sync"
@@ -11,15 +10,11 @@ import (
 	"google.golang.org/grpc/codes"
 	gstatus "google.golang.org/grpc/status"
 
+	"github.com/cenkalti/backoff/v4"
 	log "github.com/sirupsen/logrus"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/connectivity"
-	"google.golang.org/grpc/credentials"
-	"google.golang.org/grpc/credentials/insecure"
-	"google.golang.org/grpc/keepalive"
-
-	"github.com/cenkalti/backoff/v4"
 
 	"github.com/netbirdio/netbird/client/system"
 	"github.com/netbirdio/netbird/encryption"
@@ -51,26 +46,21 @@ type GrpcClient struct {
 
 // NewClient creates a new client to Management service
 func NewClient(ctx context.Context, addr string, ourPrivateKey wgtypes.Key, tlsEnabled bool) (*GrpcClient, error) {
-	transportOption := grpc.WithTransportCredentials(insecure.NewCredentials())
+	var conn *grpc.ClientConn
 
-	if tlsEnabled {
-		transportOption = grpc.WithTransportCredentials(credentials.NewTLS(&tls.Config{}))
+	operation := func() error {
+		var err error
+		conn, err = nbgrpc.CreateConnection(addr, tlsEnabled)
+		if err != nil {
+			log.Printf("createConnection error: %v", err)
+			return err
+		}
+		return nil
 	}
 
-	mgmCtx, cancel := context.WithTimeout(ctx, ConnectTimeout)
-	defer cancel()
-	conn, err := grpc.DialContext(
-		mgmCtx,
-		addr,
-		transportOption,
-		nbgrpc.WithCustomDialer(),
-		grpc.WithBlock(),
-		grpc.WithKeepaliveParams(keepalive.ClientParameters{
-			Time:    30 * time.Second,
-			Timeout: 10 * time.Second,
-		}))
+	err := backoff.Retry(operation, nbgrpc.Backoff(ctx))
 	if err != nil {
-		log.Errorf("failed creating connection to Management Service %v", err)
+		log.Errorf("failed creating connection to Management Service: %v", err)
 		return nil, err
 	}
 
@@ -326,25 +316,44 @@ func (c *GrpcClient) login(serverKey wgtypes.Key, req *proto.LoginRequest) (*pro
 	if !c.ready() {
 		return nil, fmt.Errorf(errMsgNoMgmtConnection)
 	}
+
 	loginReq, err := encryption.EncryptMessage(serverKey, c.key, req)
 	if err != nil {
 		log.Errorf("failed to encrypt message: %s", err)
 		return nil, err
 	}
-	mgmCtx, cancel := context.WithTimeout(c.ctx, ConnectTimeout)
-	defer cancel()
-	resp, err := c.realClient.Login(mgmCtx, &proto.EncryptedMessage{
-		WgPubKey: c.key.PublicKey().String(),
-		Body:     loginReq,
-	})
+
+	var resp *proto.EncryptedMessage
+	operation := func() error {
+		mgmCtx, cancel := context.WithTimeout(context.Background(), ConnectTimeout)
+		defer cancel()
+
+		var err error
+		resp, err = c.realClient.Login(mgmCtx, &proto.EncryptedMessage{
+			WgPubKey: c.key.PublicKey().String(),
+			Body:     loginReq,
+		})
+		if err != nil {
+			// retry only on context canceled
+			if s, ok := gstatus.FromError(err); ok && s.Code() == codes.Canceled {
+				return err
+			}
+			return backoff.Permanent(err)
+		}
+
+		return nil
+	}
+
+	err = backoff.Retry(operation, nbgrpc.Backoff(c.ctx))
 	if err != nil {
+		log.Errorf("failed to login to Management Service: %v", err)
 		return nil, err
 	}
 
 	loginResp := &proto.LoginResponse{}
 	err = encryption.DecryptMessage(serverKey, c.key, resp.Body, loginResp)
 	if err != nil {
-		log.Errorf("failed to decrypt registration message: %s", err)
+		log.Errorf("failed to decrypt login response: %s", err)
 		return nil, err
 	}
 

management/cmd/management.go

@@ -190,7 +190,7 @@ var (
 				return fmt.Errorf("failed to initialize integrated peer validator: %v", err)
 			}
 			accountManager, err := server.BuildManager(ctx, store, peersUpdateManager, idpManager, mgmtSingleAccModeDomain,
-				dnsDomain, eventStore, geo, userDeleteFromIDPEnabled, integratedPeerValidator)
+				dnsDomain, eventStore, geo, userDeleteFromIDPEnabled, integratedPeerValidator, appMetrics)
 			if err != nil {
 				return fmt.Errorf("failed to build default manager: %v", err)
 			}

management/server/account.go

@@ -18,6 +18,8 @@ import (
 
 	"github.com/eko/gocache/v3/cache"
 	cacheStore "github.com/eko/gocache/v3/store"
+	"github.com/hashicorp/go-multierror"
+	"github.com/miekg/dns"
 	gocache "github.com/patrickmn/go-cache"
 	"github.com/rs/xid"
 	log "github.com/sirupsen/logrus"
@@ -37,6 +39,7 @@ import (
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/netbirdio/netbird/management/server/status"
+	"github.com/netbirdio/netbird/management/server/telemetry"
 	"github.com/netbirdio/netbird/route"
 )
 
@@ -65,6 +68,7 @@ type AccountManager interface {
 	SaveSetupKey(ctx context.Context, accountID string, key *SetupKey, userID string) (*SetupKey, error)
 	CreateUser(ctx context.Context, accountID, initiatorUserID string, key *UserInfo) (*UserInfo, error)
 	DeleteUser(ctx context.Context, accountID, initiatorUserID string, targetUserID string) error
+	DeleteRegularUsers(ctx context.Context, accountID, initiatorUserID string, targetUserIDs []string) error
 	InviteUser(ctx context.Context, accountID string, initiatorUserID string, targetUserID string) error
 	ListSetupKeys(ctx context.Context, accountID, userID string) ([]*SetupKey, error)
 	SaveUser(ctx context.Context, accountID, initiatorUserID string, update *User) (*UserInfo, error)
@@ -97,6 +101,7 @@ type AccountManager interface {
 	SaveGroup(ctx context.Context, accountID, userID string, group *nbgroup.Group) error
 	SaveGroups(ctx context.Context, accountID, userID string, newGroups []*nbgroup.Group) error
 	DeleteGroup(ctx context.Context, accountId, userId, groupID string) error
+	DeleteGroups(ctx context.Context, accountId, userId string, groupIDs []string) error
 	ListGroups(ctx context.Context, accountId string) ([]*nbgroup.Group, error)
 	GroupAddPeer(ctx context.Context, accountId, groupID, peerID string) error
 	GroupDeletePeer(ctx context.Context, accountId, groupID, peerID string) error
@@ -134,8 +139,8 @@ type AccountManager interface {
 	UpdateIntegratedValidatorGroups(ctx context.Context, accountID string, userID string, groups []string) error
 	GroupValidation(ctx context.Context, accountId string, groups []string) (bool, error)
 	GetValidatedPeers(account *Account) (map[string]struct{}, error)
-	SyncAndMarkPeer(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *NetworkMap, []*posture.Checks, error)
-	CancelPeerRoutines(ctx context.Context, peer *nbpeer.Peer) error
+	SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *NetworkMap, []*posture.Checks, error)
+	OnPeerDisconnected(ctx context.Context, accountID string, peerPubKey string) error
 	SyncPeerMeta(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta) error
 	FindExistingPostureCheck(accountID string, checks *posture.ChecksDefinition) (*posture.Checks, error)
 	GetAccountIDForPeerKey(ctx context.Context, peerKey string) (string, error)
@@ -169,6 +174,8 @@ type DefaultAccountManager struct {
 	userDeleteFromIDPEnabled bool
 
 	integratedPeerValidator integrated_validator.IntegratedValidator
+
+	metrics telemetry.AppMetrics
 }
 
 // Settings represents Account settings structure that can be modified via API and Dashboard
@@ -400,8 +407,16 @@ func (a *Account) GetGroup(groupID string) *nbgroup.Group {
 	return a.Groups[groupID]
 }
 
-// GetPeerNetworkMap returns a group by ID if exists, nil otherwise
-func (a *Account) GetPeerNetworkMap(ctx context.Context, peerID, dnsDomain string, validatedPeersMap map[string]struct{}) *NetworkMap {
+// GetPeerNetworkMap returns the networkmap for the given peer ID.
+func (a *Account) GetPeerNetworkMap(
+	ctx context.Context,
+	peerID string,
+	peersCustomZone nbdns.CustomZone,
+	validatedPeersMap map[string]struct{},
+	metrics *telemetry.AccountManagerMetrics,
+) *NetworkMap {
+	start := time.Now()
+
 	peer := a.Peers[peerID]
 	if peer == nil {
 		return &NetworkMap{
@@ -437,7 +452,7 @@ func (a *Account) GetPeerNetworkMap(ctx context.Context, peerID, dnsDomain strin
 
 	if dnsManagementStatus {
 		var zones []nbdns.CustomZone
-		peersCustomZone := getPeersCustomZone(ctx, a, dnsDomain)
+
 		if peersCustomZone.Domain != "" {
 			zones = append(zones, peersCustomZone)
 		}
@@ -445,7 +460,7 @@ func (a *Account) GetPeerNetworkMap(ctx context.Context, peerID, dnsDomain strin
 		dnsUpdate.NameServerGroups = getPeerNSGroups(a, peerID)
 	}
 
-	return &NetworkMap{
+	nm := &NetworkMap{
 		Peers:         peersToConnect,
 		Network:       a.Network.Copy(),
 		Routes:        routesUpdate,
@@ -453,6 +468,60 @@ func (a *Account) GetPeerNetworkMap(ctx context.Context, peerID, dnsDomain strin
 		OfflinePeers:  expiredPeers,
 		FirewallRules: firewallRules,
 	}
+
+	if metrics != nil {
+		objectCount := int64(len(peersToConnect) + len(expiredPeers) + len(routesUpdate) + len(firewallRules))
+		metrics.CountNetworkMapObjects(objectCount)
+		metrics.CountGetPeerNetworkMapDuration(time.Since(start))
+	}
+
+	return nm
+}
+
+func (a *Account) GetPeersCustomZone(ctx context.Context, dnsDomain string) nbdns.CustomZone {
+	var merr *multierror.Error
+
+	if dnsDomain == "" {
+		log.WithContext(ctx).Error("no dns domain is set, returning empty zone")
+		return nbdns.CustomZone{}
+	}
+
+	customZone := nbdns.CustomZone{
+		Domain:  dns.Fqdn(dnsDomain),
+		Records: make([]nbdns.SimpleRecord, 0, len(a.Peers)),
+	}
+
+	domainSuffix := "." + dnsDomain
+
+	var sb strings.Builder
+	for _, peer := range a.Peers {
+		if peer.DNSLabel == "" {
+			merr = multierror.Append(merr, fmt.Errorf("peer %s has an empty DNS label", peer.Name))
+			continue
+		}
+
+		sb.Grow(len(peer.DNSLabel) + len(domainSuffix))
+		sb.WriteString(peer.DNSLabel)
+		sb.WriteString(domainSuffix)
+
+		customZone.Records = append(customZone.Records, nbdns.SimpleRecord{
+			Name:  sb.String(),
+			Type:  int(dns.TypeA),
+			Class: nbdns.DefaultClass,
+			TTL:   defaultTTL,
+			RData: peer.IP.String(),
+		})
+
+		sb.Reset()
+	}
+
+	go func() {
+		if merr != nil {
+			log.WithContext(ctx).Errorf("error generating custom zone for account %s: %v", a.Id, merr)
+		}
+	}()
+
+	return customZone
 }
 
 // GetExpiredPeers returns peers that have been expired
@@ -769,10 +838,6 @@ func (a *Account) GetPeer(peerID string) *nbpeer.Peer {
 // SetJWTGroups updates the user's auto groups by synchronizing JWT groups.
 // Returns true if there are changes in the JWT group membership.
 func (a *Account) SetJWTGroups(userID string, groupsNames []string) bool {
-	if len(groupsNames) == 0 {
-		return false
-	}
-
 	user, ok := a.Users[userID]
 	if !ok {
 		return false
@@ -856,7 +921,7 @@ func (a *Account) UserGroupsAddToPeers(userID string, groups ...string) {
 func (a *Account) UserGroupsRemoveFromPeers(userID string, groups ...string) {
 	for _, gid := range groups {
 		group, ok := a.Groups[gid]
-		if !ok {
+		if !ok || group.Name == "All" {
 			continue
 		}
 		update := make([]string, 0, len(group.Peers))
@@ -874,10 +939,18 @@ func (a *Account) UserGroupsRemoveFromPeers(userID string, groups ...string) {
 }
 
 // BuildManager creates a new DefaultAccountManager with a provided Store
-func BuildManager(ctx context.Context, store Store, peersUpdateManager *PeersUpdateManager, idpManager idp.Manager,
-	singleAccountModeDomain string, dnsDomain string, eventStore activity.Store, geo *geolocation.Geolocation,
+func BuildManager(
+	ctx context.Context,
+	store Store,
+	peersUpdateManager *PeersUpdateManager,
+	idpManager idp.Manager,
+	singleAccountModeDomain string,
+	dnsDomain string,
+	eventStore activity.Store,
+	geo *geolocation.Geolocation,
 	userDeleteFromIDPEnabled bool,
 	integratedPeerValidator integrated_validator.IntegratedValidator,
+	metrics telemetry.AppMetrics,
 ) (*DefaultAccountManager, error) {
 	am := &DefaultAccountManager{
 		Store:                    store,
@@ -892,6 +965,7 @@ func BuildManager(ctx context.Context, store Store, peersUpdateManager *PeersUpd
 		peerLoginExpiry:          NewDefaultScheduler(),
 		userDeleteFromIDPEnabled: userDeleteFromIDPEnabled,
 		integratedPeerValidator:  integratedPeerValidator,
+		metrics:                  metrics,
 	}
 	allAccounts := store.GetAllAccounts(ctx)
 	// enable single account mode only if configured by user and number of existing accounts is not grater than 1
@@ -977,7 +1051,7 @@ func (am *DefaultAccountManager) UpdateAccountSettings(ctx context.Context, acco
 		return nil, status.Errorf(status.InvalidArgument, "peer login expiration can't be smaller than one hour")
 	}
 
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -1028,7 +1102,7 @@ func (am *DefaultAccountManager) UpdateAccountSettings(ctx context.Context, acco
 
 func (am *DefaultAccountManager) peerLoginExpirationJob(ctx context.Context, accountID string) func() (time.Duration, bool) {
 	return func() (time.Duration, bool) {
-		unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+		unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 		defer unlock()
 
 		account, err := am.Store.GetAccount(ctx, accountID)
@@ -1127,7 +1201,7 @@ func (am *DefaultAccountManager) warmupIDPCache(ctx context.Context) error {
 
 // DeleteAccount deletes an account and all its users from local store and from the remote IDP if the requester is an admin and account owner
 func (am *DefaultAccountManager) DeleteAccount(ctx context.Context, accountID, userID string) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
@@ -1587,7 +1661,7 @@ func (am *DefaultAccountManager) MarkPATUsed(ctx context.Context, tokenID string
 		return err
 	}
 
-	unlock := am.Store.AcquireAccountWriteLock(ctx, account.Id)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, account.Id)
 	defer unlock()
 
 	account, err = am.Store.GetAccountByUser(ctx, user.Id)
@@ -1670,7 +1744,7 @@ func (am *DefaultAccountManager) GetAccountFromToken(ctx context.Context, claims
 	if err != nil {
 		return nil, nil, err
 	}
-	unlock := am.Store.AcquireAccountWriteLock(ctx, newAcc.Id)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, newAcc.Id)
 	alreadyUnlocked := false
 	defer func() {
 		if !alreadyUnlocked {
@@ -1831,7 +1905,7 @@ func (am *DefaultAccountManager) getAccountWithAuthorizationClaims(ctx context.C
 
 	account, err := am.Store.GetAccountByUser(ctx, claims.UserId)
 	if err == nil {
-		unlockAccount := am.Store.AcquireAccountWriteLock(ctx, account.Id)
+		unlockAccount := am.Store.AcquireWriteLockByUID(ctx, account.Id)
 		defer unlockAccount()
 		account, err = am.Store.GetAccountByUser(ctx, claims.UserId)
 		if err != nil {
@@ -1851,7 +1925,7 @@ func (am *DefaultAccountManager) getAccountWithAuthorizationClaims(ctx context.C
 		return account, nil
 	} else if s, ok := status.FromError(err); ok && s.Type() == status.NotFound {
 		if domainAccount != nil {
-			unlockAccount := am.Store.AcquireAccountWriteLock(ctx, domainAccount.Id)
+			unlockAccount := am.Store.AcquireWriteLockByUID(ctx, domainAccount.Id)
 			defer unlockAccount()
 			domainAccount, err = am.Store.GetAccountByPrivateDomain(ctx, claims.Domain)
 			if err != nil {
@@ -1865,17 +1939,11 @@ func (am *DefaultAccountManager) getAccountWithAuthorizationClaims(ctx context.C
 	}
 }
 
-func (am *DefaultAccountManager) SyncAndMarkPeer(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *NetworkMap, []*posture.Checks, error) {
-	accountID, err := am.Store.GetAccountIDByPeerPubKey(ctx, peerPubKey)
-	if err != nil {
-		if errStatus, ok := status.FromError(err); ok && errStatus.Type() == status.NotFound {
-			return nil, nil, nil, status.Errorf(status.Unauthenticated, "peer not registered")
-		}
-		return nil, nil, nil, err
-	}
-
-	unlock := am.Store.AcquireAccountReadLock(ctx, accountID)
-	defer unlock()
+func (am *DefaultAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *NetworkMap, []*posture.Checks, error) {
+	accountUnlock := am.Store.AcquireReadLockByUID(ctx, accountID)
+	defer accountUnlock()
+	peerUnlock := am.Store.AcquireWriteLockByUID(ctx, peerPubKey)
+	defer peerUnlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
@@ -1895,26 +1963,20 @@ func (am *DefaultAccountManager) SyncAndMarkPeer(ctx context.Context, peerPubKey
 	return peer, netMap, postureChecks, nil
 }
 
-func (am *DefaultAccountManager) CancelPeerRoutines(ctx context.Context, peer *nbpeer.Peer) error {
-	accountID, err := am.Store.GetAccountIDByPeerPubKey(ctx, peer.Key)
-	if err != nil {
-		if errStatus, ok := status.FromError(err); ok && errStatus.Type() == status.NotFound {
-			return status.Errorf(status.Unauthenticated, "peer not registered")
-		}
-		return err
-	}
-
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
-	defer unlock()
+func (am *DefaultAccountManager) OnPeerDisconnected(ctx context.Context, accountID string, peerPubKey string) error {
+	accountUnlock := am.Store.AcquireReadLockByUID(ctx, accountID)
+	defer accountUnlock()
+	peerUnlock := am.Store.AcquireWriteLockByUID(ctx, peerPubKey)
+	defer peerUnlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
 		return err
 	}
 
-	err = am.MarkPeerConnected(ctx, peer.Key, false, nil, account)
+	err = am.MarkPeerConnected(ctx, peerPubKey, false, nil, account)
 	if err != nil {
-		log.WithContext(ctx).Warnf("failed marking peer as connected %s %v", peer.Key, err)
+		log.WithContext(ctx).Warnf("failed marking peer as connected %s %v", peerPubKey, err)
 	}
 
 	return nil
@@ -1927,7 +1989,7 @@ func (am *DefaultAccountManager) SyncPeerMeta(ctx context.Context, peerPubKey st
 		return err
 	}
 
-	unlock := am.Store.AcquireAccountReadLock(ctx, accountID)
+	unlock := am.Store.AcquireReadLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)

management/server/account_test.go

@@ -24,6 +24,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/jwtclaims"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/posture"
+	"github.com/netbirdio/netbird/management/server/telemetry"
 	"github.com/netbirdio/netbird/route"
 )
 
@@ -410,7 +411,8 @@ func TestAccount_GetPeerNetworkMap(t *testing.T) {
 			validatedPeers[p] = struct{}{}
 		}
 
-		networkMap := account.GetPeerNetworkMap(context.Background(), testCase.peerID, "netbird.io", validatedPeers)
+		customZone := account.GetPeersCustomZone(context.Background(), "netbird.io")
+		networkMap := account.GetPeerNetworkMap(context.Background(), testCase.peerID, customZone, validatedPeers, nil)
 		assert.Len(t, networkMap.Peers, len(testCase.expectedPeers))
 		assert.Len(t, networkMap.OfflinePeers, len(testCase.expectedOfflinePeers))
 	}
@@ -2238,6 +2240,13 @@ func TestAccount_SetJWTGroups(t *testing.T) {
 		assert.Len(t, account.Users["user2"].AutoGroups, 1, "new group should be added")
 		assert.Contains(t, account.Groups, account.Users["user2"].AutoGroups[0], "groups must contain group3 from user groups")
 	})
+
+	t.Run("remove all JWT groups", func(t *testing.T) {
+		updated := account.SetJWTGroups("user1", []string{})
+		assert.True(t, updated, "account should be updated")
+		assert.Len(t, account.Users["user1"].AutoGroups, 1, "only non-JWT groups should remain")
+		assert.Contains(t, account.Users["user1"].AutoGroups, "group1", " group1 should still be present")
+	})
 }
 
 func TestAccount_UserGroupsAddToPeers(t *testing.T) {
@@ -2305,7 +2314,13 @@ func TestAccount_UserGroupsRemoveFromPeers(t *testing.T) {
 	})
 }
 
-func createManager(t *testing.T) (*DefaultAccountManager, error) {
+type TB interface {
+	Cleanup(func())
+	Helper()
+	TempDir() string
+}
+
+func createManager(t TB) (*DefaultAccountManager, error) {
 	t.Helper()
 
 	store, err := createStore(t)
@@ -2314,7 +2329,12 @@ func createManager(t *testing.T) (*DefaultAccountManager, error) {
 	}
 	eventStore := &activity.InMemoryEventStore{}
 
-	manager, err := BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.cloud", eventStore, nil, false, MocIntegratedValidator{})
+	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
+	if err != nil {
+		return nil, err
+	}
+
+	manager, err := BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.cloud", eventStore, nil, false, MocIntegratedValidator{}, metrics)
 	if err != nil {
 		return nil, err
 	}
@@ -2322,7 +2342,7 @@ func createManager(t *testing.T) (*DefaultAccountManager, error) {
 	return manager, nil
 }
 
-func createStore(t *testing.T) (Store, error) {
+func createStore(t TB) (Store, error) {
 	t.Helper()
 	dataDir := t.TempDir()
 	store, cleanUp, err := NewTestStoreFromJson(context.Background(), dataDir)

management/server/config.go

@@ -56,6 +56,10 @@ type Config struct {
 func (c Config) GetAuthAudiences() []string {
 	audiences := []string{c.HttpConfig.AuthAudience}
 
+	if c.HttpConfig.ExtraAuthAudience != "" {
+		audiences = append(audiences, c.HttpConfig.ExtraAuthAudience)
+	}
+
 	if c.DeviceAuthorizationFlow != nil && c.DeviceAuthorizationFlow.ProviderConfig.Audience != "" {
 		audiences = append(audiences, c.DeviceAuthorizationFlow.ProviderConfig.Audience)
 	}
@@ -90,6 +94,8 @@ type HttpServerConfig struct {
 	OIDCConfigEndpoint string
 	// IdpSignKeyRefreshEnabled identifies the signing key is currently being rotated or not
 	IdpSignKeyRefreshEnabled bool
+	// Extra audience
+	ExtraAuthAudience string
 }
 
 // Host represents a Wiretrustee host (e.g. STUN, TURN, Signal)

management/server/dns.go

@@ -4,8 +4,8 @@ import (
 	"context"
 	"fmt"
 	"strconv"
+	"sync"
 
-	"github.com/miekg/dns"
 	log "github.com/sirupsen/logrus"
 
 	nbdns "github.com/netbirdio/netbird/dns"
@@ -17,6 +17,50 @@ import (
 
 const defaultTTL = 300
 
+// DNSConfigCache is a thread-safe cache for DNS configuration components
+type DNSConfigCache struct {
+	CustomZones      sync.Map
+	NameServerGroups sync.Map
+}
+
+// GetCustomZone retrieves a cached custom zone
+func (c *DNSConfigCache) GetCustomZone(key string) (*proto.CustomZone, bool) {
+	if c == nil {
+		return nil, false
+	}
+	if value, ok := c.CustomZones.Load(key); ok {
+		return value.(*proto.CustomZone), true
+	}
+	return nil, false
+}
+
+// SetCustomZone stores a custom zone in the cache
+func (c *DNSConfigCache) SetCustomZone(key string, value *proto.CustomZone) {
+	if c == nil {
+		return
+	}
+	c.CustomZones.Store(key, value)
+}
+
+// GetNameServerGroup retrieves a cached name server group
+func (c *DNSConfigCache) GetNameServerGroup(key string) (*proto.NameServerGroup, bool) {
+	if c == nil {
+		return nil, false
+	}
+	if value, ok := c.NameServerGroups.Load(key); ok {
+		return value.(*proto.NameServerGroup), true
+	}
+	return nil, false
+}
+
+// SetNameServerGroup stores a name server group in the cache
+func (c *DNSConfigCache) SetNameServerGroup(key string, value *proto.NameServerGroup) {
+	if c == nil {
+		return
+	}
+	c.NameServerGroups.Store(key, value)
+}
+
 type lookupMap map[string]struct{}
 
 // DNSSettings defines dns settings at the account level
@@ -36,7 +80,7 @@ func (d DNSSettings) Copy() DNSSettings {
 
 // GetDNSSettings validates a user role and returns the DNS settings for the provided account ID
 func (am *DefaultAccountManager) GetDNSSettings(ctx context.Context, accountID string, userID string) (*DNSSettings, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -58,7 +102,7 @@ func (am *DefaultAccountManager) GetDNSSettings(ctx context.Context, accountID s
 
 // SaveDNSSettings validates a user role and updates the account's DNS settings
 func (am *DefaultAccountManager) SaveDNSSettings(ctx context.Context, accountID string, userID string, dnsSettingsToSave *DNSSettings) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -118,69 +162,73 @@ func (am *DefaultAccountManager) SaveDNSSettings(ctx context.Context, accountID
 	return nil
 }
 
-func toProtocolDNSConfig(update nbdns.Config) *proto.DNSConfig {
-	protoUpdate := &proto.DNSConfig{ServiceEnable: update.ServiceEnable}
+// toProtocolDNSConfig converts nbdns.Config to proto.DNSConfig using the cache
+func toProtocolDNSConfig(update nbdns.Config, cache *DNSConfigCache) *proto.DNSConfig {
+	protoUpdate := &proto.DNSConfig{
+		ServiceEnable:    update.ServiceEnable,
+		CustomZones:      make([]*proto.CustomZone, 0, len(update.CustomZones)),
+		NameServerGroups: make([]*proto.NameServerGroup, 0, len(update.NameServerGroups)),
+	}
 
 	for _, zone := range update.CustomZones {
-		protoZone := &proto.CustomZone{Domain: zone.Domain}
-		for _, record := range zone.Records {
-			protoZone.Records = append(protoZone.Records, &proto.SimpleRecord{
-				Name:  record.Name,
-				Type:  int64(record.Type),
-				Class: record.Class,
-				TTL:   int64(record.TTL),
-				RData: record.RData,
-			})
+		cacheKey := zone.Domain
+		if cachedZone, exists := cache.GetCustomZone(cacheKey); exists {
+			protoUpdate.CustomZones = append(protoUpdate.CustomZones, cachedZone)
+		} else {
+			protoZone := convertToProtoCustomZone(zone)
+			cache.SetCustomZone(cacheKey, protoZone)
+			protoUpdate.CustomZones = append(protoUpdate.CustomZones, protoZone)
 		}
-		protoUpdate.CustomZones = append(protoUpdate.CustomZones, protoZone)
 	}
 
 	for _, nsGroup := range update.NameServerGroups {
-		protoGroup := &proto.NameServerGroup{
-			Primary:              nsGroup.Primary,
-			Domains:              nsGroup.Domains,
-			SearchDomainsEnabled: nsGroup.SearchDomainsEnabled,
+		cacheKey := nsGroup.ID
+		if cachedGroup, exists := cache.GetNameServerGroup(cacheKey); exists {
+			protoUpdate.NameServerGroups = append(protoUpdate.NameServerGroups, cachedGroup)
+		} else {
+			protoGroup := convertToProtoNameServerGroup(nsGroup)
+			cache.SetNameServerGroup(cacheKey, protoGroup)
+			protoUpdate.NameServerGroups = append(protoUpdate.NameServerGroups, protoGroup)
 		}
-		for _, ns := range nsGroup.NameServers {
-			protoNS := &proto.NameServer{
-				IP:     ns.IP.String(),
-				Port:   int64(ns.Port),
-				NSType: int64(ns.NSType),
-			}
-			protoGroup.NameServers = append(protoGroup.NameServers, protoNS)
-		}
-		protoUpdate.NameServerGroups = append(protoUpdate.NameServerGroups, protoGroup)
 	}
 
 	return protoUpdate
 }
 
-func getPeersCustomZone(ctx context.Context, account *Account, dnsDomain string) nbdns.CustomZone {
-	if dnsDomain == "" {
-		log.WithContext(ctx).Errorf("no dns domain is set, returning empty zone")
-		return nbdns.CustomZone{}
+// Helper function to convert nbdns.CustomZone to proto.CustomZone
+func convertToProtoCustomZone(zone nbdns.CustomZone) *proto.CustomZone {
+	protoZone := &proto.CustomZone{
+		Domain:  zone.Domain,
+		Records: make([]*proto.SimpleRecord, 0, len(zone.Records)),
 	}
-
-	customZone := nbdns.CustomZone{
-		Domain: dns.Fqdn(dnsDomain),
-	}
-
-	for _, peer := range account.Peers {
-		if peer.DNSLabel == "" {
-			log.WithContext(ctx).Errorf("found a peer with empty dns label. It was probably caused by a invalid character in its name. Peer Name: %s", peer.Name)
-			continue
-		}
-
-		customZone.Records = append(customZone.Records, nbdns.SimpleRecord{
-			Name:  dns.Fqdn(peer.DNSLabel + "." + dnsDomain),
-			Type:  int(dns.TypeA),
-			Class: nbdns.DefaultClass,
-			TTL:   defaultTTL,
-			RData: peer.IP.String(),
+	for _, record := range zone.Records {
+		protoZone.Records = append(protoZone.Records, &proto.SimpleRecord{
+			Name:  record.Name,
+			Type:  int64(record.Type),
+			Class: record.Class,
+			TTL:   int64(record.TTL),
+			RData: record.RData,
 		})
 	}
+	return protoZone
+}
 
-	return customZone
+// Helper function to convert nbdns.NameServerGroup to proto.NameServerGroup
+func convertToProtoNameServerGroup(nsGroup *nbdns.NameServerGroup) *proto.NameServerGroup {
+	protoGroup := &proto.NameServerGroup{
+		Primary:              nsGroup.Primary,
+		Domains:              nsGroup.Domains,
+		SearchDomainsEnabled: nsGroup.SearchDomainsEnabled,
+		NameServers:          make([]*proto.NameServer, 0, len(nsGroup.NameServers)),
+	}
+	for _, ns := range nsGroup.NameServers {
+		protoGroup.NameServers = append(protoGroup.NameServers, &proto.NameServer{
+			IP:     ns.IP.String(),
+			Port:   int64(ns.Port),
+			NSType: int64(ns.NSType),
+		})
+	}
+	return protoGroup
 }
 
 func getPeerNSGroups(account *Account, peerID string) []*nbdns.NameServerGroup {

management/server/dns_test.go

@@ -2,11 +2,14 @@ package server
 
 import (
 	"context"
+	"fmt"
 	"net/netip"
+	"reflect"
 	"testing"
-	"time"
 
-	"github.com/stretchr/testify/assert"
+	nbdns "github.com/netbirdio/netbird/dns"
+	"github.com/netbirdio/netbird/management/server/telemetry"
+
 	"github.com/stretchr/testify/require"
 
 	"github.com/netbirdio/netbird/dns"
@@ -197,7 +200,11 @@ func createDNSManager(t *testing.T) (*DefaultAccountManager, error) {
 		return nil, err
 	}
 	eventStore := &activity.InMemoryEventStore{}
-	return BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.test", eventStore, nil, false, MocIntegratedValidator{})
+
+	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
+	require.NoError(t, err)
+
+	return BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.test", eventStore, nil, false, MocIntegratedValidator{}, metrics)
 }
 
 func createDNSStore(t *testing.T) (Store, error) {
@@ -323,91 +330,149 @@ func initTestDNSAccount(t *testing.T, am *DefaultAccountManager) (*Account, erro
 	return am.Store.GetAccount(context.Background(), account.Id)
 }
 
-func TestDNSAccountPeerUpdate(t *testing.T) {
-	manager, account, peer1, peer2, peer3 := setupNetworkMapTest(t)
+func generateTestData(size int) nbdns.Config {
+	config := nbdns.Config{
+		ServiceEnable:    true,
+		CustomZones:      make([]nbdns.CustomZone, size),
+		NameServerGroups: make([]*nbdns.NameServerGroup, size),
+	}
 
-	err := manager.SaveGroup(context.Background(), account.Id, userID, &group.Group{
-		ID:    "group-id",
-		Name:  "GroupA",
-		Peers: []string{peer1.ID, peer2.ID, peer3.ID},
-	})
-	assert.NoError(t, err)
-
-	updMsg := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID)
-	t.Cleanup(func() {
-		manager.peersUpdateManager.CloseChannel(context.Background(), peer1.ID)
-	})
-
-	// Saving DNS settings with unused groups should not update account peers and not send peer update
-	t.Run("saving dns setting with unused groups", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		err := manager.SaveDNSSettings(context.Background(), account.Id, userID, &DNSSettings{
-			DisabledManagementGroups: []string{"group-id"},
-		})
-		assert.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(200 * time.Millisecond):
-			t.Error("timeout waiting for peerShouldNotReceiveUpdate")
+	for i := 0; i < size; i++ {
+		config.CustomZones[i] = nbdns.CustomZone{
+			Domain: fmt.Sprintf("domain%d.com", i),
+			Records: []nbdns.SimpleRecord{
+				{
+					Name:  fmt.Sprintf("record%d", i),
+					Type:  1,
+					Class: "IN",
+					TTL:   3600,
+					RData: "192.168.1.1",
+				},
+			},
 		}
-	})
 
-	_, err = manager.CreateNameServerGroup(
-		context.Background(), account.Id, "ns-group-1", "ns-group-1", []dns.NameServer{{
-			IP:     netip.MustParseAddr(peer1.IP.String()),
-			NSType: dns.UDPNameServerType,
-			Port:   dns.DefaultDNSPort,
-		}},
-		[]string{"group-id"},
-		true, []string{}, true, userID, false,
-	)
-	assert.NoError(t, err)
-
-	// Saving DNS settings with used groups should update account peers and send peer update
-	t.Run("saving dns setting with used groups", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		err := manager.SaveDNSSettings(context.Background(), account.Id, userID, &DNSSettings{
-			DisabledManagementGroups: []string{"group-id"},
-		})
-		assert.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(200 * time.Millisecond):
-			t.Error("timeout waiting for peerShouldReceiveUpdate")
+		config.NameServerGroups[i] = &nbdns.NameServerGroup{
+			ID:                   fmt.Sprintf("group%d", i),
+			Primary:              i == 0,
+			Domains:              []string{fmt.Sprintf("domain%d.com", i)},
+			SearchDomainsEnabled: true,
+			NameServers: []nbdns.NameServer{
+				{
+					IP:     netip.MustParseAddr("8.8.8.8"),
+					Port:   53,
+					NSType: 1,
+				},
+			},
 		}
-	})
-
-	// Saving unchanged DNS settings with used groups should update account peers and not send peer update
-	// since there is no change in the network map
-	t.Run("saving unchanged dns setting with used groups", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		err := manager.SaveDNSSettings(context.Background(), account.Id, userID, &DNSSettings{
-			DisabledManagementGroups: []string{"group-id"},
-		})
-		assert.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(200 * time.Millisecond):
-			t.Error("timeout waiting for peerShouldNotReceiveUpdate")
-		}
-	})
+	}
 
+	return config
+}
+
+func BenchmarkToProtocolDNSConfig(b *testing.B) {
+	sizes := []int{10, 100, 1000}
+
+	for _, size := range sizes {
+		testData := generateTestData(size)
+
+		b.Run(fmt.Sprintf("WithCache-Size%d", size), func(b *testing.B) {
+			cache := &DNSConfigCache{}
+
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				toProtocolDNSConfig(testData, cache)
+			}
+		})
+
+		b.Run(fmt.Sprintf("WithoutCache-Size%d", size), func(b *testing.B) {
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				cache := &DNSConfigCache{}
+				toProtocolDNSConfig(testData, cache)
+			}
+		})
+	}
+}
+
+func TestToProtocolDNSConfigWithCache(t *testing.T) {
+	var cache DNSConfigCache
+
+	// Create two different configs
+	config1 := nbdns.Config{
+		ServiceEnable: true,
+		CustomZones: []nbdns.CustomZone{
+			{
+				Domain: "example.com",
+				Records: []nbdns.SimpleRecord{
+					{Name: "www", Type: 1, Class: "IN", TTL: 300, RData: "192.168.1.1"},
+				},
+			},
+		},
+		NameServerGroups: []*nbdns.NameServerGroup{
+			{
+				ID:   "group1",
+				Name: "Group 1",
+				NameServers: []nbdns.NameServer{
+					{IP: netip.MustParseAddr("8.8.8.8"), Port: 53},
+				},
+			},
+		},
+	}
+
+	config2 := nbdns.Config{
+		ServiceEnable: true,
+		CustomZones: []nbdns.CustomZone{
+			{
+				Domain: "example.org",
+				Records: []nbdns.SimpleRecord{
+					{Name: "mail", Type: 1, Class: "IN", TTL: 300, RData: "192.168.1.2"},
+				},
+			},
+		},
+		NameServerGroups: []*nbdns.NameServerGroup{
+			{
+				ID:   "group2",
+				Name: "Group 2",
+				NameServers: []nbdns.NameServer{
+					{IP: netip.MustParseAddr("8.8.4.4"), Port: 53},
+				},
+			},
+		},
+	}
+
+	// First run with config1
+	result1 := toProtocolDNSConfig(config1, &cache)
+
+	// Second run with config2
+	result2 := toProtocolDNSConfig(config2, &cache)
+
+	// Third run with config1 again
+	result3 := toProtocolDNSConfig(config1, &cache)
+
+	// Verify that result1 and result3 are identical
+	if !reflect.DeepEqual(result1, result3) {
+		t.Errorf("Results are not identical when run with the same input. Expected %v, got %v", result1, result3)
+	}
+
+	// Verify that result2 is different from result1 and result3
+	if reflect.DeepEqual(result1, result2) || reflect.DeepEqual(result2, result3) {
+		t.Errorf("Results should be different for different inputs")
+	}
+
+	// Verify that the cache contains elements from both configs
+	if _, exists := cache.GetCustomZone("example.com"); !exists {
+		t.Errorf("Cache should contain custom zone for example.com")
+	}
+
+	if _, exists := cache.GetCustomZone("example.org"); !exists {
+		t.Errorf("Cache should contain custom zone for example.org")
+	}
+
+	if _, exists := cache.GetNameServerGroup("group1"); !exists {
+		t.Errorf("Cache should contain name server group 'group1'")
+	}
+
+	if _, exists := cache.GetNameServerGroup("group2"); !exists {
+		t.Errorf("Cache should contain name server group 'group2'")
+	}
 }

management/server/event.go

@@ -13,7 +13,7 @@ import (
 
 // GetEvents returns a list of activity events of an account
 func (am *DefaultAccountManager) GetEvents(ctx context.Context, accountID, userID string) ([]*activity.Event, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)

management/server/file_store.go

@@ -39,8 +39,8 @@ type FileStore struct {
 	mux       sync.Mutex `json:"-"`
 	storeFile string     `json:"-"`
 
-	// sync.Mutex indexed by accountID
-	accountLocks      sync.Map   `json:"-"`
+	// sync.Mutex indexed by resource ID
+	resourceLocks     sync.Map   `json:"-"`
 	globalAccountLock sync.Mutex `json:"-"`
 
 	metrics telemetry.AppMetrics `json:"-"`
@@ -281,26 +281,26 @@ func (s *FileStore) AcquireGlobalLock(ctx context.Context) (unlock func()) {
 	return unlock
 }
 
-// AcquireAccountWriteLock acquires account lock for writing to a resource and returns a function that releases the lock
-func (s *FileStore) AcquireAccountWriteLock(ctx context.Context, accountID string) (unlock func()) {
-	log.WithContext(ctx).Debugf("acquiring lock for account %s", accountID)
+// AcquireWriteLockByUID acquires an ID lock for writing to a resource and returns a function that releases the lock
+func (s *FileStore) AcquireWriteLockByUID(ctx context.Context, uniqueID string) (unlock func()) {
+	log.WithContext(ctx).Debugf("acquiring lock for ID %s", uniqueID)
 	start := time.Now()
-	value, _ := s.accountLocks.LoadOrStore(accountID, &sync.Mutex{})
+	value, _ := s.resourceLocks.LoadOrStore(uniqueID, &sync.Mutex{})
 	mtx := value.(*sync.Mutex)
 	mtx.Lock()
 
 	unlock = func() {
 		mtx.Unlock()
-		log.WithContext(ctx).Debugf("released lock for account %s in %v", accountID, time.Since(start))
+		log.WithContext(ctx).Debugf("released lock for ID %s in %v", uniqueID, time.Since(start))
 	}
 
 	return unlock
 }
 
-// AcquireAccountReadLock AcquireAccountWriteLock acquires account lock for reading a resource and returns a function that releases the lock
+// AcquireReadLockByUID acquires an ID lock for reading a resource and returns a function that releases the lock
 // This method is still returns a write lock as file store can't handle read locks
-func (s *FileStore) AcquireAccountReadLock(ctx context.Context, accountID string) (unlock func()) {
-	return s.AcquireAccountWriteLock(ctx, accountID)
+func (s *FileStore) AcquireReadLockByUID(ctx context.Context, uniqueID string) (unlock func()) {
+	return s.AcquireWriteLockByUID(ctx, uniqueID)
 }
 
 func (s *FileStore) SaveAccount(ctx context.Context, account *Account) error {
@@ -666,6 +666,26 @@ func (s *FileStore) SaveInstallationID(ctx context.Context, ID string) error {
 	return s.persist(ctx, s.storeFile)
 }
 
+// SavePeer saves the peer in the account
+func (s *FileStore) SavePeer(_ context.Context, accountID string, peer *nbpeer.Peer) error {
+	s.mux.Lock()
+	defer s.mux.Unlock()
+
+	account, err := s.getAccount(accountID)
+	if err != nil {
+		return err
+	}
+
+	newPeer := peer.Copy()
+
+	account.Peers[peer.ID] = newPeer
+
+	s.PeerKeyID2AccountID[peer.Key] = accountID
+	s.PeerID2AccountID[peer.ID] = accountID
+
+	return nil
+}
+
 // SavePeerStatus stores the PeerStatus in memory. It doesn't attempt to persist data to speed up things.
 // PeerStatus will be saved eventually when some other changes occur.
 func (s *FileStore) SavePeerStatus(accountID, peerID string, peerStatus nbpeer.PeerStatus) error {

management/server/geolocation/database.go

@@ -9,6 +9,7 @@ import (
 	"path"
 	"strconv"
 
+	log "github.com/sirupsen/logrus"
 	"gorm.io/driver/sqlite"
 	"gorm.io/gorm"
 	"gorm.io/gorm/logger"
@@ -30,6 +31,8 @@ func loadGeolocationDatabases(dataDir string) error {
 			continue
 		}
 
+		log.Infof("geo location file %s not found , file will be downloaded", file)
+
 		switch file {
 		case MMDBFileName:
 			extractFunc := func(src string, dst string) error {

management/server/group.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"slices"
 
@@ -26,7 +27,7 @@ func (e *GroupLinkError) Error() string {
 
 // GetGroup object of the peers
 func (am *DefaultAccountManager) GetGroup(ctx context.Context, accountID, groupID, userID string) (*nbgroup.Group, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -53,7 +54,7 @@ func (am *DefaultAccountManager) GetGroup(ctx context.Context, accountID, groupI
 
 // GetAllGroups returns all groups in an account
 func (am *DefaultAccountManager) GetAllGroups(ctx context.Context, accountID string, userID string) ([]*nbgroup.Group, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -80,7 +81,7 @@ func (am *DefaultAccountManager) GetAllGroups(ctx context.Context, accountID str
 
 // GetGroupByName filters all groups in an account by name and returns the one with the most peers
 func (am *DefaultAccountManager) GetGroupByName(ctx context.Context, groupName, accountID string) (*nbgroup.Group, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -113,7 +114,7 @@ func (am *DefaultAccountManager) GetGroupByName(ctx context.Context, groupName,
 
 // SaveGroup object of the peers
 func (am *DefaultAccountManager) SaveGroup(ctx context.Context, accountID, userID string, newGroup *nbgroup.Group) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 	return am.SaveGroups(ctx, accountID, userID, []*nbgroup.Group{newGroup})
 }
@@ -165,19 +166,12 @@ func (am *DefaultAccountManager) SaveGroups(ctx context.Context, accountID, user
 		eventsToStore = append(eventsToStore, events...)
 	}
 
-	newGroupIDs := make([]string, 0, len(newGroups))
-	for _, newGroup := range newGroups {
-		newGroupIDs = append(newGroupIDs, newGroup.ID)
-	}
-
 	account.Network.IncSerial()
-	if err = am.Store.SaveGroups(account.Id, account.Groups); err != nil {
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
 		return err
 	}
 
-	if areGroupChangesAffectPeers(account, newGroupIDs) {
-		am.updateAccountPeers(ctx, account)
-	}
+	am.updateAccountPeers(ctx, account)
 
 	for _, storeEvent := range eventsToStore {
 		storeEvent()
@@ -253,12 +247,12 @@ func difference(a, b []string) []string {
 	return diff
 }
 
-// DeleteGroup object of the peers
-func (am *DefaultAccountManager) DeleteGroup(ctx context.Context, accountID, userID, groupID string) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+// DeleteGroup object of the peers.
+func (am *DefaultAccountManager) DeleteGroup(ctx context.Context, accountId, userId, groupID string) error {
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountId)
 	defer unlock()
 
-	account, err := am.Store.GetAccount(ctx, accountID)
+	account, err := am.Store.GetAccount(ctx, accountId)
 	if err != nil {
 		return err
 	}
@@ -268,22 +262,70 @@ func (am *DefaultAccountManager) DeleteGroup(ctx context.Context, accountID, use
 		return nil
 	}
 
-	if err := validateDeleteGroup(account, group, userID); err != nil {
+	if err = validateDeleteGroup(account, group, userId); err != nil {
 		return err
 	}
-
 	delete(account.Groups, groupID)
+
+	account.Network.IncSerial()
 	if err = am.Store.SaveAccount(ctx, account); err != nil {
 		return err
 	}
-	am.StoreEvent(ctx, userID, groupID, accountID, activity.GroupDeleted, group.EventMeta())
+
+	am.StoreEvent(ctx, userId, groupID, accountId, activity.GroupDeleted, group.EventMeta())
+
+	am.updateAccountPeers(ctx, account)
 
 	return nil
 }
 
+// DeleteGroups deletes groups from an account.
+// Note: This function does not acquire the global lock.
+// It is the caller's responsibility to ensure proper locking is in place before invoking this method.
+//
+// If an error occurs while deleting a group, the function skips it and continues deleting other groups.
+// Errors are collected and returned at the end.
+func (am *DefaultAccountManager) DeleteGroups(ctx context.Context, accountId, userId string, groupIDs []string) error {
+	account, err := am.Store.GetAccount(ctx, accountId)
+	if err != nil {
+		return err
+	}
+
+	var allErrors error
+
+	deletedGroups := make([]*nbgroup.Group, 0, len(groupIDs))
+	for _, groupID := range groupIDs {
+		group, ok := account.Groups[groupID]
+		if !ok {
+			continue
+		}
+
+		if err := validateDeleteGroup(account, group, userId); err != nil {
+			allErrors = errors.Join(allErrors, fmt.Errorf("failed to delete group %s: %w", groupID, err))
+			continue
+		}
+
+		delete(account.Groups, groupID)
+		deletedGroups = append(deletedGroups, group)
+	}
+
+	account.Network.IncSerial()
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
+		return err
+	}
+
+	for _, g := range deletedGroups {
+		am.StoreEvent(ctx, userId, g.ID, accountId, activity.GroupDeleted, g.EventMeta())
+	}
+
+	am.updateAccountPeers(ctx, account)
+
+	return allErrors
+}
+
 // ListGroups objects of the peers
 func (am *DefaultAccountManager) ListGroups(ctx context.Context, accountID string) ([]*nbgroup.Group, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -301,7 +343,7 @@ func (am *DefaultAccountManager) ListGroups(ctx context.Context, accountID strin
 
 // GroupAddPeer appends peer to the group
 func (am *DefaultAccountManager) GroupAddPeer(ctx context.Context, accountID, groupID, peerID string) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -330,16 +372,14 @@ func (am *DefaultAccountManager) GroupAddPeer(ctx context.Context, accountID, gr
 		return err
 	}
 
-	if areGroupChangesAffectPeers(account, []string{group.ID}) {
-		am.updateAccountPeers(ctx, account)
-	}
+	am.updateAccountPeers(ctx, account)
 
 	return nil
 }
 
 // GroupDeletePeer removes peer from the group
 func (am *DefaultAccountManager) GroupDeletePeer(ctx context.Context, accountID, groupID, peerID string) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -362,29 +402,11 @@ func (am *DefaultAccountManager) GroupDeletePeer(ctx context.Context, accountID,
 		}
 	}
 
-	if areGroupChangesAffectPeers(account, []string{group.ID}) {
-		am.updateAccountPeers(ctx, account)
-	}
+	am.updateAccountPeers(ctx, account)
 
 	return nil
 }
 
-func areGroupChangesAffectPeers(account *Account, groupIDs []string) bool {
-	for _, groupID := range groupIDs {
-		if linked, _ := isGroupLinkedToDns(account.NameServerGroups, groupID); linked {
-			return true
-		}
-		if linked, _ := isGroupLinkedToPolicy(account.Policies, groupID); linked {
-			return true
-		}
-		if linked, _ := isGroupLinkedToRoute(account.Routes, groupID); linked {
-			return true
-		}
-	}
-
-	return false
-}
-
 func validateDeleteGroup(account *Account, group *nbgroup.Group, userID string) error {
 	// disable a deleting integration group if the initiator is not an admin service user
 	if group.Issued == nbgroup.GroupIssuedIntegration {
@@ -478,18 +500,8 @@ func isGroupLinkedToSetupKey(setupKeys map[string]*SetupKey, groupID string) (bo
 func isGroupLinkedToUser(users map[string]*User, groupID string) (bool, *User) {
 	for _, user := range users {
 		if slices.Contains(user.AutoGroups, groupID) {
-			return false, user
+			return true, user
 		}
 	}
 	return false, nil
 }
-
-// anyGroupHasPeers checks if any of the given groups in the account have peers.
-func anyGroupHasPeers(account *Account, groupIDs []string) bool {
-	for _, groupID := range groupIDs {
-		if group, exists := account.Groups[groupID]; exists && group.HasPeers() {
-			return true
-		}
-	}
-	return false
-}

management/server/group_test.go

@@ -3,6 +3,7 @@ package server
 import (
 	"context"
 	"errors"
+	"fmt"
 	"testing"
 	"time"
 
@@ -23,7 +24,7 @@ func TestDefaultAccountManager_CreateGroup(t *testing.T) {
 		t.Error("failed to create account manager")
 	}
 
-	account, err := initTestGroupAccount(am)
+	_, account, err := initTestGroupAccount(am)
 	if err != nil {
 		t.Error("failed to init testing account")
 	}
@@ -58,7 +59,7 @@ func TestDefaultAccountManager_DeleteGroup(t *testing.T) {
 		t.Error("failed to create account manager")
 	}
 
-	account, err := initTestGroupAccount(am)
+	_, account, err := initTestGroupAccount(am)
 	if err != nil {
 		t.Error("failed to init testing account")
 	}
@@ -134,7 +135,136 @@ func TestDefaultAccountManager_DeleteGroup(t *testing.T) {
 	}
 }
 
-func initTestGroupAccount(am *DefaultAccountManager) (*Account, error) {
+func TestDefaultAccountManager_DeleteGroups(t *testing.T) {
+	am, err := createManager(t)
+	assert.NoError(t, err, "Failed to create account manager")
+
+	manager, account, err := initTestGroupAccount(am)
+	assert.NoError(t, err, "Failed to init testing account")
+
+	groups := make([]*nbgroup.Group, 10)
+	for i := 0; i < 10; i++ {
+		groups[i] = &nbgroup.Group{
+			ID:        fmt.Sprintf("group-%d", i+1),
+			AccountID: account.Id,
+			Name:      fmt.Sprintf("group-%d", i+1),
+			Issued:    nbgroup.GroupIssuedAPI,
+		}
+	}
+
+	err = manager.SaveGroups(context.Background(), account.Id, groupAdminUserID, groups)
+	assert.NoError(t, err, "Failed to save test groups")
+
+	testCases := []struct {
+		name               string
+		groupIDs           []string
+		expectedReasons    []string
+		expectedDeleted    []string
+		expectedNotDeleted []string
+	}{
+		{
+			name:            "route",
+			groupIDs:        []string{"grp-for-route"},
+			expectedReasons: []string{"route"},
+		},
+		{
+			name:            "route with peer groups",
+			groupIDs:        []string{"grp-for-route2"},
+			expectedReasons: []string{"route"},
+		},
+		{
+			name:            "name server groups",
+			groupIDs:        []string{"grp-for-name-server-grp"},
+			expectedReasons: []string{"name server groups"},
+		},
+		{
+			name:            "policy",
+			groupIDs:        []string{"grp-for-policies"},
+			expectedReasons: []string{"policy"},
+		},
+		{
+			name:            "setup keys",
+			groupIDs:        []string{"grp-for-keys"},
+			expectedReasons: []string{"setup key"},
+		},
+		{
+			name:            "users",
+			groupIDs:        []string{"grp-for-users"},
+			expectedReasons: []string{"user"},
+		},
+		{
+			name:            "integration",
+			groupIDs:        []string{"grp-for-integration"},
+			expectedReasons: []string{"only service users with admin power can delete integration group"},
+		},
+		{
+			name:            "successfully delete multiple groups",
+			groupIDs:        []string{"group-1", "group-2"},
+			expectedDeleted: []string{"group-1", "group-2"},
+		},
+		{
+			name:            "delete non-existent group",
+			groupIDs:        []string{"non-existent-group"},
+			expectedDeleted: []string{"non-existent-group"},
+		},
+		{
+			name:               "delete multiple groups with mixed results",
+			groupIDs:           []string{"group-3", "grp-for-policies", "group-4", "grp-for-users"},
+			expectedReasons:    []string{"policy", "user"},
+			expectedDeleted:    []string{"group-3", "group-4"},
+			expectedNotDeleted: []string{"grp-for-policies", "grp-for-users"},
+		},
+		{
+			name:               "delete groups with multiple errors",
+			groupIDs:           []string{"grp-for-policies", "grp-for-users"},
+			expectedReasons:    []string{"policy", "user"},
+			expectedNotDeleted: []string{"grp-for-policies", "grp-for-users"},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			err = am.DeleteGroups(context.Background(), account.Id, groupAdminUserID, tc.groupIDs)
+			if len(tc.expectedReasons) > 0 {
+				assert.Error(t, err)
+				var foundExpectedErrors int
+
+				wrappedErr, ok := err.(interface{ Unwrap() []error })
+				assert.Equal(t, ok, true)
+
+				for _, e := range wrappedErr.Unwrap() {
+					var sErr *status.Error
+					if errors.As(e, &sErr) {
+						assert.Contains(t, tc.expectedReasons, sErr.Message, "unexpected error message")
+						foundExpectedErrors++
+					}
+
+					var gErr *GroupLinkError
+					if errors.As(e, &gErr) {
+						assert.Contains(t, tc.expectedReasons, gErr.Resource, "unexpected error resource")
+						foundExpectedErrors++
+					}
+				}
+				assert.Equal(t, len(tc.expectedReasons), foundExpectedErrors, "not all expected errors were found")
+			} else {
+				assert.NoError(t, err)
+			}
+
+			for _, groupID := range tc.expectedDeleted {
+				_, err := am.GetGroup(context.Background(), account.Id, groupID, groupAdminUserID)
+				assert.Error(t, err, "group should have been deleted: %s", groupID)
+			}
+
+			for _, groupID := range tc.expectedNotDeleted {
+				group, err := am.GetGroup(context.Background(), account.Id, groupID, groupAdminUserID)
+				assert.NoError(t, err, "group should not have been deleted: %s", groupID)
+				assert.NotNil(t, group, "group should exist: %s", groupID)
+			}
+		})
+	}
+}
+
+func initTestGroupAccount(am *DefaultAccountManager) (*DefaultAccountManager, *Account, error) {
 	accountID := "testingAcc"
 	domain := "example.com"
 
@@ -238,7 +368,7 @@ func initTestGroupAccount(am *DefaultAccountManager) (*Account, error) {
 
 	err := am.Store.SaveAccount(context.Background(), account)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	_ = am.SaveGroup(context.Background(), accountID, groupAdminUserID, groupForRoute)
@@ -249,7 +379,11 @@ func initTestGroupAccount(am *DefaultAccountManager) (*Account, error) {
 	_ = am.SaveGroup(context.Background(), accountID, groupAdminUserID, groupForUsers)
 	_ = am.SaveGroup(context.Background(), accountID, groupAdminUserID, groupForIntegration)
 
-	return am.Store.GetAccount(context.Background(), account.Id)
+	acc, err := am.Store.GetAccount(context.Background(), account.Id)
+	if err != nil {
+		return nil, nil, err
+	}
+	return am, acc, nil
 }
 
 func TestGroupAccountPeerUpdate(t *testing.T) {

management/server/grpcserver.go

@@ -156,7 +156,7 @@ func (s *GRPCServer) Sync(req *proto.EncryptedMessage, srv proto.ManagementServi
 		log.WithContext(ctx).Tracef("peer system meta has to be provided on sync. Peer %s, remote addr %s", peerKey.String(), realIP)
 	}
 
-	peer, netMap, postureChecks, err := s.accountManager.SyncAndMarkPeer(ctx, peerKey.String(), extractPeerMeta(ctx, syncReq.GetMeta()), realIP)
+	peer, netMap, postureChecks, err := s.accountManager.SyncAndMarkPeer(ctx, accountID, peerKey.String(), extractPeerMeta(ctx, syncReq.GetMeta()), realIP)
 	if err != nil {
 		return mapError(ctx, err)
 	}
@@ -179,11 +179,11 @@ func (s *GRPCServer) Sync(req *proto.EncryptedMessage, srv proto.ManagementServi
 		s.appMetrics.GRPCMetrics().CountSyncRequestDuration(time.Since(reqStart))
 	}
 
-	return s.handleUpdates(ctx, peerKey, peer, updates, srv)
+	return s.handleUpdates(ctx, accountID, peerKey, peer, updates, srv)
 }
 
 // handleUpdates sends updates to the connected peer until the updates channel is closed.
-func (s *GRPCServer) handleUpdates(ctx context.Context, peerKey wgtypes.Key, peer *nbpeer.Peer, updates chan *UpdateMessage, srv proto.ManagementService_SyncServer) error {
+func (s *GRPCServer) handleUpdates(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, updates chan *UpdateMessage, srv proto.ManagementService_SyncServer) error {
 	for {
 		select {
 		// condition when there are some updates
@@ -194,12 +194,12 @@ func (s *GRPCServer) handleUpdates(ctx context.Context, peerKey wgtypes.Key, pee
 
 			if !open {
 				log.WithContext(ctx).Debugf("updates channel for peer %s was closed", peerKey.String())
-				s.cancelPeerRoutines(ctx, peer)
+				s.cancelPeerRoutines(ctx, accountID, peer)
 				return nil
 			}
 			log.WithContext(ctx).Debugf("received an update for peer %s", peerKey.String())
 
-			if err := s.sendUpdate(ctx, peerKey, peer, update, srv); err != nil {
+			if err := s.sendUpdate(ctx, accountID, peerKey, peer, update, srv); err != nil {
 				return err
 			}
 
@@ -207,7 +207,7 @@ func (s *GRPCServer) handleUpdates(ctx context.Context, peerKey wgtypes.Key, pee
 		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, peer)
+			s.cancelPeerRoutines(ctx, accountID, peer)
 			return srv.Context().Err()
 		}
 	}
@@ -215,10 +215,10 @@ func (s *GRPCServer) handleUpdates(ctx context.Context, peerKey wgtypes.Key, pee
 
 // 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 *GRPCServer) sendUpdate(ctx context.Context, peerKey wgtypes.Key, peer *nbpeer.Peer, update *UpdateMessage, srv proto.ManagementService_SyncServer) error {
+func (s *GRPCServer) sendUpdate(ctx context.Context, accountID string, peerKey wgtypes.Key, peer *nbpeer.Peer, update *UpdateMessage, srv proto.ManagementService_SyncServer) error {
 	encryptedResp, err := encryption.EncryptMessage(peerKey, s.wgKey, update.Update)
 	if err != nil {
-		s.cancelPeerRoutines(ctx, peer)
+		s.cancelPeerRoutines(ctx, accountID, peer)
 		return status.Errorf(codes.Internal, "failed processing update message")
 	}
 	err = srv.SendMsg(&proto.EncryptedMessage{
@@ -226,17 +226,17 @@ func (s *GRPCServer) sendUpdate(ctx context.Context, peerKey wgtypes.Key, peer *
 		Body:     encryptedResp,
 	})
 	if err != nil {
-		s.cancelPeerRoutines(ctx, peer)
+		s.cancelPeerRoutines(ctx, accountID, peer)
 		return status.Errorf(codes.Internal, "failed sending update message")
 	}
 	log.WithContext(ctx).Debugf("sent an update to peer %s", peerKey.String())
 	return nil
 }
 
-func (s *GRPCServer) cancelPeerRoutines(ctx context.Context, peer *nbpeer.Peer) {
+func (s *GRPCServer) cancelPeerRoutines(ctx context.Context, accountID string, peer *nbpeer.Peer) {
 	s.peersUpdateManager.CloseChannel(ctx, peer.ID)
 	s.turnCredentialsManager.CancelRefresh(peer.ID)
-	_ = s.accountManager.CancelPeerRoutines(ctx, peer)
+	_ = s.accountManager.OnPeerDisconnected(ctx, accountID, peer.Key)
 	s.ephemeralManager.OnPeerDisconnected(ctx, peer)
 }
 
@@ -533,53 +533,46 @@ func toPeerConfig(peer *nbpeer.Peer, network *Network, dnsName string) *proto.Pe
 	}
 }
 
-func toRemotePeerConfig(peers []*nbpeer.Peer, dnsName string) []*proto.RemotePeerConfig {
-	remotePeers := []*proto.RemotePeerConfig{}
-	for _, rPeer := range peers {
-		fqdn := rPeer.FQDN(dnsName)
-		remotePeers = append(remotePeers, &proto.RemotePeerConfig{
-			WgPubKey:   rPeer.Key,
-			AllowedIps: []string{fmt.Sprintf(AllowedIPsFormat, rPeer.IP)},
-			SshConfig:  &proto.SSHConfig{SshPubKey: []byte(rPeer.SSHKey)},
-			Fqdn:       fqdn,
-		})
-	}
-	return remotePeers
-}
-
-func toSyncResponse(ctx context.Context, config *Config, peer *nbpeer.Peer, turnCredentials *TURNCredentials, networkMap *NetworkMap, dnsName string, checks []*posture.Checks) *proto.SyncResponse {
-	wtConfig := toWiretrusteeConfig(config, turnCredentials)
-
-	pConfig := toPeerConfig(peer, networkMap.Network, dnsName)
-
-	remotePeers := toRemotePeerConfig(networkMap.Peers, dnsName)
-
-	routesUpdate := toProtocolRoutes(networkMap.Routes)
-
-	dnsUpdate := toProtocolDNSConfig(networkMap.DNSConfig)
-
-	offlinePeers := toRemotePeerConfig(networkMap.OfflinePeers, dnsName)
-
-	firewallRules := toProtocolFirewallRules(networkMap.FirewallRules)
-
-	return &proto.SyncResponse{
-		WiretrusteeConfig:  wtConfig,
-		PeerConfig:         pConfig,
-		RemotePeers:        remotePeers,
-		RemotePeersIsEmpty: len(remotePeers) == 0,
+func toSyncResponse(ctx context.Context, config *Config, peer *nbpeer.Peer, turnCredentials *TURNCredentials, networkMap *NetworkMap, dnsName string, checks []*posture.Checks, dnsCache *DNSConfigCache) *proto.SyncResponse {
+	response := &proto.SyncResponse{
+		WiretrusteeConfig: toWiretrusteeConfig(config, turnCredentials),
+		PeerConfig:        toPeerConfig(peer, networkMap.Network, dnsName),
 		NetworkMap: &proto.NetworkMap{
-			Serial:               networkMap.Network.CurrentSerial(),
-			PeerConfig:           pConfig,
-			RemotePeers:          remotePeers,
-			OfflinePeers:         offlinePeers,
-			RemotePeersIsEmpty:   len(remotePeers) == 0,
-			Routes:               routesUpdate,
-			DNSConfig:            dnsUpdate,
-			FirewallRules:        firewallRules,
-			FirewallRulesIsEmpty: len(firewallRules) == 0,
+			Serial:    networkMap.Network.CurrentSerial(),
+			Routes:    toProtocolRoutes(networkMap.Routes),
+			DNSConfig: toProtocolDNSConfig(networkMap.DNSConfig, dnsCache),
 		},
 		Checks: toProtocolChecks(ctx, checks),
 	}
+
+	response.NetworkMap.PeerConfig = response.PeerConfig
+
+	allPeers := make([]*proto.RemotePeerConfig, 0, len(networkMap.Peers)+len(networkMap.OfflinePeers))
+	allPeers = appendRemotePeerConfig(allPeers, networkMap.Peers, dnsName)
+	response.RemotePeers = allPeers
+	response.NetworkMap.RemotePeers = allPeers
+	response.RemotePeersIsEmpty = len(allPeers) == 0
+	response.NetworkMap.RemotePeersIsEmpty = response.RemotePeersIsEmpty
+
+	response.NetworkMap.OfflinePeers = appendRemotePeerConfig(nil, networkMap.OfflinePeers, dnsName)
+
+	firewallRules := toProtocolFirewallRules(networkMap.FirewallRules)
+	response.NetworkMap.FirewallRules = firewallRules
+	response.NetworkMap.FirewallRulesIsEmpty = len(firewallRules) == 0
+
+	return response
+}
+
+func appendRemotePeerConfig(dst []*proto.RemotePeerConfig, peers []*nbpeer.Peer, dnsName string) []*proto.RemotePeerConfig {
+	for _, rPeer := range peers {
+		dst = append(dst, &proto.RemotePeerConfig{
+			WgPubKey:   rPeer.Key,
+			AllowedIps: []string{rPeer.IP.String() + "/32"},
+			SshConfig:  &proto.SSHConfig{SshPubKey: []byte(rPeer.SSHKey)},
+			Fqdn:       rPeer.FQDN(dnsName),
+		})
+	}
+	return dst
 }
 
 // IsHealthy indicates whether the service is healthy
@@ -597,7 +590,7 @@ func (s *GRPCServer) sendInitialSync(ctx context.Context, peerKey wgtypes.Key, p
 	} else {
 		turnCredentials = nil
 	}
-	plainResp := toSyncResponse(ctx, s.config, peer, turnCredentials, networkMap, s.accountManager.GetDNSDomain(), postureChecks)
+	plainResp := toSyncResponse(ctx, s.config, peer, turnCredentials, networkMap, s.accountManager.GetDNSDomain(), postureChecks, nil)
 
 	encryptedResp, err := encryption.EncryptMessage(peerKey, s.wgKey, plainResp)
 	if err != nil {

management/server/http/api/openapi.yml

@@ -526,6 +526,43 @@ components:
         - revoked
         - auto_groups
         - usage_limit
+    CreateSetupKeyRequest:
+      type: object
+      properties:
+        name:
+          description: Setup Key name
+          type: string
+          example: Default key
+        type:
+          description: Setup key type, one-off for single time usage and reusable
+          type: string
+          example: reusable
+        expires_in:
+          description: Expiration time in seconds
+          type: integer
+          minimum: 86400
+          maximum: 31536000
+          example: 86400
+        auto_groups:
+          description: List of group IDs to auto-assign to peers registered with this key
+          type: array
+          items:
+            type: string
+            example: "ch8i4ug6lnn4g9hqv7m0"
+        usage_limit:
+          description: A number of times this key can be used. The value of 0 indicates the unlimited usage.
+          type: integer
+          example: 0
+        ephemeral:
+          description: Indicate that the peer will be ephemeral or not
+          type: boolean
+          example: true
+      required:
+        - name
+        - type
+        - expires_in
+        - auto_groups
+        - usage_limit
     PersonalAccessToken:
       type: object
       properties:
@@ -1806,7 +1843,7 @@ paths:
         content:
           'application/json':
             schema:
-              $ref: '#/components/schemas/SetupKeyRequest'
+              $ref: '#/components/schemas/CreateSetupKeyRequest'
       responses:
         '200':
           description: A Setup Keys Object

management/server/http/api/types.gen.go

@@ -254,6 +254,27 @@ type Country struct {
 // CountryCode 2-letter ISO 3166-1 alpha-2 code that represents the country
 type CountryCode = string
 
+// CreateSetupKeyRequest defines model for CreateSetupKeyRequest.
+type CreateSetupKeyRequest struct {
+	// AutoGroups List of group IDs to auto-assign to peers registered with this key
+	AutoGroups []string `json:"auto_groups"`
+
+	// Ephemeral Indicate that the peer will be ephemeral or not
+	Ephemeral *bool `json:"ephemeral,omitempty"`
+
+	// ExpiresIn Expiration time in seconds
+	ExpiresIn int `json:"expires_in"`
+
+	// Name Setup Key name
+	Name string `json:"name"`
+
+	// Type Setup key type, one-off for single time usage and reusable
+	Type string `json:"type"`
+
+	// UsageLimit A number of times this key can be used. The value of 0 indicates the unlimited usage.
+	UsageLimit int `json:"usage_limit"`
+}
+
 // DNSSettings defines model for DNSSettings.
 type DNSSettings struct {
 	// DisabledManagementGroups Groups whose DNS management is disabled
@@ -1241,7 +1262,7 @@ type PostApiRoutesJSONRequestBody = RouteRequest
 type PutApiRoutesRouteIdJSONRequestBody = RouteRequest
 
 // PostApiSetupKeysJSONRequestBody defines body for PostApiSetupKeys for application/json ContentType.
-type PostApiSetupKeysJSONRequestBody = SetupKeyRequest
+type PostApiSetupKeysJSONRequestBody = CreateSetupKeyRequest
 
 // PutApiSetupKeysKeyIdJSONRequestBody defines body for PutApiSetupKeysKeyId for application/json ContentType.
 type PutApiSetupKeysKeyIdJSONRequestBody = SetupKeyRequest

management/server/http/peers_handler.go

@@ -71,7 +71,8 @@ func (h *PeersHandler) getPeer(ctx context.Context, account *server.Account, pee
 		return
 	}
 
-	netMap := account.GetPeerNetworkMap(ctx, peerID, h.accountManager.GetDNSDomain(), validPeers)
+	customZone := account.GetPeersCustomZone(ctx, h.accountManager.GetDNSDomain())
+	netMap := account.GetPeerNetworkMap(ctx, peerID, customZone, validPeers, nil)
 	accessiblePeers := toAccessiblePeers(netMap, dnsDomain)
 
 	_, valid := validPeers[peer.ID]
@@ -115,7 +116,9 @@ func (h *PeersHandler) updatePeer(ctx context.Context, account *server.Account,
 		util.WriteError(ctx, fmt.Errorf("internal error"), w)
 		return
 	}
-	netMap := account.GetPeerNetworkMap(ctx, peerID, h.accountManager.GetDNSDomain(), validPeers)
+
+	customZone := account.GetPeersCustomZone(ctx, h.accountManager.GetDNSDomain())
+	netMap := account.GetPeerNetworkMap(ctx, peerID, customZone, validPeers, nil)
 	accessiblePeers := toAccessiblePeers(netMap, dnsDomain)
 
 	_, valid := validPeers[peer.ID]
@@ -194,9 +197,7 @@ func (h *PeersHandler) GetAllPeers(w http.ResponseWriter, r *http.Request) {
 		}
 		groupMinimumInfo := toGroupsInfo(account.Groups, peer.ID)
 
-		accessiblePeerNumbers, _ := h.accessiblePeersNumber(r.Context(), account, peer.ID)
-
-		respBody = append(respBody, toPeerListItemResponse(peerToReturn, groupMinimumInfo, dnsDomain, accessiblePeerNumbers))
+		respBody = append(respBody, toPeerListItemResponse(peerToReturn, groupMinimumInfo, dnsDomain, 0))
 	}
 
 	validPeersMap, err := h.accountManager.GetValidatedPeers(account)
@@ -210,16 +211,6 @@ func (h *PeersHandler) GetAllPeers(w http.ResponseWriter, r *http.Request) {
 	util.WriteJSONObject(r.Context(), w, respBody)
 }
 
-func (h *PeersHandler) accessiblePeersNumber(ctx context.Context, account *server.Account, peerID string) (int, error) {
-	validatedPeersMap, err := h.accountManager.GetValidatedPeers(account)
-	if err != nil {
-		return 0, err
-	}
-
-	netMap := account.GetPeerNetworkMap(ctx, peerID, h.accountManager.GetDNSDomain(), validatedPeersMap)
-	return len(netMap.Peers) + len(netMap.OfflinePeers), nil
-}
-
 func (h *PeersHandler) setApprovalRequiredFlag(respBody []*api.PeerBatch, approvedPeersMap map[string]struct{}) {
 	for _, peer := range respBody {
 		_, ok := approvedPeersMap[peer.Id]

management/server/integrated_validator.go

@@ -32,7 +32,7 @@ func (am *DefaultAccountManager) UpdateIntegratedValidatorGroups(ctx context.Con
 		return errors.New("invalid groups")
 	}
 
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	a, err := am.Store.GetAccountByUser(ctx, userID)

management/server/management_proto_test.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+	"fmt"
 	"net"
 	"os"
 	"path/filepath"
@@ -16,8 +17,10 @@ import (
 	"google.golang.org/grpc/keepalive"
 
 	"github.com/netbirdio/netbird/encryption"
+	"github.com/netbirdio/netbird/formatter"
 	mgmtProto "github.com/netbirdio/netbird/management/proto"
 	"github.com/netbirdio/netbird/management/server/activity"
+	"github.com/netbirdio/netbird/management/server/telemetry"
 	"github.com/netbirdio/netbird/util"
 )
 
@@ -83,7 +86,7 @@ func Test_SyncProtocol(t *testing.T) {
 	defer func() {
 		os.Remove(filepath.Join(dir, "store.json")) //nolint
 	}()
-	mgmtServer, mgmtAddr, err := startManagement(t, &Config{
+	mgmtServer, _, mgmtAddr, err := startManagement(t, &Config{
 		Stuns: []*Host{{
 			Proto: "udp",
 			URI:   "stun:stun.wiretrustee.com:3468",
@@ -399,32 +402,39 @@ func TestServer_GetDeviceAuthorizationFlow(t *testing.T) {
 	}
 }
 
-func startManagement(t *testing.T, config *Config) (*grpc.Server, string, error) {
+func startManagement(t *testing.T, config *Config) (*grpc.Server, *DefaultAccountManager, string, error) {
 	t.Helper()
 	lis, err := net.Listen("tcp", "localhost:0")
 	if err != nil {
-		return nil, "", err
+		return nil, nil, "", err
 	}
 	s := grpc.NewServer(grpc.KeepaliveEnforcementPolicy(kaep), grpc.KeepaliveParams(kasp))
 	store, cleanUp, err := NewTestStoreFromJson(context.Background(), config.Datadir)
 	if err != nil {
-		return nil, "", err
+		return nil, nil, "", err
 	}
 	t.Cleanup(cleanUp)
 
 	peersUpdateManager := NewPeersUpdateManager(nil)
 	eventStore := &activity.InMemoryEventStore{}
-	accountManager, err := BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted",
-		eventStore, nil, false, MocIntegratedValidator{})
+
+	ctx := context.WithValue(context.Background(), formatter.ExecutionContextKey, formatter.SystemSource) //nolint:staticcheck
+
+	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
+	require.NoError(t, err)
+
+	accountManager, err := BuildManager(ctx, store, peersUpdateManager, nil, "", "netbird.selfhosted",
+		eventStore, nil, false, MocIntegratedValidator{}, metrics)
+
 	if err != nil {
-		return nil, "", err
+		return nil, nil, "", err
 	}
 	turnManager := NewTimeBasedAuthSecretsManager(peersUpdateManager, config.TURNConfig)
 
 	ephemeralMgr := NewEphemeralManager(store, accountManager)
 	mgmtServer, err := NewServer(context.Background(), config, accountManager, peersUpdateManager, turnManager, nil, ephemeralMgr)
 	if err != nil {
-		return nil, "", err
+		return nil, nil, "", err
 	}
 	mgmtProto.RegisterManagementServiceServer(s, mgmtServer)
 
@@ -434,7 +444,7 @@ func startManagement(t *testing.T, config *Config) (*grpc.Server, string, error)
 		}
 	}()
 
-	return s, lis.Addr().String(), nil
+	return s, accountManager, lis.Addr().String(), nil
 }
 
 func createRawClient(addr string) (mgmtProto.ManagementServiceClient, *grpc.ClientConn, error) {
@@ -454,3 +464,165 @@ func createRawClient(addr string) (mgmtProto.ManagementServiceClient, *grpc.Clie
 
 	return mgmtProto.NewManagementServiceClient(conn), conn, nil
 }
+func Test_SyncStatusRace(t *testing.T) {
+	if os.Getenv("CI") == "true" && os.Getenv("NETBIRD_STORE_ENGINE") == "postgres" {
+		t.Skip("Skipping on CI and Postgres store")
+	}
+	for i := 0; i < 500; i++ {
+		t.Run(fmt.Sprintf("TestRun-%d", i), func(t *testing.T) {
+			testSyncStatusRace(t)
+		})
+	}
+}
+func testSyncStatusRace(t *testing.T) {
+	t.Helper()
+	dir := t.TempDir()
+	err := util.CopyFileContents("testdata/store_with_expired_peers.json", filepath.Join(dir, "store.json"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		os.Remove(filepath.Join(dir, "store.json")) //nolint
+	}()
+
+	mgmtServer, am, mgmtAddr, err := startManagement(t, &Config{
+		Stuns: []*Host{{
+			Proto: "udp",
+			URI:   "stun:stun.wiretrustee.com:3468",
+		}},
+		TURNConfig: &TURNConfig{
+			TimeBasedCredentials: false,
+			CredentialsTTL:       util.Duration{},
+			Secret:               "whatever",
+			Turns: []*Host{{
+				Proto: "udp",
+				URI:   "turn:stun.wiretrustee.com:3468",
+			}},
+		},
+		Signal: &Host{
+			Proto: "http",
+			URI:   "signal.wiretrustee.com:10000",
+		},
+		Datadir:    dir,
+		HttpConfig: nil,
+	})
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+	defer mgmtServer.GracefulStop()
+
+	client, clientConn, err := createRawClient(mgmtAddr)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	defer clientConn.Close()
+
+	// there are two peers already in the store, add two more
+	peers, err := registerPeers(2, client)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	serverKey, err := getServerKey(client)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	concurrentPeerKey2 := peers[1]
+	t.Log("Public key of concurrent peer: ", concurrentPeerKey2.PublicKey().String())
+
+	syncReq2 := &mgmtProto.SyncRequest{Meta: &mgmtProto.PeerSystemMeta{}}
+	message2, err := encryption.EncryptMessage(*serverKey, *concurrentPeerKey2, syncReq2)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	ctx2, cancelFunc2 := context.WithCancel(context.Background())
+
+	//client.
+	sync2, err := client.Sync(ctx2, &mgmtProto.EncryptedMessage{
+		WgPubKey: concurrentPeerKey2.PublicKey().String(),
+		Body:     message2,
+	})
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	resp2 := &mgmtProto.EncryptedMessage{}
+	err = sync2.RecvMsg(resp2)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	peerWithInvalidStatus := peers[0]
+	t.Log("Public key of peer with invalid status: ", peerWithInvalidStatus.PublicKey().String())
+
+	syncReq := &mgmtProto.SyncRequest{Meta: &mgmtProto.PeerSystemMeta{}}
+	message, err := encryption.EncryptMessage(*serverKey, *peerWithInvalidStatus, syncReq)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	ctx, cancelFunc := context.WithCancel(context.Background())
+
+	//client.
+	sync, err := client.Sync(ctx, &mgmtProto.EncryptedMessage{
+		WgPubKey: peerWithInvalidStatus.PublicKey().String(),
+		Body:     message,
+	})
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	// take the first registered peer as a base for the test. Total four.
+
+	resp := &mgmtProto.EncryptedMessage{}
+	err = sync.RecvMsg(resp)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	cancelFunc2()
+	time.Sleep(1 * time.Millisecond)
+	cancelFunc()
+	time.Sleep(10 * time.Millisecond)
+
+	ctx, cancelFunc = context.WithCancel(context.Background())
+	defer cancelFunc()
+	sync, err = client.Sync(ctx, &mgmtProto.EncryptedMessage{
+		WgPubKey: peerWithInvalidStatus.PublicKey().String(),
+		Body:     message,
+	})
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	resp = &mgmtProto.EncryptedMessage{}
+	err = sync.RecvMsg(resp)
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+
+	time.Sleep(10 * time.Millisecond)
+	peer, err := am.Store.GetPeerByPeerPubKey(context.Background(), peerWithInvalidStatus.PublicKey().String())
+	if err != nil {
+		t.Fatal(err)
+		return
+	}
+	if !peer.Status.Connected {
+		t.Fatal("Peer should be connected")
+	}
+}

management/server/management_test.go

@@ -26,6 +26,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/group"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/telemetry"
 	"github.com/netbirdio/netbird/util"
 )
 
@@ -541,8 +542,13 @@ func startServer(config *server.Config) (*grpc.Server, net.Listener) {
 
 	peersUpdateManager := server.NewPeersUpdateManager(nil)
 	eventStore := &activity.InMemoryEventStore{}
-	accountManager, err := server.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted",
-		eventStore, nil, false, MocIntegratedValidator{})
+
+	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
+	if err != nil {
+		log.Fatalf("failed creating metrics: %v", err)
+	}
+
+	accountManager, err := server.BuildManager(context.Background(), store, peersUpdateManager, nil, "", "netbird.selfhosted", eventStore, nil, false, MocIntegratedValidator{}, metrics)
 	if err != nil {
 		log.Fatalf("failed creating a manager: %v", err)
 	}

management/server/mock_server/account_mock.go

@@ -31,7 +31,7 @@ type MockAccountManager struct {
 	ListUsersFunc                       func(ctx context.Context, accountID string) ([]*server.User, error)
 	GetPeersFunc                        func(ctx context.Context, accountID, userID string) ([]*nbpeer.Peer, error)
 	MarkPeerConnectedFunc               func(ctx context.Context, peerKey string, connected bool, realIP net.IP) error
-	SyncAndMarkPeerFunc                 func(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error)
+	SyncAndMarkPeerFunc                 func(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error)
 	DeletePeerFunc                      func(ctx context.Context, accountID, peerKey, userID string) error
 	GetNetworkMapFunc                   func(ctx context.Context, peerKey string) (*server.NetworkMap, error)
 	GetPeerNetworkFunc                  func(ctx context.Context, peerKey string) (*server.Network, error)
@@ -42,6 +42,7 @@ type MockAccountManager struct {
 	SaveGroupFunc                       func(ctx context.Context, accountID, userID string, group *group.Group) error
 	SaveGroupsFunc                      func(ctx context.Context, accountID, userID string, groups []*group.Group) error
 	DeleteGroupFunc                     func(ctx context.Context, accountID, userId, groupID string) error
+	DeleteGroupsFunc                    func(ctx context.Context, accountId, userId string, groupIDs []string) error
 	ListGroupsFunc                      func(ctx context.Context, accountID string) ([]*group.Group, error)
 	GroupAddPeerFunc                    func(ctx context.Context, accountID, groupID, peerID string) error
 	GroupDeletePeerFunc                 func(ctx context.Context, accountID, groupID, peerID string) error
@@ -66,6 +67,7 @@ type MockAccountManager struct {
 	SaveOrAddUserFunc                   func(ctx context.Context, accountID, userID string, user *server.User, addIfNotExists bool) (*server.UserInfo, error)
 	SaveOrAddUsersFunc                  func(ctx context.Context, accountID, initiatorUserID string, update []*server.User, addIfNotExists bool) ([]*server.UserInfo, error)
 	DeleteUserFunc                      func(ctx context.Context, accountID string, initiatorUserID string, targetUserID string) error
+	DeleteRegularUsersFunc              func(ctx context.Context, accountID, initiatorUserID string, targetUserIDs []string) error
 	CreatePATFunc                       func(ctx context.Context, accountID string, initiatorUserID string, targetUserId string, tokenName string, expiresIn int) (*server.PersonalAccessTokenGenerated, error)
 	DeletePATFunc                       func(ctx context.Context, accountID string, initiatorUserID string, targetUserId string, tokenID string) error
 	GetPATFunc                          func(ctx context.Context, accountID string, initiatorUserID string, targetUserId string, tokenID string) (*server.PersonalAccessToken, error)
@@ -104,14 +106,14 @@ type MockAccountManager struct {
 	GetAccountIDForPeerKeyFunc          func(ctx context.Context, peerKey string) (string, error)
 }
 
-func (am *MockAccountManager) SyncAndMarkPeer(ctx context.Context, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error) {
+func (am *MockAccountManager) SyncAndMarkPeer(ctx context.Context, accountID string, peerPubKey string, meta nbpeer.PeerSystemMeta, realIP net.IP) (*nbpeer.Peer, *server.NetworkMap, []*posture.Checks, error) {
 	if am.SyncAndMarkPeerFunc != nil {
-		return am.SyncAndMarkPeerFunc(ctx, peerPubKey, meta, realIP)
+		return am.SyncAndMarkPeerFunc(ctx, accountID, peerPubKey, meta, realIP)
 	}
 	return nil, nil, nil, status.Errorf(codes.Unimplemented, "method MarkPeerConnected is not implemented")
 }
 
-func (am *MockAccountManager) CancelPeerRoutines(_ context.Context, peer *nbpeer.Peer) error {
+func (am *MockAccountManager) OnPeerDisconnected(_ context.Context, accountID string, peerPubKey string) error {
 	// TODO implement me
 	panic("implement me")
 }
@@ -325,6 +327,14 @@ func (am *MockAccountManager) DeleteGroup(ctx context.Context, accountId, userId
 	return status.Errorf(codes.Unimplemented, "method DeleteGroup is not implemented")
 }
 
+// DeleteGroups mock implementation of DeleteGroups from server.AccountManager interface
+func (am *MockAccountManager) DeleteGroups(ctx context.Context, accountId, userId string, groupIDs []string) error {
+	if am.DeleteGroupsFunc != nil {
+		return am.DeleteGroupsFunc(ctx, accountId, userId, groupIDs)
+	}
+	return status.Errorf(codes.Unimplemented, "method DeleteGroups is not implemented")
+}
+
 // ListGroups mock implementation of ListGroups from server.AccountManager interface
 func (am *MockAccountManager) ListGroups(ctx context.Context, accountID string) ([]*group.Group, error) {
 	if am.ListGroupsFunc != nil {
@@ -519,6 +529,14 @@ func (am *MockAccountManager) DeleteUser(ctx context.Context, accountID string,
 	return status.Errorf(codes.Unimplemented, "method DeleteUser is not implemented")
 }
 
+// DeleteRegularUsers mocks DeleteRegularUsers of the AccountManager interface
+func (am *MockAccountManager) DeleteRegularUsers(ctx context.Context, accountID string, initiatorUserID string, targetUserIDs []string) error {
+	if am.DeleteRegularUsersFunc != nil {
+		return am.DeleteRegularUsersFunc(ctx, accountID, initiatorUserID, targetUserIDs)
+	}
+	return status.Errorf(codes.Unimplemented, "method DeleteRegularUsers is not implemented")
+}
+
 func (am *MockAccountManager) InviteUser(ctx context.Context, accountID string, initiatorUserID string, targetUserID string) error {
 	if am.InviteUserFunc != nil {
 		return am.InviteUserFunc(ctx, accountID, initiatorUserID, targetUserID)

management/server/nameserver.go

@@ -20,7 +20,7 @@ const domainPattern = `^(?i)[a-z0-9]+([\-\.]{1}[a-z0-9]+)*\.[a-z]{2,}$`
 // GetNameServerGroup gets a nameserver group object from account and nameserver group IDs
 func (am *DefaultAccountManager) GetNameServerGroup(ctx context.Context, accountID, userID, nsGroupID string) (*nbdns.NameServerGroup, error) {
 
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -48,7 +48,7 @@ func (am *DefaultAccountManager) GetNameServerGroup(ctx context.Context, account
 // CreateNameServerGroup creates and saves a new nameserver group
 func (am *DefaultAccountManager) CreateNameServerGroup(ctx context.Context, accountID string, name, description string, nameServerList []nbdns.NameServer, groups []string, primary bool, domains []string, enabled bool, userID string, searchDomainEnabled bool) (*nbdns.NameServerGroup, error) {
 
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -80,13 +80,13 @@ func (am *DefaultAccountManager) CreateNameServerGroup(ctx context.Context, acco
 	account.NameServerGroups[newNSGroup.ID] = newNSGroup
 
 	account.Network.IncSerial()
-	if err := am.Store.SaveAccount(ctx, account); err != nil {
+	err = am.Store.SaveAccount(ctx, account)
+	if err != nil {
 		return nil, err
 	}
 
-	if anyGroupHasPeers(account, newNSGroup.Groups) {
-		am.updateAccountPeers(ctx, account)
-	}
+	am.updateAccountPeers(ctx, account)
+
 	am.StoreEvent(ctx, userID, newNSGroup.ID, accountID, activity.NameserverGroupCreated, newNSGroup.EventMeta())
 
 	return newNSGroup.Copy(), nil
@@ -94,7 +94,8 @@ func (am *DefaultAccountManager) CreateNameServerGroup(ctx context.Context, acco
 
 // SaveNameServerGroup saves nameserver group
 func (am *DefaultAccountManager) SaveNameServerGroup(ctx context.Context, accountID, userID string, nsGroupToSave *nbdns.NameServerGroup) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	if nsGroupToSave == nil {
@@ -111,17 +112,16 @@ func (am *DefaultAccountManager) SaveNameServerGroup(ctx context.Context, accoun
 		return err
 	}
 
-	oldNSGroup := account.NameServerGroups[nsGroupToSave.ID]
 	account.NameServerGroups[nsGroupToSave.ID] = nsGroupToSave
 
 	account.Network.IncSerial()
-	if err = am.Store.SaveAccount(ctx, account); err != nil {
+	err = am.Store.SaveAccount(ctx, account)
+	if err != nil {
 		return err
 	}
 
-	if anyGroupHasPeers(account, nsGroupToSave.Groups) || anyGroupHasPeers(account, oldNSGroup.Groups) {
-		am.updateAccountPeers(ctx, account)
-	}
+	am.updateAccountPeers(ctx, account)
+
 	am.StoreEvent(ctx, userID, nsGroupToSave.ID, accountID, activity.NameserverGroupUpdated, nsGroupToSave.EventMeta())
 
 	return nil
@@ -130,7 +130,7 @@ func (am *DefaultAccountManager) SaveNameServerGroup(ctx context.Context, accoun
 // DeleteNameServerGroup deletes nameserver group with nsGroupID
 func (am *DefaultAccountManager) DeleteNameServerGroup(ctx context.Context, accountID, nsGroupID, userID string) error {
 
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -145,13 +145,13 @@ func (am *DefaultAccountManager) DeleteNameServerGroup(ctx context.Context, acco
 	delete(account.NameServerGroups, nsGroupID)
 
 	account.Network.IncSerial()
-	if err := am.Store.SaveAccount(ctx, account); err != nil {
+	err = am.Store.SaveAccount(ctx, account)
+	if err != nil {
 		return err
 	}
 
-	if anyGroupHasPeers(account, nsGroup.Groups) {
-		am.updateAccountPeers(ctx, account)
-	}
+	am.updateAccountPeers(ctx, account)
+
 	am.StoreEvent(ctx, userID, nsGroup.ID, accountID, activity.NameserverGroupDeleted, nsGroup.EventMeta())
 
 	return nil
@@ -160,7 +160,7 @@ func (am *DefaultAccountManager) DeleteNameServerGroup(ctx context.Context, acco
 // ListNameServerGroups returns a list of nameserver groups from account
 func (am *DefaultAccountManager) ListNameServerGroups(ctx context.Context, accountID string, userID string) ([]*nbdns.NameServerGroup, error) {
 
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)

management/server/nameserver_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/activity"
 	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/telemetry"
 )
 
 const (
@@ -764,7 +765,11 @@ func createNSManager(t *testing.T) (*DefaultAccountManager, error) {
 		return nil, err
 	}
 	eventStore := &activity.InMemoryEventStore{}
-	return BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.selfhosted", eventStore, nil, false, MocIntegratedValidator{})
+
+	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
+	require.NoError(t, err)
+
+	return BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.selfhosted", eventStore, nil, false, MocIntegratedValidator{}, metrics)
 }
 
 func createNSStore(t *testing.T) (Store, error) {

management/server/peer.go

@@ -4,14 +4,15 @@ import (
 	"context"
 	"fmt"
 	"net"
-	"slices"
 	"strings"
+	"sync"
 	"time"
 
-	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/rs/xid"
 	log "github.com/sirupsen/logrus"
 
+	"github.com/netbirdio/netbird/management/server/posture"
+
 	"github.com/netbirdio/netbird/management/proto"
 	"github.com/netbirdio/netbird/management/server/activity"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
@@ -65,12 +66,14 @@ func (am *DefaultAccountManager) GetPeers(ctx context.Context, accountID, userID
 	peers := make([]*nbpeer.Peer, 0)
 	peersMap := make(map[string]*nbpeer.Peer)
 
-	if !user.HasAdminPower() && !user.IsServiceUser && account.Settings.RegularUsersViewBlocked {
+	regularUser := !user.HasAdminPower() && !user.IsServiceUser
+
+	if regularUser && account.Settings.RegularUsersViewBlocked {
 		return peers, nil
 	}
 
 	for _, peer := range account.Peers {
-		if !(user.HasAdminPower() || user.IsServiceUser) && user.Id != peer.UserID {
+		if regularUser && user.Id != peer.UserID {
 			// only display peers that belong to the current user if the current user is not an admin
 			continue
 		}
@@ -79,6 +82,10 @@ func (am *DefaultAccountManager) GetPeers(ctx context.Context, accountID, userID
 		peersMap[peer.ID] = p
 	}
 
+	if !regularUser {
+		return peers, nil
+	}
+
 	// fetch all the peers that have access to the user's peers
 	for _, peer := range peers {
 		aclPeers, _ := account.getPeerConnectionResources(ctx, peer.ID, approvedPeersMap)
@@ -150,7 +157,7 @@ func (am *DefaultAccountManager) MarkPeerConnected(ctx context.Context, peerPubK
 
 // UpdatePeer updates peer. Only Peer.Name, Peer.SSHEnabled, and Peer.LoginExpirationEnabled can be updated.
 func (am *DefaultAccountManager) UpdatePeer(ctx context.Context, accountID, userID string, update *nbpeer.Peer) (*nbpeer.Peer, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -212,17 +219,13 @@ func (am *DefaultAccountManager) UpdatePeer(ctx context.Context, accountID, user
 	}
 
 	account.UpdatePeer(peer)
-	account.Network.IncSerial()
 
 	err = am.Store.SaveAccount(ctx, account)
 	if err != nil {
 		return nil, err
 	}
 
-	expired, _ := peer.LoginExpired(account.Settings.PeerLoginExpiration)
-	if expired && peer.LoginExpirationEnabled {
-		am.updateAccountPeers(ctx, account)
-	}
+	am.updateAccountPeers(ctx, account)
 
 	return peer, nil
 }
@@ -278,7 +281,7 @@ func (am *DefaultAccountManager) deletePeers(ctx context.Context, account *Accou
 
 // DeletePeer removes peer from the account by its IP
 func (am *DefaultAccountManager) DeletePeer(ctx context.Context, accountID, peerID, userID string) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -286,7 +289,6 @@ func (am *DefaultAccountManager) DeletePeer(ctx context.Context, accountID, peer
 		return err
 	}
 
-	updateAccountPeers := isPeerInActiveGroup(account, peerID)
 	err = am.deletePeers(ctx, account, []string{peerID}, userID)
 	if err != nil {
 		return err
@@ -297,9 +299,7 @@ func (am *DefaultAccountManager) DeletePeer(ctx context.Context, accountID, peer
 		return err
 	}
 
-	if updateAccountPeers {
-		am.updateAccountPeers(ctx, account)
-	}
+	am.updateAccountPeers(ctx, account)
 
 	return nil
 }
@@ -325,7 +325,8 @@ func (am *DefaultAccountManager) GetNetworkMap(ctx context.Context, peerID strin
 	if err != nil {
 		return nil, err
 	}
-	return account.GetPeerNetworkMap(ctx, peer.ID, am.dnsDomain, validatedPeers), nil
+	customZone := account.GetPeersCustomZone(ctx, am.dnsDomain)
+	return account.GetPeerNetworkMap(ctx, peer.ID, customZone, validatedPeers, nil), nil
 }
 
 // GetPeerNetwork returns the Network for a given peer
@@ -365,7 +366,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 		return nil, nil, nil, status.Errorf(status.NotFound, "failed adding new peer: account not found")
 	}
 
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer func() {
 		if unlock != nil {
 			unlock()
@@ -389,7 +390,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 	}
 
 	// This is a handling for the case when the same machine (with the same WireGuard pub key) tries to register twice.
-	// Such case is possible when AddPeer function takes long time to finish after AcquireAccountWriteLock (e.g., database is slow)
+	// Such case is possible when AddPeer function takes long time to finish after AcquireWriteLockByUID (e.g., database is slow)
 	// and the peer disconnects with a timeout and tries to register again.
 	// We just check if this machine has been registered before and reject the second registration.
 	// The connecting peer should be able to recover with a retry.
@@ -462,6 +463,17 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 		Location:               peer.Location,
 	}
 
+	if am.geo != nil && newPeer.Location.ConnectionIP != nil {
+		location, err := am.geo.Lookup(newPeer.Location.ConnectionIP)
+		if err != nil {
+			log.WithContext(ctx).Warnf("failed to get location for new peer realip: [%s]: %v", newPeer.Location.ConnectionIP.String(), err)
+		} else {
+			newPeer.Location.CountryCode = location.Country.ISOCode
+			newPeer.Location.CityName = location.City.Names.En
+			newPeer.Location.GeoNameID = location.City.GeonameID
+		}
+	}
+
 	// add peer to 'All' group
 	group, err := account.GetGroupAll()
 	if err != nil {
@@ -501,7 +513,6 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 	}
 
 	account.Peers[newPeer.ID] = newPeer
-
 	account.Network.IncSerial()
 	err = am.Store.SaveAccount(ctx, account)
 	if err != nil {
@@ -520,9 +531,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 
 	am.StoreEvent(ctx, opEvent.InitiatorID, opEvent.TargetID, opEvent.AccountID, opEvent.Activity, opEvent.Meta)
 
-	if areGroupChangesAffectPeers(account, groupsToAdd) {
-		am.updateAccountPeers(ctx, account)
-	}
+	am.updateAccountPeers(ctx, account)
 
 	approvedPeersMap, err := am.GetValidatedPeers(account)
 	if err != nil {
@@ -530,7 +539,8 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 	}
 
 	postureChecks := am.getPeerPostureChecks(account, peer)
-	networkMap := account.GetPeerNetworkMap(ctx, newPeer.ID, am.dnsDomain, approvedPeersMap)
+	customZone := account.GetPeersCustomZone(ctx, am.dnsDomain)
+	networkMap := account.GetPeerNetworkMap(ctx, newPeer.ID, customZone, approvedPeersMap, am.metrics.AccountManagerMetrics())
 	return newPeer, networkMap, postureChecks, nil
 }
 
@@ -547,7 +557,19 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
 	}
 
 	if peerLoginExpired(ctx, peer, account.Settings) {
-		return nil, nil, nil, status.Errorf(status.PermissionDenied, "peer login has expired, please log in once more")
+		return nil, nil, nil, status.NewPeerLoginExpiredError()
+	}
+
+	peer, updated := updatePeerMeta(peer, sync.Meta, account)
+	if updated {
+		err = am.Store.SavePeer(ctx, account.Id, peer)
+		if err != nil {
+			return nil, nil, nil, err
+		}
+
+		if sync.UpdateAccountPeers {
+			am.updateAccountPeers(ctx, account)
+		}
 	}
 
 	peerNotValid, isStatusChanged, err := am.integratedPeerValidator.IsNotValidPeer(ctx, account.Id, peer, account.GetPeerGroupsList(peer.ID), account.Settings.Extra)
@@ -555,22 +577,16 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
 		return nil, nil, nil, err
 	}
 
+	var postureChecks []*posture.Checks
+
 	if peerNotValid {
 		emptyMap := &NetworkMap{
 			Network: account.Network.Copy(),
 		}
-		return peer, emptyMap, nil, nil
+		return peer, emptyMap, postureChecks, nil
 	}
 
-	peer, peerMetaUpdated := updatePeerMeta(peer, sync.Meta, account)
-	if peerMetaUpdated {
-		err = am.Store.SaveAccount(ctx, account)
-		if err != nil {
-			return nil, nil, nil, err
-		}
-	}
-
-	if isStatusChanged || (peerMetaUpdated && sync.UpdateAccountPeers) {
+	if isStatusChanged {
 		am.updateAccountPeers(ctx, account)
 	}
 
@@ -578,9 +594,10 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
 	if err != nil {
 		return nil, nil, nil, err
 	}
-	postureChecks := am.getPeerPostureChecks(account, peer)
+	postureChecks = am.getPeerPostureChecks(account, peer)
 
-	return peer, account.GetPeerNetworkMap(ctx, peer.ID, am.dnsDomain, validPeersMap), postureChecks, nil
+	customZone := account.GetPeersCustomZone(ctx, am.dnsDomain)
+	return peer, account.GetPeerNetworkMap(ctx, peer.ID, customZone, validPeersMap, am.metrics.AccountManagerMetrics()), postureChecks, nil
 }
 
 // LoginPeer logs in or registers a peer.
@@ -592,21 +609,10 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 			// we couldn't find this peer by its public key which can mean that peer hasn't been registered yet.
 			// Try registering it.
 			newPeer := &nbpeer.Peer{
-				Key:    login.WireGuardPubKey,
-				Meta:   login.Meta,
-				SSHKey: login.SSHKey,
-			}
-			if am.geo != nil && login.ConnectionIP != nil {
-				location, err := am.geo.Lookup(login.ConnectionIP)
-				if err != nil {
-					log.WithContext(ctx).Warnf("failed to get location for new peer realip: [%s]: %v", login.ConnectionIP.String(), err)
-				} else {
-					newPeer.Location.ConnectionIP = login.ConnectionIP
-					newPeer.Location.CountryCode = location.Country.ISOCode
-					newPeer.Location.CityName = location.City.Names.En
-					newPeer.Location.GeoNameID = location.City.GeonameID
-
-				}
+				Key:      login.WireGuardPubKey,
+				Meta:     login.Meta,
+				SSHKey:   login.SSHKey,
+				Location: nbpeer.Location{ConnectionIP: login.ConnectionIP},
 			}
 
 			return am.AddPeer(ctx, login.SetupKey, login.UserID, newPeer)
@@ -616,44 +622,17 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 		return nil, nil, nil, status.Errorf(status.Internal, "failed while logging in peer")
 	}
 
-	peer, err := am.Store.GetPeerByPeerPubKey(ctx, login.WireGuardPubKey)
-	if err != nil {
-		return nil, nil, nil, status.NewPeerNotRegisteredError()
-	}
-
-	accSettings, err := am.Store.GetAccountSettings(ctx, accountID)
-	if err != nil {
-		return nil, nil, nil, status.Errorf(status.Internal, "failed to get account settings: %s", err)
-	}
-
-	var isWriteLock bool
-
-	// duplicated logic from after the lock to have an early exit
-	expired := peerLoginExpired(ctx, peer, accSettings)
-	switch {
-	case expired:
-		if err := checkAuth(ctx, login.UserID, peer); err != nil {
+	// when the client sends a login request with a JWT which is used to get the user ID,
+	// it means that the client has already checked if it needs login and had been through the SSO flow
+	// so, we can skip this check and directly proceed with the login
+	if login.UserID == "" {
+		err = am.checkIFPeerNeedsLoginWithoutLock(ctx, accountID, login)
+		if err != nil {
 			return nil, nil, nil, err
 		}
-		isWriteLock = true
-		log.WithContext(ctx).Debugf("peer login expired, acquiring write lock")
-
-	case peer.UpdateMetaIfNew(login.Meta):
-		isWriteLock = true
-		log.WithContext(ctx).Debugf("peer changed meta, acquiring write lock")
-
-	default:
-		isWriteLock = false
-		log.WithContext(ctx).Debugf("peer meta is the same, acquiring read lock")
 	}
 
-	var unlock func()
-
-	if isWriteLock {
-		unlock = am.Store.AcquireAccountWriteLock(ctx, accountID)
-	} else {
-		unlock = am.Store.AcquireAccountReadLock(ctx, accountID)
-	}
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer func() {
 		if unlock != nil {
 			unlock()
@@ -666,7 +645,7 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 		return nil, nil, nil, err
 	}
 
-	peer, err = account.FindPeerByPubKey(login.WireGuardPubKey)
+	peer, err := account.FindPeerByPubKey(login.WireGuardPubKey)
 	if err != nil {
 		return nil, nil, nil, status.NewPeerNotRegisteredError()
 	}
@@ -677,53 +656,39 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 	}
 
 	// this flag prevents unnecessary calls to the persistent store.
-	shouldStoreAccount := false
+	shouldStorePeer := false
 	updateRemotePeers := false
 	if peerLoginExpired(ctx, peer, account.Settings) {
-		err = checkAuth(ctx, login.UserID, peer)
+		err = am.handleExpiredPeer(ctx, login, account, peer)
 		if err != nil {
 			return nil, nil, nil, err
 		}
-		// If peer was expired before and if it reached this point, it is re-authenticated.
-		// UserID is present, meaning that JWT validation passed successfully in the API layer.
-		updatePeerLastLogin(peer, account)
 		updateRemotePeers = true
-		shouldStoreAccount = true
-
-		// sync user last login with peer last login
-		user, err := account.FindUser(login.UserID)
-		if err != nil {
-			return nil, nil, nil, status.Errorf(status.Internal, "couldn't find user")
-		}
-		user.updateLastLogin(peer.LastLogin)
-
-		am.StoreEvent(ctx, login.UserID, peer.ID, account.Id, activity.UserLoggedInPeer, peer.EventMeta(am.GetDNSDomain()))
+		shouldStorePeer = true
 	}
 
 	isRequiresApproval, isStatusChanged, err := am.integratedPeerValidator.IsNotValidPeer(ctx, account.Id, peer, account.GetPeerGroupsList(peer.ID), account.Settings.Extra)
 	if err != nil {
 		return nil, nil, nil, err
 	}
+
 	peer, updated := updatePeerMeta(peer, login.Meta, account)
 	if updated {
-		shouldStoreAccount = true
+		shouldStorePeer = true
 	}
 
-	peer, err = am.checkAndUpdatePeerSSHKey(ctx, peer, account, login.SSHKey)
-	if err != nil {
-		return nil, nil, nil, err
+	if peer.SSHKey != login.SSHKey {
+		peer.SSHKey = login.SSHKey
+		shouldStorePeer = true
 	}
 
-	if shouldStoreAccount {
-		if !isWriteLock {
-			log.WithContext(ctx).Errorf("account %s should be stored but is not write locked", accountID)
-			return nil, nil, nil, status.Errorf(status.Internal, "account should be stored but is not write locked")
-		}
-		err = am.Store.SaveAccount(ctx, account)
+	if shouldStorePeer {
+		err = am.Store.SavePeer(ctx, accountID, peer)
 		if err != nil {
 			return nil, nil, nil, err
 		}
 	}
+
 	unlock()
 	unlock = nil
 
@@ -731,13 +696,46 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 		am.updateAccountPeers(ctx, account)
 	}
 
+	return am.getValidatedPeerWithMap(ctx, isRequiresApproval, account, peer)
+}
+
+// checkIFPeerNeedsLoginWithoutLock checks if the peer needs login without acquiring the account lock. The check validate if the peer was not added via SSO
+// and if the peer login is expired.
+// The NetBird client doesn't have a way to check if the peer needs login besides sending a login request
+// with no JWT token and usually no setup-key. As the client can send up to two login request to check if it is expired
+// and before starting the engine, we do the checks without an account lock to avoid piling up requests.
+func (am *DefaultAccountManager) checkIFPeerNeedsLoginWithoutLock(ctx context.Context, accountID string, login PeerLogin) error {
+	peer, err := am.Store.GetPeerByPeerPubKey(ctx, login.WireGuardPubKey)
+	if err != nil {
+		return err
+	}
+
+	// if the peer was not added with SSO login we can exit early because peers activated with setup-key
+	// doesn't expire, and we avoid extra databases calls.
+	if !peer.AddedWithSSOLogin() {
+		return nil
+	}
+
+	settings, err := am.Store.GetAccountSettings(ctx, accountID)
+	if err != nil {
+		return err
+	}
+
+	if peerLoginExpired(ctx, peer, settings) {
+		return status.NewPeerLoginExpiredError()
+	}
+
+	return nil
+}
+
+func (am *DefaultAccountManager) getValidatedPeerWithMap(ctx context.Context, isRequiresApproval bool, account *Account, peer *nbpeer.Peer) (*nbpeer.Peer, *NetworkMap, []*posture.Checks, error) {
 	var postureChecks []*posture.Checks
 
 	if isRequiresApproval {
 		emptyMap := &NetworkMap{
 			Network: account.Network.Copy(),
 		}
-		return peer, emptyMap, postureChecks, nil
+		return peer, emptyMap, nil, nil
 	}
 
 	approvedPeersMap, err := am.GetValidatedPeers(account)
@@ -746,7 +744,32 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 	}
 	postureChecks = am.getPeerPostureChecks(account, peer)
 
-	return peer, account.GetPeerNetworkMap(ctx, peer.ID, am.dnsDomain, approvedPeersMap), postureChecks, nil
+	customZone := account.GetPeersCustomZone(ctx, am.dnsDomain)
+	return peer, account.GetPeerNetworkMap(ctx, peer.ID, customZone, approvedPeersMap, am.metrics.AccountManagerMetrics()), postureChecks, nil
+}
+
+func (am *DefaultAccountManager) handleExpiredPeer(ctx context.Context, login PeerLogin, account *Account, peer *nbpeer.Peer) error {
+	err := checkAuth(ctx, login.UserID, peer)
+	if err != nil {
+		return err
+	}
+	// If peer was expired before and if it reached this point, it is re-authenticated.
+	// UserID is present, meaning that JWT validation passed successfully in the API layer.
+	updatePeerLastLogin(peer, account)
+
+	// sync user last login with peer last login
+	user, err := account.FindUser(login.UserID)
+	if err != nil {
+		return status.Errorf(status.Internal, "couldn't find user")
+	}
+
+	err = am.Store.SaveUserLastLogin(account.Id, user.Id, peer.LastLogin)
+	if err != nil {
+		return err
+	}
+
+	am.StoreEvent(ctx, login.UserID, peer.ID, account.Id, activity.UserLoggedInPeer, peer.EventMeta(am.GetDNSDomain()))
+	return nil
 }
 
 func checkIfPeerOwnerIsBlocked(peer *nbpeer.Peer, account *Account) error {
@@ -765,11 +788,11 @@ func checkIfPeerOwnerIsBlocked(peer *nbpeer.Peer, account *Account) error {
 func checkAuth(ctx context.Context, loginUserID string, peer *nbpeer.Peer) error {
 	if loginUserID == "" {
 		// absence of a user ID indicates that JWT wasn't provided.
-		return status.Errorf(status.PermissionDenied, "peer login has expired, please log in once more")
+		return status.NewPeerLoginExpiredError()
 	}
 	if peer.UserID != loginUserID {
 		log.WithContext(ctx).Warnf("user mismatch when logging in peer %s: peer user %s, login user %s ", peer.ID, peer.UserID, loginUserID)
-		return status.Errorf(status.Unauthenticated, "can't login")
+		return status.Errorf(status.Unauthenticated, "can't login with this credentials")
 	}
 	return nil
 }
@@ -789,31 +812,54 @@ func updatePeerLastLogin(peer *nbpeer.Peer, account *Account) {
 	account.UpdatePeer(peer)
 }
 
-func (am *DefaultAccountManager) checkAndUpdatePeerSSHKey(ctx context.Context, peer *nbpeer.Peer, account *Account, newSSHKey string) (*nbpeer.Peer, error) {
-	if len(newSSHKey) == 0 {
-		log.WithContext(ctx).Debugf("no new SSH key provided for peer %s, skipping update", peer.ID)
-		return peer, nil
+// UpdatePeerSSHKey updates peer's public SSH key
+func (am *DefaultAccountManager) UpdatePeerSSHKey(ctx context.Context, peerID string, sshKey string) error {
+	if sshKey == "" {
+		log.WithContext(ctx).Debugf("empty SSH key provided for peer %s, skipping update", peerID)
+		return nil
 	}
 
-	if peer.SSHKey == newSSHKey {
-		log.WithContext(ctx).Debugf("same SSH key provided for peer %s, skipping update", peer.ID)
-		return peer, nil
+	account, err := am.Store.GetAccountByPeerID(ctx, peerID)
+	if err != nil {
+		return err
 	}
 
-	peer.SSHKey = newSSHKey
+	unlock := am.Store.AcquireWriteLockByUID(ctx, account.Id)
+	defer unlock()
+
+	// ensure that we consider modification happened meanwhile (because we were outside the account lock when we fetched the account)
+	account, err = am.Store.GetAccount(ctx, account.Id)
+	if err != nil {
+		return err
+	}
+
+	peer := account.GetPeer(peerID)
+	if peer == nil {
+		return status.Errorf(status.NotFound, "peer with ID %s not found", peerID)
+	}
+
+	if peer.SSHKey == sshKey {
+		log.WithContext(ctx).Debugf("same SSH key provided for peer %s, skipping update", peerID)
+		return nil
+	}
+
+	peer.SSHKey = sshKey
 	account.UpdatePeer(peer)
 
-	err := am.Store.SaveAccount(ctx, account)
+	err = am.Store.SaveAccount(ctx, account)
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	return peer, nil
+	// trigger network map update
+	am.updateAccountPeers(ctx, account)
+
+	return nil
 }
 
 // GetPeer for a given accountID, peerID and userID error if not found.
 func (am *DefaultAccountManager) GetPeer(ctx context.Context, accountID, peerID, userID string) (*nbpeer.Peer, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -875,34 +921,45 @@ func updatePeerMeta(peer *nbpeer.Peer, meta nbpeer.PeerSystemMeta, account *Acco
 // updateAccountPeers updates all peers that belong to an account.
 // Should be called when changes have to be synced to peers.
 func (am *DefaultAccountManager) updateAccountPeers(ctx context.Context, account *Account) {
+	start := time.Now()
+	defer func() {
+		if am.metrics != nil {
+			am.metrics.AccountManagerMetrics().CountUpdateAccountPeersDuration(time.Since(start))
+		}
+	}()
+
 	peers := account.GetPeers()
 
 	approvedPeersMap, err := am.GetValidatedPeers(account)
 	if err != nil {
-		log.WithContext(ctx).Errorf("failed send out updates to peers, failed to validate peer: %v", err)
+		log.WithContext(ctx).Errorf("failed to send out updates to peers, failed to validate peer: %v", err)
 		return
 	}
+
+	var wg sync.WaitGroup
+	semaphore := make(chan struct{}, 10)
+
+	dnsCache := &DNSConfigCache{}
+	customZone := account.GetPeersCustomZone(ctx, am.dnsDomain)
+
 	for _, peer := range peers {
 		if !am.peersUpdateManager.HasChannel(peer.ID) {
 			log.WithContext(ctx).Tracef("peer %s doesn't have a channel, skipping network map update", peer.ID)
 			continue
 		}
 
-		postureChecks := am.getPeerPostureChecks(account, peer)
-		remotePeerNetworkMap := account.GetPeerNetworkMap(ctx, peer.ID, am.dnsDomain, approvedPeersMap)
-		update := toSyncResponse(ctx, nil, peer, nil, remotePeerNetworkMap, am.GetDNSDomain(), postureChecks)
-		go am.peersUpdateManager.SendUpdate(ctx, peer.ID, &UpdateMessage{Update: update, NetworkMap: remotePeerNetworkMap, Checks: postureChecks})
-	}
-}
+		wg.Add(1)
+		semaphore <- struct{}{}
+		go func(p *nbpeer.Peer) {
+			defer wg.Done()
+			defer func() { <-semaphore }()
 
-// IsPeerInActiveGroup checks if the given peer is part of a group that is used
-// in an active DNS, route, or ACL configuration.
-func isPeerInActiveGroup(account *Account, peerID string) bool {
-	peerGroupIDs := make([]string, 0)
-	for _, group := range account.Groups {
-		if slices.Contains(group.Peers, peerID) {
-			peerGroupIDs = append(peerGroupIDs, group.ID)
-		}
+			postureChecks := am.getPeerPostureChecks(account, p)
+			remotePeerNetworkMap := account.GetPeerNetworkMap(ctx, p.ID, customZone, approvedPeersMap, am.metrics.AccountManagerMetrics())
+			update := toSyncResponse(ctx, nil, p, nil, remotePeerNetworkMap, am.GetDNSDomain(), postureChecks, dnsCache)
+			am.peersUpdateManager.SendUpdate(ctx, p.ID, &UpdateMessage{Update: update, NetworkMap: remotePeerNetworkMap, Checks: postureChecks})
+		}(peer)
 	}
-	return areGroupChangesAffectPeers(account, peerGroupIDs)
+
+	wg.Wait()
 }

management/server/peer/peer.go

@@ -1,7 +1,6 @@
 package peer
 
 import (
-	"fmt"
 	"net"
 	"net/netip"
 	"slices"
@@ -241,7 +240,7 @@ func (p *Peer) FQDN(dnsDomain string) string {
 	if dnsDomain == "" {
 		return ""
 	}
-	return fmt.Sprintf("%s.%s", p.DNSLabel, dnsDomain)
+	return p.DNSLabel + "." + dnsDomain
 }
 
 // EventMeta returns activity event meta related to the peer

management/server/peer/peer_test.go

@@ -0,0 +1,31 @@
+package peer
+
+import (
+	"fmt"
+	"testing"
+)
+
+// FQDNOld is the original implementation for benchmarking purposes
+func (p *Peer) FQDNOld(dnsDomain string) string {
+	if dnsDomain == "" {
+		return ""
+	}
+	return fmt.Sprintf("%s.%s", p.DNSLabel, dnsDomain)
+}
+
+func BenchmarkFQDN(b *testing.B) {
+	p := &Peer{DNSLabel: "test-peer"}
+	dnsDomain := "example.com"
+
+	b.Run("Old", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			p.FQDNOld(dnsDomain)
+		}
+	})
+
+	b.Run("New", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			p.FQDN(dnsDomain)
+		}
+	})
+}

management/server/peer_test.go

@@ -2,16 +2,26 @@ package server
 
 import (
 	"context"
+	"fmt"
+	"io"
+	"net"
+	"net/netip"
+	"os"
 	"testing"
 	"time"
 
 	"github.com/rs/xid"
+	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 	"golang.zx2c4.com/wireguard/wgctrl/wgtypes"
 
+	nbdns "github.com/netbirdio/netbird/dns"
+	"github.com/netbirdio/netbird/management/domain"
+	"github.com/netbirdio/netbird/management/proto"
 	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/posture"
+	nbroute "github.com/netbirdio/netbird/route"
 )
 
 func TestPeer_LoginExpired(t *testing.T) {
@@ -635,155 +645,353 @@ func TestDefaultAccountManager_GetPeers(t *testing.T) {
 
 }
 
-func TestPeerAccountPeerUpdate(t *testing.T) {
-	manager, account, peer1, peer2, peer3 := setupNetworkMapTest(t)
+func setupTestAccountManager(b *testing.B, peers int, groups int) (*DefaultAccountManager, string, string, error) {
+	b.Helper()
 
-	err := manager.DeletePolicy(context.Background(), account.Id, account.Policies[0].ID, userID)
-	require.NoError(t, err)
+	manager, err := createManager(b)
+	if err != nil {
+		return nil, "", "", err
+	}
 
-	err = manager.SaveGroup(context.Background(), account.Id, userID, &nbgroup.Group{
-		ID:    "group-id",
-		Name:  "GroupA",
-		Peers: []string{peer1.ID, peer2.ID, peer3.ID},
-	})
-	require.NoError(t, err)
+	accountID := "test_account"
+	adminUser := "account_creator"
+	regularUser := "regular_user"
 
-	updMsg := manager.peersUpdateManager.CreateChannel(context.Background(), peer1.ID)
-	t.Cleanup(func() {
-		manager.peersUpdateManager.CloseChannel(context.Background(), peer1.ID)
-	})
+	account := newAccountWithId(context.Background(), accountID, adminUser, "")
+	account.Users[regularUser] = &User{
+		Id:   regularUser,
+		Role: UserRoleUser,
+	}
 
-	// create a user with auto groups
-	_, err = manager.SaveOrAddUser(context.Background(), account.Id, userID, &User{
-		Id:         "regularUser1",
-		AccountID:  account.Id,
-		Role:       UserRoleAdmin,
-		Issued:     UserIssuedAPI,
-		AutoGroups: []string{"group-id"},
-	}, true)
-	require.NoError(t, err)
-
-	var peer4 *nbpeer.Peer
-
-	// Updating not expired peer and peer expiration is enabled should not update account peers and not send peer update
-	t.Run("updating not expired peer and peer expiration is enabled", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		_, err := manager.UpdatePeer(context.Background(), account.Id, userID, peer2)
-		require.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(200 * time.Millisecond):
-			t.Error("timeout waiting for peerShouldNotReceiveUpdate")
+	// Create peers
+	for i := 0; i < peers; i++ {
+		peerKey, _ := wgtypes.GeneratePrivateKey()
+		peer := &nbpeer.Peer{
+			ID:       fmt.Sprintf("peer-%d", i),
+			DNSLabel: fmt.Sprintf("peer-%d", i),
+			Key:      peerKey.PublicKey().String(),
+			IP:       net.ParseIP(fmt.Sprintf("100.64.%d.%d", i/256, i%256)),
+			Status:   &nbpeer.PeerStatus{},
+			UserID:   regularUser,
 		}
-	})
+		account.Peers[peer.ID] = peer
+	}
 
-	// Adding peer with an unused group in active dns, route, acl should not update account peers and not send peer update
-	t.Run("adding peer with unused group", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		key, err := wgtypes.GeneratePrivateKey()
-		require.NoError(t, err)
-
-		expectedPeerKey := key.PublicKey().String()
-		peer4, _, _, err = manager.AddPeer(context.Background(), "", "regularUser1", &nbpeer.Peer{
-			Key:  expectedPeerKey,
-			Meta: nbpeer.PeerSystemMeta{Hostname: expectedPeerKey},
-		})
-		require.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(200 * time.Millisecond):
-			t.Error("timeout waiting for peerShouldNotReceiveUpdate")
+	// Create groups and policies
+	account.Policies = make([]*Policy, 0, groups)
+	for i := 0; i < groups; i++ {
+		groupID := fmt.Sprintf("group-%d", i)
+		group := &nbgroup.Group{
+			ID:   groupID,
+			Name: fmt.Sprintf("Group %d", i),
 		}
-	})
-
-	// Deleting peer with an unused group in active dns, route, acl should not update account peers and not send peer update
-	t.Run("deleting peer with unused group", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldNotReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		err = manager.DeletePeer(context.Background(), account.Id, peer4.ID, userID)
-		require.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(200 * time.Millisecond):
-			t.Error("timeout waiting for peerShouldNotReceiveUpdate")
+		for j := 0; j < peers/groups; j++ {
+			peerIndex := i*(peers/groups) + j
+			group.Peers = append(group.Peers, fmt.Sprintf("peer-%d", peerIndex))
 		}
-	})
+		account.Groups[groupID] = group
 
-	// use the group-id in policy
-	err = manager.SavePolicy(context.Background(), account.Id, userID, &Policy{
-		ID:      "policy",
-		Enabled: true,
-		Rules: []*PolicyRule{
-			{
-				Enabled:       true,
-				Sources:       []string{"group-id"},
-				Destinations:  []string{"group-id"},
-				Bidirectional: true,
-				Action:        PolicyTrafficActionAccept,
+		// Create a policy for this group
+		policy := &Policy{
+			ID:      fmt.Sprintf("policy-%d", i),
+			Name:    fmt.Sprintf("Policy for Group %d", i),
+			Enabled: true,
+			Rules: []*PolicyRule{
+				{
+					ID:            fmt.Sprintf("rule-%d", i),
+					Name:          fmt.Sprintf("Rule for Group %d", i),
+					Enabled:       true,
+					Sources:       []string{groupID},
+					Destinations:  []string{groupID},
+					Bidirectional: true,
+					Protocol:      PolicyRuleProtocolALL,
+					Action:        PolicyTrafficActionAccept,
+				},
+			},
+		}
+		account.Policies = append(account.Policies, policy)
+	}
+
+	account.PostureChecks = []*posture.Checks{
+		{
+			ID:   "PostureChecksAll",
+			Name: "All",
+			Checks: posture.ChecksDefinition{
+				NBVersionCheck: &posture.NBVersionCheck{
+					MinVersion: "0.0.1",
+				},
 			},
 		},
-	})
-	require.NoError(t, err)
+	}
 
-	// Adding peer with a used group in active dns, route or policy should update account peers and send peer update
-	t.Run("adding peer with used group", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		key, err := wgtypes.GeneratePrivateKey()
-		require.NoError(t, err)
-
-		expectedPeerKey := key.PublicKey().String()
-		peer4, _, _, err = manager.AddPeer(context.Background(), "", "regularUser1", &nbpeer.Peer{
-			Key:                    expectedPeerKey,
-			LoginExpirationEnabled: true,
-			Meta:                   nbpeer.PeerSystemMeta{Hostname: expectedPeerKey},
-		})
-		require.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(200 * time.Millisecond):
-			t.Error("timeout waiting for peerShouldReceiveUpdate")
-		}
-	})
-
-	//Deleting peer with a used group in active dns, route or acl should update account peers and send peer update
-	t.Run("deleting peer with used group", func(t *testing.T) {
-		done := make(chan struct{})
-		go func() {
-			peerShouldReceiveUpdate(t, updMsg)
-			close(done)
-		}()
-
-		err = manager.DeletePeer(context.Background(), account.Id, peer4.ID, userID)
-		require.NoError(t, err)
-
-		select {
-		case <-done:
-		case <-time.After(200 * time.Millisecond):
-			t.Error("timeout waiting for peerShouldReceiveUpdate")
-		}
-	})
+	err = manager.Store.SaveAccount(context.Background(), account)
+	if err != nil {
+		return nil, "", "", err
+	}
 
+	return manager, accountID, regularUser, nil
+}
+
+func BenchmarkGetPeers(b *testing.B) {
+	benchCases := []struct {
+		name   string
+		peers  int
+		groups int
+	}{
+		{"Small", 50, 5},
+		{"Medium", 500, 10},
+		{"Large", 5000, 20},
+		{"Small single", 50, 1},
+		{"Medium single", 500, 1},
+		{"Large 5", 5000, 5},
+	}
+
+	log.SetOutput(io.Discard)
+	defer log.SetOutput(os.Stderr)
+	for _, bc := range benchCases {
+		b.Run(bc.name, func(b *testing.B) {
+			manager, accountID, userID, err := setupTestAccountManager(b, bc.peers, bc.groups)
+			if err != nil {
+				b.Fatalf("Failed to setup test account manager: %v", err)
+			}
+
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				_, err := manager.GetPeers(context.Background(), accountID, userID)
+				if err != nil {
+					b.Fatalf("GetPeers failed: %v", err)
+				}
+			}
+		})
+	}
+}
+
+func BenchmarkUpdateAccountPeers(b *testing.B) {
+	benchCases := []struct {
+		name   string
+		peers  int
+		groups int
+	}{
+		{"Small", 50, 5},
+		{"Medium", 500, 10},
+		{"Large", 5000, 20},
+		{"Small single", 50, 1},
+		{"Medium single", 500, 1},
+		{"Large 5", 5000, 5},
+	}
+
+	log.SetOutput(io.Discard)
+	defer log.SetOutput(os.Stderr)
+
+	for _, bc := range benchCases {
+		b.Run(bc.name, func(b *testing.B) {
+			manager, accountID, _, err := setupTestAccountManager(b, bc.peers, bc.groups)
+			if err != nil {
+				b.Fatalf("Failed to setup test account manager: %v", err)
+			}
+
+			ctx := context.Background()
+
+			account, err := manager.Store.GetAccount(ctx, accountID)
+			if err != nil {
+				b.Fatalf("Failed to get account: %v", err)
+			}
+
+			peerChannels := make(map[string]chan *UpdateMessage)
+
+			for peerID := range account.Peers {
+				peerChannels[peerID] = make(chan *UpdateMessage, channelBufferSize)
+			}
+
+			manager.peersUpdateManager.peerChannels = peerChannels
+
+			b.ResetTimer()
+			start := time.Now()
+
+			for i := 0; i < b.N; i++ {
+				manager.updateAccountPeers(ctx, account)
+			}
+
+			duration := time.Since(start)
+			b.ReportMetric(float64(duration.Nanoseconds())/float64(b.N)/1e6, "ms/op")
+			b.ReportMetric(0, "ns/op")
+		})
+	}
+}
+
+func TestToSyncResponse(t *testing.T) {
+	_, ipnet, err := net.ParseCIDR("192.168.1.0/24")
+	if err != nil {
+		t.Fatal(err)
+	}
+	domainList, err := domain.FromStringList([]string{"example.com"})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	config := &Config{
+		Signal: &Host{
+			Proto:    "https",
+			URI:      "signal.uri",
+			Username: "",
+			Password: "",
+		},
+		Stuns: []*Host{{URI: "stun.uri", Proto: UDP}},
+		TURNConfig: &TURNConfig{
+			Turns: []*Host{{URI: "turn.uri", Proto: UDP, Username: "turn-user", Password: "turn-pass"}},
+		},
+	}
+	peer := &nbpeer.Peer{
+		IP:         net.ParseIP("192.168.1.1"),
+		SSHEnabled: true,
+		Key:        "peer-key",
+		DNSLabel:   "peer1",
+		SSHKey:     "peer1-ssh-key",
+	}
+	turnCredentials := &TURNCredentials{
+		Username: "turn-user",
+		Password: "turn-pass",
+	}
+	networkMap := &NetworkMap{
+		Network:      &Network{Net: *ipnet, Serial: 1000},
+		Peers:        []*nbpeer.Peer{{IP: net.ParseIP("192.168.1.2"), Key: "peer2-key", DNSLabel: "peer2", SSHEnabled: true, SSHKey: "peer2-ssh-key"}},
+		OfflinePeers: []*nbpeer.Peer{{IP: net.ParseIP("192.168.1.3"), Key: "peer3-key", DNSLabel: "peer3", SSHEnabled: true, SSHKey: "peer3-ssh-key"}},
+		Routes: []*nbroute.Route{
+			{
+				ID:          "route1",
+				Network:     netip.MustParsePrefix("10.0.0.0/24"),
+				Domains:     domainList,
+				KeepRoute:   true,
+				NetID:       "route1",
+				Peer:        "peer1",
+				NetworkType: 1,
+				Masquerade:  true,
+				Metric:      9999,
+				Enabled:     true,
+			},
+		},
+		DNSConfig: nbdns.Config{
+			ServiceEnable: true,
+			NameServerGroups: []*nbdns.NameServerGroup{
+				{
+					NameServers: []nbdns.NameServer{{
+						IP:     netip.MustParseAddr("8.8.8.8"),
+						NSType: nbdns.UDPNameServerType,
+						Port:   nbdns.DefaultDNSPort,
+					}},
+					Primary:              true,
+					Domains:              []string{"example.com"},
+					Enabled:              true,
+					SearchDomainsEnabled: true,
+				},
+				{
+					ID: "ns1",
+					NameServers: []nbdns.NameServer{{
+						IP:     netip.MustParseAddr("1.1.1.1"),
+						NSType: nbdns.UDPNameServerType,
+						Port:   nbdns.DefaultDNSPort,
+					}},
+					Groups:               []string{"group1"},
+					Primary:              true,
+					Domains:              []string{"example.com"},
+					Enabled:              true,
+					SearchDomainsEnabled: true,
+				},
+			},
+			CustomZones: []nbdns.CustomZone{{Domain: "example.com", Records: []nbdns.SimpleRecord{{Name: "example.com", Type: 1, Class: "IN", TTL: 60, RData: "100.64.0.1"}}}},
+		},
+		FirewallRules: []*FirewallRule{
+			{PeerIP: "192.168.1.2", Direction: firewallRuleDirectionIN, Action: string(PolicyTrafficActionAccept), Protocol: string(PolicyRuleProtocolTCP), Port: "80"},
+		},
+	}
+	dnsName := "example.com"
+	checks := []*posture.Checks{
+		{
+			Checks: posture.ChecksDefinition{
+				ProcessCheck: &posture.ProcessCheck{
+					Processes: []posture.Process{{LinuxPath: "/usr/bin/netbird"}},
+				},
+			},
+		},
+	}
+	dnsCache := &DNSConfigCache{}
+
+	response := toSyncResponse(context.Background(), config, peer, turnCredentials, networkMap, dnsName, checks, dnsCache)
+
+	assert.NotNil(t, response)
+	// assert peer config
+	assert.Equal(t, "192.168.1.1/24", response.PeerConfig.Address)
+	assert.Equal(t, "peer1.example.com", response.PeerConfig.Fqdn)
+	assert.Equal(t, true, response.PeerConfig.SshConfig.SshEnabled)
+	// assert wiretrustee config
+	assert.Equal(t, "signal.uri", response.WiretrusteeConfig.Signal.Uri)
+	assert.Equal(t, proto.HostConfig_HTTPS, response.WiretrusteeConfig.Signal.GetProtocol())
+	assert.Equal(t, "stun.uri", response.WiretrusteeConfig.Stuns[0].Uri)
+	assert.Equal(t, "turn.uri", response.WiretrusteeConfig.Turns[0].HostConfig.GetUri())
+	assert.Equal(t, "turn-user", response.WiretrusteeConfig.Turns[0].User)
+	assert.Equal(t, "turn-pass", response.WiretrusteeConfig.Turns[0].Password)
+	// assert RemotePeers
+	assert.Equal(t, 1, len(response.RemotePeers))
+	assert.Equal(t, "192.168.1.2/32", response.RemotePeers[0].AllowedIps[0])
+	assert.Equal(t, "peer2-key", response.RemotePeers[0].WgPubKey)
+	assert.Equal(t, "peer2.example.com", response.RemotePeers[0].GetFqdn())
+	assert.Equal(t, false, response.RemotePeers[0].GetSshConfig().GetSshEnabled())
+	assert.Equal(t, []byte("peer2-ssh-key"), response.RemotePeers[0].GetSshConfig().GetSshPubKey())
+	// assert network map
+	assert.Equal(t, uint64(1000), response.NetworkMap.Serial)
+	assert.Equal(t, "192.168.1.1/24", response.NetworkMap.PeerConfig.Address)
+	assert.Equal(t, "peer1.example.com", response.NetworkMap.PeerConfig.Fqdn)
+	assert.Equal(t, true, response.NetworkMap.PeerConfig.SshConfig.SshEnabled)
+	// assert network map RemotePeers
+	assert.Equal(t, 1, len(response.NetworkMap.RemotePeers))
+	assert.Equal(t, "192.168.1.2/32", response.NetworkMap.RemotePeers[0].AllowedIps[0])
+	assert.Equal(t, "peer2-key", response.NetworkMap.RemotePeers[0].WgPubKey)
+	assert.Equal(t, "peer2.example.com", response.NetworkMap.RemotePeers[0].GetFqdn())
+	assert.Equal(t, []byte("peer2-ssh-key"), response.NetworkMap.RemotePeers[0].GetSshConfig().GetSshPubKey())
+	// assert network map OfflinePeers
+	assert.Equal(t, 1, len(response.NetworkMap.OfflinePeers))
+	assert.Equal(t, "192.168.1.3/32", response.NetworkMap.OfflinePeers[0].AllowedIps[0])
+	assert.Equal(t, "peer3-key", response.NetworkMap.OfflinePeers[0].WgPubKey)
+	assert.Equal(t, "peer3.example.com", response.NetworkMap.OfflinePeers[0].GetFqdn())
+	assert.Equal(t, []byte("peer3-ssh-key"), response.NetworkMap.OfflinePeers[0].GetSshConfig().GetSshPubKey())
+	// assert network map Routes
+	assert.Equal(t, 1, len(response.NetworkMap.Routes))
+	assert.Equal(t, "10.0.0.0/24", response.NetworkMap.Routes[0].Network)
+	assert.Equal(t, "route1", response.NetworkMap.Routes[0].ID)
+	assert.Equal(t, "peer1", response.NetworkMap.Routes[0].Peer)
+	assert.Equal(t, "example.com", response.NetworkMap.Routes[0].Domains[0])
+	assert.Equal(t, true, response.NetworkMap.Routes[0].KeepRoute)
+	assert.Equal(t, true, response.NetworkMap.Routes[0].Masquerade)
+	assert.Equal(t, int64(9999), response.NetworkMap.Routes[0].Metric)
+	assert.Equal(t, int64(1), response.NetworkMap.Routes[0].NetworkType)
+	assert.Equal(t, "route1", response.NetworkMap.Routes[0].NetID)
+	// assert network map DNSConfig
+	assert.Equal(t, true, response.NetworkMap.DNSConfig.ServiceEnable)
+	assert.Equal(t, 1, len(response.NetworkMap.DNSConfig.CustomZones))
+	assert.Equal(t, 2, len(response.NetworkMap.DNSConfig.NameServerGroups))
+	// assert network map DNSConfig.CustomZones
+	assert.Equal(t, "example.com", response.NetworkMap.DNSConfig.CustomZones[0].Domain)
+	assert.Equal(t, 1, len(response.NetworkMap.DNSConfig.CustomZones[0].Records))
+	assert.Equal(t, "example.com", response.NetworkMap.DNSConfig.CustomZones[0].Records[0].Name)
+	assert.Equal(t, int64(1), response.NetworkMap.DNSConfig.CustomZones[0].Records[0].Type)
+	assert.Equal(t, "IN", response.NetworkMap.DNSConfig.CustomZones[0].Records[0].Class)
+	assert.Equal(t, int64(60), response.NetworkMap.DNSConfig.CustomZones[0].Records[0].TTL)
+	assert.Equal(t, "100.64.0.1", response.NetworkMap.DNSConfig.CustomZones[0].Records[0].RData)
+	// assert network map DNSConfig.NameServerGroups
+	assert.Equal(t, true, response.NetworkMap.DNSConfig.NameServerGroups[0].Primary)
+	assert.Equal(t, true, response.NetworkMap.DNSConfig.NameServerGroups[0].SearchDomainsEnabled)
+	assert.Equal(t, "example.com", response.NetworkMap.DNSConfig.NameServerGroups[0].Domains[0])
+	assert.Equal(t, "8.8.8.8", response.NetworkMap.DNSConfig.NameServerGroups[0].NameServers[0].GetIP())
+	assert.Equal(t, int64(1), response.NetworkMap.DNSConfig.NameServerGroups[0].NameServers[0].GetNSType())
+	assert.Equal(t, int64(53), response.NetworkMap.DNSConfig.NameServerGroups[0].NameServers[0].GetPort())
+	// assert network map Firewall
+	assert.Equal(t, 1, len(response.NetworkMap.FirewallRules))
+	assert.Equal(t, "192.168.1.2", response.NetworkMap.FirewallRules[0].PeerIP)
+	assert.Equal(t, proto.FirewallRule_IN, response.NetworkMap.FirewallRules[0].Direction)
+	assert.Equal(t, proto.FirewallRule_ACCEPT, response.NetworkMap.FirewallRules[0].Action)
+	assert.Equal(t, proto.FirewallRule_TCP, response.NetworkMap.FirewallRules[0].Protocol)
+	assert.Equal(t, "80", response.NetworkMap.FirewallRules[0].Port)
+	// assert posture checks
+	assert.Equal(t, 1, len(response.Checks))
+	assert.Equal(t, "/usr/bin/netbird", response.Checks[0].Files[0])
 }

management/server/policy.go

@@ -223,7 +223,6 @@ type FirewallRule struct {
 //
 // This function returns the list of peers and firewall rules that are applicable to a given peer.
 func (a *Account) getPeerConnectionResources(ctx context.Context, peerID string, validatedPeersMap map[string]struct{}) ([]*nbpeer.Peer, []*FirewallRule) {
-
 	generateResources, getAccumulatedResources := a.connResourcesGenerator(ctx)
 	for _, policy := range a.Policies {
 		if !policy.Enabled {
@@ -235,8 +234,8 @@ func (a *Account) getPeerConnectionResources(ctx context.Context, peerID string,
 				continue
 			}
 
-			sourcePeers, peerInSources := getAllPeersFromGroups(ctx, a, rule.Sources, peerID, policy.SourcePostureChecks, validatedPeersMap)
-			destinationPeers, peerInDestinations := getAllPeersFromGroups(ctx, a, rule.Destinations, peerID, nil, validatedPeersMap)
+			sourcePeers, peerInSources := a.getAllPeersFromGroups(ctx, rule.Sources, peerID, policy.SourcePostureChecks, validatedPeersMap)
+			destinationPeers, peerInDestinations := a.getAllPeersFromGroups(ctx, rule.Destinations, peerID, nil, validatedPeersMap)
 
 			if rule.Bidirectional {
 				if peerInSources {
@@ -300,8 +299,8 @@ func (a *Account) connResourcesGenerator(ctx context.Context) (func(*PolicyRule,
 					fr.PeerIP = "0.0.0.0"
 				}
 
-				ruleID := (rule.ID + fr.PeerIP + strconv.Itoa(direction) +
-					fr.Protocol + fr.Action + strings.Join(rule.Ports, ","))
+				ruleID := rule.ID + fr.PeerIP + strconv.Itoa(direction) +
+					fr.Protocol + fr.Action + strings.Join(rule.Ports, ",")
 				if _, ok := rulesExists[ruleID]; ok {
 					continue
 				}
@@ -325,7 +324,7 @@ func (a *Account) connResourcesGenerator(ctx context.Context) (func(*PolicyRule,
 
 // GetPolicy from the store
 func (am *DefaultAccountManager) GetPolicy(ctx context.Context, accountID, policyID, userID string) (*Policy, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -353,7 +352,7 @@ func (am *DefaultAccountManager) GetPolicy(ctx context.Context, accountID, polic
 
 // SavePolicy in the store
 func (am *DefaultAccountManager) SavePolicy(ctx context.Context, accountID, userID string, policy *Policy) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -383,7 +382,7 @@ func (am *DefaultAccountManager) SavePolicy(ctx context.Context, accountID, user
 
 // DeletePolicy from the store
 func (am *DefaultAccountManager) DeletePolicy(ctx context.Context, accountID, policyID, userID string) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -412,7 +411,7 @@ func (am *DefaultAccountManager) DeletePolicy(ctx context.Context, accountID, po
 
 // ListPolicies from the store
 func (am *DefaultAccountManager) ListPolicies(ctx context.Context, accountID, userID string) ([]*Policy, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -505,23 +504,23 @@ func toProtocolFirewallRules(update []*FirewallRule) []*proto.FirewallRule {
 //
 // Important: Posture checks are applicable only to source group peers,
 // for destination group peers, call this method with an empty list of sourcePostureChecksIDs
-func getAllPeersFromGroups(ctx context.Context, account *Account, groups []string, peerID string, sourcePostureChecksIDs []string, validatedPeersMap map[string]struct{}) ([]*nbpeer.Peer, bool) {
+func (a *Account) getAllPeersFromGroups(ctx context.Context, groups []string, peerID string, sourcePostureChecksIDs []string, validatedPeersMap map[string]struct{}) ([]*nbpeer.Peer, bool) {
 	peerInGroups := false
 	filteredPeers := make([]*nbpeer.Peer, 0, len(groups))
 	for _, g := range groups {
-		group, ok := account.Groups[g]
+		group, ok := a.Groups[g]
 		if !ok {
 			continue
 		}
 
 		for _, p := range group.Peers {
-			peer, ok := account.Peers[p]
+			peer, ok := a.Peers[p]
 			if !ok || peer == nil {
 				continue
 			}
 
 			// validate the peer based on policy posture checks applied
-			isValid := account.validatePostureChecksOnPeer(ctx, sourcePostureChecksIDs, peer.ID)
+			isValid := a.validatePostureChecksOnPeer(ctx, sourcePostureChecksIDs, peer.ID)
 			if !isValid {
 				continue
 			}
@@ -549,7 +548,7 @@ func (a *Account) validatePostureChecksOnPeer(ctx context.Context, sourcePosture
 	}
 
 	for _, postureChecksID := range sourcePostureChecksID {
-		postureChecks := getPostureChecks(a, postureChecksID)
+		postureChecks := a.getPostureChecks(postureChecksID)
 		if postureChecks == nil {
 			continue
 		}
@@ -567,8 +566,8 @@ func (a *Account) validatePostureChecksOnPeer(ctx context.Context, sourcePosture
 	return true
 }
 
-func getPostureChecks(account *Account, postureChecksID string) *posture.Checks {
-	for _, postureChecks := range account.PostureChecks {
+func (a *Account) getPostureChecks(postureChecksID string) *posture.Checks {
+	for _, postureChecks := range a.PostureChecks {
 		if postureChecks.ID == postureChecksID {
 			return postureChecks
 		}

management/server/posture_checks.go

@@ -15,7 +15,7 @@ const (
 )
 
 func (am *DefaultAccountManager) GetPostureChecks(ctx context.Context, accountID, postureChecksID, userID string) (*posture.Checks, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -42,7 +42,7 @@ func (am *DefaultAccountManager) GetPostureChecks(ctx context.Context, accountID
 }
 
 func (am *DefaultAccountManager) SavePostureChecks(ctx context.Context, accountID, userID string, postureChecks *posture.Checks) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -91,7 +91,7 @@ func (am *DefaultAccountManager) SavePostureChecks(ctx context.Context, accountI
 }
 
 func (am *DefaultAccountManager) DeletePostureChecks(ctx context.Context, accountID, postureChecksID, userID string) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -123,7 +123,7 @@ func (am *DefaultAccountManager) DeletePostureChecks(ctx context.Context, accoun
 }
 
 func (am *DefaultAccountManager) ListPostureChecks(ctx context.Context, accountID, userID string) ([]*posture.Checks, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)

management/server/route.go

@@ -17,7 +17,7 @@ import (
 
 // GetRoute gets a route object from account and route IDs
 func (am *DefaultAccountManager) GetRoute(ctx context.Context, accountID string, routeID route.ID, userID string) (*route.Route, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -126,7 +126,7 @@ func getRouteDescriptor(prefix netip.Prefix, domains domain.List) string {
 
 // CreateRoute creates and saves a new route
 func (am *DefaultAccountManager) CreateRoute(ctx context.Context, accountID string, prefix netip.Prefix, networkType route.NetworkType, domains domain.List, peerID string, peerGroupIDs []string, description string, netID route.NetID, masquerade bool, metric int, groups []string, enabled bool, userID string, keepRoute bool) (*route.Route, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -216,7 +216,7 @@ func (am *DefaultAccountManager) CreateRoute(ctx context.Context, accountID stri
 
 // SaveRoute saves route
 func (am *DefaultAccountManager) SaveRoute(ctx context.Context, accountID, userID string, routeToSave *route.Route) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	if routeToSave == nil {
@@ -288,7 +288,7 @@ func (am *DefaultAccountManager) SaveRoute(ctx context.Context, accountID, userI
 
 // DeleteRoute deletes route with routeID
 func (am *DefaultAccountManager) DeleteRoute(ctx context.Context, accountID string, routeID route.ID, userID string) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -318,7 +318,7 @@ func (am *DefaultAccountManager) DeleteRoute(ctx context.Context, accountID stri
 
 // ListRoutes returns a list of routes from account
 func (am *DefaultAccountManager) ListRoutes(ctx context.Context, accountID, userID string) ([]*route.Route, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)

management/server/route_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/netbirdio/netbird/management/server/activity"
 	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	nbpeer "github.com/netbirdio/netbird/management/server/peer"
+	"github.com/netbirdio/netbird/management/server/telemetry"
 	"github.com/netbirdio/netbird/route"
 )
 
@@ -1234,7 +1235,11 @@ func createRouterManager(t *testing.T) (*DefaultAccountManager, error) {
 		return nil, err
 	}
 	eventStore := &activity.InMemoryEventStore{}
-	return BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.selfhosted", eventStore, nil, false, MocIntegratedValidator{})
+
+	metrics, err := telemetry.NewDefaultAppMetrics(context.Background())
+	require.NoError(t, err)
+
+	return BuildManager(context.Background(), store, NewPeersUpdateManager(nil), nil, "", "netbird.selfhosted", eventStore, nil, false, MocIntegratedValidator{}, metrics)
 }
 
 func createRouterStore(t *testing.T) (Store, error) {

management/server/setupkey.go

@@ -210,7 +210,7 @@ func Hash(s string) uint32 {
 // and adds it to the specified account. A list of autoGroups IDs can be empty.
 func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID string, keyName string, keyType SetupKeyType,
 	expiresIn time.Duration, autoGroups []string, usageLimit int, userID string, ephemeral bool) (*SetupKey, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	keyDuration := DefaultSetupKeyDuration
@@ -223,10 +223,8 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 		return nil, err
 	}
 
-	for _, group := range autoGroups {
-		if _, ok := account.Groups[group]; !ok {
-			return nil, status.Errorf(status.NotFound, "group %s doesn't exist", group)
-		}
+	if err := validateSetupKeyAutoGroups(account, autoGroups); err != nil {
+		return nil, err
 	}
 
 	setupKey := GenerateSetupKey(keyName, keyType, keyDuration, autoGroups, usageLimit, ephemeral)
@@ -256,7 +254,7 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 // (e.g. the key itself, creation date, ID, etc).
 // These properties are overwritten: Name, AutoGroups, Revoked. The rest is copied from the existing key.
 func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID string, keyToSave *SetupKey, userID string) (*SetupKey, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	if keyToSave == nil {
@@ -279,6 +277,10 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 		return nil, status.Errorf(status.NotFound, "setup key not found")
 	}
 
+	if err := validateSetupKeyAutoGroups(account, keyToSave.AutoGroups); err != nil {
+		return nil, err
+	}
+
 	// only auto groups, revoked status, and name can be updated for now
 	newKey := oldKey.Copy()
 	newKey.Name = keyToSave.Name
@@ -326,7 +328,7 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 
 // ListSetupKeys returns a list of all setup keys of the account
 func (am *DefaultAccountManager) ListSetupKeys(ctx context.Context, accountID, userID string) ([]*SetupKey, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 	account, err := am.Store.GetAccount(ctx, accountID)
 	if err != nil {
@@ -358,7 +360,7 @@ func (am *DefaultAccountManager) ListSetupKeys(ctx context.Context, accountID, u
 
 // GetSetupKey looks up a SetupKey by KeyID, returns NotFound error if not found.
 func (am *DefaultAccountManager) GetSetupKey(ctx context.Context, accountID, userID, keyID string) (*SetupKey, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -397,3 +399,16 @@ func (am *DefaultAccountManager) GetSetupKey(ctx context.Context, accountID, use
 
 	return foundKey, nil
 }
+
+func validateSetupKeyAutoGroups(account *Account, autoGroups []string) error {
+	for _, group := range autoGroups {
+		g, ok := account.Groups[group]
+		if !ok {
+			return status.Errorf(status.NotFound, "group %s doesn't exist", group)
+		}
+		if g.Name == "All" {
+			return status.Errorf(status.InvalidArgument, "can't add All group to the setup key")
+		}
+	}
+	return nil
+}

management/server/setupkey_test.go

@@ -27,10 +27,17 @@ func TestDefaultAccountManager_SaveSetupKey(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	err = manager.SaveGroup(context.Background(), account.Id, userID, &nbgroup.Group{
-		ID:    "group_1",
-		Name:  "group_name_1",
-		Peers: []string{},
+	err = manager.SaveGroups(context.Background(), account.Id, userID, []*nbgroup.Group{
+		{
+			ID:    "group_1",
+			Name:  "group_name_1",
+			Peers: []string{},
+		},
+		{
+			ID:    "group_2",
+			Name:  "group_name_2",
+			Peers: []string{},
+		},
 	})
 	if err != nil {
 		t.Fatal(err)
@@ -71,6 +78,19 @@ func TestDefaultAccountManager_SaveSetupKey(t *testing.T) {
 	assert.NotEmpty(t, ev.Meta["key"])
 	assert.Equal(t, userID, ev.InitiatorID)
 	assert.Equal(t, key.Id, ev.TargetID)
+
+	groupAll, err := account.GetGroupAll()
+	assert.NoError(t, err)
+
+	// saving setup key with All group assigned to auto groups should return error
+	autoGroups = append(autoGroups, groupAll.ID)
+	_, err = manager.SaveSetupKey(context.Background(), account.Id, &SetupKey{
+		Id:         key.Id,
+		Name:       newKeyName,
+		Revoked:    revoked,
+		AutoGroups: autoGroups,
+	}, userID)
+	assert.Error(t, err, "should not save setup key with All group assigned in auto groups")
 }
 
 func TestDefaultAccountManager_CreateSetupKey(t *testing.T) {
@@ -103,6 +123,9 @@ func TestDefaultAccountManager_CreateSetupKey(t *testing.T) {
 		t.Fatal(err)
 	}
 
+	groupAll, err := account.GetGroupAll()
+	assert.NoError(t, err)
+
 	type testCase struct {
 		name string
 
@@ -135,8 +158,14 @@ func TestDefaultAccountManager_CreateSetupKey(t *testing.T) {
 		expectedGroups:  []string{"FAKE"},
 		expectedFailure: true,
 	}
+	testCase3 := testCase{
+		name:            "Create Setup Key should fail because of All group",
+		expectedKeyName: "my-test-key",
+		expectedGroups:  []string{groupAll.ID},
+		expectedFailure: true,
+	}
 
-	for _, tCase := range []testCase{testCase1, testCase2} {
+	for _, tCase := range []testCase{testCase1, testCase2, testCase3} {
 		t.Run(tCase.name, func(t *testing.T) {
 			key, err := manager.CreateSetupKey(context.Background(), account.Id, tCase.expectedKeyName, SetupKeyReusable, expiresIn,
 				tCase.expectedGroups, SetupKeyUnlimitedUsage, userID, false)

management/server/sql_store.go

@@ -31,14 +31,16 @@ import (
 )
 
 const (
-	storeSqliteFileName = "store.db"
-	idQueryCondition    = "id = ?"
+	storeSqliteFileName        = "store.db"
+	idQueryCondition           = "id = ?"
+	accountAndIDQueryCondition = "account_id = ? and id = ?"
+	peerNotFoundFMT            = "peer %s not found"
 )
 
 // SqlStore represents an account storage backed by a Sql DB persisted to disk
 type SqlStore struct {
 	db                *gorm.DB
-	accountLocks      sync.Map
+	resourceLocks     sync.Map
 	globalAccountLock sync.Mutex
 	metrics           telemetry.AppMetrics
 	installationPK    int
@@ -96,33 +98,35 @@ func (s *SqlStore) AcquireGlobalLock(ctx context.Context) (unlock func()) {
 	return unlock
 }
 
-func (s *SqlStore) AcquireAccountWriteLock(ctx context.Context, accountID string) (unlock func()) {
-	log.WithContext(ctx).Tracef("acquiring write lock for account %s", accountID)
+// AcquireWriteLockByUID acquires an ID lock for writing to a resource and returns a function that releases the lock
+func (s *SqlStore) AcquireWriteLockByUID(ctx context.Context, uniqueID string) (unlock func()) {
+	log.WithContext(ctx).Tracef("acquiring write lock for ID %s", uniqueID)
 
 	start := time.Now()
-	value, _ := s.accountLocks.LoadOrStore(accountID, &sync.RWMutex{})
+	value, _ := s.resourceLocks.LoadOrStore(uniqueID, &sync.RWMutex{})
 	mtx := value.(*sync.RWMutex)
 	mtx.Lock()
 
 	unlock = func() {
 		mtx.Unlock()
-		log.WithContext(ctx).Tracef("released write lock for account %s in %v", accountID, time.Since(start))
+		log.WithContext(ctx).Tracef("released write lock for ID %s in %v", uniqueID, time.Since(start))
 	}
 
 	return unlock
 }
 
-func (s *SqlStore) AcquireAccountReadLock(ctx context.Context, accountID string) (unlock func()) {
-	log.WithContext(ctx).Tracef("acquiring read lock for account %s", accountID)
+// AcquireReadLockByUID acquires an ID lock for writing to a resource and returns a function that releases the lock
+func (s *SqlStore) AcquireReadLockByUID(ctx context.Context, uniqueID string) (unlock func()) {
+	log.WithContext(ctx).Tracef("acquiring read lock for ID %s", uniqueID)
 
 	start := time.Now()
-	value, _ := s.accountLocks.LoadOrStore(accountID, &sync.RWMutex{})
+	value, _ := s.resourceLocks.LoadOrStore(uniqueID, &sync.RWMutex{})
 	mtx := value.(*sync.RWMutex)
 	mtx.RLock()
 
 	unlock = func() {
 		mtx.RUnlock()
-		log.WithContext(ctx).Tracef("released read lock for account %s in %v", accountID, time.Since(start))
+		log.WithContext(ctx).Tracef("released read lock for ID %s in %v", uniqueID, time.Since(start))
 	}
 
 	return unlock
@@ -271,6 +275,38 @@ func (s *SqlStore) GetInstallationID() string {
 	return installation.InstallationIDValue
 }
 
+func (s *SqlStore) SavePeer(ctx context.Context, accountID string, peer *nbpeer.Peer) error {
+	// To maintain data integrity, we create a copy of the peer's to prevent unintended updates to other fields.
+	peerCopy := peer.Copy()
+	peerCopy.AccountID = accountID
+
+	err := s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
+		// check if peer exists before saving
+		var peerID string
+		result := tx.Model(&nbpeer.Peer{}).Select("id").Find(&peerID, accountAndIDQueryCondition, accountID, peer.ID)
+		if result.Error != nil {
+			return result.Error
+		}
+
+		if peerID == "" {
+			return status.Errorf(status.NotFound, peerNotFoundFMT, peer.ID)
+		}
+
+		result = tx.Model(&nbpeer.Peer{}).Where(accountAndIDQueryCondition, accountID, peer.ID).Save(peerCopy)
+		if result.Error != nil {
+			return result.Error
+		}
+
+		return nil
+	})
+
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
 func (s *SqlStore) SavePeerStatus(accountID, peerID string, peerStatus nbpeer.PeerStatus) error {
 	var peerCopy nbpeer.Peer
 	peerCopy.Status = &peerStatus
@@ -281,14 +317,14 @@ func (s *SqlStore) SavePeerStatus(accountID, peerID string, peerStatus nbpeer.Pe
 	}
 	result := s.db.Model(&nbpeer.Peer{}).
 		Select(fieldsToUpdate).
-		Where("account_id = ? AND id = ?", accountID, peerID).
+		Where(accountAndIDQueryCondition, accountID, peerID).
 		Updates(&peerCopy)
 	if result.Error != nil {
 		return result.Error
 	}
 
 	if result.RowsAffected == 0 {
-		return status.Errorf(status.NotFound, "peer %s not found", peerID)
+		return status.Errorf(status.NotFound, peerNotFoundFMT, peerID)
 	}
 
 	return nil
@@ -302,7 +338,7 @@ func (s *SqlStore) SavePeerLocation(accountID string, peerWithLocation *nbpeer.P
 	peerCopy.Location = peerWithLocation.Location
 
 	result := s.db.Model(&nbpeer.Peer{}).
-		Where("account_id = ? and id = ?", accountID, peerWithLocation.ID).
+		Where(accountAndIDQueryCondition, accountID, peerWithLocation.ID).
 		Updates(peerCopy)
 
 	if result.Error != nil {
@@ -310,7 +346,7 @@ func (s *SqlStore) SavePeerLocation(accountID string, peerWithLocation *nbpeer.P
 	}
 
 	if result.RowsAffected == 0 {
-		return status.Errorf(status.NotFound, "peer %s not found", peerWithLocation.ID)
+		return status.Errorf(status.NotFound, peerNotFoundFMT, peerWithLocation.ID)
 	}
 
 	return nil
@@ -644,7 +680,7 @@ func (s *SqlStore) GetAccountSettings(ctx context.Context, accountID string) (*S
 func (s *SqlStore) SaveUserLastLogin(accountID, userID string, lastLogin time.Time) error {
 	var user User
 
-	result := s.db.First(&user, "account_id = ? and id = ?", accountID, userID)
+	result := s.db.First(&user, accountAndIDQueryCondition, accountID, userID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return status.Errorf(status.NotFound, "user %s not found", userID)

management/server/sql_store_test.go

@@ -362,6 +362,54 @@ func TestSqlite_GetAccount(t *testing.T) {
 	require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
 }
 
+func TestSqlite_SavePeer(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skip("The SQLite store is not properly supported by Windows yet")
+	}
+
+	store := newSqliteStoreFromFile(t, "testdata/store.json")
+
+	account, err := store.GetAccount(context.Background(), "bf1c8084-ba50-4ce7-9439-34653001fc3b")
+	require.NoError(t, err)
+
+	// save status of non-existing peer
+	peer := &nbpeer.Peer{
+		Key:      "peerkey",
+		ID:       "testpeer",
+		SetupKey: "peerkeysetupkey",
+		IP:       net.IP{127, 0, 0, 1},
+		Meta:     nbpeer.PeerSystemMeta{Hostname: "testingpeer"},
+		Name:     "peer name",
+		Status:   &nbpeer.PeerStatus{Connected: true, LastSeen: time.Now().UTC()},
+	}
+	ctx := context.Background()
+	err = store.SavePeer(ctx, account.Id, peer)
+	assert.Error(t, err)
+	parsedErr, ok := status.FromError(err)
+	require.True(t, ok)
+	require.Equal(t, status.NotFound, parsedErr.Type(), "should return not found error")
+
+	// save new status of existing peer
+	account.Peers[peer.ID] = peer
+
+	err = store.SaveAccount(context.Background(), account)
+	require.NoError(t, err)
+
+	updatedPeer := peer.Copy()
+	updatedPeer.Status.Connected = false
+	updatedPeer.Meta.Hostname = "updatedpeer"
+
+	err = store.SavePeer(ctx, account.Id, updatedPeer)
+	require.NoError(t, err)
+
+	account, err = store.GetAccount(context.Background(), account.Id)
+	require.NoError(t, err)
+
+	actual := account.Peers[peer.ID]
+	assert.Equal(t, updatedPeer.Status, actual.Status)
+	assert.Equal(t, updatedPeer.Meta, actual.Meta)
+}
+
 func TestSqlite_SavePeerStatus(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("The SQLite store is not properly supported by Windows yet")
@@ -402,7 +450,19 @@ func TestSqlite_SavePeerStatus(t *testing.T) {
 
 	actual := account.Peers["testpeer"].Status
 	assert.Equal(t, newStatus, *actual)
+
+	newStatus.Connected = true
+
+	err = store.SavePeerStatus(account.Id, "testpeer", newStatus)
+	require.NoError(t, err)
+
+	account, err = store.GetAccount(context.Background(), account.Id)
+	require.NoError(t, err)
+
+	actual = account.Peers["testpeer"].Status
+	assert.Equal(t, newStatus, *actual)
 }
+
 func TestSqlite_SavePeerLocation(t *testing.T) {
 	if runtime.GOOS == "windows" {
 		t.Skip("The SQLite store is not properly supported by Windows yet")

management/server/status/error.go

@@ -95,3 +95,8 @@ func NewUserNotFoundError(userKey string) error {
 func NewPeerNotRegisteredError() error {
 	return Errorf(Unauthenticated, "peer is not registered")
 }
+
+// NewPeerLoginExpiredError creates a new Error with PermissionDenied type for an expired peer
+func NewPeerLoginExpiredError() error {
+	return Errorf(PermissionDenied, "peer login has expired, please log in once more")
+}

management/server/store.go

@@ -12,10 +12,11 @@ import (
 	"strings"
 	"time"
 
-	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	log "github.com/sirupsen/logrus"
 	"gorm.io/gorm"
 
+	nbgroup "github.com/netbirdio/netbird/management/server/group"
+
 	"github.com/netbirdio/netbird/management/server/telemetry"
 	"github.com/netbirdio/netbird/util"
 
@@ -48,12 +49,13 @@ type Store interface {
 	DeleteTokenID2UserIDIndex(tokenID string) error
 	GetInstallationID() string
 	SaveInstallationID(ctx context.Context, ID string) error
-	// AcquireAccountWriteLock should attempt to acquire account lock for write purposes and return a function that releases the lock
-	AcquireAccountWriteLock(ctx context.Context, accountID string) func()
-	// AcquireAccountReadLock should attempt to acquire account lock for read purposes and return a function that releases the lock
-	AcquireAccountReadLock(ctx context.Context, accountID string) func()
+	// AcquireWriteLockByUID should attempt to acquire a lock for write purposes and return a function that releases the lock
+	AcquireWriteLockByUID(ctx context.Context, uniqueID string) func()
+	// AcquireReadLockByUID should attempt to acquire lock for read purposes and return a function that releases the lock
+	AcquireReadLockByUID(ctx context.Context, uniqueID string) func()
 	// AcquireGlobalLock should attempt to acquire a global lock and return a function that releases the lock
 	AcquireGlobalLock(ctx context.Context) func()
+	SavePeer(ctx context.Context, accountID string, peer *nbpeer.Peer) error
 	SavePeerStatus(accountID, peerID string, status nbpeer.PeerStatus) error
 	SavePeerLocation(accountID string, peer *nbpeer.Peer) error
 	SaveUserLastLogin(accountID, userID string, lastLogin time.Time) error

management/server/telemetry/accountmanager_metrics.go

@@ -0,0 +1,69 @@
+package telemetry
+
+import (
+	"context"
+	"time"
+
+	"go.opentelemetry.io/otel/metric"
+)
+
+// AccountManagerMetrics represents all metrics related to the AccountManager
+type AccountManagerMetrics struct {
+	ctx                          context.Context
+	updateAccountPeersDurationMs metric.Float64Histogram
+	getPeerNetworkMapDurationMs  metric.Float64Histogram
+	networkMapObjectCount        metric.Int64Histogram
+}
+
+// NewAccountManagerMetrics creates an instance of AccountManagerMetrics
+func NewAccountManagerMetrics(ctx context.Context, meter metric.Meter) (*AccountManagerMetrics, error) {
+	updateAccountPeersDurationMs, err := meter.Float64Histogram("management.account.update.account.peers.duration.ms",
+		metric.WithUnit("milliseconds"),
+		metric.WithExplicitBucketBoundaries(
+			0.5, 1, 2.5, 5, 10, 25, 50, 100, 250, 500, 1000, 2500, 5000, 10000, 30000,
+		))
+	if err != nil {
+		return nil, err
+	}
+
+	getPeerNetworkMapDurationMs, err := meter.Float64Histogram("management.account.get.peer.network.map.duration.ms",
+		metric.WithUnit("milliseconds"),
+		metric.WithExplicitBucketBoundaries(
+			0.1, 0.5, 1, 2.5, 5, 10, 25, 50, 100, 250, 500, 1000,
+		))
+	if err != nil {
+		return nil, err
+	}
+
+	networkMapObjectCount, err := meter.Int64Histogram("management.account.network.map.object.count",
+		metric.WithUnit("objects"),
+		metric.WithExplicitBucketBoundaries(
+			50, 100, 200, 500, 1000, 2500, 5000, 10000,
+		))
+	if err != nil {
+		return nil, err
+	}
+
+	return &AccountManagerMetrics{
+		ctx:                          ctx,
+		getPeerNetworkMapDurationMs:  getPeerNetworkMapDurationMs,
+		updateAccountPeersDurationMs: updateAccountPeersDurationMs,
+		networkMapObjectCount:        networkMapObjectCount,
+	}, nil
+
+}
+
+// CountUpdateAccountPeersDuration counts the duration of updating account peers
+func (metrics *AccountManagerMetrics) CountUpdateAccountPeersDuration(duration time.Duration) {
+	metrics.updateAccountPeersDurationMs.Record(metrics.ctx, float64(duration.Nanoseconds())/1e6)
+}
+
+// CountGetPeerNetworkMapDuration counts the duration of getting the peer network map
+func (metrics *AccountManagerMetrics) CountGetPeerNetworkMapDuration(duration time.Duration) {
+	metrics.getPeerNetworkMapDurationMs.Record(metrics.ctx, float64(duration.Nanoseconds())/1e6)
+}
+
+// CountNetworkMapObjects counts the number of network map objects
+func (metrics *AccountManagerMetrics) CountNetworkMapObjects(count int64) {
+	metrics.networkMapObjectCount.Record(metrics.ctx, count)
+}

management/server/telemetry/app_metrics.go

@@ -20,14 +20,15 @@ const defaultEndpoint = "/metrics"
 
 // MockAppMetrics mocks the AppMetrics interface
 type MockAppMetrics struct {
-	GetMeterFunc             func() metric2.Meter
-	CloseFunc                func() error
-	ExposeFunc               func(ctx context.Context, port int, endpoint string) error
-	IDPMetricsFunc           func() *IDPMetrics
-	HTTPMiddlewareFunc       func() *HTTPMiddleware
-	GRPCMetricsFunc          func() *GRPCMetrics
-	StoreMetricsFunc         func() *StoreMetrics
-	UpdateChannelMetricsFunc func() *UpdateChannelMetrics
+	GetMeterFunc                 func() metric2.Meter
+	CloseFunc                    func() error
+	ExposeFunc                   func(ctx context.Context, port int, endpoint string) error
+	IDPMetricsFunc               func() *IDPMetrics
+	HTTPMiddlewareFunc           func() *HTTPMiddleware
+	GRPCMetricsFunc              func() *GRPCMetrics
+	StoreMetricsFunc             func() *StoreMetrics
+	UpdateChannelMetricsFunc     func() *UpdateChannelMetrics
+	AddAccountManagerMetricsFunc func() *AccountManagerMetrics
 }
 
 // GetMeter mocks the GetMeter function of the AppMetrics interface
@@ -94,6 +95,14 @@ func (mock *MockAppMetrics) UpdateChannelMetrics() *UpdateChannelMetrics {
 	return nil
 }
 
+// AccountManagerMetrics mocks the MockAppMetrics function of the AccountManagerMetrics interface
+func (mock *MockAppMetrics) AccountManagerMetrics() *AccountManagerMetrics {
+	if mock.AddAccountManagerMetricsFunc != nil {
+		return mock.AddAccountManagerMetricsFunc()
+	}
+	return nil
+}
+
 // AppMetrics is metrics interface
 type AppMetrics interface {
 	GetMeter() metric2.Meter
@@ -104,19 +113,21 @@ type AppMetrics interface {
 	GRPCMetrics() *GRPCMetrics
 	StoreMetrics() *StoreMetrics
 	UpdateChannelMetrics() *UpdateChannelMetrics
+	AccountManagerMetrics() *AccountManagerMetrics
 }
 
 // defaultAppMetrics are core application metrics based on OpenTelemetry https://opentelemetry.io/
 type defaultAppMetrics struct {
 	// Meter can be used by different application parts to create counters and measure things
-	Meter                metric2.Meter
-	listener             net.Listener
-	ctx                  context.Context
-	idpMetrics           *IDPMetrics
-	httpMiddleware       *HTTPMiddleware
-	grpcMetrics          *GRPCMetrics
-	storeMetrics         *StoreMetrics
-	updateChannelMetrics *UpdateChannelMetrics
+	Meter                 metric2.Meter
+	listener              net.Listener
+	ctx                   context.Context
+	idpMetrics            *IDPMetrics
+	httpMiddleware        *HTTPMiddleware
+	grpcMetrics           *GRPCMetrics
+	storeMetrics          *StoreMetrics
+	updateChannelMetrics  *UpdateChannelMetrics
+	accountManagerMetrics *AccountManagerMetrics
 }
 
 // IDPMetrics returns metrics for the idp package
@@ -144,6 +155,11 @@ func (appMetrics *defaultAppMetrics) UpdateChannelMetrics() *UpdateChannelMetric
 	return appMetrics.updateChannelMetrics
 }
 
+// AccountManagerMetrics returns metrics for the account manager
+func (appMetrics *defaultAppMetrics) AccountManagerMetrics() *AccountManagerMetrics {
+	return appMetrics.accountManagerMetrics
+}
+
 // Close stop application metrics HTTP handler and closes listener.
 func (appMetrics *defaultAppMetrics) Close() error {
 	if appMetrics.listener == nil {
@@ -220,13 +236,19 @@ func NewDefaultAppMetrics(ctx context.Context) (AppMetrics, error) {
 		return nil, err
 	}
 
+	accountManagerMetrics, err := NewAccountManagerMetrics(ctx, meter)
+	if err != nil {
+		return nil, err
+	}
+
 	return &defaultAppMetrics{
-		Meter:                meter,
-		ctx:                  ctx,
-		idpMetrics:           idpMetrics,
-		httpMiddleware:       middleware,
-		grpcMetrics:          grpcMetrics,
-		storeMetrics:         storeMetrics,
-		updateChannelMetrics: updateChannelMetrics,
+		Meter:                 meter,
+		ctx:                   ctx,
+		idpMetrics:            idpMetrics,
+		httpMiddleware:        middleware,
+		grpcMetrics:           grpcMetrics,
+		storeMetrics:          storeMetrics,
+		updateChannelMetrics:  updateChannelMetrics,
+		accountManagerMetrics: accountManagerMetrics,
 	}, nil
 }

management/server/user.go

@@ -2,8 +2,8 @@ package server
 
 import (
 	"context"
+	"errors"
 	"fmt"
-	"slices"
 	"strings"
 	"time"
 
@@ -212,7 +212,7 @@ func NewOwnerUser(id string) *User {
 
 // createServiceUser creates a new service user under the given account.
 func (am *DefaultAccountManager) createServiceUser(ctx context.Context, accountID string, initiatorUserID string, role UserRole, serviceUserName string, nonDeletable bool, autoGroups []string) (*UserInfo, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -268,7 +268,7 @@ func (am *DefaultAccountManager) CreateUser(ctx context.Context, accountID, user
 
 // inviteNewUser Invites a USer to a given account and creates reference in datastore
 func (am *DefaultAccountManager) inviteNewUser(ctx context.Context, accountID, userID string, invite *UserInfo) (*UserInfo, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	if am.idpManager == nil {
@@ -369,7 +369,7 @@ func (am *DefaultAccountManager) GetUser(ctx context.Context, claims jwtclaims.A
 		return nil, fmt.Errorf("failed to get account with token claims %v", err)
 	}
 
-	unlock := am.Store.AcquireAccountWriteLock(ctx, account.Id)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, account.Id)
 	defer unlock()
 
 	account, err = am.Store.GetAccount(ctx, account.Id)
@@ -402,7 +402,7 @@ func (am *DefaultAccountManager) GetUser(ctx context.Context, claims jwtclaims.A
 // ListUsers returns lists of all users under the account.
 // It doesn't populate user information such as email or name.
 func (am *DefaultAccountManager) ListUsers(ctx context.Context, accountID string) ([]*User, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -429,7 +429,7 @@ func (am *DefaultAccountManager) DeleteUser(ctx context.Context, accountID, init
 	if initiatorUserID == targetUserID {
 		return status.Errorf(status.InvalidArgument, "self deletion is not allowed")
 	}
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -473,68 +473,27 @@ func (am *DefaultAccountManager) DeleteUser(ctx context.Context, accountID, init
 }
 
 func (am *DefaultAccountManager) deleteRegularUser(ctx context.Context, account *Account, initiatorUserID, targetUserID string) error {
-	tuEmail, tuName, err := am.getEmailAndNameOfTargetUser(ctx, account.Id, initiatorUserID, targetUserID)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to resolve email address: %s", err)
-		return err
-	}
-
-	if !isNil(am.idpManager) {
-		// Delete if the user already exists in the IdP.Necessary in cases where a user account
-		// was created where a user account was provisioned but the user did not sign in
-		_, err = am.idpManager.GetUserDataByID(ctx, targetUserID, idp.AppMetadata{WTAccountID: account.Id})
-		if err == nil {
-			err = am.deleteUserFromIDP(ctx, targetUserID, account.Id)
-			if err != nil {
-				log.WithContext(ctx).Debugf("failed to delete user from IDP: %s", targetUserID)
-				return err
-			}
-		} else {
-			log.WithContext(ctx).Debugf("skipped deleting user %s from IDP, error: %v", targetUserID, err)
-		}
-	}
-
-	userHasPeers, err := am.deleteUserPeers(ctx, initiatorUserID, targetUserID, account)
+	meta, err := am.prepareUserDeletion(ctx, account, initiatorUserID, targetUserID)
 	if err != nil {
 		return err
 	}
 
-	u, err := account.FindUser(targetUserID)
-	if err != nil {
-		log.WithContext(ctx).Errorf("failed to find user %s for deletion, this should never happen: %s", targetUserID, err)
-	}
-
-	var tuCreatedAt time.Time
-	if u != nil {
-		tuCreatedAt = u.CreatedAt
-	}
-
 	delete(account.Users, targetUserID)
 	err = am.Store.SaveAccount(ctx, account)
 	if err != nil {
 		return err
 	}
 
-	meta := map[string]any{"name": tuName, "email": tuEmail, "created_at": tuCreatedAt}
 	am.StoreEvent(ctx, initiatorUserID, targetUserID, account.Id, activity.UserDeleted, meta)
-
-	if userHasPeers && account.Settings.GroupsPropagationEnabled {
-		am.updateAccountPeers(ctx, account)
-	}
+	am.updateAccountPeers(ctx, account)
 
 	return nil
 }
 
-// deleteUserPeers deletes all peers associated with the target user in the specified account.
-func (am *DefaultAccountManager) deleteUserPeers(ctx context.Context, initiatorUserID string, targetUserID string, account *Account) (bool, error) {
+func (am *DefaultAccountManager) deleteUserPeers(ctx context.Context, initiatorUserID string, targetUserID string, account *Account) error {
 	peers, err := account.FindUserPeers(targetUserID)
 	if err != nil {
-		return false, status.Errorf(status.Internal, "failed to find user peers")
-	}
-
-	hadPeers := len(peers) > 0
-	if !hadPeers {
-		return false, nil
+		return status.Errorf(status.Internal, "failed to find user peers")
 	}
 
 	peerIDs := make([]string, 0, len(peers))
@@ -542,12 +501,12 @@ func (am *DefaultAccountManager) deleteUserPeers(ctx context.Context, initiatorU
 		peerIDs = append(peerIDs, peer.ID)
 	}
 
-	return hadPeers, am.deletePeers(ctx, account, peerIDs, initiatorUserID)
+	return am.deletePeers(ctx, account, peerIDs, initiatorUserID)
 }
 
 // InviteUser resend invitations to users who haven't activated their accounts prior to the expiration period.
 func (am *DefaultAccountManager) InviteUser(ctx context.Context, accountID string, initiatorUserID string, targetUserID string) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	if am.idpManager == nil {
@@ -587,7 +546,7 @@ func (am *DefaultAccountManager) InviteUser(ctx context.Context, accountID strin
 
 // CreatePAT creates a new PAT for the given user
 func (am *DefaultAccountManager) CreatePAT(ctx context.Context, accountID string, initiatorUserID string, targetUserID string, tokenName string, expiresIn int) (*PersonalAccessTokenGenerated, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	if tokenName == "" {
@@ -637,7 +596,7 @@ func (am *DefaultAccountManager) CreatePAT(ctx context.Context, accountID string
 
 // DeletePAT deletes a specific PAT from a user
 func (am *DefaultAccountManager) DeletePAT(ctx context.Context, accountID string, initiatorUserID string, targetUserID string, tokenID string) error {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -687,7 +646,7 @@ func (am *DefaultAccountManager) DeletePAT(ctx context.Context, accountID string
 
 // GetPAT returns a specific PAT from a user
 func (am *DefaultAccountManager) GetPAT(ctx context.Context, accountID string, initiatorUserID string, targetUserID string, tokenID string) (*PersonalAccessToken, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -719,7 +678,7 @@ func (am *DefaultAccountManager) GetPAT(ctx context.Context, accountID string, i
 
 // GetAllPATs returns all PATs for a user
 func (am *DefaultAccountManager) GetAllPATs(ctx context.Context, accountID string, initiatorUserID string, targetUserID string) ([]*PersonalAccessToken, error) {
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	account, err := am.Store.GetAccount(ctx, accountID)
@@ -761,7 +720,7 @@ func (am *DefaultAccountManager) SaveOrAddUser(ctx context.Context, accountID, i
 		return nil, status.Errorf(status.InvalidArgument, "provided user update is nil")
 	}
 
-	unlock := am.Store.AcquireAccountWriteLock(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()
 
 	updatedUsers, err := am.SaveOrAddUsers(ctx, accountID, initiatorUserID, []*User{update}, addIfNotExists)
@@ -801,7 +760,6 @@ func (am *DefaultAccountManager) SaveOrAddUsers(ctx context.Context, accountID,
 	updatedUsers := make([]*UserInfo, 0, len(updates))
 	var (
 		expiredPeers  []*nbpeer.Peer
-		userIDs       []string
 		eventsToStore []func()
 	)
 
@@ -810,8 +768,6 @@ func (am *DefaultAccountManager) SaveOrAddUsers(ctx context.Context, accountID,
 			return nil, status.Errorf(status.InvalidArgument, "provided user update is nil")
 		}
 
-		userIDs = append(userIDs, update.Id)
-
 		oldUser := account.Users[update.Id]
 		if oldUser == nil {
 			if !addIfNotExists {
@@ -871,11 +827,11 @@ func (am *DefaultAccountManager) SaveOrAddUsers(ctx context.Context, accountID,
 	}
 
 	account.Network.IncSerial()
-	if err = am.Store.SaveUsers(account.Id, account.Users); err != nil {
+	if err = am.Store.SaveAccount(ctx, account); err != nil {
 		return nil, err
 	}
 
-	if areUsersLinkedToPeers(account, userIDs) && account.Settings.GroupsPropagationEnabled {
+	if account.Settings.GroupsPropagationEnabled {
 		am.updateAccountPeers(ctx, account)
 	}
 
@@ -988,10 +944,14 @@ func validateUserUpdate(account *Account, initiatorUser, oldUser, update *User)
 	}
 
 	for _, newGroupID := range update.AutoGroups {
-		if _, ok := account.Groups[newGroupID]; !ok {
+		group, ok := account.Groups[newGroupID]
+		if !ok {
 			return status.Errorf(status.InvalidArgument, "provided group ID %s in the user %s update doesn't exist",
 				newGroupID, update.Id)
 		}
+		if group.Name == "All" {
+			return status.Errorf(status.InvalidArgument, "can't add All group to the user")
+		}
 	}
 
 	return nil
@@ -1202,14 +1162,114 @@ func (am *DefaultAccountManager) getEmailAndNameOfTargetUser(ctx context.Context
 	return "", "", fmt.Errorf("user info not found for user: %s", targetId)
 }
 
-// areUsersLinkedToPeers checks if any of the given userIDs are linked to any of the peers in the account.
-func areUsersLinkedToPeers(account *Account, userIDs []string) bool {
-	for _, peer := range account.Peers {
-		if slices.Contains(userIDs, peer.UserID) {
-			return true
+// DeleteRegularUsers deletes regular users from an account.
+// Note: This function does not acquire the global lock.
+// It is the caller's responsibility to ensure proper locking is in place before invoking this method.
+//
+// If an error occurs while deleting the user, the function skips it and continues deleting other users.
+// Errors are collected and returned at the end.
+func (am *DefaultAccountManager) DeleteRegularUsers(ctx context.Context, accountID, initiatorUserID string, targetUserIDs []string) error {
+	account, err := am.Store.GetAccount(ctx, accountID)
+	if err != nil {
+		return err
+	}
+
+	executingUser := account.Users[initiatorUserID]
+	if executingUser == nil {
+		return status.Errorf(status.NotFound, "user not found")
+	}
+	if !executingUser.HasAdminPower() {
+		return status.Errorf(status.PermissionDenied, "only users with admin power can delete users")
+	}
+
+	var allErrors error
+
+	deletedUsersMeta := make(map[string]map[string]any)
+	for _, targetUserID := range targetUserIDs {
+		if initiatorUserID == targetUserID {
+			allErrors = errors.Join(allErrors, errors.New("self deletion is not allowed"))
+			continue
+		}
+
+		targetUser := account.Users[targetUserID]
+		if targetUser == nil {
+			allErrors = errors.Join(allErrors, fmt.Errorf("target user: %s not found", targetUserID))
+			continue
+		}
+
+		if targetUser.Role == UserRoleOwner {
+			allErrors = errors.Join(allErrors, fmt.Errorf("unable to delete a user: %s with owner role", targetUserID))
+			continue
+		}
+
+		// disable deleting integration user if the initiator is not admin service user
+		if targetUser.Issued == UserIssuedIntegration && !executingUser.IsServiceUser {
+			allErrors = errors.Join(allErrors, errors.New("only integration service user can delete this user"))
+			continue
+		}
+
+		meta, err := am.prepareUserDeletion(ctx, account, initiatorUserID, targetUserID)
+		if err != nil {
+			allErrors = errors.Join(allErrors, fmt.Errorf("failed to delete user %s: %s", targetUserID, err))
+			continue
+		}
+
+		delete(account.Users, targetUserID)
+		deletedUsersMeta[targetUserID] = meta
+	}
+
+	err = am.Store.SaveAccount(ctx, account)
+	if err != nil {
+		return fmt.Errorf("failed to delete users: %w", err)
+	}
+
+	am.updateAccountPeers(ctx, account)
+
+	for targetUserID, meta := range deletedUsersMeta {
+		am.StoreEvent(ctx, initiatorUserID, targetUserID, account.Id, activity.UserDeleted, meta)
+	}
+
+	return allErrors
+}
+
+func (am *DefaultAccountManager) prepareUserDeletion(ctx context.Context, account *Account, initiatorUserID, targetUserID string) (map[string]any, error) {
+	tuEmail, tuName, err := am.getEmailAndNameOfTargetUser(ctx, account.Id, initiatorUserID, targetUserID)
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to resolve email address: %s", err)
+		return nil, err
+	}
+
+	if !isNil(am.idpManager) {
+		// Delete if the user already exists in the IdP. Necessary in cases where a user account
+		// was created where a user account was provisioned but the user did not sign in
+		_, err = am.idpManager.GetUserDataByID(ctx, targetUserID, idp.AppMetadata{WTAccountID: account.Id})
+		if err == nil {
+			err = am.deleteUserFromIDP(ctx, targetUserID, account.Id)
+			if err != nil {
+				log.WithContext(ctx).Debugf("failed to delete user from IDP: %s", targetUserID)
+				return nil, err
+			}
+		} else {
+			log.WithContext(ctx).Debugf("skipped deleting user %s from IDP, error: %v", targetUserID, err)
 		}
 	}
-	return false
+
+	err = am.deleteUserPeers(ctx, initiatorUserID, targetUserID, account)
+	if err != nil {
+		return nil, err
+	}
+
+	u, err := account.FindUser(targetUserID)
+	if err != nil {
+		log.WithContext(ctx).Errorf("failed to find user %s for deletion, this should never happen: %s", targetUserID, err)
+	}
+
+	var tuCreatedAt time.Time
+	if u != nil {
+		tuCreatedAt = u.CreatedAt
+	}
+
+	return map[string]any{"name": tuName, "email": tuEmail, "created_at": tuCreatedAt}, nil
 }
 
 func findUserInIDPUserdata(userID string, userData []*idp.UserData) (*idp.UserData, bool) {

management/server/user_test.go

@@ -665,6 +665,157 @@ func TestUser_DeleteUser_regularUser(t *testing.T) {
 
 }
 
+func TestUser_DeleteUser_RegularUsers(t *testing.T) {
+	store := newStore(t)
+	defer store.Close(context.Background())
+	account := newAccountWithId(context.Background(), mockAccountID, mockUserID, "")
+
+	targetId := "user2"
+	account.Users[targetId] = &User{
+		Id:              targetId,
+		IsServiceUser:   true,
+		ServiceUserName: "user2username",
+	}
+	targetId = "user3"
+	account.Users[targetId] = &User{
+		Id:            targetId,
+		IsServiceUser: false,
+		Issued:        UserIssuedAPI,
+	}
+	targetId = "user4"
+	account.Users[targetId] = &User{
+		Id:            targetId,
+		IsServiceUser: false,
+		Issued:        UserIssuedIntegration,
+	}
+
+	targetId = "user5"
+	account.Users[targetId] = &User{
+		Id:            targetId,
+		IsServiceUser: false,
+		Issued:        UserIssuedAPI,
+		Role:          UserRoleOwner,
+	}
+	account.Users["user6"] = &User{
+		Id:            "user6",
+		IsServiceUser: false,
+		Issued:        UserIssuedAPI,
+	}
+	account.Users["user7"] = &User{
+		Id:            "user7",
+		IsServiceUser: false,
+		Issued:        UserIssuedAPI,
+	}
+	account.Users["user8"] = &User{
+		Id:            "user8",
+		IsServiceUser: false,
+		Issued:        UserIssuedAPI,
+		Role:          UserRoleAdmin,
+	}
+	account.Users["user9"] = &User{
+		Id:            "user9",
+		IsServiceUser: false,
+		Issued:        UserIssuedAPI,
+		Role:          UserRoleAdmin,
+	}
+
+	err := store.SaveAccount(context.Background(), account)
+	if err != nil {
+		t.Fatalf("Error when saving account: %s", err)
+	}
+
+	am := DefaultAccountManager{
+		Store:                   store,
+		eventStore:              &activity.InMemoryEventStore{},
+		integratedPeerValidator: MocIntegratedValidator{},
+	}
+
+	testCases := []struct {
+		name               string
+		userIDs            []string
+		expectedReasons    []string
+		expectedDeleted    []string
+		expectedNotDeleted []string
+	}{
+		{
+			name:            "Delete service user successfully ",
+			userIDs:         []string{"user2"},
+			expectedDeleted: []string{"user2"},
+		},
+		{
+			name:            "Delete regular user successfully",
+			userIDs:         []string{"user3"},
+			expectedDeleted: []string{"user3"},
+		},
+		{
+			name:               "Delete integration regular user permission denied",
+			userIDs:            []string{"user4"},
+			expectedReasons:    []string{"only integration service user can delete this user"},
+			expectedNotDeleted: []string{"user4"},
+		},
+		{
+			name:               "Delete user with owner role should return permission denied",
+			userIDs:            []string{"user5"},
+			expectedReasons:    []string{"unable to delete a user: user5 with owner role"},
+			expectedNotDeleted: []string{"user5"},
+		},
+		{
+			name:               "Delete multiple users with mixed results",
+			userIDs:            []string{"user5", "user5", "user6", "user7"},
+			expectedReasons:    []string{"only integration service user can delete this user", "unable to delete a user: user5 with owner role"},
+			expectedDeleted:    []string{"user6", "user7"},
+			expectedNotDeleted: []string{"user4", "user5"},
+		},
+		{
+			name:               "Delete non-existent user",
+			userIDs:            []string{"non-existent-user"},
+			expectedReasons:    []string{"target user: non-existent-user not found"},
+			expectedNotDeleted: []string{},
+		},
+		{
+			name:            "Delete multiple regular users successfully",
+			userIDs:         []string{"user8", "user9"},
+			expectedDeleted: []string{"user8", "user9"},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			err = am.DeleteRegularUsers(context.Background(), mockAccountID, mockUserID, tc.userIDs)
+			if len(tc.expectedReasons) > 0 {
+				assert.Error(t, err)
+				var foundExpectedErrors int
+
+				wrappedErr, ok := err.(interface{ Unwrap() []error })
+				assert.Equal(t, ok, true)
+
+				for _, e := range wrappedErr.Unwrap() {
+					assert.Contains(t, tc.expectedReasons, e.Error(), "unexpected error message")
+					foundExpectedErrors++
+				}
+
+				assert.Equal(t, len(tc.expectedReasons), foundExpectedErrors, "not all expected errors were found")
+			} else {
+				assert.NoError(t, err)
+			}
+
+			acc, err := am.GetAccountByUserOrAccountID(context.Background(), "", account.Id, "")
+			assert.NoError(t, err)
+
+			for _, id := range tc.expectedDeleted {
+				_, exists := acc.Users[id]
+				assert.False(t, exists, "user should have been deleted: %s", id)
+			}
+
+			for _, id := range tc.expectedNotDeleted {
+				user, exists := acc.Users[id]
+				assert.True(t, exists, "user should not have been deleted: %s", id)
+				assert.NotNil(t, user, "user should exist: %s", id)
+			}
+		})
+	}
+}
+
 func TestDefaultAccountManager_GetUser(t *testing.T) {
 	store := newStore(t)
 	defer store.Close(context.Background())