refactor setup keys

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

management/server/file_store.go

@@ -1051,6 +1051,14 @@ func (s *FileStore) GetSetupKeyByID(_ context.Context, _ LockingStrength, _ stri
 	return nil, status.Errorf(status.Internal, "GetSetupKeyByID is not implemented")
 }
 
+func (s *FileStore) SaveSetupKey(_ context.Context, _ LockingStrength, _ *SetupKey) error {
+	return status.Errorf(status.Internal, "GetSetupKeyByID is not implemented")
+}
+
+func (s *FileStore) DeleteSetupKey(_ context.Context, _ LockingStrength, _ string) error {
+	return status.Errorf(status.Internal, "DeleteSetupKey is not implemented")
+}
+
 func (s *FileStore) GetAccountNameServerGroups(_ context.Context, _ LockingStrength, _ string) ([]*dns.NameServerGroup, error) {
 	return nil, status.Errorf(status.Internal, "GetAccountNameServerGroups is not implemented")
 }

management/server/setupkey.go

@@ -9,6 +9,7 @@ import (
 	"unicode/utf8"
 
 	"github.com/google/uuid"
+	nbgroup "github.com/netbirdio/netbird/management/server/group"
 	log "github.com/sirupsen/logrus"
 
 	"github.com/netbirdio/netbird/management/server/activity"
@@ -210,39 +211,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)
+	user, err := am.Store.GetUserByUserID(ctx, LockingStrengthShare, userID)
-	defer unlock()
+	if err != nil {
+		return nil, err
+	}
+
+	if user.AccountID != accountID {
+		return nil, status.Errorf(status.PermissionDenied, "only users with admin power can update setup keys")
+	}
 
 	keyDuration := DefaultSetupKeyDuration
 	if expiresIn != 0 {
 		keyDuration = expiresIn
 	}
 
-	account, err := am.Store.GetAccount(ctx, accountID)
+	groups, err := am.Store.GetAccountGroups(ctx, accountID)
 	if err != nil {
 		return nil, err
 	}
 
-	if err := validateSetupKeyAutoGroups(account, autoGroups); err != nil {
+	if err = validateSetupKeyAutoGroups(groups, autoGroups); err != nil {
 		return nil, err
 	}
 
 	setupKey := GenerateSetupKey(keyName, keyType, keyDuration, autoGroups, usageLimit, ephemeral)
-	account.SetupKeys[setupKey.Key] = setupKey
+	setupKey.AccountID = accountID
-	err = am.Store.SaveAccount(ctx, account)
+
-	if err != nil {
+	if err = am.Store.SaveSetupKey(ctx, LockingStrengthUpdate, setupKey); 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(groups))
+	for _, g := range groups {
+		groupMap[g.ID] = g
+	}
 
 	for _, g := range setupKey.AutoGroups {
-		group := account.GetGroup(g)
+		group := groupMap[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)
+			log.WithContext(ctx).Errorf("group %s not found while saving setup key activity event of account %s", g, accountID)
 		}
 	}
 
@@ -254,30 +265,30 @@ 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
 	}
 
-	var oldKey *SetupKey
+	if user.AccountID != accountID {
-	for _, key := range account.SetupKeys {
+		return nil, status.Errorf(status.PermissionDenied, "only users with admin power can update setup keys")
-		if key.Id == keyToSave.Id {
-			oldKey = key.Copy()
-			break
-		}
-	}
-	if oldKey == nil {
-		return nil, status.Errorf(status.NotFound, "setup key not found")
 	}
 
-	if err := validateSetupKeyAutoGroups(account, keyToSave.AutoGroups); err != nil {
+	groups, err := am.Store.GetAccountGroups(ctx, accountID)
+	if err != nil {
+		return nil, err
+	}
+
+	if err = validateSetupKeyAutoGroups(groups, keyToSave.AutoGroups); err != nil {
+		return nil, err
+	}
+
+	oldKey, err := am.Store.GetSetupKeyByID(ctx, LockingStrengthShare, keyToSave.Id, accountID)
+	if err != nil {
 		return nil, err
 	}
 
@@ -288,9 +299,7 @@ func (am *DefaultAccountManager) SaveSetupKey(ctx context.Context, accountID str
 	newKey.Revoked = keyToSave.Revoked
 	newKey.UpdatedAt = time.Now().UTC()
 
-	account.SetupKeys[newKey.Key] = newKey
+	if err = am.Store.SaveSetupKey(ctx, LockingStrengthUpdate, newKey); err != nil {
-
-	if err = am.Store.SaveAccount(ctx, account); err != nil {
 		return nil, err
 	}
 
@@ -301,30 +310,34 @@ 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(groups))
+		for _, g := range groups {
+			groupMap[g.ID] = g
+		}
+
 		for _, g := range removedGroups {
-			group := account.GetGroup(g)
+			group := groupMap[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)
+				log.WithContext(ctx).Errorf("group %s not found while saving setup key activity event of account %s", g, accountID)
 			}
 
 		}
 
 		for _, g := range addedGroups {
-			group := account.GetGroup(g)
+			group := groupMap[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)
+				log.WithContext(ctx).Errorf("group %s not found while saving setup key activity event of account %s", g, accountID)
 			}
 		}
 	}()
 
-	am.updateAccountPeers(ctx, account)
-
 	return newKey, nil
 }
 
@@ -386,15 +399,22 @@ func (am *DefaultAccountManager) GetSetupKey(ctx context.Context, accountID, use
 	return setupKey, nil
 }
 
-func validateSetupKeyAutoGroups(account *Account, autoGroups []string) error {
+func validateSetupKeyAutoGroups(groups []*nbgroup.Group, autoGroups []string) error {
-	for _, group := range autoGroups {
+	groupMap := make(map[string]*nbgroup.Group, len(groups))
-		g, ok := account.Groups[group]
+	for _, g := range groups {
-		if !ok {
+		groupMap[g.ID] = g
-			return status.Errorf(status.NotFound, "group %s doesn't exist", group)
+	}
+
+	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
 }

management/server/sql_store.go

@@ -1180,6 +1180,18 @@ func (s *SqlStore) GetSetupKeyByID(ctx context.Context, lockStrength LockingStre
 	return getRecordByID[SetupKey](s.db.WithContext(ctx), lockStrength, setupKeyID, accountID)
 }
 
+// SaveSetupKey saves a setup key to the database.
+func (s *SqlStore) SaveSetupKey(ctx context.Context, lockStrength LockingStrength, setupKey *SetupKey) error {
+	return s.db.WithContext(ctx).Session(&gorm.Session{FullSaveAssociations: true}).
+		Clauses(clause.Locking{Strength: string(lockStrength)}).Save(&setupKey).Error
+}
+
+// DeleteSetupKey deletes a setup key from the database.
+func (s *SqlStore) DeleteSetupKey(ctx context.Context, lockStrength LockingStrength, setupKeyID string) error {
+	return s.db.WithContext(ctx).Clauses(clause.Locking{Strength: string(lockStrength)}).
+		Delete(&posture.Checks{}, idQueryCondition, setupKeyID).Error
+}
+
 // GetAccountNameServerGroups retrieves name server groups for an account.
 func (s *SqlStore) GetAccountNameServerGroups(ctx context.Context, lockStrength LockingStrength, accountID string) ([]*nbdns.NameServerGroup, error) {
 	return getRecords[*nbdns.NameServerGroup](s.db.WithContext(ctx), lockStrength, accountID)

management/server/store.go

@@ -93,6 +93,8 @@ type Store interface {
 	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)
+	SaveSetupKey(ctx context.Context, lockStrength LockingStrength, setupKey *SetupKey) error
+	DeleteSetupKey(ctx context.Context, lockStrength LockingStrength, setupKeyID 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)