netbird/management/server/migration/migration_test.go

package migration_test

import (
	"context"
	"encoding/gob"
	"net"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"gorm.io/driver/mysql"
	"gorm.io/driver/postgres"
	"gorm.io/driver/sqlite"
	"gorm.io/gorm"

	"github.com/netbirdio/netbird/management/server/migration"
	nbpeer "github.com/netbirdio/netbird/management/server/peer"
	"github.com/netbirdio/netbird/management/server/testutil"
	"github.com/netbirdio/netbird/management/server/types"
	"github.com/netbirdio/netbird/route"
)

func setupDatabase(t *testing.T) *gorm.DB {
	t.Helper()

	var db *gorm.DB
	var err error
	var dsn string
	var cleanup func()
	switch os.Getenv("NETBIRD_STORE_ENGINE") {
	case "mysql":
		cleanup, dsn, err = testutil.CreateMysqlTestContainer()
		if err != nil {
			t.Fatalf("Failed to create MySQL test container: %v", err)
		}

		if dsn == "" {
			t.Fatal("MySQL connection string is empty, ensure the test container is running")
		}

		db, err = gorm.Open(mysql.Open(dsn+"?charset=utf8&parseTime=True&loc=Local"), &gorm.Config{})
	case "postgres":
		cleanup, dsn, err = testutil.CreatePostgresTestContainer()
		if err != nil {
			t.Fatalf("Failed to create PostgreSQL test container: %v", err)
		}

		if dsn == "" {
			t.Fatalf("PostgreSQL connection string is empty, ensure the test container is running")
		}

		db, err = gorm.Open(postgres.Open(dsn), &gorm.Config{})
	case "sqlite":
		db, err = gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
	default:
		db, err = gorm.Open(sqlite.Open("file::memory:?cache=shared"), &gorm.Config{})
	}
	if cleanup != nil {
		t.Cleanup(cleanup)
	}

	require.NoError(t, err, "Failed to open database")
	return db
}

func TestMigrateFieldFromGobToJSON_EmptyDB(t *testing.T) {
	db := setupDatabase(t)
	err := migration.MigrateFieldFromGobToJSON[types.Account, net.IPNet](context.Background(), db, "network_net")
	require.NoError(t, err, "Migration should not fail for an empty database")
}

func TestMigrateFieldFromGobToJSON_WithGobData(t *testing.T) {
	t.Setenv("NETBIRD_STORE_ENGINE", "sqlite")
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.Account{}, &route.Route{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	_, ipnet, err := net.ParseCIDR("10.0.0.0/24")
	require.NoError(t, err, "Failed to parse CIDR")

	type network struct {
		types.Network
		Net net.IPNet `gorm:"serializer:gob"`
	}

	type account struct {
		types.Account
		Network *network `gorm:"embedded;embeddedPrefix:network_"`
	}

	err = db.Save(&account{Account: types.Account{Id: "123"}, Network: &network{Net: *ipnet}}).Error
	require.NoError(t, err, "Failed to insert Gob data")

	var gobStr string
	err = db.Model(&types.Account{}).Select("network_net").First(&gobStr).Error
	assert.NoError(t, err, "Failed to fetch Gob data")

	err = gob.NewDecoder(strings.NewReader(gobStr)).Decode(&ipnet)
	require.NoError(t, err, "Failed to decode Gob data")

	err = migration.MigrateFieldFromGobToJSON[types.Account, net.IPNet](context.Background(), db, "network_net")
	require.NoError(t, err, "Migration should not fail with Gob data")

	var jsonStr string
	db.Model(&types.Account{}).Select("network_net").First(&jsonStr)
	assert.JSONEq(t, `{"IP":"10.0.0.0","Mask":"////AA=="}`, jsonStr, "Data should be migrated")
}

func TestMigrateFieldFromGobToJSON_WithJSONData(t *testing.T) {
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.Account{}, &route.Route{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	_, ipnet, err := net.ParseCIDR("10.0.0.0/24")
	require.NoError(t, err, "Failed to parse CIDR")

	err = db.Save(&types.Account{Network: &types.Network{Net: *ipnet}}).Error
	require.NoError(t, err, "Failed to insert JSON data")

	err = migration.MigrateFieldFromGobToJSON[types.Account, net.IPNet](context.Background(), db, "network_net")
	require.NoError(t, err, "Migration should not fail with JSON data")

	var jsonStr string
	db.Model(&types.Account{}).Select("network_net").First(&jsonStr)
	assert.JSONEq(t, `{"IP":"10.0.0.0","Mask":"////AA=="}`, jsonStr, "Data should be unchanged")
}

func TestMigrateNetIPFieldFromBlobToJSON_EmptyDB(t *testing.T) {
	db := setupDatabase(t)
	err := migration.MigrateNetIPFieldFromBlobToJSON[nbpeer.Peer](context.Background(), db, "ip", "idx_peers_account_id_ip")
	require.NoError(t, err, "Migration should not fail for an empty database")
}

func TestMigrateNetIPFieldFromBlobToJSON_WithBlobData(t *testing.T) {
	t.Setenv("NETBIRD_STORE_ENGINE", "sqlite")
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.Account{}, &nbpeer.Peer{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	type location struct {
		nbpeer.Location
		ConnectionIP net.IP
	}

	type peer struct {
		nbpeer.Peer
		Location location `gorm:"embedded;embeddedPrefix:location_"`
	}

	type account struct {
		types.Account
		Peers []peer `gorm:"foreignKey:AccountID;references:id"`
	}

	a := &account{
		Account: types.Account{Id: "123"},
	}

	err = db.Save(a).Error
	require.NoError(t, err, "Failed to insert account")

	a.Peers = []peer{
		{Location: location{ConnectionIP: net.IP{10, 0, 0, 1}}},
	}

	err = db.Save(a).Error
	require.NoError(t, err, "Failed to insert blob data")

	var blobValue string
	err = db.Model(&nbpeer.Peer{}).Select("location_connection_ip").First(&blobValue).Error
	assert.NoError(t, err, "Failed to fetch blob data")

	err = migration.MigrateNetIPFieldFromBlobToJSON[nbpeer.Peer](context.Background(), db, "location_connection_ip", "")
	require.NoError(t, err, "Migration should not fail with net.IP blob data")

	var jsonStr string
	db.Model(&nbpeer.Peer{}).Select("location_connection_ip").First(&jsonStr)
	assert.JSONEq(t, `"10.0.0.1"`, jsonStr, "Data should be migrated")
}

func TestMigrateNetIPFieldFromBlobToJSON_WithJSONData(t *testing.T) {
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.Account{}, &nbpeer.Peer{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	account := &types.Account{
		Id: "1234",
	}

	err = db.Save(account).Error
	require.NoError(t, err, "Failed to insert account")

	account.PeersG = []nbpeer.Peer{
		{AccountID: "1234", Location: nbpeer.Location{ConnectionIP: net.IP{10, 0, 0, 1}}},
	}

	err = db.Save(account).Error
	require.NoError(t, err, "Failed to insert JSON data")

	err = migration.MigrateNetIPFieldFromBlobToJSON[nbpeer.Peer](context.Background(), db, "location_connection_ip", "")
	require.NoError(t, err, "Migration should not fail with net.IP JSON data")

	var jsonStr string
	db.Model(&nbpeer.Peer{}).Select("location_connection_ip").First(&jsonStr)
	assert.JSONEq(t, `"10.0.0.1"`, jsonStr, "Data should be unchanged")
}

func TestMigrateSetupKeyToHashedSetupKey_ForPlainKey(t *testing.T) {
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.SetupKey{}, &nbpeer.Peer{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	err = db.Save(&types.SetupKey{
		Id:        "1",
		Key:       "EEFDAB47-C1A5-4472-8C05-71DE9A1E8382",
		UpdatedAt: time.Now(),
	}).Error
	require.NoError(t, err, "Failed to insert setup key")

	err = migration.MigrateSetupKeyToHashedSetupKey[types.SetupKey](context.Background(), db)
	require.NoError(t, err, "Migration should not fail to migrate setup key")

	var key types.SetupKey
	err = db.Model(&types.SetupKey{}).First(&key).Error
	assert.NoError(t, err, "Failed to fetch setup key")

	assert.Equal(t, "EEFDA****", key.KeySecret, "Key should be secret")
	assert.Equal(t, "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=", key.Key, "Key should be hashed")
}

func TestMigrateSetupKeyToHashedSetupKey_ForAlreadyMigratedKey_Case1(t *testing.T) {
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.SetupKey{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	err = db.Save(&types.SetupKey{
		Id:        "1",
		Key:       "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=",
		KeySecret: "EEFDA****",
		UpdatedAt: time.Now(),
	}).Error
	require.NoError(t, err, "Failed to insert setup key")

	err = migration.MigrateSetupKeyToHashedSetupKey[types.SetupKey](context.Background(), db)
	require.NoError(t, err, "Migration should not fail to migrate setup key")

	var key types.SetupKey
	err = db.Model(&types.SetupKey{}).First(&key).Error
	assert.NoError(t, err, "Failed to fetch setup key")

	assert.Equal(t, "EEFDA****", key.KeySecret, "Key should be secret")
	assert.Equal(t, "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=", key.Key, "Key should be hashed")
}

func TestMigrateSetupKeyToHashedSetupKey_ForAlreadyMigratedKey_Case2(t *testing.T) {
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.SetupKey{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	err = db.Save(&types.SetupKey{
		Id:        "1",
		Key:       "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=",
		UpdatedAt: time.Now(),
	}).Error
	require.NoError(t, err, "Failed to insert setup key")

	err = migration.MigrateSetupKeyToHashedSetupKey[types.SetupKey](context.Background(), db)
	require.NoError(t, err, "Migration should not fail to migrate setup key")

	var key types.SetupKey
	err = db.Model(&types.SetupKey{}).First(&key).Error
	assert.NoError(t, err, "Failed to fetch setup key")

	assert.Equal(t, "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=", key.Key, "Key should be hashed")
}

func TestDropIndex(t *testing.T) {
	db := setupDatabase(t)

	err := db.AutoMigrate(&types.SetupKey{})
	require.NoError(t, err, "Failed to auto-migrate tables")

	err = db.Save(&types.SetupKey{
		Id:        "1",
		Key:       "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=",
		UpdatedAt: time.Now(),
	}).Error
	require.NoError(t, err, "Failed to insert setup key")

	exist := db.Migrator().HasIndex(&types.SetupKey{}, "idx_setup_keys_account_id")
	assert.True(t, exist, "Should have the index")

	err = migration.DropIndex[types.SetupKey](context.Background(), db, "idx_setup_keys_account_id")
	require.NoError(t, err, "Migration should not fail to remove index")

	exist = db.Migrator().HasIndex(&types.SetupKey{}, "idx_setup_keys_account_id")
	assert.False(t, exist, "Should not have the index")
}

func TestCreateIndex(t *testing.T) {
	db := setupDatabase(t)
	err := db.AutoMigrate(&nbpeer.Peer{})
	assert.NoError(t, err, "Failed to auto-migrate tables")

	indexName := "idx_account_ip"

	err = migration.CreateIndexIfNotExists[nbpeer.Peer](context.Background(), db, indexName, "account_id", "ip")
	assert.NoError(t, err, "Migration should not fail to create index")

	exist := db.Migrator().HasIndex(&nbpeer.Peer{}, indexName)
	assert.True(t, exist, "Should have the index")
}

func TestCreateIndexIfExists(t *testing.T) {
	db := setupDatabase(t)
	err := db.AutoMigrate(&nbpeer.Peer{})
	assert.NoError(t, err, "Failed to auto-migrate tables")

	indexName := "idx_account_ip"

	err = migration.CreateIndexIfNotExists[nbpeer.Peer](context.Background(), db, indexName, "account_id", "ip")
	assert.NoError(t, err, "Migration should not fail to create index")

	exist := db.Migrator().HasIndex(&nbpeer.Peer{}, indexName)
	assert.True(t, exist, "Should have the index")

	err = migration.CreateIndexIfNotExists[nbpeer.Peer](context.Background(), db, indexName, "account_id", "ip")
	assert.NoError(t, err, "Create index should not fail if index exists")

	exist = db.Migrator().HasIndex(&nbpeer.Peer{}, indexName)
	assert.True(t, exist, "Should have the index")
}