Files
frpc-console/db.go
T

618 lines
17 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package main
import (
"crypto/rand"
"database/sql"
"encoding/hex"
"fmt"
"io"
"log"
"os"
"sort"
"strings"
"time"
_ "modernc.org/sqlite"
)
var DB *sql.DB
// ============================================================
// 数据库版本常量
// ============================================================
const (
SchemaVersion = "v2" // 当前数据库 Schema 版本(与 schemaHistory 中的版本对应)
)
// ============================================================
// 数据模型
// ============================================================
// GlobalConfig 全局配置表
type GlobalConfig struct {
ID int `json:"id"`
ServerAddr string `json:"serverAddr"`
ServerPort int `json:"serverPort"`
Token string `json:"token"`
LogLevel string `json:"logLevel"`
LogMaxDays int `json:"logMaxDays"`
TcpMux bool `json:"tcpMux"`
TcpMuxKeepalive int `json:"tcpMuxKeepalive"`
HeartbeatInterval int `json:"heartbeatInterval"`
HeartbeatTimeout int `json:"heartbeatTimeout"`
PoolCount int `json:"poolCount"`
WireProtocolV2 bool `json:"wireProtocolV2"` // v2 协议全局开关
}
// Proxy 隧道表(不含 wire_protocol_v2
type Proxy struct {
ID int `json:"id"`
Name string `json:"name"`
Type string `json:"type"`
LocalIP string `json:"localIP"`
LocalPort int `json:"localPort"`
RemotePort int `json:"remotePort"`
Enabled bool `json:"enabled"`
}
// User 用户表
type User struct {
ID int `json:"id"`
Username string `json:"username"`
PasswordHash string `json:"-"`
CreatedAt string `json:"createdAt"`
}
// ============================================================
// 数据库初始化
// ============================================================
func InitDB() error {
var err error
DB, err = sql.Open("sqlite", "./frpc-console.db")
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
}
// ============================================================
// 迁移引擎
// ============================================================
// getCurrentSchemaVersion 读取当前数据库的 Schema 版本
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
}
// setSchemaVersion 更新 Schema 版本
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
}
// backupDatabase 备份数据库文件
func backupDatabase() (string, error) {
src := "./frpc-console.db"
if _, err := os.Stat(src); os.IsNotExist(err) {
return "", nil
}
timestamp := time.Now().Format("20060102_150405")
dst := fmt.Sprintf("./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
}
// restoreDatabase 从备份恢复数据库
func restoreDatabase(backupPath string) error {
srcFile, err := os.Open(backupPath)
if err != nil {
return err
}
defer srcFile.Close()
dstFile, err := os.Create("./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
}
// runMigrations 执行迁移
func runMigrations() error {
currentVer := getCurrentSchemaVersion()
targetVer := SchemaVersion
log.Printf("📌 当前数据库 Schema: %s, 目标版本: %s", currentVer, targetVer)
if currentVer == targetVer {
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 未在 schemaHistory 中定义", targetVer)
}
if currentSchema == nil || schemaVersionsEqual(currentSchema, targetSchema) {
log.Println(" 迁移类型: 轻量复制(Schema 无变更)")
} 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
}
// heavyMigration 重型迁移
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
}
// buildCreateTableSQL 根据 Schema 定义生成 CREATE TABLE 语句
func buildCreateTableSQL(tableName string, def *SchemaVersionDef) string {
var cols []string
var primaryKey string
names := make([]string, 0, len(def.Columns))
for name := range def.Columns {
names = append(names, name)
}
sort.Strings(names)
for _, name := range names {
col := def.Columns[name]
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 "))
}
// buildInsertSQL 构建 INSERT INTO new_table SELECT ... FROM old_table
func buildInsertSQL(oldTable, newTable string, oldDef, newDef *SchemaVersionDef) (string, error) {
newCols := make([]string, 0, len(newDef.Columns))
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
}
// ============================================================
// 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
}
// ============================================================
// 全局配置 CRUD
// ============================================================
func GetGlobalConfig() (*GlobalConfig, error) {
var cfg GlobalConfig
err := DB.QueryRow(`
SELECT 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
FROM global_config WHERE id = 1
`).Scan(
&cfg.ID, &cfg.ServerAddr, &cfg.ServerPort, &cfg.Token,
&cfg.LogLevel, &cfg.LogMaxDays, &cfg.TcpMux, &cfg.TcpMuxKeepalive,
&cfg.HeartbeatInterval, &cfg.HeartbeatTimeout, &cfg.PoolCount,
&cfg.WireProtocolV2,
)
if err != nil {
return nil, err
}
cfg.TcpMux = true
return &cfg, nil
}
func UpdateGlobalConfig(cfg *GlobalConfig) error {
_, err := DB.Exec(`
UPDATE global_config SET
server_addr = ?, server_port = ?, token = ?, log_level = ?, log_max_days = ?,
tcp_mux = 1,
tcp_mux_keepalive = ?, heartbeat_interval = ?, heartbeat_timeout = ?, pool_count = ?,
wire_protocol_v2 = ?,
updated_at = CURRENT_TIMESTAMP
WHERE id = 1
`, cfg.ServerAddr, cfg.ServerPort, cfg.Token, cfg.LogLevel, cfg.LogMaxDays,
cfg.TcpMuxKeepalive, cfg.HeartbeatInterval, cfg.HeartbeatTimeout, cfg.PoolCount,
cfg.WireProtocolV2)
return err
}
// ============================================================
// 隧道 CRUD
// ============================================================
func GetProxies() ([]Proxy, error) {
rows, err := DB.Query(`
SELECT id, name, type, local_ip, local_port, remote_port, enabled
FROM proxies ORDER BY id
`)
if err != nil {
return nil, err
}
defer rows.Close()
var proxies []Proxy
for rows.Next() {
var p Proxy
err := rows.Scan(&p.ID, &p.Name, &p.Type, &p.LocalIP, &p.LocalPort, &p.RemotePort, &p.Enabled)
if err != nil {
return nil, err
}
proxies = append(proxies, p)
}
return proxies, rows.Err()
}
func GetProxy(id int) (*Proxy, error) {
var p Proxy
err := DB.QueryRow(`
SELECT id, name, type, local_ip, local_port, remote_port, enabled
FROM proxies WHERE id = ?
`, id).Scan(&p.ID, &p.Name, &p.Type, &p.LocalIP, &p.LocalPort, &p.RemotePort, &p.Enabled)
if err != nil {
return nil, err
}
return &p, nil
}
func CreateProxy(p *Proxy) error {
result, err := DB.Exec(`
INSERT INTO proxies (name, type, local_ip, local_port, remote_port, enabled)
VALUES (?, ?, ?, ?, ?, ?)
`, p.Name, p.Type, p.LocalIP, p.LocalPort, p.RemotePort, p.Enabled)
if err != nil {
return err
}
id, _ := result.LastInsertId()
p.ID = int(id)
return nil
}
func UpdateProxy(p *Proxy) error {
_, err := DB.Exec(`
UPDATE proxies SET
name = ?, type = ?, local_ip = ?, local_port = ?, remote_port = ?, enabled = ?,
updated_at = CURRENT_TIMESTAMP
WHERE id = ?
`, p.Name, p.Type, p.LocalIP, p.LocalPort, p.RemotePort, p.Enabled, p.ID)
return err
}
func DeleteProxy(id int) error {
_, err := DB.Exec("DELETE FROM proxies WHERE id = ?", id)
return err
}
// ============================================================
// 用户 CRUD
// ============================================================
func GetUserByUsername(username string) (*User, error) {
var u User
err := DB.QueryRow(`
SELECT id, username, password_hash, created_at
FROM users WHERE username = ?
`, username).Scan(&u.ID, &u.Username, &u.PasswordHash, &u.CreatedAt)
if err != nil {
return nil, err
}
return &u, nil
}
func CreateUser(username, passwordHash string) error {
_, err := DB.Exec(`
INSERT INTO users (username, password_hash) VALUES (?, ?)
`, username, passwordHash)
return err
}
func UpdatePassword(username, passwordHash string) error {
_, err := DB.Exec(`
UPDATE users SET password_hash = ?, updated_at = CURRENT_TIMESTAMP
WHERE username = ?
`, passwordHash, username)
return err
}
func CountUsers() (int, error) {
var count int
err := DB.QueryRow("SELECT COUNT(*) FROM users").Scan(&count)
return count, err
}