[management] Setup key improvements (#2775)

management/server/migration/migration.go

@@ -2,13 +2,16 @@ package migration
 
 import (
 	"context"
+	"crypto/sha256"
 	"database/sql"
+	b64 "encoding/base64"
 	"encoding/gob"
 	"encoding/json"
 	"errors"
 	"fmt"
 	"net"
 	"strings"
+	"unicode/utf8"
 
 	log "github.com/sirupsen/logrus"
 	"gorm.io/gorm"
@@ -205,3 +208,90 @@ func MigrateNetIPFieldFromBlobToJSON[T any](ctx context.Context, db *gorm.DB, fi
 
 	return nil
 }
+
+func MigrateSetupKeyToHashedSetupKey[T any](ctx context.Context, db *gorm.DB) error {
+	oldColumnName := "key"
+	newColumnName := "key_secret"
+
+	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 err := db.Transaction(func(tx *gorm.DB) error {
+		if !tx.Migrator().HasColumn(&model, newColumnName) {
+			log.WithContext(ctx).Infof("Column %s does not exist in table %s, adding it", newColumnName, tableName)
+			if err := tx.Migrator().AddColumn(&model, newColumnName); err != nil {
+				return fmt.Errorf("add column %s: %w", newColumnName, err)
+			}
+		}
+
+		var rows []map[string]any
+		if err := tx.Table(tableName).
+			Select("id", oldColumnName, newColumnName).
+			Where(newColumnName + " IS NULL OR " + newColumnName + " = ''").
+			Where("SUBSTR(" + oldColumnName + ", 9, 1) = '-'").
+			Find(&rows).Error; err != nil {
+			return fmt.Errorf("find rows with empty secret key and matching pattern: %w", err)
+		}
+
+		if len(rows) == 0 {
+			log.WithContext(ctx).Infof("No plain setup keys found in table %s, no migration needed", tableName)
+			return nil
+		}
+
+		for _, row := range rows {
+			var plainKey string
+			if columnValue := row[oldColumnName]; columnValue != nil {
+				value, ok := columnValue.(string)
+				if !ok {
+					return fmt.Errorf("type assertion failed")
+				}
+				plainKey = value
+			}
+
+			secretKey := hiddenKey(plainKey, 4)
+
+			hashedKey := sha256.Sum256([]byte(plainKey))
+			encodedHashedKey := b64.StdEncoding.EncodeToString(hashedKey[:])
+
+			if err := tx.Table(tableName).Where("id = ?", row["id"]).Update(newColumnName, secretKey).Error; err != nil {
+				return fmt.Errorf("update row with secret key: %w", err)
+			}
+
+			if err := tx.Table(tableName).Where("id = ?", row["id"]).Update(oldColumnName, encodedHashedKey).Error; err != nil {
+				return fmt.Errorf("update row with hashed key: %w", err)
+			}
+		}
+
+		if err := tx.Exec(fmt.Sprintf("ALTER TABLE %s DROP COLUMN %s", "peers", "setup_key")).Error; err != nil {
+			log.WithContext(ctx).Errorf("Failed to drop column %s: %v", "setup_key", err)
+		}
+
+		return nil
+	}); err != nil {
+		return err
+	}
+
+	log.Printf("Migration of plain setup key to hashed setup key completed")
+	return nil
+}
+
+// hiddenKey returns the Key value hidden with "*" and a 5 character prefix.
+// E.g., "831F6*******************************"
+func hiddenKey(key string, length int) string {
+	prefix := key[0:5]
+	if length > utf8.RuneCountInString(key) {
+		length = utf8.RuneCountInString(key) - len(prefix)
+	}
+	return prefix + strings.Repeat("*", length)
+}

management/server/migration/migration_test.go

@@ -160,3 +160,72 @@ func TestMigrateNetIPFieldFromBlobToJSON_WithJSONData(t *testing.T) {
 	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(&server.SetupKey{})
+	require.NoError(t, err, "Failed to auto-migrate tables")
+
+	err = db.Save(&server.SetupKey{
+		Id:  "1",
+		Key: "EEFDAB47-C1A5-4472-8C05-71DE9A1E8382",
+	}).Error
+	require.NoError(t, err, "Failed to insert setup key")
+
+	err = migration.MigrateSetupKeyToHashedSetupKey[server.SetupKey](context.Background(), db)
+	require.NoError(t, err, "Migration should not fail to migrate setup key")
+
+	var key server.SetupKey
+	err = db.Model(&server.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(&server.SetupKey{})
+	require.NoError(t, err, "Failed to auto-migrate tables")
+
+	err = db.Save(&server.SetupKey{
+		Id:        "1",
+		Key:       "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=",
+		KeySecret: "EEFDA****",
+	}).Error
+	require.NoError(t, err, "Failed to insert setup key")
+
+	err = migration.MigrateSetupKeyToHashedSetupKey[server.SetupKey](context.Background(), db)
+	require.NoError(t, err, "Migration should not fail to migrate setup key")
+
+	var key server.SetupKey
+	err = db.Model(&server.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(&server.SetupKey{})
+	require.NoError(t, err, "Failed to auto-migrate tables")
+
+	err = db.Save(&server.SetupKey{
+		Id:  "1",
+		Key: "9+FQcmNd2GCxIK+SvHmtp6PPGV4MKEicDS+xuSQmvlE=",
+	}).Error
+	require.NoError(t, err, "Failed to insert setup key")
+
+	err = migration.MigrateSetupKeyToHashedSetupKey[server.SetupKey](context.Background(), db)
+	require.NoError(t, err, "Migration should not fail to migrate setup key")
+
+	var key server.SetupKey
+	err = db.Model(&server.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")
+}