diff --git a/management/server/account.go b/management/server/account.go index 8a80aefb6..fdaa6ddef 100644 --- a/management/server/account.go +++ b/management/server/account.go @@ -1314,7 +1314,7 @@ func (am *DefaultAccountManager) SyncUserJWTGroups(ctx context.Context, userAuth // Propagate changes to peers if group propagation is enabled if settings.GroupsPropagationEnabled { - groups, err = transaction.GetAccountGroups(ctx, store.LockingStrengthShare, userAuth.AccountId) + groups, err = transaction.GetAccountGroups(ctx, store.LockingStrengthUpdate, userAuth.AccountId) if err != nil { return fmt.Errorf("error getting account groups: %w", err) } @@ -1899,7 +1899,7 @@ func (am *DefaultAccountManager) UpdateToPrimaryAccount(ctx context.Context, acc // propagateUserGroupMemberships propagates all account users' group memberships to their peers. // Returns true if any groups were modified, true if those updates affect peers and an error. func propagateUserGroupMemberships(ctx context.Context, transaction store.Store, accountID string) (groupsUpdated bool, peersAffected bool, err error) { - groups, err := transaction.GetAccountGroups(ctx, store.LockingStrengthShare, accountID) + groups, err := transaction.GetAccountGroups(ctx, store.LockingStrengthUpdate, accountID) if err != nil { return false, false, err } diff --git a/management/server/group.go b/management/server/group.go index 130a67145..c18c34240 100644 --- a/management/server/group.go +++ b/management/server/group.go @@ -265,20 +265,10 @@ func (am *DefaultAccountManager) GroupAddPeer(ctx context.Context, accountID, gr unlock := am.Store.AcquireWriteLockByUID(ctx, accountID) defer unlock() - var group *types.Group var updateAccountPeers bool var err error err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error { - group, err = transaction.GetGroupByID(context.Background(), store.LockingStrengthUpdate, accountID, groupID) - if err != nil { - return err - } - - if updated := group.AddPeer(peerID); !updated { - return nil - } - updateAccountPeers, err = areGroupChangesAffectPeers(ctx, transaction, accountID, []string{groupID}) if err != nil { return err @@ -288,7 +278,7 @@ func (am *DefaultAccountManager) GroupAddPeer(ctx context.Context, accountID, gr return err } - return transaction.SaveGroup(ctx, store.LockingStrengthUpdate, group) + return transaction.AddPeerToGroup(ctx, peerID, groupID) }) if err != nil { return err @@ -347,20 +337,10 @@ func (am *DefaultAccountManager) GroupDeletePeer(ctx context.Context, accountID, unlock := am.Store.AcquireWriteLockByUID(ctx, accountID) defer unlock() - var group *types.Group var updateAccountPeers bool var err error err = am.Store.ExecuteInTransaction(ctx, func(transaction store.Store) error { - group, err = transaction.GetGroupByID(context.Background(), store.LockingStrengthUpdate, accountID, groupID) - if err != nil { - return err - } - - if updated := group.RemovePeer(peerID); !updated { - return nil - } - updateAccountPeers, err = areGroupChangesAffectPeers(ctx, transaction, accountID, []string{groupID}) if err != nil { return err @@ -370,7 +350,7 @@ func (am *DefaultAccountManager) GroupDeletePeer(ctx context.Context, accountID, return err } - return transaction.SaveGroup(ctx, store.LockingStrengthUpdate, group) + return transaction.RemovePeerFromGroup(ctx, peerID, groupID) }) if err != nil { return err diff --git a/management/server/group_test.go b/management/server/group_test.go index 631fe3a71..306175c1d 100644 --- a/management/server/group_test.go +++ b/management/server/group_test.go @@ -2,14 +2,19 @@ package server import ( "context" + "encoding/binary" "errors" "fmt" + "net" "net/netip" + "strconv" + "sync" "testing" "time" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "golang.org/x/exp/maps" nbdns "github.com/netbirdio/netbird/dns" "github.com/netbirdio/netbird/management/server/groups" @@ -18,8 +23,10 @@ import ( "github.com/netbirdio/netbird/management/server/networks/routers" routerTypes "github.com/netbirdio/netbird/management/server/networks/routers/types" networkTypes "github.com/netbirdio/netbird/management/server/networks/types" + peer2 "github.com/netbirdio/netbird/management/server/peer" "github.com/netbirdio/netbird/management/server/permissions" "github.com/netbirdio/netbird/management/server/status" + "github.com/netbirdio/netbird/management/server/store" "github.com/netbirdio/netbird/management/server/types" "github.com/netbirdio/netbird/route" ) @@ -733,3 +740,265 @@ func TestGroupAccountPeersUpdate(t *testing.T) { } }) } + +func Test_AddPeerToGroup(t *testing.T) { + t.Setenv("NETBIRD_STORE_ENGINE", string(types.PostgresStoreEngine)) + manager, err := createManager(t) + if err != nil { + t.Fatal(err) + return + } + + accountID := "testaccount" + userID := "testuser" + + acc, err := createAccount(manager, accountID, userID, "domain.com") + if err != nil { + t.Fatal("error creating account") + return + } + + const totalPeers = 10000 // totalPeers / differentHostnames should be less than 10 (due to concurrent retries) + const differentHostnames = 50 + + var wg sync.WaitGroup + errs := make(chan error, totalPeers+differentHostnames) + start := make(chan struct{}) + for i := 0; i < totalPeers; i++ { + wg.Add(1) + + go func(i int) { + defer wg.Done() + + <-start + + err = manager.Store.AddPeerToGroup(context.Background(), strconv.Itoa(i), acc.GroupsG[0].ID) + if err != nil { + errs <- fmt.Errorf("AddPeer failed for peer %d: %w", i, err) + return + } + + }(i) + } + startTime := time.Now() + + close(start) + wg.Wait() + close(errs) + + t.Logf("time since start: %s", time.Since(startTime)) + + for err := range errs { + t.Fatal(err) + } + + account, err := manager.Store.GetAccount(context.Background(), accountID) + if err != nil { + t.Fatalf("Failed to get account %s: %v", accountID, err) + } + + assert.Equal(t, totalPeers, len(maps.Values(account.Groups)[0].Peers), "Expected %d peers in account %s, got %d", totalPeers, accountID, len(account.Peers)) +} + +func Test_AddPeerToAll(t *testing.T) { + t.Setenv("NETBIRD_STORE_ENGINE", string(types.PostgresStoreEngine)) + manager, err := createManager(t) + if err != nil { + t.Fatal(err) + return + } + + accountID := "testaccount" + userID := "testuser" + + _, err = createAccount(manager, accountID, userID, "domain.com") + if err != nil { + t.Fatal("error creating account") + return + } + + const totalPeers = 10000 // totalPeers / differentHostnames should be less than 10 (due to concurrent retries) + const differentHostnames = 50 + + var wg sync.WaitGroup + errs := make(chan error, totalPeers+differentHostnames) + start := make(chan struct{}) + for i := 0; i < totalPeers; i++ { + wg.Add(1) + + go func(i int) { + defer wg.Done() + + <-start + + err = manager.Store.AddPeerToAllGroup(context.Background(), accountID, strconv.Itoa(i)) + if err != nil { + errs <- fmt.Errorf("AddPeer failed for peer %d: %w", i, err) + return + } + + }(i) + } + startTime := time.Now() + + close(start) + wg.Wait() + close(errs) + + t.Logf("time since start: %s", time.Since(startTime)) + + for err := range errs { + t.Fatal(err) + } + + account, err := manager.Store.GetAccount(context.Background(), accountID) + if err != nil { + t.Fatalf("Failed to get account %s: %v", accountID, err) + } + + assert.Equal(t, totalPeers, len(maps.Values(account.Groups)[0].Peers), "Expected %d peers in account %s, got %d", totalPeers, accountID, len(account.Peers)) +} + +func Test_AddPeerAndAddToAll(t *testing.T) { + t.Setenv("NETBIRD_STORE_ENGINE", string(types.PostgresStoreEngine)) + manager, err := createManager(t) + if err != nil { + t.Fatal(err) + return + } + + accountID := "testaccount" + userID := "testuser" + + _, err = createAccount(manager, accountID, userID, "domain.com") + if err != nil { + t.Fatal("error creating account") + return + } + + const totalPeers = 10000 // totalPeers / differentHostnames should be less than 10 (due to concurrent retries) + const differentHostnames = 50 + + var wg sync.WaitGroup + errs := make(chan error, totalPeers+differentHostnames) + start := make(chan struct{}) + for i := 0; i < totalPeers; i++ { + wg.Add(1) + + go func(i int) { + defer wg.Done() + + <-start + + peer := &peer2.Peer{ + ID: strconv.Itoa(i), + AccountID: accountID, + Meta: peer2.PeerSystemMeta{Hostname: "peer" + strconv.Itoa(i)}, + IP: uint32ToIP(uint32(i)), + } + + err = manager.Store.ExecuteInTransaction(context.Background(), func(transaction store.Store) error { + err = manager.Store.AddPeerToAccount(context.Background(), store.LockingStrengthNone, peer) + if err != nil { + return fmt.Errorf("AddPeer failed for peer %d: %w", i, err) + } + err = manager.Store.AddPeerToAllGroup(context.Background(), accountID, strconv.Itoa(i)) + if err != nil { + return fmt.Errorf("AddPeer failed for peer %d: %w", i, err) + } + return nil + }) + if err != nil { + t.Errorf("AddPeer failed for peer %d: %v", i, err) + return + } + }(i) + } + startTime := time.Now() + + close(start) + wg.Wait() + close(errs) + + t.Logf("time since start: %s", time.Since(startTime)) + + for err := range errs { + t.Fatal(err) + } + + account, err := manager.Store.GetAccount(context.Background(), accountID) + if err != nil { + t.Fatalf("Failed to get account %s: %v", accountID, err) + } + + assert.Equal(t, totalPeers, len(maps.Values(account.Groups)[0].Peers), "Expected %d peers in account %s, got %d", totalPeers, accountID, len(account.Peers)) +} + +func uint32ToIP(n uint32) net.IP { + ip := make(net.IP, 4) + binary.BigEndian.PutUint32(ip, n) + return ip +} + +func Test_IncrementNetworkSerial(t *testing.T) { + t.Setenv("NETBIRD_STORE_ENGINE", string(types.PostgresStoreEngine)) + manager, err := createManager(t) + if err != nil { + t.Fatal(err) + return + } + + accountID := "testaccount" + userID := "testuser" + + _, err = createAccount(manager, accountID, userID, "domain.com") + if err != nil { + t.Fatal("error creating account") + return + } + + const totalPeers = 3000 + + var wg sync.WaitGroup + errs := make(chan error, totalPeers) + start := make(chan struct{}) + for i := 0; i < totalPeers; i++ { + wg.Add(1) + + go func(i int) { + defer wg.Done() + + <-start + + err = manager.Store.ExecuteInTransaction(context.Background(), func(transaction store.Store) error { + err = transaction.IncrementNetworkSerial(context.Background(), store.LockingStrengthNone, accountID) + if err != nil { + t.Fatalf("Failed to get account %s: %v", accountID, err) + } + return nil + }) + if err != nil { + t.Errorf("AddPeer failed for peer %d: %v", i, err) + return + } + }(i) + } + startTime := time.Now() + + close(start) + wg.Wait() + close(errs) + + t.Logf("time since start: %s", time.Since(startTime)) + + for err := range errs { + t.Fatal(err) + } + + account, err := manager.Store.GetAccount(context.Background(), accountID) + if err != nil { + t.Fatalf("Failed to get account %s: %v", accountID, err) + } + + assert.Equal(t, totalPeers, int(account.Network.Serial), "Expected %d peers in account %s, got %d", totalPeers, accountID, account.Network.Serial) +} diff --git a/management/server/migration/migration.go b/management/server/migration/migration.go index ab11be731..cd81dbc7d 100644 --- a/management/server/migration/migration.go +++ b/management/server/migration/migration.go @@ -412,3 +412,62 @@ func CreateIndexIfNotExists[T any](ctx context.Context, db *gorm.DB, indexName s log.WithContext(ctx).Infof("successfully created index %s on table %s", indexName, tableName) return nil } + +func MigrateJsonToTable[T any](ctx context.Context, db *gorm.DB, columnName string, mapperFunc func(id string, value string) any) error { + var model T + + if !db.Migrator().HasTable(&model) { + log.WithContext(ctx).Debugf("table for %T does not exist, no migration needed", model) + return nil + } + + stmt := &gorm.Statement{DB: db} + err := stmt.Parse(&model) + if err != nil { + return fmt.Errorf("parse model: %w", err) + } + tableName := stmt.Schema.Table + + if !db.Migrator().HasColumn(&model, columnName) { + log.WithContext(ctx).Debugf("column %s does not exist in table %s, no migration needed", columnName, tableName) + return nil + } + + if err := db.Transaction(func(tx *gorm.DB) error { + var rows []map[string]any + if err := tx.Table(tableName).Select("id", columnName).Find(&rows).Error; err != nil { + return fmt.Errorf("find rows: %w", err) + } + + for _, row := range rows { + jsonValue, ok := row[columnName].(string) + if !ok || jsonValue == "" { + continue + } + + var data []string + if err := json.Unmarshal([]byte(jsonValue), &data); err != nil { + return fmt.Errorf("unmarshal json: %w", err) + } + + for _, value := range data { + if err := tx.Create( + mapperFunc(row["id"].(string), value), + ).Error; err != nil { + return fmt.Errorf("failed to insert id %v: %w", row["id"], err) + } + } + } + + if err := tx.Migrator().DropColumn(&model, columnName); err != nil { + return fmt.Errorf("drop column %s: %w", columnName, err) + } + + return nil + }); err != nil { + return err + } + + log.WithContext(ctx).Infof("Migration of JSON field %s from table %s into seperte table completed", columnName, tableName) + return nil +} diff --git a/management/server/peer.go b/management/server/peer.go index 2c1d8f64c..77d4b024c 100644 --- a/management/server/peer.go +++ b/management/server/peer.go @@ -360,25 +360,20 @@ func (am *DefaultAccountManager) DeletePeer(ctx context.Context, accountID, peer return err } - if err = transaction.IncrementNetworkSerial(ctx, store.LockingStrengthUpdate, accountID); err != nil { - return err - } - - groups, err := transaction.GetPeerGroups(ctx, store.LockingStrengthUpdate, accountID, peerID) - if err != nil { - return fmt.Errorf("failed to get peer groups: %w", err) - } - - for _, group := range groups { - group.RemovePeer(peerID) - err = transaction.SaveGroup(ctx, store.LockingStrengthUpdate, group) - if err != nil { - return fmt.Errorf("failed to save group: %w", err) - } + if err = transaction.RemovePeerFromAllGroups(ctx, peer.ID); err != nil { + return fmt.Errorf("failed to remove peer from groups: %w", err) } eventsToStore, err = deletePeers(ctx, am, transaction, accountID, userID, []*nbpeer.Peer{peer}) - return err + if err != nil { + return fmt.Errorf("failed to delete peer: %w", err) + } + + if err = transaction.IncrementNetworkSerial(ctx, store.LockingStrengthUpdate, accountID); err != nil { + return fmt.Errorf("failed to increment network serial: %w", err) + } + + return nil }) if err != nil { return err @@ -477,7 +472,6 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s } var newPeer *nbpeer.Peer - var updateAccountPeers bool var setupKeyID string var setupKeyName string @@ -607,20 +601,20 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s return err } - err = transaction.AddPeerToAllGroup(ctx, store.LockingStrengthUpdate, accountID, newPeer.ID) - if err != nil { - return fmt.Errorf("failed adding peer to All group: %w", err) - } - if len(groupsToAdd) > 0 { for _, g := range groupsToAdd { - err = transaction.AddPeerToGroup(ctx, store.LockingStrengthUpdate, accountID, newPeer.ID, g) + err = transaction.AddPeerToGroup(ctx, newPeer.ID, g) if err != nil { return err } } } + err = transaction.AddPeerToAllGroup(ctx, accountID, newPeer.ID) + if err != nil { + return fmt.Errorf("failed adding peer to All group: %w", err) + } + if addedByUser { err := transaction.SaveUserLastLogin(ctx, accountID, userID, newPeer.GetLastLogin()) if err != nil { @@ -670,11 +664,6 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s return nil, nil, nil, fmt.Errorf("failed to add peer to database after %d attempts: %w", maxAttempts, err) } - updateAccountPeers, err = isPeerInActiveGroup(ctx, am.Store, accountID, newPeer.ID) - if err != nil { - updateAccountPeers = true - } - if newPeer == nil { return nil, nil, nil, fmt.Errorf("new peer is nil") } @@ -687,9 +676,7 @@ func (am *DefaultAccountManager) AddPeer(ctx context.Context, setupKey, userID s am.StoreEvent(ctx, opEvent.InitiatorID, opEvent.TargetID, opEvent.AccountID, opEvent.Activity, opEvent.Meta) - if updateAccountPeers { - am.BufferUpdateAccountPeers(ctx, accountID) - } + am.BufferUpdateAccountPeers(ctx, accountID) return am.getValidatedPeerWithMap(ctx, false, accountID, newPeer) } @@ -1019,7 +1006,7 @@ func (am *DefaultAccountManager) getValidatedPeerWithMap(ctx context.Context, is }() if isRequiresApproval { - network, err := am.Store.GetAccountNetwork(ctx, store.LockingStrengthShare, accountID) + network, err := am.Store.GetAccountNetwork(ctx, store.LockingStrengthNone, accountID) if err != nil { return nil, nil, nil, err } diff --git a/management/server/peer_test.go b/management/server/peer_test.go index 3edf7e82c..67cc6abeb 100644 --- a/management/server/peer_test.go +++ b/management/server/peer_test.go @@ -2144,6 +2144,7 @@ func Test_IsUniqueConstraintError(t *testing.T) { func Test_AddPeer(t *testing.T) { t.Setenv("NETBIRD_STORE_ENGINE", string(types.PostgresStoreEngine)) + t.Setenv("NB_GET_ACCOUNT_BUFFER_INTERVAL", "300ms") manager, err := createManager(t) if err != nil { t.Fatal(err) @@ -2155,7 +2156,7 @@ func Test_AddPeer(t *testing.T) { _, err = createAccount(manager, accountID, userID, "domain.com") if err != nil { - t.Fatal("error creating account") + t.Fatalf("error creating account: %v", err) return } @@ -2165,22 +2166,21 @@ func Test_AddPeer(t *testing.T) { return } - const totalPeers = 300 // totalPeers / differentHostnames should be less than 10 (due to concurrent retries) - const differentHostnames = 50 + const totalPeers = 10000 var wg sync.WaitGroup - errs := make(chan error, totalPeers+differentHostnames) + errs := make(chan error, totalPeers) start := make(chan struct{}) for i := 0; i < totalPeers; i++ { wg.Add(1) - hostNameID := i % differentHostnames go func(i int) { defer wg.Done() newPeer := &nbpeer.Peer{ - Key: "key" + strconv.Itoa(i), - Meta: nbpeer.PeerSystemMeta{Hostname: "peer" + strconv.Itoa(hostNameID), GoOS: "linux"}, + AccountID: accountID, + Key: "key" + strconv.Itoa(i), + Meta: nbpeer.PeerSystemMeta{Hostname: "peer" + strconv.Itoa(i), GoOS: "linux"}, } <-start diff --git a/management/server/store/sql_store.go b/management/server/store/sql_store.go index 197255ab6..2b1073c33 100644 --- a/management/server/store/sql_store.go +++ b/management/server/store/sql_store.go @@ -96,7 +96,7 @@ func NewSqlStore(ctx context.Context, db *gorm.DB, storeEngine types.Engine, met return nil, fmt.Errorf("migratePreAuto: %w", err) } err = db.AutoMigrate( - &types.SetupKey{}, &nbpeer.Peer{}, &types.User{}, &types.PersonalAccessToken{}, &types.Group{}, + &types.SetupKey{}, &nbpeer.Peer{}, &types.User{}, &types.PersonalAccessToken{}, &types.Group{}, &types.GroupPeer{}, &types.Account{}, &types.Policy{}, &types.PolicyRule{}, &route.Route{}, &nbdns.NameServerGroup{}, &installation{}, &types.ExtraSettings{}, &posture.Checks{}, &nbpeer.NetworkAddress{}, &networkTypes.Network{}, &routerTypes.NetworkRouter{}, &resourceTypes.NetworkResource{}, @@ -455,19 +455,34 @@ func (s *SqlStore) SaveGroups(ctx context.Context, lockStrength LockingStrength, return nil } - result := s.db. - Clauses( - clause.Locking{Strength: string(lockStrength)}, - clause.OnConflict{ - Where: clause.Where{Exprs: []clause.Expression{clause.Eq{Column: "groups.account_id", Value: accountID}}}, - UpdateAll: true, - }, - ). - Create(&groups) - if result.Error != nil { - return status.Errorf(status.Internal, "failed to save groups to store: %v", result.Error) - } - return nil + return s.db.Transaction(func(tx *gorm.DB) error { + for _, g := range groups { + g.StoreGroupPeers() + + if err := tx.Model(&g). + Association("GroupPeers"). + Replace(g.GroupPeers); err != nil { + log.WithContext(ctx).Errorf("failed to save group peers to store: %s", err) + return status.Errorf(status.Internal, "failed to save group peers to store") + } + } + + result := tx.Session(&gorm.Session{FullSaveAssociations: true}). + Clauses( + clause.Locking{Strength: string(lockStrength)}, + clause.OnConflict{ + Where: clause.Where{Exprs: []clause.Expression{clause.Eq{Column: "groups.account_id", Value: accountID}}}, + UpdateAll: true, + }, + ). + Save(&groups) + if result.Error != nil { + log.WithContext(ctx).Errorf("failed to save groups to store: %v", result.Error) + return status.Errorf(status.Internal, "failed to save groups to store") + } + + return nil + }) } // DeleteHashedPAT2TokenIDIndex is noop in SqlStore @@ -646,7 +661,7 @@ func (s *SqlStore) GetAccountGroups(ctx context.Context, lockStrength LockingStr } var groups []*types.Group - result := tx.Find(&groups, accountIDCondition, accountID) + result := tx.Preload(clause.Associations).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") @@ -655,6 +670,10 @@ func (s *SqlStore) GetAccountGroups(ctx context.Context, lockStrength LockingStr return nil, status.Errorf(status.Internal, "failed to get account groups from the store") } + for _, g := range groups { + g.LoadGroupPeers() + } + return groups, nil } @@ -669,6 +688,7 @@ func (s *SqlStore) GetResourceGroups(ctx context.Context, lockStrength LockingSt likePattern := `%"ID":"` + resourceID + `"%` result := tx. + Preload(clause.Associations). Where("resources LIKE ?", likePattern). Find(&groups) @@ -679,6 +699,10 @@ func (s *SqlStore) GetResourceGroups(ctx context.Context, lockStrength LockingSt return nil, result.Error } + for _, g := range groups { + g.LoadGroupPeers() + } + return groups, nil } @@ -739,7 +763,8 @@ func (s *SqlStore) GetAccount(ctx context.Context, accountID string) (*types.Acc var account types.Account result := s.db.Model(&account). - Preload("UsersG.PATsG"). // have to be specifies as this is nester reference + Preload("UsersG.PATsG"). // have to be specifies as this is nester reference + Preload("GroupsG.GroupPeers"). // have to be specifies as this is nester reference Preload(clause.Associations). First(&account, idQueryCondition, accountID) if result.Error != nil { @@ -784,6 +809,7 @@ func (s *SqlStore) GetAccount(ctx context.Context, accountID string) (*types.Acc account.Groups = make(map[string]*types.Group, len(account.GroupsG)) for _, group := range account.GroupsG { + group.LoadGroupPeers() account.Groups[group.ID] = group.Copy() } account.GroupsG = nil @@ -1285,55 +1311,71 @@ func (s *SqlStore) IncrementSetupKeyUsage(ctx context.Context, setupKeyID string } // AddPeerToAllGroup adds a peer to the 'All' group. Method always needs to run in a transaction -func (s *SqlStore) AddPeerToAllGroup(ctx context.Context, lockStrength LockingStrength, accountID string, peerID string) error { - var group types.Group - result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}). - First(&group, "account_id = ? AND name = ?", accountID, "All") - if result.Error != nil { - if errors.Is(result.Error, gorm.ErrRecordNotFound) { - return status.Errorf(status.NotFound, "group 'All' not found for account") - } - return status.Errorf(status.Internal, "issue finding group 'All': %s", result.Error) +func (s *SqlStore) AddPeerToAllGroup(ctx context.Context, accountID string, peerID string) error { + var groupID string + _ = s.db.Model(types.Group{}). + Select("id"). + Where("account_id = ? AND name = ?", accountID, "All"). + Limit(1). + Scan(&groupID) + + if groupID == "" { + return status.Errorf(status.NotFound, "group 'All' not found for account %s", accountID) } - for _, existingPeerID := range group.Peers { - if existingPeerID == peerID { - return nil - } - } + err := s.db.Create(&types.GroupPeer{ + GroupID: groupID, + PeerID: peerID, + }).Error - group.Peers = append(group.Peers, peerID) - - if err := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(&group).Error; err != nil { - return status.Errorf(status.Internal, "issue updating group 'All': %s", err) + if err != nil { + return status.Errorf(status.Internal, "error adding peer to group 'All': %v", err) } return nil } -// AddPeerToGroup adds a peer to a group. Method always needs to run in a transaction -func (s *SqlStore) AddPeerToGroup(ctx context.Context, lockStrength LockingStrength, accountId string, peerId string, groupID string) error { - var group types.Group - result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Where(accountAndIDQueryCondition, accountId, groupID). - First(&group) - if result.Error != nil { - if errors.Is(result.Error, gorm.ErrRecordNotFound) { - return status.NewGroupNotFoundError(groupID) - } - - return status.Errorf(status.Internal, "issue finding group: %s", result.Error) +// AddPeerToGroup adds a peer to a group +func (s *SqlStore) AddPeerToGroup(ctx context.Context, peerID string, groupID string) error { + peer := &types.GroupPeer{ + GroupID: groupID, + PeerID: peerID, } - for _, existingPeerID := range group.Peers { - if existingPeerID == peerId { - return nil - } + err := s.db.WithContext(ctx).Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "group_id"}, {Name: "peer_id"}}, + DoNothing: true, + }).Create(peer).Error + + if err != nil { + log.WithContext(ctx).Errorf("failed to add peer %s to group %s: %v", peerID, groupID, err) + return status.Errorf(status.Internal, "failed to add peer to group") } - group.Peers = append(group.Peers, peerId) + return nil +} - if err := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(&group).Error; err != nil { - return status.Errorf(status.Internal, "issue updating group: %s", err) +// RemovePeerFromGroup removes a peer from a group +func (s *SqlStore) RemovePeerFromGroup(ctx context.Context, peerID string, groupID string) error { + err := s.db.WithContext(ctx). + Delete(&types.GroupPeer{}, "group_id = ? AND peer_id = ?", groupID, peerID).Error + + if err != nil { + log.WithContext(ctx).Errorf("failed to remove peer %s from group %s: %v", peerID, groupID, err) + return status.Errorf(status.Internal, "failed to remove peer from group") + } + + return nil +} + +// RemovePeerFromAllGroups removes a peer from all groups +func (s *SqlStore) RemovePeerFromAllGroups(ctx context.Context, peerID string) error { + err := s.db.WithContext(ctx). + Delete(&types.GroupPeer{}, "peer_id = ?", peerID).Error + + if err != nil { + log.WithContext(ctx).Errorf("failed to remove peer %s from all groups: %v", peerID, err) + return status.Errorf(status.Internal, "failed to remove peer from all groups") } return nil @@ -1401,12 +1443,19 @@ func (s *SqlStore) GetPeerGroups(ctx context.Context, lockStrength LockingStreng var groups []*types.Group query := tx. - Find(&groups, "account_id = ? AND peers LIKE ?", accountId, fmt.Sprintf(`%%"%s"%%`, peerId)) + Joins("JOIN group_peers ON group_peers.group_id = groups.id"). + Where("group_peers.peer_id = ?", peerId). + Preload(clause.Associations). + Find(&groups) if query.Error != nil { return nil, query.Error } + for _, group := range groups { + group.LoadGroupPeers() + } + return groups, nil } @@ -1455,7 +1504,7 @@ func (s *SqlStore) GetUserPeers(ctx context.Context, lockStrength LockingStrengt } func (s *SqlStore) AddPeerToAccount(ctx context.Context, lockStrength LockingStrength, peer *nbpeer.Peer) error { - if err := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Create(peer).Error; err != nil { + if err := s.db.Create(peer).Error; err != nil { return status.Errorf(status.Internal, "issue adding peer to account: %s", err) } @@ -1692,7 +1741,7 @@ func (s *SqlStore) GetGroupByID(ctx context.Context, lockStrength LockingStrengt } var group *types.Group - result := tx.First(&group, accountAndIDQueryCondition, accountID, groupID) + result := tx.Preload(clause.Associations).First(&group, accountAndIDQueryCondition, accountID, groupID) if err := result.Error; err != nil { if errors.Is(err, gorm.ErrRecordNotFound) { return nil, status.NewGroupNotFoundError(groupID) @@ -1701,6 +1750,8 @@ func (s *SqlStore) GetGroupByID(ctx context.Context, lockStrength LockingStrengt return nil, status.Errorf(status.Internal, "failed to get group from store") } + group.LoadGroupPeers() + return group, nil } @@ -1717,15 +1768,6 @@ func (s *SqlStore) GetGroupByName(ctx context.Context, lockStrength LockingStren // we may need to reconsider changing the types. query := tx.Preload(clause.Associations) - switch s.storeEngine { - case types.PostgresStoreEngine: - query = query.Order("json_array_length(peers::json) DESC") - case types.MysqlStoreEngine: - query = query.Order("JSON_LENGTH(JSON_EXTRACT(peers, \"$\")) DESC") - default: - query = query.Order("json_array_length(peers) DESC") - } - result := query.First(&group, "account_id = ? AND name = ?", accountID, groupName) if err := result.Error; err != nil { if errors.Is(result.Error, gorm.ErrRecordNotFound) { @@ -1745,7 +1787,7 @@ func (s *SqlStore) GetGroupsByIDs(ctx context.Context, lockStrength LockingStren } var groups []*types.Group - result := tx.Find(&groups, accountAndIDsQueryCondition, accountID, groupIDs) + result := tx.Preload(clause.Associations).Find(&groups, accountAndIDsQueryCondition, accountID, groupIDs) if result.Error != nil { log.WithContext(ctx).Errorf("failed to get groups by ID's from store: %s", result.Error) return nil, status.Errorf(status.Internal, "failed to get groups by ID's from store") @@ -1753,6 +1795,7 @@ func (s *SqlStore) GetGroupsByIDs(ctx context.Context, lockStrength LockingStren groupsMap := make(map[string]*types.Group) for _, group := range groups { + group.LoadGroupPeers() groupsMap[group.ID] = group } @@ -1761,17 +1804,34 @@ func (s *SqlStore) GetGroupsByIDs(ctx context.Context, lockStrength LockingStren // SaveGroup saves a group to the store. func (s *SqlStore) SaveGroup(ctx context.Context, lockStrength LockingStrength, group *types.Group) error { - result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}).Save(group) - if result.Error != nil { - log.WithContext(ctx).Errorf("failed to save group to store: %v", result.Error) + if group == nil { + return status.Errorf(status.InvalidArgument, "group is nil") + } + + group.StoreGroupPeers() + + tx := s.db + if lockStrength != LockingStrengthNone { + tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)}) + } + + if err := tx.Model(group).Association("Peers").Replace(group.Peers); err != nil { + log.WithContext(ctx).Errorf("failed to replace peers for group %s: %v", group.ID, err) + return status.Errorf(status.Internal, "failed to sync group peers") + } + + if err := tx.Save(group).Error; err != nil { + log.WithContext(ctx).Errorf("failed to save group to store: %v", err) return status.Errorf(status.Internal, "failed to save group to store") } + return nil } // DeleteGroup deletes a group from the database. func (s *SqlStore) DeleteGroup(ctx context.Context, lockStrength LockingStrength, accountID, groupID string) error { result := s.db.Clauses(clause.Locking{Strength: string(lockStrength)}). + Select(clause.Associations). Delete(&types.Group{}, accountAndIDQueryCondition, accountID, groupID) if err := result.Error; err != nil { log.WithContext(ctx).Errorf("failed to delete group from store: %s", result.Error) @@ -1788,6 +1848,7 @@ func (s *SqlStore) DeleteGroup(ctx context.Context, lockStrength LockingStrength // DeleteGroups deletes groups from the database. func (s *SqlStore) DeleteGroups(ctx context.Context, strength LockingStrength, accountID string, groupIDs []string) error { result := s.db.Clauses(clause.Locking{Strength: string(strength)}). + Select(clause.Associations). Delete(&types.Group{}, accountAndIDsQueryCondition, accountID, groupIDs) if result.Error != nil { log.WithContext(ctx).Errorf("failed to delete groups from store: %v", result.Error) diff --git a/management/server/store/sql_store_test.go b/management/server/store/sql_store_test.go index 928486ab4..9b5101c79 100644 --- a/management/server/store/sql_store_test.go +++ b/management/server/store/sql_store_test.go @@ -2506,7 +2506,7 @@ func TestSqlStore_AddPeerToGroup(t *testing.T) { require.NoError(t, err, "failed to get group") require.Len(t, group.Peers, 0, "group should have 0 peers") - err = store.AddPeerToGroup(context.Background(), LockingStrengthUpdate, accountID, peerID, groupID) + err = store.AddPeerToGroup(context.Background(), peerID, groupID) require.NoError(t, err, "failed to add peer to group") group, err = store.GetGroupByID(context.Background(), LockingStrengthShare, accountID, groupID) @@ -2537,7 +2537,7 @@ func TestSqlStore_AddPeerToAllGroup(t *testing.T) { err = store.AddPeerToAccount(context.Background(), LockingStrengthUpdate, peer) require.NoError(t, err, "failed to add peer to account") - err = store.AddPeerToAllGroup(context.Background(), LockingStrengthUpdate, accountID, peer.ID) + err = store.AddPeerToAllGroup(context.Background(), accountID, peer.ID) require.NoError(t, err, "failed to add peer to all group") group, err = store.GetGroupByID(context.Background(), LockingStrengthShare, accountID, groupID) @@ -2623,7 +2623,7 @@ func TestSqlStore_GetPeerGroups(t *testing.T) { assert.Len(t, groups, 1) assert.Equal(t, groups[0].Name, "All") - err = store.AddPeerToGroup(context.Background(), LockingStrengthUpdate, accountID, peerID, "cfefqs706sqkneg59g4h") + err = store.AddPeerToGroup(context.Background(), peerID, "cfefqs706sqkneg59g4h") require.NoError(t, err) groups, err = store.GetPeerGroups(context.Background(), LockingStrengthShare, accountID, peerID) diff --git a/management/server/store/store.go b/management/server/store/store.go index 30ff1549d..fad12ce77 100644 --- a/management/server/store/store.go +++ b/management/server/store/store.go @@ -118,8 +118,10 @@ type Store interface { DeletePostureChecks(ctx context.Context, lockStrength LockingStrength, accountID, postureChecksID string) error GetPeerLabelsInAccount(ctx context.Context, lockStrength LockingStrength, accountId string, hostname string) ([]string, error) - AddPeerToAllGroup(ctx context.Context, lockStrength LockingStrength, accountID string, peerID string) error - AddPeerToGroup(ctx context.Context, lockStrength LockingStrength, accountId string, peerId string, groupID string) error + AddPeerToAllGroup(ctx context.Context, accountID string, peerID string) error + AddPeerToGroup(ctx context.Context, peerId string, groupID string) error + RemovePeerFromGroup(ctx context.Context, peerID string, groupID string) error + RemovePeerFromAllGroups(ctx context.Context, peerID string) error GetPeerGroups(ctx context.Context, lockStrength LockingStrength, accountId string, peerId string) ([]*types.Group, error) AddResourceToGroup(ctx context.Context, accountId string, groupID string, resource *types.Resource) error RemoveResourceFromGroup(ctx context.Context, accountId string, groupID string, resourceID string) error @@ -351,6 +353,14 @@ func getMigrationsPostAuto(ctx context.Context) []migrationFunc { func(db *gorm.DB) error { return migration.CreateIndexIfNotExists[nbpeer.Peer](ctx, db, "idx_account_dnslabel", "account_id", "dns_label") }, + func(db *gorm.DB) error { + return migration.MigrateJsonToTable[types.Group](ctx, db, "peers", func(id, value string) any { + return &types.GroupPeer{ + GroupID: id, + PeerID: value, + } + }) + }, } } diff --git a/management/server/testdata/extended-store.sql b/management/server/testdata/extended-store.sql index 0393d1ade..650ecbec4 100644 --- a/management/server/testdata/extended-store.sql +++ b/management/server/testdata/extended-store.sql @@ -33,8 +33,8 @@ INSERT INTO users VALUES('edafee4e-63fb-11ec-90d6-0242ac120003','bf1c8084-ba50-4 INSERT INTO users VALUES('f4f6d672-63fb-11ec-90d6-0242ac120003','bf1c8084-ba50-4ce7-9439-34653001fc3b','user',0,0,'','[]',0,NULL,'2024-10-02 16:01:38.210678+02:00','api',0,''); INSERT INTO personal_access_tokens VALUES('9dj38s35-63fb-11ec-90d6-0242ac120003','f4f6d672-63fb-11ec-90d6-0242ac120003','','SoMeHaShEdToKeN','2023-02-27 00:00:00+00:00','user','2023-01-01 00:00:00+00:00','2023-02-01 00:00:00+00:00'); INSERT INTO "groups" VALUES('cfefqs706sqkneg59g4g','bf1c8084-ba50-4ce7-9439-34653001fc3b','All','api','[]',0,''); -INSERT INTO "groups" VALUES('cfefqs706sqkneg59g3g','bf1c8084-ba50-4ce7-9439-34653001fc3b','AwesomeGroup1','api','[]',0,''); -INSERT INTO "groups" VALUES('cfefqs706sqkneg59g2g','bf1c8084-ba50-4ce7-9439-34653001fc3b','AwesomeGroup2','api','[]',0,''); +INSERT INTO "groups" VALUES('cfefqs706sqkneg59g3g','bf1c8084-ba50-4ce7-9439-34653001fc3b','AwesomeGroup1','api','["peer1"]',0,''); +INSERT INTO "groups" VALUES('cfefqs706sqkneg59g2g','bf1c8084-ba50-4ce7-9439-34653001fc3b','AwesomeGroup2','api','["peer1","peer2","peer3"]',0,''); INSERT INTO posture_checks VALUES('csplshq7qv948l48f7t0','NetBird Version > 0.32.0','','bf1c8084-ba50-4ce7-9439-34653001fc3b','{"NBVersionCheck":{"MinVersion":"0.31.0"}}'); INSERT INTO posture_checks VALUES('cspnllq7qv95uq1r4k90','Allow Berlin and Deny local network 172.16.1.0/24','','bf1c8084-ba50-4ce7-9439-34653001fc3b','{"GeoLocationCheck":{"Locations":[{"CountryCode":"DE","CityName":"Berlin"}],"Action":"allow"},"PeerNetworkRangeCheck":{"Action":"deny","Ranges":["172.16.1.0/24"]}}'); INSERT INTO name_server_groups VALUES('csqdelq7qv97ncu7d9t0','bf1c8084-ba50-4ce7-9439-34653001fc3b','Google DNS','Google DNS Servers','[{"IP":"8.8.8.8","NSType":1,"Port":53},{"IP":"8.8.4.4","NSType":1,"Port":53}]','["cfefqs706sqkneg59g2g"]',1,'[]',1,0); diff --git a/management/server/types/group.go b/management/server/types/group.go index 1b321387c..b8284efbb 100644 --- a/management/server/types/group.go +++ b/management/server/types/group.go @@ -26,7 +26,8 @@ type Group struct { Issued string // Peers list of the group - Peers []string `gorm:"serializer:json"` + Peers []string `gorm:"-"` + GroupPeers []GroupPeer `gorm:"foreignKey:GroupID;references:id;constraint:OnDelete:CASCADE;"` // Resources contains a list of resources in that group Resources []Resource `gorm:"serializer:json"` @@ -34,6 +35,29 @@ type Group struct { IntegrationReference integration_reference.IntegrationReference `gorm:"embedded;embeddedPrefix:integration_ref_"` } +type GroupPeer struct { + GroupID string `gorm:"primaryKey"` + PeerID string `gorm:"primaryKey"` +} + +func (g *Group) LoadGroupPeers() { + g.Peers = make([]string, len(g.GroupPeers)) + for i, peer := range g.GroupPeers { + g.Peers[i] = peer.PeerID + } + // g.GroupPeers = nil +} +func (g *Group) StoreGroupPeers() { + g.GroupPeers = make([]GroupPeer, len(g.Peers)) + for i, peer := range g.Peers { + g.GroupPeers[i] = GroupPeer{ + GroupID: g.ID, + PeerID: peer, + } + } + // g.Peers = nil +} + // EventMeta returns activity event meta related to the group func (g *Group) EventMeta() map[string]any { return map[string]any{"name": g.Name} diff --git a/management/server/types/setupkey.go b/management/server/types/setupkey.go index 69b381ae5..3d421342d 100644 --- a/management/server/types/setupkey.go +++ b/management/server/types/setupkey.go @@ -35,7 +35,7 @@ type SetupKey struct { // AccountID is a reference to Account that this object belongs AccountID string `json:"-" gorm:"index"` Key string - KeySecret string + KeySecret string `gorm:"index"` Name string Type SetupKeyType CreatedAt time.Time