From 389c9619afe8b6d129a12d137bf332491503bb83 Mon Sep 17 00:00:00 2001 From: bcmmbaga Date: Fri, 8 Nov 2024 00:31:41 +0300 Subject: [PATCH] Refactor setup key handling to use store methods Signed-off-by: bcmmbaga --- management/server/setupkey.go | 183 +++++++++++++++++----------- 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, 180 insertions(+), 112 deletions(-) diff --git a/management/server/setupkey.go b/management/server/setupkey.go index 43b6e02c9..f54eafdc1 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 user.IsRegularUser() { + return status.NewAdminPermissionError() + } + + 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 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 +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 646184578..a11370e4f 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 b371e2313..3f3b2a453 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 a145edf80..5a75c94b1 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 087c98847..73c9ef6a6 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