Files
mailgo/internal/db/db.go
T
kevin 6ee2dd8184 fix(security): 修复 P4 封禁记录数据错位 + P5 枚举爆破宽限(方案 A)
P4 #17 手动封禁 Create 非 upsert 导致数据错位:
- BanStore 新增 BanIP(ip, reason, duration):事务内清理该 IP 全部
  既有记录(兼容历史脏数据)后插入单条封禁记录,计数清零;
  DisconnectConnection 改用(原裸 Create 为全仓库唯一调用点)
- BanEntry.IPAddress 升级 uniqueIndex;InitDB 在 AutoMigrate 前
  dedupeBanEntries 清理历史重复行(MySQL 1093 兼容写法)
- IncrementFail 原子化:SQL 侧 fail_count+1,miss 时 OnConflict
  DoNothing 插入兜底并发竞态,回读计数

P5 #18 方案 A(按失败性质区分宽限):
- RecordAuthFailure 新增 knownUser 参数:用户名不存在(枚举型
  爆破)跳过 3 次宽限、首次触发即封第 1 档;用户名存在(真实
  用户输错)保留宽限防误封
- 新增 UserStore.LoginExists(邮箱/裸用户名);五个失败调用点
  接线(Web 登录查存在性;SMTP/IMAP/POP3 用登录名;LDAP 侧
  存在性不可判定,保守按已知用户处理)
- 封禁原因注明「未知用户名,跳过宽限」便于后台审计

新增 8 个测试(-race 通过):BanIP 单行 upsert、唯一索引约束、
16 协程并发计数精确、dedupe 清理/表不存在静默、未知用户即时
封禁、已知用户宽限回归、LoginExists 矩阵。
2026-08-20 11:25:10 +08:00

90 lines
3.2 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 db
import (
"fmt"
"log"
"os"
"path/filepath"
"strings"
"mail_go/config"
"gorm.io/driver/mysql"
"gorm.io/driver/sqlite"
"gorm.io/gorm"
"gorm.io/gorm/logger"
)
// InitDB initializes the database connection and performs auto-migration.
// It selects the appropriate driver based on cfg.Driver and resolves
// the DSN path for SQLite relative to the storage base directory.
func InitDB(cfg config.DatabaseConfig, storageCfg config.StorageConfig) (*gorm.DB, error) {
var dialector gorm.Dialector
switch cfg.Driver {
case "sqlite":
dsn := cfg.DSN
// If the DSN is the default relative path, prepend the storage base directory
if dsn == config.DefaultDSNWin || dsn == config.DefaultDSNLinux {
dsn = filepath.Join(storageCfg.BaseDir, "mail.db")
}
// Ensure the parent directory exists for SQLite
dir := filepath.Dir(dsn)
if err := os.MkdirAll(dir, 0755); err != nil {
return nil, fmt.Errorf("创建数据库目录失败 %s: %w", dir, err)
}
// 多连接并发(SMTP/IMAP/POP3/Web/外发 worker/推送)下:
// - WAL 模式:读不阻塞写,消除瞬时 SQLITE_BUSY 导致写失败被吞
// - busy_timeout=5000ms:写竞争时等待而非立刻失败
// - synchronous=NORMALWAL 下安全且写入更快
sep := "?"
if strings.Contains(dsn, "?") {
sep = "&"
}
dsn = dsn + sep + "_busy_timeout=5000&_journal_mode=WAL&_synchronous=NORMAL"
dialector = sqlite.Open(dsn)
case "mysql":
dialector = mysql.Open(cfg.DSN)
default:
return nil, fmt.Errorf("不支持的数据库驱动: %s", cfg.Driver)
}
db, err := gorm.Open(dialector, &gorm.Config{
Logger: logger.Default.LogMode(logger.Warn),
})
if err != nil {
return nil, fmt.Errorf("连接数据库失败: %w", err)
}
// AutoMigrate 前清理 ban_entries 的历史重复行(保留每 IP 最大 id):
// ip_address 将升级为唯一索引,重复行会使索引创建失败。
// 首次安装表不存在时忽略错误(AutoMigrate 会建新表)。
dedupeBanEntries(db)
// Auto-migrate all models
if err := db.AutoMigrate(&User{}, &Domain{}, &Message{}, &Attachment{}, &BanEntry{}, &OutboundMessage{}, &ProtocolLog{}, &MailboxState{}, &Mailbox{}); err != nil {
return nil, fmt.Errorf("数据库迁移失败: %w", err)
}
return db, nil
}
// dedupeBanEntries 删除 ban_entries 中同一 ip_address 的重复行(保留
// 每组 id 最大的一条),为 ip_address 唯一索引的 AutoMigrate 扫清障碍。
// 表不存在(首次安装)时静默跳过;清理失败仅告警,不阻断启动
// (索引创建失败会在 AutoMigrate 中显式报错)。
func dedupeBanEntries(db *gorm.DB) {
// SQLite 与 MySQL 均支持;MySQL 不允许 DELETE 子查询直接引用同表
// (1093),因此用派生表包一层。
sql := "DELETE FROM ban_entries WHERE id NOT IN (" +
"SELECT mid FROM (SELECT MAX(id) AS mid FROM ban_entries GROUP BY ip_address) AS t)"
if err := db.Exec(sql).Error; err != nil {
// 表不存在(首次安装)为预期情况
msg := err.Error()
if strings.Contains(msg, "no such table") || strings.Contains(msg, "doesn't exist") {
return
}
log.Printf("清理 ban_entries 重复行失败(唯一索引可能无法创建): %v", err)
}
}