[management] add uniqueness constraint for peer ip and label and optimize generation (#4042)

management/server/store/file_store.go

@@ -156,7 +156,7 @@ func restore(ctx context.Context, file string) (*FileStore, error) {
 
 		allGroup, err := account.GetGroupAll()
 		if err != nil {
-			log.WithContext(ctx).Errorf("unable to find the All group, this should happen only when migrate from a version that didn't support groups. Error: %v", err)
+			log.WithContext(ctx).Errorf("unable to find the All group, this should happen only when migratePreAuto from a version that didn't support groups. Error: %v", err)
 			// if the All group didn't exist we probably don't have routes to update
 			continue
 		}

management/server/store/sql_store.go

@@ -92,8 +92,8 @@ func NewSqlStore(ctx context.Context, db *gorm.DB, storeEngine types.Engine, met
 		return &SqlStore{db: db, storeEngine: storeEngine, metrics: metrics, installationPK: 1}, nil
 	}
 
-	if err := migrate(ctx, db); err != nil {
-		return nil, fmt.Errorf("migrate: %w", err)
+	if err := migratePreAuto(ctx, db); err != nil {
+		return nil, fmt.Errorf("migratePreAuto: %w", err)
 	}
 	err = db.AutoMigrate(
 		&types.SetupKey{}, &nbpeer.Peer{}, &types.User{}, &types.PersonalAccessToken{}, &types.Group{},
@@ -102,7 +102,10 @@ func NewSqlStore(ctx context.Context, db *gorm.DB, storeEngine types.Engine, met
 		&networkTypes.Network{}, &routerTypes.NetworkRouter{}, &resourceTypes.NetworkResource{},
 	)
 	if err != nil {
-		return nil, fmt.Errorf("auto migrate: %w", err)
+		return nil, fmt.Errorf("auto migratePreAuto: %w", err)
+	}
+	if err := migratePostAuto(ctx, db); err != nil {
+		return nil, fmt.Errorf("migratePostAuto: %w", err)
 	}
 
 	return &SqlStore{db: db, storeEngine: storeEngine, metrics: metrics, installationPK: 1}, nil
@@ -967,7 +970,7 @@ func (s *SqlStore) GetTakenIPs(ctx context.Context, lockStrength LockingStrength
 	return ips, nil
 }
 
-func (s *SqlStore) GetPeerLabelsInAccount(ctx context.Context, lockStrength LockingStrength, accountID string) ([]string, error) {
+func (s *SqlStore) GetPeerLabelsInAccount(ctx context.Context, lockStrength LockingStrength, accountID string, dnsLabel string) ([]string, error) {
 	tx := s.db
 	if lockStrength != LockingStrengthNone {
 		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
@@ -975,7 +978,7 @@ func (s *SqlStore) GetPeerLabelsInAccount(ctx context.Context, lockStrength Lock
 
 	var labels []string
 	result := tx.Model(&nbpeer.Peer{}).
-		Where("account_id = ?", accountID).
+		Where("account_id = ? AND dns_label LIKE ?", accountID, dnsLabel+"%").
 		Pluck("dns_label", &labels)
 
 	if result.Error != nil {
@@ -1254,7 +1257,7 @@ func (s *SqlStore) GetSetupKeyBySecret(ctx context.Context, lockStrength Locking
 
 	if result.Error != nil {
 		if errors.Is(result.Error, gorm.ErrRecordNotFound) {
-			return nil, status.NewSetupKeyNotFoundError(key)
+			return nil, status.Errorf(status.PreconditionFailed, "setup key not found")
 		}
 		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")
@@ -2546,6 +2549,27 @@ func (s *SqlStore) GetPeerByIP(ctx context.Context, lockStrength LockingStrength
 	return &peer, nil
 }
 
+func (s *SqlStore) GetPeerIdByLabel(ctx context.Context, lockStrength LockingStrength, accountID string, hostname string) (string, error) {
+	tx := s.db.WithContext(ctx)
+	if lockStrength != LockingStrengthNone {
+		tx = tx.Clauses(clause.Locking{Strength: string(lockStrength)})
+	}
+
+	var peerID string
+	result := tx.Model(&nbpeer.Peer{}).
+		Select("id").
+		// Where(" = ?", hostname).
+		Where("account_id = ? AND dns_label = ?", accountID, hostname).
+		Limit(1).
+		Scan(&peerID)
+
+	if peerID == "" {
+		return "", gorm.ErrRecordNotFound
+	}
+
+	return peerID, result.Error
+}
+
 func (s *SqlStore) CountAccountsByPrivateDomain(ctx context.Context, domain string) (int64, error) {
 	var count int64
 	result := s.db.Model(&types.Account{}).

management/server/store/sql_store_test.go

@@ -10,6 +10,7 @@ import (
 	"net/netip"
 	"os"
 	"runtime"
+	"sort"
 	"sync"
 	"testing"
 	"time"
@@ -630,7 +631,7 @@ func TestMigrate(t *testing.T) {
 	t.Cleanup(cleanUp)
 	assert.NoError(t, err)
 
-	err = migrate(context.Background(), store.(*SqlStore).db)
+	err = migratePreAuto(context.Background(), store.(*SqlStore).db)
 	require.NoError(t, err, "Migration should not fail on empty db")
 
 	_, ipnet, err := net.ParseCIDR("10.0.0.0/24")
@@ -685,10 +686,10 @@ func TestMigrate(t *testing.T) {
 	err = store.(*SqlStore).db.Save(rt).Error
 	require.NoError(t, err, "Failed to insert Gob data")
 
-	err = migrate(context.Background(), store.(*SqlStore).db)
+	err = migratePreAuto(context.Background(), store.(*SqlStore).db)
 	require.NoError(t, err, "Migration should not fail on gob populated db")
 
-	err = migrate(context.Background(), store.(*SqlStore).db)
+	err = migratePreAuto(context.Background(), store.(*SqlStore).db)
 	require.NoError(t, err, "Migration should not fail on migrated db")
 
 	err = store.(*SqlStore).db.Delete(rt).Where("id = ?", "route1").Error
@@ -704,10 +705,10 @@ func TestMigrate(t *testing.T) {
 	err = store.(*SqlStore).db.Save(nRT).Error
 	require.NoError(t, err, "Failed to insert json nil slice data")
 
-	err = migrate(context.Background(), store.(*SqlStore).db)
+	err = migratePreAuto(context.Background(), store.(*SqlStore).db)
 	require.NoError(t, err, "Migration should not fail on json nil slice populated db")
 
-	err = migrate(context.Background(), store.(*SqlStore).db)
+	err = migratePreAuto(context.Background(), store.(*SqlStore).db)
 	require.NoError(t, err, "Migration should not fail on migrated db")
 
 }
@@ -950,6 +951,7 @@ func TestSqlite_GetTakenIPs(t *testing.T) {
 	peer1 := &nbpeer.Peer{
 		ID:        "peer1",
 		AccountID: existingAccountID,
+		DNSLabel:  "peer1",
 		IP:        net.IP{1, 1, 1, 1},
 	}
 	err = store.AddPeerToAccount(context.Background(), LockingStrengthUpdate, peer1)
@@ -961,8 +963,9 @@ func TestSqlite_GetTakenIPs(t *testing.T) {
 	assert.Equal(t, []net.IP{ip1}, takenIPs)
 
 	peer2 := &nbpeer.Peer{
-		ID:        "peer2",
+		ID:        "peer1second",
 		AccountID: existingAccountID,
+		DNSLabel:  "peer1-1",
 		IP:        net.IP{2, 2, 2, 2},
 	}
 	err = store.AddPeerToAccount(context.Background(), LockingStrengthUpdate, peer2)
@@ -972,49 +975,100 @@ func TestSqlite_GetTakenIPs(t *testing.T) {
 	require.NoError(t, err)
 	ip2 := net.IP{2, 2, 2, 2}.To16()
 	assert.Equal(t, []net.IP{ip1, ip2}, takenIPs)
-
 }
 
 func TestSqlite_GetPeerLabelsInAccount(t *testing.T) {
-	t.Setenv("NETBIRD_STORE_ENGINE", string(types.SqliteStoreEngine))
-	store, cleanup, err := NewTestStoreFromSQL(context.Background(), "../testdata/extended-store.sql", t.TempDir())
-	if err != nil {
-		return
-	}
-	t.Cleanup(cleanup)
+	runTestForAllEngines(t, "../testdata/extended-store.sql", func(t *testing.T, store Store) {
+		existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+		peerHostname := "peer1"
 
-	existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+		_, err := store.GetAccount(context.Background(), existingAccountID)
+		require.NoError(t, err)
 
-	_, err = store.GetAccount(context.Background(), existingAccountID)
-	require.NoError(t, err)
+		labels, err := store.GetPeerLabelsInAccount(context.Background(), LockingStrengthShare, existingAccountID, peerHostname)
+		require.NoError(t, err)
+		assert.Equal(t, []string{}, labels)
 
-	labels, err := store.GetPeerLabelsInAccount(context.Background(), LockingStrengthShare, existingAccountID)
-	require.NoError(t, err)
-	assert.Equal(t, []string{}, labels)
+		peer1 := &nbpeer.Peer{
+			ID:        "peer1",
+			AccountID: existingAccountID,
+			DNSLabel:  "peer1",
+			IP:        net.IP{1, 1, 1, 1},
+		}
+		err = store.AddPeerToAccount(context.Background(), LockingStrengthUpdate, peer1)
+		require.NoError(t, err)
 
-	peer1 := &nbpeer.Peer{
-		ID:        "peer1",
-		AccountID: existingAccountID,
-		DNSLabel:  "peer1.domain.test",
-	}
-	err = store.AddPeerToAccount(context.Background(), LockingStrengthUpdate, peer1)
-	require.NoError(t, err)
+		labels, err = store.GetPeerLabelsInAccount(context.Background(), LockingStrengthShare, existingAccountID, peerHostname)
+		require.NoError(t, err)
+		assert.Equal(t, []string{"peer1"}, labels)
 
-	labels, err = store.GetPeerLabelsInAccount(context.Background(), LockingStrengthShare, existingAccountID)
-	require.NoError(t, err)
-	assert.Equal(t, []string{"peer1.domain.test"}, labels)
+		peer2 := &nbpeer.Peer{
+			ID:        "peer1second",
+			AccountID: existingAccountID,
+			DNSLabel:  "peer1-1",
+			IP:        net.IP{2, 2, 2, 2},
+		}
+		err = store.AddPeerToAccount(context.Background(), LockingStrengthUpdate, peer2)
+		require.NoError(t, err)
 
-	peer2 := &nbpeer.Peer{
-		ID:        "peer2",
-		AccountID: existingAccountID,
-		DNSLabel:  "peer2.domain.test",
-	}
-	err = store.AddPeerToAccount(context.Background(), LockingStrengthUpdate, peer2)
-	require.NoError(t, err)
+		labels, err = store.GetPeerLabelsInAccount(context.Background(), LockingStrengthShare, existingAccountID, peerHostname)
+		require.NoError(t, err)
 
-	labels, err = store.GetPeerLabelsInAccount(context.Background(), LockingStrengthShare, existingAccountID)
-	require.NoError(t, err)
-	assert.Equal(t, []string{"peer1.domain.test", "peer2.domain.test"}, labels)
+		expected := []string{"peer1", "peer1-1"}
+		sort.Strings(expected)
+		sort.Strings(labels)
+		assert.Equal(t, expected, labels)
+	})
+}
+
+func Test_AddPeerWithSameDnsLabel(t *testing.T) {
+	runTestForAllEngines(t, "../testdata/extended-store.sql", func(t *testing.T, store Store) {
+		existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+		_, err := store.GetAccount(context.Background(), existingAccountID)
+		require.NoError(t, err)
+
+		peer1 := &nbpeer.Peer{
+			ID:        "peer1",
+			AccountID: existingAccountID,
+			DNSLabel:  "peer1.domain.test",
+		}
+		err = store.AddPeerToAccount(context.Background(), LockingStrengthUpdate, peer1)
+		require.NoError(t, err)
+
+		peer2 := &nbpeer.Peer{
+			ID:        "peer1second",
+			AccountID: existingAccountID,
+			DNSLabel:  "peer1.domain.test",
+		}
+		err = store.AddPeerToAccount(context.Background(), LockingStrengthUpdate, peer2)
+		require.Error(t, err)
+	})
+}
+
+func Test_AddPeerWithSameIP(t *testing.T) {
+	runTestForAllEngines(t, "../testdata/extended-store.sql", func(t *testing.T, store Store) {
+		existingAccountID := "bf1c8084-ba50-4ce7-9439-34653001fc3b"
+
+		_, err := store.GetAccount(context.Background(), existingAccountID)
+		require.NoError(t, err)
+
+		peer1 := &nbpeer.Peer{
+			ID:        "peer1",
+			AccountID: existingAccountID,
+			IP:        net.IP{1, 1, 1, 1},
+		}
+		err = store.AddPeerToAccount(context.Background(), LockingStrengthUpdate, peer1)
+		require.NoError(t, err)
+
+		peer2 := &nbpeer.Peer{
+			ID:        "peer1second",
+			AccountID: existingAccountID,
+			IP:        net.IP{1, 1, 1, 1},
+		}
+		err = store.AddPeerToAccount(context.Background(), LockingStrengthUpdate, peer2)
+		require.Error(t, err)
+	})
 }
 
 func TestSqlite_GetAccountNetwork(t *testing.T) {

management/server/store/store.go

@@ -117,7 +117,7 @@ type Store interface {
 	SavePostureChecks(ctx context.Context, lockStrength LockingStrength, postureCheck *posture.Checks) error
 	DeletePostureChecks(ctx context.Context, lockStrength LockingStrength, accountID, postureChecksID string) error
 
-	GetPeerLabelsInAccount(ctx context.Context, lockStrength LockingStrength, accountId string) ([]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
 	GetPeerGroups(ctx context.Context, lockStrength LockingStrength, accountId string, peerId string) ([]*types.Group, error)
@@ -193,6 +193,7 @@ type Store interface {
 	SaveNetworkResource(ctx context.Context, lockStrength LockingStrength, resource *resourceTypes.NetworkResource) error
 	DeleteNetworkResource(ctx context.Context, lockStrength LockingStrength, accountID, resourceID string) error
 	GetPeerByIP(ctx context.Context, lockStrength LockingStrength, accountID string, ip net.IP) (*nbpeer.Peer, error)
+	GetPeerIdByLabel(ctx context.Context, lockStrength LockingStrength, accountID string, hostname string) (string, error)
 }
 
 const (
@@ -234,9 +235,9 @@ func getStoreEngine(ctx context.Context, dataDir string, kind types.Engine) type
 			if util.FileExists(jsonStoreFile) && !util.FileExists(sqliteStoreFile) {
 				log.WithContext(ctx).Warnf("unsupported store engine specified, but found %s. Automatically migrating to SQLite.", jsonStoreFile)
 
-				// Attempt to migrate from JSON store to SQLite
+				// Attempt to migratePreAuto from JSON store to SQLite
 				if err := MigrateFileStoreToSqlite(ctx, dataDir); err != nil {
-					log.WithContext(ctx).Errorf("failed to migrate filestore to SQLite: %v", err)
+					log.WithContext(ctx).Errorf("failed to migratePreAuto filestore to SQLite: %v", err)
 					kind = types.FileStoreEngine
 				}
 			}
@@ -280,9 +281,9 @@ func checkFileStoreEngine(kind types.Engine, dataDir string) error {
 	return nil
 }
 
-// migrate migrates the SQLite database to the latest schema
-func migrate(ctx context.Context, db *gorm.DB) error {
-	migrations := getMigrations(ctx)
+// migratePreAuto migrates the SQLite database to the latest schema
+func migratePreAuto(ctx context.Context, db *gorm.DB) error {
+	migrations := getMigrationsPreAuto(ctx)
 
 	for _, m := range migrations {
 		if err := m(db); err != nil {
@@ -293,7 +294,7 @@ func migrate(ctx context.Context, db *gorm.DB) error {
 	return nil
 }
 
-func getMigrations(ctx context.Context) []migrationFunc {
+func getMigrationsPreAuto(ctx context.Context) []migrationFunc {
 	return []migrationFunc{
 		func(db *gorm.DB) error {
 			return migration.MigrateFieldFromGobToJSON[types.Account, net.IPNet](ctx, db, "network_net")
@@ -329,6 +330,28 @@ func getMigrations(ctx context.Context) []migrationFunc {
 			return migration.DropIndex[routerTypes.NetworkRouter](ctx, db, "idx_network_routers_id")
 		},
 	}
+} // migratePostAuto migrates the SQLite database to the latest schema
+func migratePostAuto(ctx context.Context, db *gorm.DB) error {
+	migrations := getMigrationsPostAuto(ctx)
+
+	for _, m := range migrations {
+		if err := m(db); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+func getMigrationsPostAuto(ctx context.Context) []migrationFunc {
+	return []migrationFunc{
+		func(db *gorm.DB) error {
+			return migration.CreateIndexIfNotExists[nbpeer.Peer](ctx, db, "idx_account_ip", "account_id", "ip")
+		},
+		func(db *gorm.DB) error {
+			return migration.CreateIndexIfNotExists[nbpeer.Peer](ctx, db, "idx_account_dnslabel", "account_id", "dns_label")
+		},
+	}
 }
 
 // NewTestStoreFromSQL is only used in tests. It will create a test database base of the store engine set in env.
@@ -577,7 +600,7 @@ func MigrateFileStoreToSqlite(ctx context.Context, dataDir string) error {
 
 	sqliteStoreAccounts := len(store.GetAllAccounts(ctx))
 	if fsStoreAccounts != sqliteStoreAccounts {
-		return fmt.Errorf("failed to migrate accounts from file to sqlite. Expected accounts: %d, got: %d",
+		return fmt.Errorf("failed to migratePreAuto accounts from file to sqlite. Expected accounts: %d, got: %d",
 			fsStoreAccounts, sqliteStoreAccounts)
 	}