[management] Refactor setup key to use store methods (#2861)

* Refactor setup key handling to use store methods Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com> * add lock to get account groups Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com> * add check for regular user Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com> * get only required groups for auto-group validation Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com> * add account lock and return auto groups map on validation Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com> * fix missing group removed from setup key activity Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com> * Remove context from DB queries Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com> * Add user permission check and add setup events into events to store slice Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com> * Retrieve all groups once during setup key auto-group validation Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com> * Fix lint Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com> * Fix sonar Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com> --------- Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>
2025-08-16 18:11:58 +02:00 · 2024-11-11 19:46:10 +03:00
parent e0bed2b0fb
commit 6cb697eed6
11 changed files with 263 additions and 143 deletions
--- a/management/server/setupkey.go
+++ b/management/server/setupkey.go
@ -4,18 +4,17 @@ import (
 	"context"
 	"crypto/sha256"
 	b64 "encoding/base64"
-	"fmt"
 	"hash/fnv"
+	"slices"
 	"strconv"
 	"strings"
 	"time"
 	"unicode/utf8"

 	"github.com/google/uuid"
-	log "github.com/sirupsen/logrus"
-
 	"github.com/netbirdio/netbird/management/server/activity"
 	"github.com/netbirdio/netbird/management/server/status"
+	log "github.com/sirupsen/logrus"
 )

 const (
@ -229,32 +228,43 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
 	defer unlock()

-	account, err := am.Store.GetAccount(ctx, accountID)
+	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
 		return nil, err
 	}

-	if err := validateSetupKeyAutoGroups(account, autoGroups); err != nil {
-		return nil, err
+	if user.AccountID != accountID {
+		return nil, status.NewUserNotPartOfAccountError()
 	}

-	setupKey, plainKey := GenerateSetupKey(keyName, keyType, expiresIn, autoGroups, usageLimit, ephemeral)
-	account.SetupKeys[setupKey.Key] = setupKey
-	err = am.Store.SaveAccount(ctx, account)
+	if user.IsRegularUser() {
+		return nil, status.NewAdminPermissionError()
+	}
+
+	var setupKey *SetupKey
+	var plainKey string
+	var eventsToStore []func()
+
+	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
+		if err = validateSetupKeyAutoGroups(ctx, transaction, accountID, autoGroups); err != nil {
+			return err
+		}
+
+		setupKey, plainKey = GenerateSetupKey(keyName, keyType, expiresIn, autoGroups, usageLimit, ephemeral)
+		setupKey.AccountID = accountID
+
+		events := am.prepareSetupKeyEvents(ctx, transaction, accountID, userID, autoGroups, nil, setupKey)
+		eventsToStore = append(eventsToStore, events...)
+
+		return transaction.SaveSetupKey(ctx, LockingStrengthUpdate, setupKey)
+	})
 	if err != nil {
-		return nil, status.Errorf(status.Internal, "failed adding account key")
+		return nil, err
 	}

 	am.StoreEvent(ctx, userID, setupKey.Id, accountID, activity.SetupKeyCreated, setupKey.EventMeta())
-
-	for _, g := range setupKey.AutoGroups {
-		group := account.GetGroup(g)
-		if group != nil {
-			am.StoreEvent(ctx, userID, setupKey.Id, accountID, activity.GroupAddedToSetupKey,
-				map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": setupKey.Name})
-		} else {
-			log.WithContext(ctx).Errorf("group %s not found while saving setup key activity event of account %s", g, account.Id)
-		}
+	for _, storeEvent := range eventsToStore {
+		storeEvent()
 	}

 	// for the creation return the plain key to the caller
@ -268,43 +278,56 @@ 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.AcquireWriteLockByUID(ctx, accountID)
-	defer unlock()
-
 	if keyToSave == nil {
 		return nil, status.Errorf(status.InvalidArgument, "provided setup key to update is nil")
 	}

-	account, err := am.Store.GetAccount(ctx, accountID)
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
+	defer unlock()
+
+	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
 		return nil, err
 	}

+	if user.AccountID != accountID {
+		return nil, status.NewUserNotPartOfAccountError()
+	}
+
+	if user.IsRegularUser() {
+		return nil, status.NewAdminPermissionError()
+	}
+
 	var oldKey *SetupKey
-	for _, key := range account.SetupKeys {
-		if key.Id == keyToSave.Id {
-			oldKey = key.Copy()
-			break
+	var newKey *SetupKey
+	var eventsToStore []func()
+
+	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
+		if err = validateSetupKeyAutoGroups(ctx, transaction, accountID, keyToSave.AutoGroups); err != nil {
+			return err
 		}
-	}
-	if oldKey == nil {
-		return nil, status.Errorf(status.NotFound, "setup key not found")
-	}

-	if err := validateSetupKeyAutoGroups(account, keyToSave.AutoGroups); err != nil {
-		return nil, err
-	}
+		oldKey, err = transaction.GetSetupKeyByID(ctx, LockingStrengthShare, accountID, keyToSave.Id)
+		if err != nil {
+			return err
+		}

-	// only auto groups, revoked status, and name can be updated for now
-	newKey := oldKey.Copy()
-	newKey.Name = keyToSave.Name
-	newKey.AutoGroups = keyToSave.AutoGroups
-	newKey.Revoked = keyToSave.Revoked
-	newKey.UpdatedAt = time.Now().UTC()
+		// only auto groups, revoked status, and name can be updated for now
+		newKey = oldKey.Copy()
+		newKey.Name = keyToSave.Name
+		newKey.AutoGroups = keyToSave.AutoGroups
+		newKey.Revoked = keyToSave.Revoked
+		newKey.UpdatedAt = time.Now().UTC()

-	account.SetupKeys[newKey.Key] = newKey
+		addedGroups := difference(newKey.AutoGroups, oldKey.AutoGroups)
+		removedGroups := difference(oldKey.AutoGroups, newKey.AutoGroups)

-	if err = am.Store.SaveAccount(ctx, account); err != nil {
+		events := am.prepareSetupKeyEvents(ctx, transaction, accountID, userID, addedGroups, removedGroups, oldKey)
+		eventsToStore = append(eventsToStore, events...)
+
+		return transaction.SaveSetupKey(ctx, LockingStrengthUpdate, newKey)
+	})
+	if err != nil {
 		return nil, err
 	}

@ -312,30 +335,9 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 		am.StoreEvent(ctx, userID, newKey.Id, accountID, activity.SetupKeyRevoked, newKey.EventMeta())
 	}

-	defer func() {
-		addedGroups := difference(newKey.AutoGroups, oldKey.AutoGroups)
-		removedGroups := difference(oldKey.AutoGroups, newKey.AutoGroups)
-		for _, g := range removedGroups {
-			group := account.GetGroup(g)
-			if group != nil {
-				am.StoreEvent(ctx, userID, oldKey.Id, accountID, activity.GroupRemovedFromSetupKey,
-					map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": newKey.Name})
-			} else {
-				log.WithContext(ctx).Errorf("group %s not found while saving setup key activity event of account %s", g, account.Id)
-			}
-
-		}
-
-		for _, g := range addedGroups {
-			group := account.GetGroup(g)
-			if group != nil {
-				am.StoreEvent(ctx, userID, oldKey.Id, accountID, activity.GroupAddedToSetupKey,
-					map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": newKey.Name})
-			} else {
-				log.WithContext(ctx).Errorf("group %s not found while saving setup key activity event of account %s", g, account.Id)
-			}
-		}
-	}()
+	for _, storeEvent := range eventsToStore {
+		storeEvent()
+	}

 	return newKey, nil
 }
@ -347,16 +349,15 @@ func (am *DefaultAccountManager) ListSetupKeys(ctx context.Context, accountID, u
 		return nil, err
 	}

-	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
-		return nil, status.NewUnauthorizedToViewSetupKeysError()
+	if user.AccountID != accountID {
+		return nil, status.NewUserNotPartOfAccountError()
 	}

-	setupKeys, err := am.Store.GetAccountSetupKeys(ctx, LockingStrengthShare, accountID)
-	if err != nil {
-		return nil, err
+	if user.IsRegularUser() {
+		return nil, status.NewAdminPermissionError()
 	}

-	return setupKeys, nil
+	return am.Store.GetAccountSetupKeys(ctx, LockingStrengthShare, accountID)
 }

 // GetSetupKey looks up a SetupKey by KeyID, returns NotFound error if not found.
@ -366,8 +367,12 @@ func (am *DefaultAccountManager) GetSetupKey(ctx context.Context, accountID, use
 		return nil, err
 	}

-	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
-		return nil, status.NewUnauthorizedToViewSetupKeysError()
+	if user.AccountID != accountID {
+		return nil, status.NewUserNotPartOfAccountError()
+	}
+
+	if user.IsRegularUser() {
+		return nil, status.NewAdminPermissionError()
 	}

 	setupKey, err := am.Store.GetSetupKeyByID(ctx, LockingStrengthShare, keyID, accountID)
@ -387,21 +392,29 @@ func (am *DefaultAccountManager) GetSetupKey(ctx context.Context, accountID, use
 func (am *DefaultAccountManager) DeleteSetupKey(ctx context.Context, accountID, userID, keyID string) error {
 	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
-		return fmt.Errorf("failed to get user: %w", err)
+		return err
 	}

-	if !user.IsAdminOrServiceUser() || user.AccountID != accountID {
-		return status.NewUnauthorizedToViewSetupKeysError()
+	if user.AccountID != accountID {
+		return status.NewUserNotPartOfAccountError()
 	}

-	deletedSetupKey, err := am.Store.GetSetupKeyByID(ctx, LockingStrengthShare, keyID, accountID)
+	if user.IsRegularUser() {
+		return status.NewAdminPermissionError()
+	}
+
+	var deletedSetupKey *SetupKey
+
+	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
+		deletedSetupKey, err = transaction.GetSetupKeyByID(ctx, LockingStrengthShare, accountID, keyID)
+		if err != nil {
+			return err
+		}
+
+		return transaction.DeleteSetupKey(ctx, LockingStrengthUpdate, accountID, keyID)
+	})
 	if err != nil {
-		return fmt.Errorf("failed to get setup key: %w", err)
-	}
-
-	err = am.Store.DeleteSetupKey(ctx, accountID, keyID)
-	if err != nil {
-		return fmt.Errorf("failed to delete setup key: %w", err)
+		return err
 	}

 	am.StoreEvent(ctx, userID, keyID, accountID, activity.SetupKeyDeleted, deletedSetupKey.EventMeta())
@ -409,15 +422,62 @@ func (am *DefaultAccountManager) DeleteSetupKey(ctx context.Context, accountID,
 	return nil
 }

-func validateSetupKeyAutoGroups(account *Account, autoGroups []string) error {
-	for _, group := range autoGroups {
-		g, ok := account.Groups[group]
+func validateSetupKeyAutoGroups(ctx context.Context, transaction Store, accountID string, autoGroupIDs []string) error {
+	groups, err := transaction.GetGroupsByIDs(ctx, LockingStrengthShare, accountID, autoGroupIDs)
+	if err != nil {
+		return err
+	}
+
+	for _, groupID := range autoGroupIDs {
+		group, ok := groups[groupID]
 		if !ok {
-			return status.Errorf(status.NotFound, "group %s doesn't exist", group)
+			return status.Errorf(status.NotFound, "group not found: %s", groupID)
 		}
-		if g.Name == "All" {
-			return status.Errorf(status.InvalidArgument, "can't add All group to the setup key")
+
+		if group.IsGroupAll() {
+			return status.Errorf(status.InvalidArgument, "can't add 'All' group to the setup key")
 		}
 	}
+
 	return nil
 }
+
+// prepareSetupKeyEvents prepares a list of event functions to be stored.
+func (am *DefaultAccountManager) prepareSetupKeyEvents(ctx context.Context, transaction Store, accountID, userID string, addedGroups, removedGroups []string, key *SetupKey) []func() {
+	var eventsToStore []func()
+
+	modifiedGroups := slices.Concat(addedGroups, removedGroups)
+	groups, err := transaction.GetGroupsByIDs(ctx, LockingStrengthShare, accountID, modifiedGroups)
+	if err != nil {
+		log.WithContext(ctx).Errorf("issue getting groups for setup key events: %v", err)
+		return nil
+	}
+
+	for _, g := range removedGroups {
+		group, ok := groups[g]
+		if !ok {
+			log.WithContext(ctx).Debugf("skipped adding group: %s GroupRemovedFromSetupKey activity: %v", g, err)
+			continue
+		}
+
+		eventsToStore = append(eventsToStore, func() {
+			meta := map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": key.Name}
+			am.StoreEvent(ctx, userID, key.Id, accountID, activity.GroupRemovedFromSetupKey, meta)
+		})
+	}
+
+	for _, g := range addedGroups {
+		group, ok := groups[g]
+		if !ok {
+			log.WithContext(ctx).Debugf("skipped adding group: %s GroupAddedToSetupKey activity: %v", g, err)
+			continue
+		}
+
+		eventsToStore = append(eventsToStore, func() {
+			meta := map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": key.Name}
+			am.StoreEvent(ctx, userID, key.Id, accountID, activity.GroupAddedToSetupKey, meta)
+		})
+	}
+
+	return eventsToStore
+}