重构进行中，preview通道尚未完全恢复全部功能

2026-08-07 19:16:08 +08:00
parent 5bce105bd6
commit 3a381c85ed
21 changed files with 889 additions and 2416 deletions
@@ -0,0 +1,424 @@
+package db
+
+import (
+	"crypto/rand"
+	"database/sql"
+	"encoding/hex"
+	"fmt"
+	"io"
+	"log"
+	"os"
+	"sort"
+	"strings"
+	"time"
+
+	_ "modernc.org/sqlite"
+)
+
+var DB *sql.DB
+
+const SchemaVersion = "v2"
+
+// ================================================================
+// 初始化
+// ================================================================
+
+func InitDB() error {
+	if err := os.MkdirAll("./data", 0755); err != nil {
+		return fmt.Errorf("创建数据目录失败: %w", err)
+	}
+
+	dbPath := "./data/frpc-console.db"
+	var err error
+	DB, err = sql.Open("sqlite", dbPath)
+	if err != nil {
+		return err
+	}
+
+	if err := createTables(); err != nil {
+		return err
+	}
+
+	if err := runMigrations(); err != nil {
+		return err
+	}
+
+	if err := ensureJwtSecret(); err != nil {
+		return err
+	}
+
+	log.Println("✅ 数据库初始化完成 (Schema " + SchemaVersion + ")")
+	return nil
+}
+
+func createTables() error {
+	_, err := DB.Exec(`
+        CREATE TABLE IF NOT EXISTS users (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            username TEXT UNIQUE NOT NULL,
+            password_hash TEXT NOT NULL,
+            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
+        )
+    `)
+	if err != nil {
+		return err
+	}
+
+	_, err = DB.Exec(`
+        CREATE TABLE IF NOT EXISTS global_config (
+            id INTEGER PRIMARY KEY CHECK (id = 1),
+            server_addr TEXT NOT NULL DEFAULT 'frp.example.com',
+            server_port INTEGER NOT NULL DEFAULT 7000,
+            token TEXT NOT NULL DEFAULT 'CHANGE_ME',
+            log_level TEXT NOT NULL DEFAULT 'info',
+            log_max_days INTEGER NOT NULL DEFAULT 3,
+            tcp_mux INTEGER NOT NULL DEFAULT 1,
+            tcp_mux_keepalive INTEGER NOT NULL DEFAULT 30,
+            heartbeat_interval INTEGER NOT NULL DEFAULT 15,
+            heartbeat_timeout INTEGER NOT NULL DEFAULT 70,
+            pool_count INTEGER NOT NULL DEFAULT 8,
+            wire_protocol_v2 INTEGER NOT NULL DEFAULT 0,
+            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
+        )
+    `)
+	if err != nil {
+		return err
+	}
+
+	_, err = DB.Exec(`
+        CREATE TABLE IF NOT EXISTS proxies (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            name TEXT UNIQUE NOT NULL,
+            type TEXT NOT NULL DEFAULT 'tcp',
+            local_ip TEXT NOT NULL,
+            local_port INTEGER NOT NULL,
+            remote_port INTEGER NOT NULL,
+            enabled INTEGER NOT NULL DEFAULT 1,
+            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
+        )
+    `)
+	if err != nil {
+		return err
+	}
+
+	_, err = DB.Exec(`
+        CREATE TABLE IF NOT EXISTS app_config (
+            key TEXT PRIMARY KEY,
+            value TEXT NOT NULL,
+            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
+        )
+    `)
+	if err != nil {
+		return err
+	}
+
+	var count int
+	DB.QueryRow("SELECT COUNT(*) FROM global_config").Scan(&count)
+	if count == 0 {
+		_, err = DB.Exec(`
+            INSERT INTO global_config (
+                id, server_addr, server_port, token, log_level, log_max_days,
+                tcp_mux, tcp_mux_keepalive, heartbeat_interval, heartbeat_timeout, pool_count,
+                wire_protocol_v2
+            ) VALUES (1, 'frp.example.com', 7000, 'CHANGE_ME', 'info', 3, 1, 30, 15, 70, 8, 0)
+        `)
+		if err != nil {
+			return err
+		}
+		log.Println("✅ 全局配置初始化完成")
+	}
+
+	return nil
+}
+
+// ================================================================
+// JWT 密钥管理
+// ================================================================
+
+func ensureJwtSecret() error {
+	var value string
+	err := DB.QueryRow("SELECT value FROM app_config WHERE key = 'jwt_secret'").Scan(&value)
+	if err == nil && value != "" {
+		return nil
+	}
+
+	bytes := make([]byte, 32)
+	if _, err := rand.Read(bytes); err != nil {
+		return err
+	}
+	secret := hex.EncodeToString(bytes)
+
+	_, err = DB.Exec(`
+        INSERT INTO app_config (key, value) VALUES ('jwt_secret', ?)
+    `, secret)
+	if err != nil {
+		return err
+	}
+	log.Printf("✅ JWT 密钥已生成")
+	return nil
+}
+
+func GetJwtSecret() (string, error) {
+	var secret string
+	err := DB.QueryRow("SELECT value FROM app_config WHERE key = 'jwt_secret'").Scan(&secret)
+	if err != nil {
+		return "", err
+	}
+	return secret, nil
+}
+
+// ================================================================
+// 迁移引擎
+// ================================================================
+
+func getCurrentSchemaVersion() string {
+	var version string
+	err := DB.QueryRow("SELECT value FROM app_config WHERE key = 'schema_version'").Scan(&version)
+	if err != nil {
+		if err == sql.ErrNoRows {
+			var count int
+			DB.QueryRow("SELECT COUNT(*) FROM users").Scan(&count)
+			if count > 0 {
+				return "v1"
+			}
+			return SchemaVersion
+		}
+		log.Printf("⚠️ 读取 Schema 版本失败: %v", err)
+		return "v1"
+	}
+	return version
+}
+
+func setSchemaVersion(version string) error {
+	_, err := DB.Exec(`
+        INSERT INTO app_config (key, value) VALUES ('schema_version', ?)
+        ON CONFLICT(key) DO UPDATE SET value = excluded.value, updated_at = CURRENT_TIMESTAMP
+    `, version)
+	return err
+}
+
+func backupDatabase() (string, error) {
+	src := "./data/frpc-console.db"
+	if _, err := os.Stat(src); os.IsNotExist(err) {
+		return "", nil
+	}
+
+	timestamp := time.Now().Format("20060102_150405")
+	dst := fmt.Sprintf("./data/frpc-console.db.pre-%s.%s", SchemaVersion, timestamp)
+
+	srcFile, err := os.Open(src)
+	if err != nil {
+		return "", fmt.Errorf("打开源数据库失败: %w", err)
+	}
+	defer srcFile.Close()
+
+	dstFile, err := os.Create(dst)
+	if err != nil {
+		return "", fmt.Errorf("创建备份文件失败: %w", err)
+	}
+	defer dstFile.Close()
+
+	if _, err := io.Copy(dstFile, srcFile); err != nil {
+		return "", fmt.Errorf("复制数据库失败: %w", err)
+	}
+
+	log.Printf("✅ 数据库备份完成: %s", dst)
+	return dst, nil
+}
+
+func restoreDatabase(backupPath string) error {
+	srcFile, err := os.Open(backupPath)
+	if err != nil {
+		return err
+	}
+	defer srcFile.Close()
+
+	dstFile, err := os.Create("./data/frpc-console.db")
+	if err != nil {
+		return err
+	}
+	defer dstFile.Close()
+
+	if _, err := io.Copy(dstFile, srcFile); err != nil {
+		return err
+	}
+
+	log.Printf("✅ 数据库已从备份恢复: %s", backupPath)
+	return nil
+}
+
+func runMigrations() error {
+	currentVer := getCurrentSchemaVersion()
+	targetVer := SchemaVersion
+
+	log.Printf("📌 当前数据库 Schema: %s, 目标版本: %s", currentVer, targetVer)
+
+	if currentVer == targetVer {
+		var userCount int
+		err := DB.QueryRow("SELECT COUNT(*) FROM users").Scan(&userCount)
+		if err != nil || userCount == 0 {
+			log.Println("  数据库为空或无效，无需迁移，直接初始化")
+			return nil
+		}
+		log.Println("✅ Schema 已是最新，数据库有效")
+		return nil
+	}
+
+	log.Printf("🔄 检测到版本变更 (%s → %s)，开始迁移...", currentVer, targetVer)
+
+	backupPath, err := backupDatabase()
+	if err != nil {
+		return fmt.Errorf("备份数据库失败: %w", err)
+	}
+	if backupPath != "" {
+		log.Printf("📦 备份文件: %s", backupPath)
+	}
+
+	currentSchema := getSchemaDef(currentVer)
+	targetSchema := getSchemaDef(targetVer)
+
+	if targetSchema == nil {
+		return fmt.Errorf("目标 Schema 版本 %s 未定义", targetVer)
+	}
+
+	if currentSchema == nil || schemaVersionsEqual(currentSchema, targetSchema) {
+		log.Println("  迁移类型: 轻量复制（Schema 无变更）")
+		var userCount int
+		err := DB.QueryRow("SELECT COUNT(*) FROM users").Scan(&userCount)
+		if err != nil || userCount == 0 {
+			log.Println("  数据库为空或无效，跳过迁移，直接初始化")
+			return nil
+		}
+		log.Println("  数据库有效，继续使用")
+	} else {
+		log.Println("  迁移类型: 重型迁移（Schema 有变更，新建表 + 搬数据）")
+		if err := heavyMigration(currentSchema, targetSchema); err != nil {
+			if backupPath != "" {
+				log.Printf("❌ 迁移失败，尝试恢复备份: %s", backupPath)
+				if restoreErr := restoreDatabase(backupPath); restoreErr != nil {
+					log.Printf("⚠️ 恢复备份失败: %v", restoreErr)
+				}
+			}
+			return fmt.Errorf("重型迁移失败: %w", err)
+		}
+	}
+
+	if err := setSchemaVersion(targetVer); err != nil {
+		return fmt.Errorf("更新 Schema 版本失败: %w", err)
+	}
+
+	log.Printf("✅ 迁移完成，当前 Schema: %s", targetVer)
+	return nil
+}
+
+func heavyMigration(oldDef, newDef *SchemaVersionDef) error {
+	if oldDef == nil {
+		return fmt.Errorf("旧 Schema 定义为空，无法执行重型迁移")
+	}
+
+	oldTable := oldDef.TableName
+	newTable := oldTable + "_new"
+
+	createSQL := buildCreateTableSQL(newTable, newDef)
+	log.Printf("  创建新表: %s", newTable)
+	if _, err := DB.Exec(createSQL); err != nil {
+		return fmt.Errorf("创建新表失败: %w", err)
+	}
+
+	insertSQL, err := buildInsertSQL(oldTable, newTable, oldDef, newDef)
+	if err != nil {
+		return fmt.Errorf("构建数据迁移 SQL 失败: %w", err)
+	}
+	log.Printf("  迁移数据: %s → %s", oldTable, newTable)
+	if _, err := DB.Exec(insertSQL); err != nil {
+		return fmt.Errorf("数据迁移失败: %w", err)
+	}
+
+	var oldCount, newCount int
+	DB.QueryRow(fmt.Sprintf("SELECT COUNT(*) FROM %s", oldTable)).Scan(&oldCount)
+	DB.QueryRow(fmt.Sprintf("SELECT COUNT(*) FROM %s", newTable)).Scan(&newCount)
+	if oldCount != newCount {
+		return fmt.Errorf("数据迁移不完整: 旧表 %d 行，新表 %d 行", oldCount, newCount)
+	}
+	log.Printf("  数据迁移验证通过: %d 行", newCount)
+
+	tempTable := oldTable + "_old_temp"
+	if _, err := DB.Exec(fmt.Sprintf("ALTER TABLE %s RENAME TO %s", oldTable, tempTable)); err != nil {
+		return fmt.Errorf("重命名旧表失败: %w", err)
+	}
+	if _, err := DB.Exec(fmt.Sprintf("ALTER TABLE %s RENAME TO %s", newTable, oldTable)); err != nil {
+		DB.Exec(fmt.Sprintf("ALTER TABLE %s RENAME TO %s", tempTable, oldTable))
+		return fmt.Errorf("重命名新表失败: %w", err)
+	}
+	if _, err := DB.Exec(fmt.Sprintf("DROP TABLE %s", tempTable)); err != nil {
+		log.Printf("⚠️ 删除临时表失败（不影响使用）: %v", err)
+	}
+
+	log.Printf("  表交换完成: %s (新表已生效)", oldTable)
+	return nil
+}
+
+func buildCreateTableSQL(tableName string, def *SchemaVersionDef) string {
+	var cols []string
+	var primaryKey string
+
+	for name, col := range def.Columns {
+		parts := []string{name, col.Type}
+		if col.NotNull {
+			parts = append(parts, "NOT NULL")
+		}
+		if col.Default != "" {
+			parts = append(parts, "DEFAULT "+col.Default)
+		}
+		if col.Primary {
+			primaryKey = "PRIMARY KEY (" + name + ")"
+		} else {
+			cols = append(cols, strings.Join(parts, " "))
+		}
+	}
+
+	if primaryKey != "" {
+		cols = append(cols, primaryKey)
+	}
+
+	return fmt.Sprintf("CREATE TABLE %s (\n    %s\n)", tableName, strings.Join(cols, ",\n    "))
+}
+
+func buildInsertSQL(oldTable, newTable string, oldDef, newDef *SchemaVersionDef) (string, error) {
+	var newCols []string
+	for name := range newDef.Columns {
+		newCols = append(newCols, name)
+	}
+	sort.Strings(newCols)
+
+	var selectParts []string
+	var colNames []string
+
+	for _, name := range newCols {
+		colNames = append(colNames, name)
+		if _, ok := oldDef.Columns[name]; ok {
+			selectParts = append(selectParts, name)
+		} else {
+			colDef := newDef.Columns[name]
+			if colDef.Default != "" {
+				selectParts = append(selectParts, colDef.Default+" AS "+name)
+			} else if colDef.Type == "INTEGER" {
+				selectParts = append(selectParts, "0 AS "+name)
+			} else if colDef.Type == "TEXT" {
+				selectParts = append(selectParts, "'' AS "+name)
+			} else {
+				selectParts = append(selectParts, "NULL AS "+name)
+			}
+		}
+	}
+
+	return fmt.Sprintf(
+		"INSERT INTO %s (%s) SELECT %s FROM %s",
+		newTable,
+		strings.Join(colNames, ", "),
+		strings.Join(selectParts, ", "),
+		oldTable,
+	), nil
+}