From 389c9619afe8b6d129a12d137bf332491503bb83 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Fri, 8 Nov 2024 00:31:41 +0300
Subject: [PATCH 01/11] Refactor setup key handling to use store methods

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>
---
 management/server/setupkey.go       | 179 +++++++++++++++++-----------
 management/server/sql_store.go      |  83 ++++++++-----
 management/server/sql_store_test.go |   4 +-
 management/server/status/error.go   |  15 ++-
 management/server/store.go          |   7 +-
 5 files changed, 178 insertions(+), 110 deletions(-)

diff --git a/management/server/setupkey.go b/management/server/setupkey.go
index 43b6e02c936..f54eafdc1fd 100644
--- a/management/server/setupkey.go
+++ b/management/server/setupkey.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"crypto/sha256"
 	b64 "encoding/base64"
-	"fmt"
 	"hash/fnv"
 	"strconv"
 	"strings"
@@ -12,9 +11,8 @@ import (
 	"unicode/utf8"
 
 	"github.com/google/uuid"
-	log "github.com/sirupsen/logrus"
-
 	"github.com/netbirdio/netbird/management/server/activity"
+	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	"github.com/netbirdio/netbird/management/server/status"
 )
 
@@ -226,34 +224,49 @@ 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.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)
+	var accountGroups []*nbgroup.Group
+	var setupKey *SetupKey
+	var plainKey string
+
+	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
+		accountGroups, err = transaction.GetAccountGroups(ctx, LockingStrengthShare, accountID)
+		if err != nil {
+			return err
+		}
+
+		if err = validateSetupKeyAutoGroups(accountGroups, autoGroups); err != nil {
+			return err
+		}
+
+		setupKey, plainKey = GenerateSetupKey(keyName, keyType, expiresIn, autoGroups, usageLimit, ephemeral)
+		setupKey.AccountID = accountID
+
+		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())
+	groupMap := make(map[string]*nbgroup.Group, len(accountGroups))
+	for _, g := range accountGroups {
+		groupMap[g.ID] = g
+	}
 
 	for _, g := range setupKey.AutoGroups {
-		group := account.GetGroup(g)
-		if group != nil {
+		group, ok := groupMap[g]
+		if ok {
 			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)
 		}
 	}
 
@@ -268,43 +281,48 @@ 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)
+	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
 		return nil, err
 	}
 
+	if user.AccountID != accountID {
+		return nil, status.NewUserNotPartOfAccountError()
+	}
+
+	var accountGroups []*nbgroup.Group
 	var oldKey *SetupKey
-	for _, key := range account.SetupKeys {
-		if key.Id == keyToSave.Id {
-			oldKey = key.Copy()
-			break
+	var newKey *SetupKey
+
+	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
+		accountGroups, err = transaction.GetAccountGroups(ctx, LockingStrengthShare, accountID)
+		if 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
-	}
+		if err = validateSetupKeyAutoGroups(accountGroups, keyToSave.AutoGroups); 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()
+		oldKey, err = transaction.GetSetupKeyByID(ctx, LockingStrengthShare, accountID, keyToSave.Id)
+		if err != nil {
+			return err
+		}
 
-	account.SetupKeys[newKey.Key] = newKey
+		// 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()
 
-	if err = am.Store.SaveAccount(ctx, account); err != nil {
+		return transaction.SaveSetupKey(ctx, LockingStrengthUpdate, newKey)
+	})
+	if err != nil {
 		return nil, err
 	}
 
@@ -315,24 +333,25 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 	defer func() {
 		addedGroups := difference(newKey.AutoGroups, oldKey.AutoGroups)
 		removedGroups := difference(oldKey.AutoGroups, newKey.AutoGroups)
+
+		groupMap := make(map[string]*nbgroup.Group, len(accountGroups))
+		for _, g := range accountGroups {
+			groupMap[g.ID] = g
+		}
+
 		for _, g := range removedGroups {
-			group := account.GetGroup(g)
-			if group != nil {
+			group, ok := groupMap[g]
+			if ok {
 				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 {
+			group, ok := groupMap[g]
+			if ok {
 				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)
 			}
 		}
 	}()
@@ -347,16 +366,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 +384,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 +409,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 err != nil {
-		return fmt.Errorf("failed to get setup key: %w", err)
+	if user.IsRegularUser() {
+		return status.NewAdminPermissionError()
 	}
 
-	err = am.Store.DeleteSetupKey(ctx, accountID, keyID)
+	var deletedSetupKey *SetupKey
+
+	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
+		deletedSetupKey, err = transaction.GetSetupKeyByID(ctx, LockingStrengthShare, keyID, accountID)
+		if err != nil {
+			return err
+		}
+
+		return transaction.DeleteSetupKey(ctx, LockingStrengthUpdate, 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 +439,22 @@ 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]
-		if !ok {
-			return status.Errorf(status.NotFound, "group %s doesn't exist", group)
+func validateSetupKeyAutoGroups(groups []*nbgroup.Group, autoGroups []string) error {
+	groupMap := make(map[string]*nbgroup.Group, len(groups))
+	for _, g := range groups {
+		groupMap[g.ID] = g
+	}
+
+	for _, groupID := range autoGroups {
+		g, exists := groupMap[groupID]
+		if !exists {
+			return status.Errorf(status.NotFound, "group %s doesn't exist", groupID)
 		}
+
 		if g.Name == "All" {
-			return status.Errorf(status.InvalidArgument, "can't add All group to the setup key")
+			return status.Errorf(status.InvalidArgument, "can't add 'All' group to the setup key")
 		}
 	}
+
 	return nil
 }
diff --git a/management/server/sql_store.go b/management/server/sql_store.go
index 646184578eb..a11370e4f9d 100644
--- a/management/server/sql_store.go
+++ b/management/server/sql_store.go
@@ -633,11 +633,11 @@ func (s *SqlStore) GetAccountUsers(ctx context.Context, accountID string) ([]*Us
 	return users, nil
 }
 
-func (s *SqlStore) GetAccountGroups(ctx context.Context, accountID string) ([]*nbgroup.Group, error) {
+func (s *SqlStore) GetAccountGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbgroup.Group, error) {
 	startTime := time.Now()
 
 	var groups []*nbgroup.Group
-	result := s.db.Find(&groups, accountIDCondition, accountID)
+	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&groups, accountIDCondition, accountID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "accountID not found: index lookup failed")
@@ -645,8 +645,8 @@ func (s *SqlStore) GetAccountGroups(ctx context.Context, accountID string) ([]*n
 		if errors.Is(result.Error, context.Canceled) {
 			return nil, status.NewStoreContextCanceledError(time.Since(startTime))
 		}
-		log.WithContext(ctx).Errorf("error when getting groups from the store: %s", result.Error)
-		return nil, status.Errorf(status.Internal, "issue getting groups from store")
+		log.WithContext(ctx).Errorf("failed to get account groups from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get account groups from the store")
 	}
 
 	return groups, nil
@@ -1404,12 +1404,59 @@ func (s *SqlStore) GetRouteByID(ctx context.Context, lockStrength LockingStrengt
 
 // GetAccountSetupKeys retrieves setup keys for an account.
 func (s *SqlStore) GetAccountSetupKeys(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*SetupKey, error) {
-	return getRecords[*SetupKey](s.db.WithContext(ctx), lockStrength, accountID)
+	var setupKeys []*SetupKey
+	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).
+		Find(&setupKeys, accountIDCondition, accountID)
+	if err := result.Error; err != nil {
+		log.WithContext(ctx).Errorf("failed to get setup keys from the store: %s", err)
+		return nil, status.Errorf(status.Internal, "failed to get setup keys from store")
+	}
+
+	return setupKeys, nil
 }
 
 // GetSetupKeyByID retrieves a setup key by its ID and account ID.
-func (s *SqlStore) GetSetupKeyByID(ctx context.Context, lockStrength LockingStrength, setupKeyID string, accountID string) (*SetupKey, error) {
-	return getRecordByID[SetupKey](s.db.WithContext(ctx), lockStrength, setupKeyID, accountID)
+func (s *SqlStore) GetSetupKeyByID(ctx context.Context, lockStrength LockingStrength, accountID, setupKeyID string) (*SetupKey, error) {
+	var setupKey *SetupKey
+	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).
+		First(&setupKey, accountAndIDQueryCondition, accountID, setupKeyID)
+	if err := result.Error; err != nil {
+		if errors.Is(err, gorm.ErrRecordNotFound) {
+			return nil, status.Errorf(status.NotFound, "setup key not found")
+		}
+		log.WithContext(ctx).Errorf("failed to get setup key from the store: %s", err)
+		return nil, status.Errorf(status.Internal, "failed to get setup key from store")
+	}
+
+	return setupKey, nil
+}
+
+// SaveSetupKey saves a setup key to the database.
+func (s *SqlStore) SaveSetupKey(ctx context.Context, lockStrength LockingStrength, setupKey *SetupKey) error {
+	result := s.db.WithContext(ctx).Session(&gorm.Session{FullSaveAssociations: true}).
+		Clauses(clause.Locking{Strength: string(lockStrength)}).Save(setupKey)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to save setup key to store: %s", result.Error)
+		return status.Errorf(status.Internal, "failed to save setup key to store")
+	}
+
+	return nil
+}
+
+// DeleteSetupKey deletes a setup key from the database.
+func (s *SqlStore) DeleteSetupKey(ctx context.Context, lockStrength LockingStrength, accountID, keyID string) error {
+	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).
+		Delete(&SetupKey{}, accountAndIDQueryCondition, accountID, keyID)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to delete setup key from store: %s", result.Error)
+		return status.Errorf(status.Internal, "failed to delete setup key from store")
+	}
+
+	if result.RowsAffected == 0 {
+		return status.Errorf(status.NotFound, "setup key not found")
+	}
+
+	return nil
 }
 
 // GetAccountNameServerGroups retrieves name server groups for an account.
@@ -1422,10 +1469,6 @@ func (s *SqlStore) GetNameServerGroupByID(ctx context.Context, lockStrength Lock
 	return getRecordByID[nbdns.NameServerGroup](s.db.WithContext(ctx), lockStrength, nsGroupID, accountID)
 }
 
-func (s *SqlStore) DeleteSetupKey(ctx context.Context, accountID, keyID string) error {
-	return deleteRecordByID[SetupKey](s.db.WithContext(ctx), LockingStrengthUpdate, keyID, accountID)
-}
-
 // getRecords retrieves records from the database based on the account ID.
 func getRecords[T any](db *gorm.DB, lockStrength LockingStrength, accountID string) ([]T, error) {
 	var record []T
@@ -1458,21 +1501,3 @@ func getRecordByID[T any](db *gorm.DB, lockStrength LockingStrength, recordID, a
 	}
 	return &record, nil
 }
-
-// deleteRecordByID deletes a record by its ID and account ID from the database.
-func deleteRecordByID[T any](db *gorm.DB, lockStrength LockingStrength, recordID, accountID string) error {
-	var record T
-	result := db.Clauses(clause.Locking{Strength: string(lockStrength)}).Delete(record, accountAndIDQueryCondition, accountID, recordID)
-	if err := result.Error; err != nil {
-		parts := strings.Split(fmt.Sprintf("%T", record), ".")
-		recordType := parts[len(parts)-1]
-
-		return status.Errorf(status.Internal, "failed to delete %s from store: %v", recordType, err)
-	}
-
-	if result.RowsAffected == 0 {
-		return status.Errorf(status.NotFound, "record not found")
-	}
-
-	return nil
-}
diff --git a/management/server/sql_store_test.go b/management/server/sql_store_test.go
index b371e231319..3f3b2a453d4 100644
--- a/management/server/sql_store_test.go
+++ b/management/server/sql_store_test.go
@@ -1274,7 +1274,7 @@ func Test_DeleteSetupKeySuccessfully(t *testing.T) {
 	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
 	setupKeyID := "A2C8E62B-38F5-4553-B31E-DD66C696CEBB"
 
-	err = store.DeleteSetupKey(context.Background(), accountID, setupKeyID)
+	err = store.DeleteSetupKey(context.Background(), LockingStrengthUpdate, accountID, setupKeyID)
 	require.NoError(t, err)
 
 	_, err = store.GetSetupKeyByID(context.Background(), LockingStrengthShare, setupKeyID, accountID)
@@ -1290,6 +1290,6 @@ func Test_DeleteSetupKeyFailsForNonExistingKey(t *testing.T) {
 	accountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
 	nonExistingKeyID := "non-existing-key-id"
 
-	err = store.DeleteSetupKey(context.Background(), accountID, nonExistingKeyID)
+	err = store.DeleteSetupKey(context.Background(), LockingStrengthUpdate, accountID, nonExistingKeyID)
 	require.Error(t, err)
 }
diff --git a/management/server/status/error.go b/management/server/status/error.go
index a145edf8002..5a75c94b1c1 100644
--- a/management/server/status/error.go
+++ b/management/server/status/error.go
@@ -111,11 +111,21 @@ func NewGetAccountFromStoreError(err error) error {
 	return Errorf(Internal, "issue getting account from store: %s", err)
 }
 
+// NewUserNotPartOfAccountError creates a new Error with PermissionDenied type for a user not being part of an account
+func NewUserNotPartOfAccountError() error {
+	return Errorf(PermissionDenied, "user is not part of this account")
+}
+
 // NewGetUserFromStoreError creates a new Error with Internal type for an issue getting user from store
 func NewGetUserFromStoreError() error {
 	return Errorf(Internal, "issue getting user from store")
 }
 
+// NewAdminPermissionError creates a new Error with PermissionDenied type for actions requiring admin role.
+func NewAdminPermissionError() error {
+	return Errorf(PermissionDenied, "admin role required to perform this action")
+}
+
 // NewStoreContextCanceledError creates a new Error with Internal type for a canceled store context
 func NewStoreContextCanceledError(duration time.Duration) error {
 	return Errorf(Internal, "store access: context canceled after %v", duration)
@@ -125,8 +135,3 @@ func NewStoreContextCanceledError(duration time.Duration) error {
 func NewInvalidKeyIDError() error {
 	return Errorf(InvalidArgument, "invalid key ID")
 }
-
-// NewUnauthorizedToViewSetupKeysError creates a new Error with Unauthorized type for an issue getting a setup key
-func NewUnauthorizedToViewSetupKeysError() error {
-	return Errorf(Unauthorized, "only users with admin power can view setup keys")
-}
diff --git a/management/server/store.go b/management/server/store.go
index 087c9884763..73c9ef6a692 100644
--- a/management/server/store.go
+++ b/management/server/store.go
@@ -70,7 +70,7 @@ type Store interface {
 	DeleteHashedPAT2TokenIDIndex(hashedToken string) error
 	DeleteTokenID2UserIDIndex(tokenID string) error
 
-	GetAccountGroups(ctx context.Context, accountID string) ([]*nbgroup.Group, error)
+	GetAccountGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbgroup.Group, error)
 	GetGroupByID(ctx context.Context, lockStrength LockingStrength, groupID, accountID string) (*nbgroup.Group, error)
 	GetGroupByName(ctx context.Context, lockStrength LockingStrength, groupName, accountID string) (*nbgroup.Group, error)
 	SaveGroups(ctx context.Context, lockStrength LockingStrength, groups []*nbgroup.Group) error
@@ -96,7 +96,9 @@ type Store interface {
 	GetSetupKeyBySecret(ctx context.Context, lockStrength LockingStrength, key string) (*SetupKey, error)
 	IncrementSetupKeyUsage(ctx context.Context, setupKeyID string) error
 	GetAccountSetupKeys(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*SetupKey, error)
-	GetSetupKeyByID(ctx context.Context, lockStrength LockingStrength, setupKeyID string, accountID string) (*SetupKey, error)
+	GetSetupKeyByID(ctx context.Context, lockStrength LockingStrength, accountID, setupKeyID string) (*SetupKey, error)
+	SaveSetupKey(ctx context.Context, lockStrength LockingStrength, setupKey *SetupKey) error
+	DeleteSetupKey(ctx context.Context, lockStrength LockingStrength, accountID, keyID string) error
 
 	GetAccountRoutes(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*route.Route, error)
 	GetRouteByID(ctx context.Context, lockStrength LockingStrength, routeID string, accountID string) (*route.Route, error)
@@ -124,7 +126,6 @@ type Store interface {
 	// This is also a method of metrics.DataSource interface.
 	GetStoreEngine() StoreEngine
 	ExecuteInTransaction(ctx context.Context, f func(store Store) error) error
-	DeleteSetupKey(ctx context.Context, accountID, keyID string) error
 }
 
 type StoreEngine string

From 78044c226d9240edcdd5bb180aaab1da86f442e4 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Fri, 8 Nov 2024 00:32:14 +0300
Subject: [PATCH 02/11] add lock to get account groups

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>
---
 management/server/account.go      | 4 ++--
 management/server/account_test.go | 2 +-
 management/server/group.go        | 2 +-
 management/server/peer.go         | 2 +-
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/management/server/account.go b/management/server/account.go
index aa7609388c0..583853f2504 100644
--- a/management/server/account.go
+++ b/management/server/account.go
@@ -2029,7 +2029,7 @@ func (am *DefaultAccountManager) syncJWTGroups(ctx context.Context, accountID st
 			return fmt.Errorf("error getting user: %w", err)
 		}
 
-		groups, err := transaction.GetAccountGroups(ctx, accountID)
+		groups, err := transaction.GetAccountGroups(ctx, LockingStrengthShare, accountID)
 		if err != nil {
 			return fmt.Errorf("error getting account groups: %w", err)
 		}
@@ -2059,7 +2059,7 @@ func (am *DefaultAccountManager) syncJWTGroups(ctx context.Context, accountID st
 
 		// Propagate changes to peers if group propagation is enabled
 		if settings.GroupsPropagationEnabled {
-			groups, err = transaction.GetAccountGroups(ctx, accountID)
+			groups, err = transaction.GetAccountGroups(ctx, LockingStrengthShare, accountID)
 			if err != nil {
 				return fmt.Errorf("error getting account groups: %w", err)
 			}
diff --git a/management/server/account_test.go b/management/server/account_test.go
index 6a2d85fe8f7..fdf004a3b8a 100644
--- a/management/server/account_test.go
+++ b/management/server/account_test.go
@@ -2773,7 +2773,7 @@ func TestAccount_SetJWTGroups(t *testing.T) {
 		err = manager.syncJWTGroups(context.Background(), "accountID", claims)
 		assert.NoError(t, err, "unable to sync jwt groups")
 
-		groups, err := manager.Store.GetAccountGroups(context.Background(), "accountID")
+		groups, err := manager.Store.GetAccountGroups(context.Background(), LockingStrengthShare, "accountID")
 		assert.NoError(t, err)
 		assert.Len(t, groups, 3, "new group3 should be added")
 
diff --git a/management/server/group.go b/management/server/group.go
index bdb569e377f..b2ec88cc0d2 100644
--- a/management/server/group.go
+++ b/management/server/group.go
@@ -59,7 +59,7 @@ func (am *DefaultAccountManager) GetAllGroups(ctx context.Context, accountID, us
 		return nil, err
 	}
 
-	return am.Store.GetAccountGroups(ctx, accountID)
+	return am.Store.GetAccountGroups(ctx, LockingStrengthShare, accountID)
 }
 
 // GetGroupByName filters all groups in an account by name and returns the one with the most peers
diff --git a/management/server/peer.go b/management/server/peer.go
index 9c5ab571bab..8ced2a1deb0 100644
--- a/management/server/peer.go
+++ b/management/server/peer.go
@@ -765,7 +765,7 @@ func (am *DefaultAccountManager) LoginPeer(ctx context.Context, login PeerLogin)
 		}
 	}
 
-	groups, err := am.Store.GetAccountGroups(ctx, accountID)
+	groups, err := am.Store.GetAccountGroups(ctx, LockingStrengthShare, accountID)
 	if err != nil {
 		return nil, nil, nil, err
 	}

From 1a5f3c653c4b78a5c52bca9bba74c966fbd7495c Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Fri, 8 Nov 2024 00:37:47 +0300
Subject: [PATCH 03/11] add check for regular user

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>
---
 management/server/user.go | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/management/server/user.go b/management/server/user.go
index 9fdd3a6eeea..1368b76b121 100644
--- a/management/server/user.go
+++ b/management/server/user.go
@@ -103,6 +103,11 @@ func (u *User) IsAdminOrServiceUser() bool {
 	return u.HasAdminPower() || u.IsServiceUser
 }
 
+// IsRegularUser checks if the user is a regular user.
+func (u *User) IsRegularUser() bool {
+	return !u.HasAdminPower() && !u.IsServiceUser
+}
+
 // ToUserInfo converts a User object to a UserInfo object.
 func (u *User) ToUserInfo(userData *idp.UserData, settings *Settings) (*UserInfo, error) {
 	autoGroups := u.AutoGroups

From 931521d505b012f45dcf5bcb5de0ee07f0c5b876 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Fri, 8 Nov 2024 00:59:37 +0300
Subject: [PATCH 04/11] get only required groups for auto-group validation

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>
---
 management/server/group/group.go |  5 ++++
 management/server/setupkey.go    | 46 +++++++++++++-------------------
 2 files changed, 23 insertions(+), 28 deletions(-)

diff --git a/management/server/group/group.go b/management/server/group/group.go
index d293e1afc6f..e98e5ecc4b5 100644
--- a/management/server/group/group.go
+++ b/management/server/group/group.go
@@ -49,3 +49,8 @@ func (g *Group) Copy() *Group {
 func (g *Group) HasPeers() bool {
 	return len(g.Peers) > 0
 }
+
+// IsGroupAll checks if the group is a default "All" group.
+func (g *Group) IsGroupAll() bool {
+	return g.Name == "All"
+}
diff --git a/management/server/setupkey.go b/management/server/setupkey.go
index f54eafdc1fd..da248be25d6 100644
--- a/management/server/setupkey.go
+++ b/management/server/setupkey.go
@@ -233,20 +233,16 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 		return nil, status.NewUserNotPartOfAccountError()
 	}
 
-	var accountGroups []*nbgroup.Group
+	var groups []*nbgroup.Group
 	var setupKey *SetupKey
 	var plainKey string
 
 	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		accountGroups, err = transaction.GetAccountGroups(ctx, LockingStrengthShare, accountID)
+		groups, err = validateSetupKeyAutoGroups(ctx, transaction, accountID, autoGroups)
 		if err != nil {
 			return err
 		}
 
-		if err = validateSetupKeyAutoGroups(accountGroups, autoGroups); err != nil {
-			return err
-		}
-
 		setupKey, plainKey = GenerateSetupKey(keyName, keyType, expiresIn, autoGroups, usageLimit, ephemeral)
 		setupKey.AccountID = accountID
 
@@ -257,8 +253,8 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 	}
 
 	am.StoreEvent(ctx, userID, setupKey.Id, accountID, activity.SetupKeyCreated, setupKey.EventMeta())
-	groupMap := make(map[string]*nbgroup.Group, len(accountGroups))
-	for _, g := range accountGroups {
+	groupMap := make(map[string]*nbgroup.Group, len(groups))
+	for _, g := range groups {
 		groupMap[g.ID] = g
 	}
 
@@ -294,20 +290,16 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 		return nil, status.NewUserNotPartOfAccountError()
 	}
 
-	var accountGroups []*nbgroup.Group
+	var groups []*nbgroup.Group
 	var oldKey *SetupKey
 	var newKey *SetupKey
 
 	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		accountGroups, err = transaction.GetAccountGroups(ctx, LockingStrengthShare, accountID)
+		groups, err = validateSetupKeyAutoGroups(ctx, transaction, accountID, keyToSave.AutoGroups)
 		if err != nil {
 			return err
 		}
 
-		if err = validateSetupKeyAutoGroups(accountGroups, keyToSave.AutoGroups); err != nil {
-			return err
-		}
-
 		oldKey, err = transaction.GetSetupKeyByID(ctx, LockingStrengthShare, accountID, keyToSave.Id)
 		if err != nil {
 			return err
@@ -334,8 +326,8 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 		addedGroups := difference(newKey.AutoGroups, oldKey.AutoGroups)
 		removedGroups := difference(oldKey.AutoGroups, newKey.AutoGroups)
 
-		groupMap := make(map[string]*nbgroup.Group, len(accountGroups))
-		for _, g := range accountGroups {
+		groupMap := make(map[string]*nbgroup.Group, len(groups))
+		for _, g := range groups {
 			groupMap[g.ID] = g
 		}
 
@@ -439,22 +431,20 @@ func (am *DefaultAccountManager) DeleteSetupKey(ctx context.Context, accountID,
 	return nil
 }
 
-func validateSetupKeyAutoGroups(groups []*nbgroup.Group, autoGroups []string) error {
-	groupMap := make(map[string]*nbgroup.Group, len(groups))
-	for _, g := range groups {
-		groupMap[g.ID] = g
-	}
+func validateSetupKeyAutoGroups(ctx context.Context, transaction Store, accountID string, autoGroupIDs []string) ([]*nbgroup.Group, error) {
+	autoGroups := make([]*nbgroup.Group, 0, len(autoGroupIDs))
 
-	for _, groupID := range autoGroups {
-		g, exists := groupMap[groupID]
-		if !exists {
-			return status.Errorf(status.NotFound, "group %s doesn't exist", groupID)
+	for _, groupID := range autoGroupIDs {
+		group, err := transaction.GetGroupByID(ctx, LockingStrengthShare, groupID, accountID)
+		if err != nil {
+			return nil, err
 		}
 
-		if g.Name == "All" {
-			return status.Errorf(status.InvalidArgument, "can't add 'All' group to the setup key")
+		if group.IsGroupAll() {
+			return nil, status.Errorf(status.InvalidArgument, "can't add 'All' group to the setup key")
 		}
+		autoGroups = append(autoGroups, group)
 	}
 
-	return nil
+	return autoGroups, nil
 }

From f8b5eedd382d8a218517cf7c7b552f3a0dd8ee3d Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Fri, 8 Nov 2024 10:14:13 +0300
Subject: [PATCH 05/11] add account lock and return auto groups map on
 validation

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>
---
 management/server/setupkey.go | 31 ++++++++++++++-----------------
 1 file changed, 14 insertions(+), 17 deletions(-)

diff --git a/management/server/setupkey.go b/management/server/setupkey.go
index da248be25d6..65d7796f1a0 100644
--- a/management/server/setupkey.go
+++ b/management/server/setupkey.go
@@ -224,6 +224,9 @@ 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.AcquireWriteLockByUID(ctx, accountID)
+	defer unlock()
+
 	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
 		return nil, err
@@ -233,7 +236,7 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 		return nil, status.NewUserNotPartOfAccountError()
 	}
 
-	var groups []*nbgroup.Group
+	var groups map[string]*nbgroup.Group
 	var setupKey *SetupKey
 	var plainKey string
 
@@ -253,13 +256,9 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 	}
 
 	am.StoreEvent(ctx, userID, setupKey.Id, accountID, activity.SetupKeyCreated, setupKey.EventMeta())
-	groupMap := make(map[string]*nbgroup.Group, len(groups))
-	for _, g := range groups {
-		groupMap[g.ID] = g
-	}
 
 	for _, g := range setupKey.AutoGroups {
-		group, ok := groupMap[g]
+		group, ok := groups[g]
 		if ok {
 			am.StoreEvent(ctx, userID, setupKey.Id, accountID, activity.GroupAddedToSetupKey,
 				map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": setupKey.Name})
@@ -281,6 +280,9 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 		return nil, status.Errorf(status.InvalidArgument, "provided setup key to update is nil")
 	}
 
+	unlock := am.Store.AcquireWriteLockByUID(ctx, accountID)
+	defer unlock()
+
 	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
 	if err != nil {
 		return nil, err
@@ -290,7 +292,7 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 		return nil, status.NewUserNotPartOfAccountError()
 	}
 
-	var groups []*nbgroup.Group
+	var groups map[string]*nbgroup.Group
 	var oldKey *SetupKey
 	var newKey *SetupKey
 
@@ -326,13 +328,8 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 		addedGroups := difference(newKey.AutoGroups, oldKey.AutoGroups)
 		removedGroups := difference(oldKey.AutoGroups, newKey.AutoGroups)
 
-		groupMap := make(map[string]*nbgroup.Group, len(groups))
-		for _, g := range groups {
-			groupMap[g.ID] = g
-		}
-
 		for _, g := range removedGroups {
-			group, ok := groupMap[g]
+			group, ok := groups[g]
 			if ok {
 				am.StoreEvent(ctx, userID, oldKey.Id, accountID, activity.GroupRemovedFromSetupKey,
 					map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": newKey.Name})
@@ -340,7 +337,7 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 		}
 
 		for _, g := range addedGroups {
-			group, ok := groupMap[g]
+			group, ok := groups[g]
 			if ok {
 				am.StoreEvent(ctx, userID, oldKey.Id, accountID, activity.GroupAddedToSetupKey,
 					map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": newKey.Name})
@@ -431,8 +428,8 @@ func (am *DefaultAccountManager) DeleteSetupKey(ctx context.Context, accountID,
 	return nil
 }
 
-func validateSetupKeyAutoGroups(ctx context.Context, transaction Store, accountID string, autoGroupIDs []string) ([]*nbgroup.Group, error) {
-	autoGroups := make([]*nbgroup.Group, 0, len(autoGroupIDs))
+func validateSetupKeyAutoGroups(ctx context.Context, transaction Store, accountID string, autoGroupIDs []string) (map[string]*nbgroup.Group, error) {
+	autoGroups := map[string]*nbgroup.Group{}
 
 	for _, groupID := range autoGroupIDs {
 		group, err := transaction.GetGroupByID(ctx, LockingStrengthShare, groupID, accountID)
@@ -443,7 +440,7 @@ func validateSetupKeyAutoGroups(ctx context.Context, transaction Store, accountI
 		if group.IsGroupAll() {
 			return nil, status.Errorf(status.InvalidArgument, "can't add 'All' group to the setup key")
 		}
-		autoGroups = append(autoGroups, group)
+		autoGroups[group.ID] = group
 	}
 
 	return autoGroups, nil

From 3ed8b9cee93e7d45f3d27210606536c06169ab06 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Sat, 9 Nov 2024 01:48:28 +0300
Subject: [PATCH 06/11] fix missing group removed from setup key activity

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>
---
 management/server/setupkey.go | 95 +++++++++++++++++++----------------
 1 file changed, 53 insertions(+), 42 deletions(-)

diff --git a/management/server/setupkey.go b/management/server/setupkey.go
index 65d7796f1a0..2e8230d1ccb 100644
--- a/management/server/setupkey.go
+++ b/management/server/setupkey.go
@@ -12,8 +12,8 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/netbirdio/netbird/management/server/activity"
-	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	"github.com/netbirdio/netbird/management/server/status"
+	log "github.com/sirupsen/logrus"
 )
 
 const (
@@ -236,19 +236,21 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 		return nil, status.NewUserNotPartOfAccountError()
 	}
 
-	var groups map[string]*nbgroup.Group
 	var setupKey *SetupKey
 	var plainKey string
+	var eventsToStore []func()
 
 	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		groups, err = validateSetupKeyAutoGroups(ctx, transaction, accountID, autoGroups)
-		if err != nil {
+		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 {
@@ -256,13 +258,8 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 	}
 
 	am.StoreEvent(ctx, userID, setupKey.Id, accountID, activity.SetupKeyCreated, setupKey.EventMeta())
-
-	for _, g := range setupKey.AutoGroups {
-		group, ok := groups[g]
-		if ok {
-			am.StoreEvent(ctx, userID, setupKey.Id, accountID, activity.GroupAddedToSetupKey,
-				map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": setupKey.Name})
-		}
+	for _, storeEvent := range eventsToStore {
+		storeEvent()
 	}
 
 	// for the creation return the plain key to the caller
@@ -292,13 +289,12 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 		return nil, status.NewUserNotPartOfAccountError()
 	}
 
-	var groups map[string]*nbgroup.Group
 	var oldKey *SetupKey
 	var newKey *SetupKey
+	var eventsToStore []func()
 
 	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		groups, err = validateSetupKeyAutoGroups(ctx, transaction, accountID, keyToSave.AutoGroups)
-		if err != nil {
+		if err = validateSetupKeyAutoGroups(ctx, transaction, accountID, keyToSave.AutoGroups); err != nil {
 			return err
 		}
 
@@ -314,6 +310,12 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 		newKey.Revoked = keyToSave.Revoked
 		newKey.UpdatedAt = time.Now().UTC()
 
+		addedGroups := difference(newKey.AutoGroups, oldKey.AutoGroups)
+		removedGroups := difference(oldKey.AutoGroups, newKey.AutoGroups)
+
+		events := am.prepareSetupKeyEvents(ctx, transaction, accountID, userID, addedGroups, removedGroups, oldKey)
+		eventsToStore = append(eventsToStore, events...)
+
 		return transaction.SaveSetupKey(ctx, LockingStrengthUpdate, newKey)
 	})
 	if err != nil {
@@ -324,26 +326,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, ok := groups[g]
-			if ok {
-				am.StoreEvent(ctx, userID, oldKey.Id, accountID, activity.GroupRemovedFromSetupKey,
-					map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": newKey.Name})
-			}
-		}
-
-		for _, g := range addedGroups {
-			group, ok := groups[g]
-			if ok {
-				am.StoreEvent(ctx, userID, oldKey.Id, accountID, activity.GroupAddedToSetupKey,
-					map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": newKey.Name})
-			}
-		}
-	}()
+	for _, storeEvent := range eventsToStore {
+		storeEvent()
+	}
 
 	return newKey, nil
 }
@@ -412,7 +397,7 @@ func (am *DefaultAccountManager) DeleteSetupKey(ctx context.Context, accountID,
 	var deletedSetupKey *SetupKey
 
 	err = am.Store.ExecuteInTransaction(ctx, func(transaction Store) error {
-		deletedSetupKey, err = transaction.GetSetupKeyByID(ctx, LockingStrengthShare, keyID, accountID)
+		deletedSetupKey, err = transaction.GetSetupKeyByID(ctx, LockingStrengthShare, accountID, keyID)
 		if err != nil {
 			return err
 		}
@@ -428,20 +413,46 @@ func (am *DefaultAccountManager) DeleteSetupKey(ctx context.Context, accountID,
 	return nil
 }
 
-func validateSetupKeyAutoGroups(ctx context.Context, transaction Store, accountID string, autoGroupIDs []string) (map[string]*nbgroup.Group, error) {
-	autoGroups := map[string]*nbgroup.Group{}
-
+func validateSetupKeyAutoGroups(ctx context.Context, transaction Store, accountID string, autoGroupIDs []string) error {
 	for _, groupID := range autoGroupIDs {
 		group, err := transaction.GetGroupByID(ctx, LockingStrengthShare, groupID, accountID)
 		if err != nil {
-			return nil, err
+			return err
 		}
 
 		if group.IsGroupAll() {
-			return nil, status.Errorf(status.InvalidArgument, "can't add 'All' group to the setup key")
+			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()
+
+	for _, g := range removedGroups {
+		group, err := transaction.GetGroupByID(ctx, LockingStrengthShare, g, accountID)
+		if err != nil {
+			log.WithContext(ctx).Debugf("skipped adding group: %s GroupRemovedFromSetupKey activity: %v", g, err)
+			continue
+		}
+
+		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, err := transaction.GetGroupByID(ctx, LockingStrengthShare, g, accountID)
+		if err != nil {
+			log.WithContext(ctx).Debugf("skipped adding group: %s GroupAddedToSetupKey activity: %v", g, err)
+			continue
 		}
-		autoGroups[group.ID] = group
+
+		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 autoGroups, nil
+	return eventsToStore
 }

From 5b4a4aab4b034b3e72d4aa69edc0977ba78b56ce Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Mon, 11 Nov 2024 17:51:56 +0300
Subject: [PATCH 07/11] Remove context from DB queries

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>
---
 management/server/sql_store.go | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/management/server/sql_store.go b/management/server/sql_store.go
index 2f028b4f70e..d6ce0022e1e 100644
--- a/management/server/sql_store.go
+++ b/management/server/sql_store.go
@@ -570,7 +570,7 @@ func (s *SqlStore) GetAccountUsers(ctx context.Context, accountID string) ([]*Us
 
 func (s *SqlStore) GetAccountGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbgroup.Group, error) {
 	var groups []*nbgroup.Group
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&groups, accountIDCondition, accountID)
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&groups, accountIDCondition, accountID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
 			return nil, status.Errorf(status.NotFound, "accountID not found: index lookup failed")
@@ -1221,7 +1221,7 @@ func (s *SqlStore) GetRouteByID(ctx context.Context, lockStrength LockingStrengt
 // GetAccountSetupKeys retrieves setup keys for an account.
 func (s *SqlStore) GetAccountSetupKeys(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*SetupKey, error) {
 	var setupKeys []*SetupKey
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
 		Find(&setupKeys, accountIDCondition, accountID)
 	if err := result.Error; err != nil {
 		log.WithContext(ctx).Errorf("failed to get setup keys from the store: %s", err)
@@ -1234,7 +1234,7 @@ func (s *SqlStore) GetAccountSetupKeys(ctx context.Context, lockStrength Locking
 // GetSetupKeyByID retrieves a setup key by its ID and account ID.
 func (s *SqlStore) GetSetupKeyByID(ctx context.Context, lockStrength LockingStrength, accountID, setupKeyID string) (*SetupKey, error) {
 	var setupKey *SetupKey
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).
 		First(&setupKey, accountAndIDQueryCondition, accountID, setupKeyID)
 	if err := result.Error; err != nil {
 		if errors.Is(err, gorm.ErrRecordNotFound) {
@@ -1249,8 +1249,7 @@ func (s *SqlStore) GetSetupKeyByID(ctx context.Context, lockStrength LockingStre
 
 // SaveSetupKey saves a setup key to the database.
 func (s *SqlStore) SaveSetupKey(ctx context.Context, lockStrength LockingStrength, setupKey *SetupKey) error {
-	result := s.db.WithContext(ctx).Session(&gorm.Session{FullSaveAssociations: true}).
-		Clauses(clause.Locking{Strength: string(lockStrength)}).Save(setupKey)
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(setupKey)
 	if result.Error != nil {
 		log.WithContext(ctx).Errorf("failed to save setup key to store: %s", result.Error)
 		return status.Errorf(status.Internal, "failed to save setup key to store")
@@ -1261,8 +1260,7 @@ func (s *SqlStore) SaveSetupKey(ctx context.Context, lockStrength LockingStrengt
 
 // DeleteSetupKey deletes a setup key from the database.
 func (s *SqlStore) DeleteSetupKey(ctx context.Context, lockStrength LockingStrength, accountID, keyID string) error {
-	result := s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).
-		Delete(&SetupKey{}, accountAndIDQueryCondition, accountID, keyID)
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Delete(&SetupKey{}, accountAndIDQueryCondition, accountID, keyID)
 	if result.Error != nil {
 		log.WithContext(ctx).Errorf("failed to delete setup key from store: %s", result.Error)
 		return status.Errorf(status.Internal, "failed to delete setup key from store")

From 69c3231f59cc301800ea8a7265d09f6ff058b389 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Mon, 11 Nov 2024 18:05:36 +0300
Subject: [PATCH 08/11] Add user permission check and add setup events into
 events to store slice

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>
---
 management/server/setupkey.go | 24 ++++++++++++++++++++----
 1 file changed, 20 insertions(+), 4 deletions(-)

diff --git a/management/server/setupkey.go b/management/server/setupkey.go
index 2e8230d1ccb..f644236c82d 100644
--- a/management/server/setupkey.go
+++ b/management/server/setupkey.go
@@ -236,6 +236,10 @@ func (am *DefaultAccountManager) CreateSetupKey(ctx context.Context, accountID s
 		return nil, status.NewUserNotPartOfAccountError()
 	}
 
+	if user.IsRegularUser() {
+		return nil, status.NewAdminPermissionError()
+	}
+
 	var setupKey *SetupKey
 	var plainKey string
 	var eventsToStore []func()
@@ -289,6 +293,10 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 		return nil, status.NewUserNotPartOfAccountError()
 	}
 
+	if user.IsRegularUser() {
+		return nil, status.NewAdminPermissionError()
+	}
+
 	var oldKey *SetupKey
 	var newKey *SetupKey
 	var eventsToStore []func()
@@ -348,6 +356,10 @@ func (am *DefaultAccountManager) ListSetupKeys(ctx context.Context, accountID, u
 		return nil, status.NewAdminPermissionError()
 	}
 
+	if user.IsRegularUser() {
+		return nil, status.NewAdminPermissionError()
+	}
+
 	return am.Store.GetAccountSetupKeys(ctx, LockingStrengthShare, accountID)
 }
 
@@ -439,8 +451,10 @@ func (am *DefaultAccountManager) prepareSetupKeyEvents(ctx context.Context, tran
 			continue
 		}
 
-		meta := map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": key.Name}
-		am.StoreEvent(ctx, userID, key.Id, accountID, activity.GroupRemovedFromSetupKey, meta)
+		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 {
@@ -450,8 +464,10 @@ func (am *DefaultAccountManager) prepareSetupKeyEvents(ctx context.Context, tran
 			continue
 		}
 
-		meta := map[string]any{"group": group.Name, "group_id": group.ID, "setupkey": key.Name}
-		am.StoreEvent(ctx, userID, key.Id, accountID, activity.GroupAddedToSetupKey, meta)
+		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

From fdfd2ada5058bb36b865c29c0f8ffd9ebcfc9d11 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Mon, 11 Nov 2024 18:44:41 +0300
Subject: [PATCH 09/11] Retrieve all groups once during setup key auto-group
 validation

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>
---
 management/server/setupkey.go  | 30 +++++++++++++++++++-----------
 management/server/sql_store.go | 17 +++++++++++++++++
 management/server/store.go     |  1 +
 3 files changed, 37 insertions(+), 11 deletions(-)

diff --git a/management/server/setupkey.go b/management/server/setupkey.go
index f644236c82d..4817b42b978 100644
--- a/management/server/setupkey.go
+++ b/management/server/setupkey.go
@@ -356,10 +356,6 @@ func (am *DefaultAccountManager) ListSetupKeys(ctx context.Context, accountID, u
 		return nil, status.NewAdminPermissionError()
 	}
 
-	if user.IsRegularUser() {
-		return nil, status.NewAdminPermissionError()
-	}
-
 	return am.Store.GetAccountSetupKeys(ctx, LockingStrengthShare, accountID)
 }
 
@@ -426,10 +422,15 @@ func (am *DefaultAccountManager) DeleteSetupKey(ctx context.Context, accountID,
 }
 
 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, err := transaction.GetGroupByID(ctx, LockingStrengthShare, groupID, accountID)
-		if err != nil {
-			return err
+		group, ok := groups[groupID]
+		if !ok {
+			return status.Errorf(status.NotFound, "group not found: %s", groupID)
 		}
 
 		if group.IsGroupAll() {
@@ -444,9 +445,16 @@ func validateSetupKeyAutoGroups(ctx context.Context, transaction Store, accountI
 func (am *DefaultAccountManager) prepareSetupKeyEvents(ctx context.Context, transaction Store, accountID, userID string, addedGroups, removedGroups []string, key *SetupKey) []func() {
 	var eventsToStore []func()
 
+	modifiedGroups := append(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, err := transaction.GetGroupByID(ctx, LockingStrengthShare, g, accountID)
-		if err != nil {
+		group, ok := groups[g]
+		if !ok {
 			log.WithContext(ctx).Debugf("skipped adding group: %s GroupRemovedFromSetupKey activity: %v", g, err)
 			continue
 		}
@@ -458,8 +466,8 @@ func (am *DefaultAccountManager) prepareSetupKeyEvents(ctx context.Context, tran
 	}
 
 	for _, g := range addedGroups {
-		group, err := transaction.GetGroupByID(ctx, LockingStrengthShare, g, accountID)
-		if err != nil {
+		group, ok := groups[g]
+		if !ok {
 			log.WithContext(ctx).Debugf("skipped adding group: %s GroupAddedToSetupKey activity: %v", g, err)
 			continue
 		}
diff --git a/management/server/sql_store.go b/management/server/sql_store.go
index d6ce0022e1e..e7b86340255 100644
--- a/management/server/sql_store.go
+++ b/management/server/sql_store.go
@@ -1179,6 +1179,23 @@ func (s *SqlStore) GetGroupByName(ctx context.Context, lockStrength LockingStren
 	return &group, nil
 }
 
+// GetGroupsByIDs retrieves groups by their IDs and account ID.
+func (s *SqlStore) GetGroupsByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, groupIDs []string) (map[string]*nbgroup.Group, error) {
+	var groups []*nbgroup.Group
+	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Find(&groups, "account_id = ? AND id in ?", accountID, groupIDs)
+	if result.Error != nil {
+		log.WithContext(ctx).Errorf("failed to get groups by ID's from the store: %s", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get groups by ID's from the store")
+	}
+
+	groupsMap := make(map[string]*nbgroup.Group)
+	for _, group := range groups {
+		groupsMap[group.ID] = group
+	}
+
+	return groupsMap, nil
+}
+
 // SaveGroup saves a group to the store.
 func (s *SqlStore) SaveGroup(ctx context.Context, lockStrength LockingStrength, group *nbgroup.Group) error {
 	result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(group)
diff --git a/management/server/store.go b/management/server/store.go
index 73c9ef6a692..c7c066629e3 100644
--- a/management/server/store.go
+++ b/management/server/store.go
@@ -73,6 +73,7 @@ type Store interface {
 	GetAccountGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbgroup.Group, error)
 	GetGroupByID(ctx context.Context, lockStrength LockingStrength, groupID, accountID string) (*nbgroup.Group, error)
 	GetGroupByName(ctx context.Context, lockStrength LockingStrength, groupName, accountID string) (*nbgroup.Group, error)
+	GetGroupsByIDs(ctx context.Context, lockStrength LockingStrength, accountID string, groupIDs []string) (map[string]*nbgroup.Group, error)
 	SaveGroups(ctx context.Context, lockStrength LockingStrength, groups []*nbgroup.Group) error
 	SaveGroup(ctx context.Context, lockStrength LockingStrength, group *nbgroup.Group) error
 

From 15194de7813e43b4b9c28beb06fbc8a61382f926 Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Mon, 11 Nov 2024 18:48:37 +0300
Subject: [PATCH 10/11] Fix lint

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>
---
 management/server/setupkey.go | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/management/server/setupkey.go b/management/server/setupkey.go
index 4817b42b978..554c66ba4fc 100644
--- a/management/server/setupkey.go
+++ b/management/server/setupkey.go
@@ -5,6 +5,7 @@ import (
 	"crypto/sha256"
 	b64 "encoding/base64"
 	"hash/fnv"
+	"slices"
 	"strconv"
 	"strings"
 	"time"
@@ -445,7 +446,7 @@ func validateSetupKeyAutoGroups(ctx context.Context, transaction Store, accountI
 func (am *DefaultAccountManager) prepareSetupKeyEvents(ctx context.Context, transaction Store, accountID, userID string, addedGroups, removedGroups []string, key *SetupKey) []func() {
 	var eventsToStore []func()
 
-	modifiedGroups := append(addedGroups, removedGroups...)
+	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)

From cf2766dc2c128debf433fe43326d9d2bc5c614ee Mon Sep 17 00:00:00 2001
From: bcmmbaga <bethuelmbaga12@gmail.com>
Date: Mon, 11 Nov 2024 19:13:23 +0300
Subject: [PATCH 11/11] Fix sonar

Signed-off-by: bcmmbaga <bethuelmbaga12@gmail.com>
---
 management/server/sql_store.go    | 21 ++++++++++++---------
 management/server/status/error.go |  4 ++--
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/management/server/sql_store.go b/management/server/sql_store.go
index e7b86340255..9dd3e778d43 100644
--- a/management/server/sql_store.go
+++ b/management/server/sql_store.go
@@ -485,9 +485,10 @@ func (s *SqlStore) GetAccountBySetupKey(ctx context.Context, setupKey string) (*
 	result := s.db.Select("account_id").First(&key, keyQueryCondition, setupKey)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.Errorf(status.NotFound, "account not found: index lookup failed")
+			return nil, status.NewSetupKeyNotFoundError(setupKey)
 		}
-		return nil, status.NewSetupKeyNotFoundError(result.Error)
+		log.WithContext(ctx).Errorf("failed to get account by setup key from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get account by setup key from store")
 	}
 
 	if key.AccountID == "" {
@@ -756,9 +757,10 @@ func (s *SqlStore) GetAccountIDBySetupKey(ctx context.Context, setupKey string)
 	result := s.db.Model(&SetupKey{}).Select("account_id").Where(keyQueryCondition, setupKey).First(&accountID)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return "", status.Errorf(status.NotFound, "account not found: index lookup failed")
+			return "", status.NewSetupKeyNotFoundError(setupKey)
 		}
-		return "", status.NewSetupKeyNotFoundError(result.Error)
+		log.WithContext(ctx).Errorf("failed to get account ID by setup key from store: %v", result.Error)
+		return "", status.Errorf(status.Internal, "failed to get account ID by setup key from store")
 	}
 
 	if accountID == "" {
@@ -986,9 +988,10 @@ func (s *SqlStore) GetSetupKeyBySecret(ctx context.Context, lockStrength Locking
 		First(&setupKey, keyQueryCondition, key)
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.Errorf(status.NotFound, "setup key not found")
+			return nil, status.NewSetupKeyNotFoundError(key)
 		}
-		return nil, status.NewSetupKeyNotFoundError(result.Error)
+		log.WithContext(ctx).Errorf("failed to get setup key by secret from store: %v", result.Error)
+		return nil, status.Errorf(status.Internal, "failed to get setup key by secret from store")
 	}
 	return &setupKey, nil
 }
@@ -1006,7 +1009,7 @@ func (s *SqlStore) IncrementSetupKeyUsage(ctx context.Context, setupKeyID string
 	}
 
 	if result.RowsAffected == 0 {
-		return status.Errorf(status.NotFound, "setup key not found")
+		return status.NewSetupKeyNotFoundError(setupKeyID)
 	}
 
 	return nil
@@ -1255,7 +1258,7 @@ func (s *SqlStore) GetSetupKeyByID(ctx context.Context, lockStrength LockingStre
 		First(&setupKey, accountAndIDQueryCondition, accountID, setupKeyID)
 	if err := result.Error; err != nil {
 		if errors.Is(err, gorm.ErrRecordNotFound) {
-			return nil, status.Errorf(status.NotFound, "setup key not found")
+			return nil, status.NewSetupKeyNotFoundError(setupKeyID)
 		}
 		log.WithContext(ctx).Errorf("failed to get setup key from the store: %s", err)
 		return nil, status.Errorf(status.Internal, "failed to get setup key from store")
@@ -1284,7 +1287,7 @@ func (s *SqlStore) DeleteSetupKey(ctx context.Context, lockStrength LockingStren
 	}
 
 	if result.RowsAffected == 0 {
-		return status.Errorf(status.NotFound, "setup key not found")
+		return status.NewSetupKeyNotFoundError(keyID)
 	}
 
 	return nil
diff --git a/management/server/status/error.go b/management/server/status/error.go
index 5a75c94b1c1..a415d5b6e2b 100644
--- a/management/server/status/error.go
+++ b/management/server/status/error.go
@@ -103,8 +103,8 @@ func NewPeerLoginExpiredError() error {
 }
 
 // NewSetupKeyNotFoundError creates a new Error with NotFound type for a missing setup key
-func NewSetupKeyNotFoundError(err error) error {
-	return Errorf(NotFound, "setup key not found: %s", err)
+func NewSetupKeyNotFoundError(setupKeyID string) error {
+	return Errorf(NotFound, "setup key: %s not found", setupKeyID)
 }
 
 func NewGetAccountFromStoreError(err error) error {