go_blog/models/db.go

package models

import (
	"crypto/rand"
	"errors"
	"log"
	"os"
	"path/filepath"

	"go_blog/config"

	"gorm.io/driver/mysql"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"github.com/glebarez/sqlite"
)

// DB 是全局数据库连接，由 InitDB 初始化。
var DB *gorm.DB

// adminPasswordAlphabet 避免了视觉上易混淆的字符（不含 l、I、O、0、1），
// 用于生成首次运行的管理员密码。
const adminPasswordAlphabet = "abcdefghjkmnpqrstuvwxyzABCDEFGHJKMNPQRSTUVWXYZ23456789" + "^$*+?%"

// randomAdminPassword 返回密码学随机的 16 位首次运行管理员密码
// （SECURITY_TODO #12：不再硬编码 admin/admin）。
func randomAdminPassword() string {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		log.Fatalf("Failed to generate admin password: %v", err)
	}
	out := make([]byte, len(b))
	for i, v := range b {
		out[i] = adminPasswordAlphabet[int(v)%len(adminPasswordAlphabet)]
	}
	return string(out)
}

// validateMySQLConfig 校验 MySQL 连接所需字段均非空。
func validateMySQLConfig(d *config.DatabaseConfig) error {
	if d.DBName == "" {
		return errors.New("'db_name' is required when type is 'mysql'")
	}
	if d.Username == "" {
		return errors.New("'username' is required when type is 'mysql'")
	}
	if d.Password == "" {
		return errors.New("'password' is required when type is 'mysql'")
	}
	if d.Host == "" {
		return errors.New("'host' is required when type is 'mysql'")
	}
	if d.Port == "" {
		return errors.New("'port' is required when type is 'mysql'")
	}
	return nil
}

// InitDB 打开数据库连接、执行迁移并初始化管理员用户。
func InitDB(cfg *config.Config) *gorm.DB {
	// 确保存储目录存在。
	if err := os.MkdirAll(cfg.Path, 0755); err != nil {
		log.Fatalf("Failed to create storage directory %s: %v", cfg.Path, err)
	}

	var dialector gorm.Dialector

	switch cfg.Database.Type {
	case "mysql":
		if err := validateMySQLConfig(&cfg.Database); err != nil {
			log.Fatalf("Invalid MySQL config: %v", err)
		}
		dialector = mysql.Open(cfg.Database.MySQLDSN())
	default:
		dbPath := filepath.Join(cfg.Path, "blog.db")
		dialector = sqlite.Open(dbPath)
	}

	db, err := gorm.Open(dialector, &gorm.Config{
		Logger: logger.Default.LogMode(logger.Warn),
	})
	if err != nil {
		log.Fatalf("Failed to connect to database: %v", err)
	}

	// 自动迁移数据表（幂等操作）。
	if err := db.AutoMigrate(&User{}, &Article{}, &SiteSetting{}, &UploadConfig{}, &UploadFileType{}, &DownloadBaseURL{}, &Attachment{}, &File{}, &Comment{}, &CommentConfig{}, &ArticleView{}, &NavLink{}, &Tag{}, &ArticleTag{}); err != nil {
		log.Fatalf("Failed to auto-migrate database: %v", err)
	}

	// 数据迁移：attachments → 全站统一 files 表（type='attachments'），幂等。
	migrateAttachmentsToFiles(db)

	// 首次运行时初始化站点平台配置。
	seedSiteSettings(db)
	seedUploadConfig(db)
	seedUploadFileTypes(db)
	seedCommentConfig(db)

	// 首次运行初始化：若不存在任何用户则创建管理员用户。
	var count int64
	db.Model(&User{}).Count(&count)
	if count == 0 {
		admin := &User{
			Username:    "admin",
			DisplayName: "Administrator",
			Gender:      "other",
			Status:      StatusNormal,
			Role:        RoleAdmin,
		}
		adminPassword := randomAdminPassword()
		if err := admin.SetPassword(adminPassword); err != nil {
			log.Fatalf("Failed to hash admin password: %v", err)
		}
		if err := db.Create(admin).Error; err != nil {
			log.Fatalf("Failed to create admin user: %v", err)
		}
		// SECURITY_TODO #12：首次运行密码为密码学随机生成，且只打印一次——
		// 请立即抄写；之后将无法找回。
		log.Println("==============================================")
		log.Println("  First run: created default admin user.")
		log.Println("  Username: admin")
		log.Println("  Password: " + adminPassword)
		log.Println("  This password is shown ONCE. Change it after login!")
		log.Println("==============================================")
	}

	// 迁移修复：始终为 admin 用户设置管理员角色。
	result := db.Model(&User{}).Where("username = ?", "admin").Update("role", RoleAdmin)
	if result.RowsAffected > 0 {
		log.Printf("Migration: set admin role for existing admin user (rows affected: %d)", result.RowsAffected)
	}

	DB = db
	return db
}

// migrateAttachmentsToFiles 将 attachments 表中的历史数据复制到统一的
// files 表，type 一律标记为 "attachments"。
//
// 幂等策略：以主键 id 对齐——files 中已存在同 id 的行视为已迁移并跳过，
// 因此 InitDB 每次启动重复执行也不会产生重复数据（含软删除行一并复制）。
func migrateAttachmentsToFiles(db *gorm.DB) {
	res := db.Exec(`
INSERT INTO files
	(id, type, article_id, session_token, uploader_id, filename, stored_name, ext, mime, size, category, created_at, updated_at, deleted_at)
SELECT
	a.id, 'attachments', a.article_id, a.session_token, a.uploader_id, a.filename,
	a.stored_name, a.ext, a.mime, a.size, a.category, a.created_at, a.updated_at, a.deleted_at
FROM attachments a
WHERE NOT EXISTS (SELECT 1 FROM files f WHERE f.id = a.id)`)
	if res.Error != nil {
		log.Printf("Migration attachments -> files failed: %v", res.Error)
		return
	}
	if res.RowsAffected > 0 {
		log.Printf("Migration: copied %d attachment(s) into files table (type=attachments)", res.RowsAffected)
	}
}