refactor get peers posture checks

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>

management/server/peer.go

@@ -583,7 +583,11 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s
 		return nil, nil, nil, err
 	}
 
-	postureChecks := am.getPeerPostureChecks(account, newPeer)
+	postureChecks, err := am.getPeerPostureChecks(ctx, account.Id, newPeer.ID)
+	if err != nil {
+		return nil, nil, nil, err
+	}
+
 	customZone := account.GetPeersCustomZone(ctx, am.dnsDomain)
 	networkMap := account.GetPeerNetworkMap(ctx, newPeer.ID, customZone, approvedPeersMap, am.metrics.AccountManagerMetrics())
 	return newPeer, networkMap, postureChecks, nil
@@ -665,7 +669,11 @@ func (am *DefaultAccountManager) SyncPeer(ctx context.Context, sync PeerSync, ac
 	if err != nil {
 		return nil, nil, nil, err
 	}
-	postureChecks = am.getPeerPostureChecks(account, peer)
+
+	postureChecks, err = am.getPeerPostureChecks(ctx, account.Id, peer.ID)
+	if err != nil {
+		return nil, nil, nil, err
+	}
 
 	customZone := account.GetPeersCustomZone(ctx, am.dnsDomain)
 	return peer, account.GetPeerNetworkMap(ctx, peer.ID, customZone, validPeersMap, am.metrics.AccountManagerMetrics()), postureChecks, nil
@@ -825,8 +833,6 @@ func (am *DefaultAccountManager) checkIFPeerNeedsLoginWithoutLock(ctx context.Co
 }
 
 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(),
@@ -838,7 +844,11 @@ func (am *DefaultAccountManager) getValidatedPeerWithMap(ctx context.Context, is
 	if err != nil {
 		return nil, nil, nil, err
 	}
-	postureChecks = am.getPeerPostureChecks(account, peer)
+
+	postureChecks, err := am.getPeerPostureChecks(ctx, account.Id, peer.ID)
+	if err != nil {
+		return nil, nil, nil, err
+	}
 
 	customZone := account.GetPeersCustomZone(ctx, am.dnsDomain)
 	return peer, account.GetPeerNetworkMap(ctx, peer.ID, customZone, approvedPeersMap, am.metrics.AccountManagerMetrics()), postureChecks, nil
@@ -1031,7 +1041,12 @@ func (am *DefaultAccountManager) updateAccountPeers(ctx context.Context, account
 			defer wg.Done()
 			defer func() { <-semaphore }()
 
-			postureChecks := am.getPeerPostureChecks(account, p)
+			postureChecks, err := am.getPeerPostureChecks(ctx, account.Id, p.ID)
+			if err != nil {
+				log.WithContext(ctx).Debugf("failed to get posture checks for peer %s: %v", peer.ID, err)
+				return
+			}
+
 			remotePeerNetworkMap := account.GetPeerNetworkMap(ctx, p.ID, customZone, approvedPeersMap, am.metrics.AccountManagerMetrics())
 			update := toSyncResponse(ctx, nil, p, nil, nil, remotePeerNetworkMap, am.GetDNSDomain(), postureChecks, dnsCache)
 			am.peersUpdateManager.SendUpdate(ctx, p.ID, &UpdateMessage{Update: update})

management/server/posture_checks.go

@@ -6,13 +6,14 @@ import (
 	"slices"
 
 	"github.com/netbirdio/netbird/management/server/activity"
-	nbpeer "github.com/netbirdio/netbird/management/server/peer"
 	"github.com/netbirdio/netbird/management/server/posture"
 	"github.com/netbirdio/netbird/management/server/status"
+	log "github.com/sirupsen/logrus"
+	"golang.org/x/exp/maps"
 )
 
 const (
-	errMsgPostureAdminOnly = "only users with admin power are allowed to view posture checks"
+	errPostureAdminOnlyMsg = "only users with admin power are allowed to view posture checks"
 )
 
 func (am *DefaultAccountManager) GetPostureChecks(ctx context.Context, accountID, postureChecksID, userID string) (*posture.Checks, error) {
@@ -21,8 +22,12 @@ func (am *DefaultAccountManager) GetPostureChecks(ctx context.Context, accountID
 		return nil, err
 	}
 
-	if !user.HasAdminPower() || user.AccountID != accountID {
+	if !user.HasAdminPower() {
-		return nil, status.Errorf(status.PermissionDenied, errMsgPostureAdminOnly)
+		return nil, status.Errorf(status.PermissionDenied, errPostureAdminOnlyMsg)
+	}
+
+	if user.AccountID != accountID {
+		return nil, status.Errorf(status.PermissionDenied, errUserNotPartOfAccountMsg)
 	}
 
 	return am.Store.GetPostureChecksByID(ctx, LockingStrengthShare, accountID, postureChecksID)
@@ -35,10 +40,14 @@ func (am *DefaultAccountManager) SavePostureChecks(ctx context.Context, accountI
 		return err
 	}
 
-	if !user.HasAdminPower() || user.AccountID != accountID {
+	if !user.HasAdminPower() {
 		return status.Errorf(status.PermissionDenied, "only admin users are allowed to update posture checks")
 	}
 
+	if user.AccountID != accountID {
+		return status.Errorf(status.PermissionDenied, errUserNotPartOfAccountMsg)
+	}
+
 	if err = am.validatePostureChecks(ctx, accountID, postureChecks); err != nil {
 		return status.Errorf(status.InvalidArgument, err.Error())
 	}
@@ -54,7 +63,7 @@ func (am *DefaultAccountManager) SavePostureChecks(ctx context.Context, accountI
 			}
 
 			if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-				return fmt.Errorf("failed to increment network serial: %w", err)
+				return fmt.Errorf(errNetworkSerialIncrementFmt, err)
 			}
 		}
 
@@ -72,7 +81,7 @@ func (am *DefaultAccountManager) SavePostureChecks(ctx context.Context, accountI
 	if isUpdate {
 		account, err := am.requestBuffer.GetAccountWithBackpressure(ctx, accountID)
 		if err != nil {
-			return fmt.Errorf("error getting account: %w", err)
+			return fmt.Errorf("failed to get account: %w", err)
 		}
 		am.updateAccountPeers(ctx, account)
 	}
@@ -106,10 +115,14 @@ func (am *DefaultAccountManager) DeletePostureChecks(ctx context.Context, accoun
 		return err
 	}
 
-	if !user.HasAdminPower() || user.AccountID != accountID {
+	if !user.HasAdminPower() {
 		return status.Errorf(status.PermissionDenied, "only admin users are allowed to delete posture checks")
 	}
 
+	if user.AccountID != accountID {
+		return status.Errorf(status.PermissionDenied, errUserNotPartOfAccountMsg)
+	}
+
 	if err = am.isPostureCheckLinkedToPolicy(ctx, postureChecksID, accountID); err != nil {
 		return err
 	}
@@ -121,7 +134,7 @@ func (am *DefaultAccountManager) DeletePostureChecks(ctx context.Context, accoun
 
 	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
 		if err = transaction.IncrementNetworkSerial(ctx, LockingStrengthUpdate, accountID); err != nil {
-			return fmt.Errorf("failed to increment network serial: %w", err)
+			return fmt.Errorf(errNetworkSerialIncrementFmt, err)
 		}
 
 		if err = transaction.DeletePostureChecks(ctx, LockingStrengthUpdate, accountID, postureChecksID); err != nil {
@@ -151,8 +164,12 @@ func (am *DefaultAccountManager) ListPostureChecks(ctx context.Context, accountI
 		return nil, err
 	}
 
-	if !user.HasAdminPower() || user.AccountID != accountID {
+	if !user.HasAdminPower() {
-		return nil, status.Errorf(status.PermissionDenied, errMsgPostureAdminOnly)
+		return nil, status.Errorf(status.PermissionDenied, errPostureAdminOnlyMsg)
+	}
+
+	if user.AccountID != accountID {
+		return nil, status.Errorf(status.PermissionDenied, errUserNotPartOfAccountMsg)
 	}
 
 	return am.Store.GetAccountPostureChecks(ctx, LockingStrengthShare, accountID)
@@ -175,56 +192,61 @@ func (am *DefaultAccountManager) isPostureCheckLinkedToPolicy(ctx context.Contex
 }
 
 // getPeerPostureChecks returns the posture checks applied for a given peer.
-func (am *DefaultAccountManager) getPeerPostureChecks(account *Account, peer *nbpeer.Peer) []*posture.Checks {
+func (am *DefaultAccountManager) getPeerPostureChecks(ctx context.Context, accountID string, peerID string) ([]*posture.Checks, error) {
-	peerPostureChecks := make(map[string]posture.Checks)
+	postureChecks, err := am.Store.GetAccountPostureChecks(ctx, LockingStrengthShare, accountID)
-
+	if err != nil || len(postureChecks) == 0 {
-	if len(account.PostureChecks) == 0 {
+		return nil, err
-		return nil
 	}
 
-	for _, policy := range account.Policies {
+	policies, err := am.Store.GetAccountPolicies(ctx, LockingStrengthShare, accountID)
+	if err != nil {
+		return nil, err
+	}
+
+	peerPostureChecks := make(map[string]*posture.Checks)
+
+	for _, policy := range policies {
 		if !policy.Enabled {
 			continue
 		}
 
-		if isPeerInPolicySourceGroups(peer.ID, account, policy) {
+		isInGroup, err := am.isPeerInPolicySourceGroups(ctx, accountID, peerID, policy)
-			addPolicyPostureChecks(account, policy, peerPostureChecks)
+		if err != nil {
+			return nil, err
+		}
+
+		if isInGroup {
+			for _, sourcePostureCheckID := range policy.SourcePostureChecks {
+				postureCheck, err := am.Store.GetPostureChecksByID(ctx, LockingStrengthShare, accountID, sourcePostureCheckID)
+				if err == nil {
+					peerPostureChecks[sourcePostureCheckID] = postureCheck
+				}
+			}
 		}
 	}
 
-	postureChecksList := make([]*posture.Checks, 0, len(peerPostureChecks))
+	return maps.Values(peerPostureChecks), nil
-	for _, check := range peerPostureChecks {
-		checkCopy := check
-		postureChecksList = append(postureChecksList, &checkCopy)
-	}
-
-	return postureChecksList
 }
 
 // isPeerInPolicySourceGroups checks if a peer is present in any of the policy rule source groups.
-func isPeerInPolicySourceGroups(peerID string, account *Account, policy *Policy) bool {
+func (am *DefaultAccountManager) isPeerInPolicySourceGroups(ctx context.Context, accountID, peerID string, policy *Policy) (bool, error) {
 	for _, rule := range policy.Rules {
 		if !rule.Enabled {
 			continue
 		}
 
 		for _, sourceGroup := range rule.Sources {
-			group, ok := account.Groups[sourceGroup]
+			group, err := am.Store.GetGroupByID(ctx, LockingStrengthShare, accountID, sourceGroup)
-			if ok && slices.Contains(group.Peers, peerID) {
+			if err != nil {
-				return true
+				log.WithContext(ctx).Debugf("failed to check peer in policy source group: %v", err)
+				return false, fmt.Errorf("failed to check peer in policy source group: %w", err)
+			}
+
+			if slices.Contains(group.Peers, peerID) {
+				return true, nil
 			}
 		}
 	}
 
-	return false
+	return false, nil
-}
-
-func addPolicyPostureChecks(account *Account, policy *Policy, peerPostureChecks map[string]posture.Checks) {
-	for _, sourcePostureCheckID := range policy.SourcePostureChecks {
-		for _, postureCheck := range account.PostureChecks {
-			if postureCheck.ID == sourcePostureCheckID {
-				peerPostureChecks[sourcePostureCheckID] = *postureCheck
-			}
-		}
-	}
 }