mailgo/internal/store/ban_store.go

package store

import (
	"fmt"
	"strings"
	"time"

	"mail_go/internal/db"

	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

// 阶段封禁档位（分钟/天），从第 4 次触发阈值开始封禁：
// 第 4 次 = ban_duration_min（默认 30 分钟）→ 第 5 次 = 3 小时 →
// 第 6 次 = 3 个月 → 第 7 次起 = 半年（上限）。
const (
	// freeTriggers 达到失败阈值但暂不封禁的触发次数（前 3 次只计数）。
	freeTriggers = 3
	banStage2Min = 3 * 60 // 3 小时
	banStage3Day = 90     // 3 个月
	banStage4Day = 180    // 半年（上限）
	banMaxDay    = banStage4Day
)

// stageDuration 返回第 banCount 次封禁（banCount 从 1 开始）的时长。
// firstBanMin 是第一次封禁的分钟数（来自配置 [ban] ban_duration_min）。
func stageDuration(banCount int, firstBanMin int) time.Duration {
	switch banCount {
	case 1:
		if firstBanMin <= 0 {
			firstBanMin = 30
		}
		return time.Duration(firstBanMin) * time.Minute
	case 2:
		return time.Duration(banStage2Min) * time.Minute
	case 3:
		return time.Duration(banStage3Day) * 24 * time.Hour
	default:
		return time.Duration(banMaxDay) * 24 * time.Hour
	}
}

// BanStore defines the interface for IP ban operations.
type BanStore interface {
	Create(entry *db.BanEntry) error
	// BanIP 手动/直接封禁某 IP 指定时长：该 IP 只保留一条记录
	// （既有观察记录一并清理，计数清零），避免产生重复行导致
	// IncrementFail 与 GetByIP 读写错位。
	BanIP(ip, reason string, duration time.Duration) error
	GetByIP(ip string) (*db.BanEntry, error)
	Update(entry *db.BanEntry) error
	Delete(id uint) error
	// List 返回已封禁或曾封禁的记录（不含仅计数未封禁的观察记录）。
	List(page, size int) ([]db.BanEntry, int64, error)
	IsBanned(ip string) (bool, *db.BanEntry)
	// IncrementFail 累加该 IP 的失败次数（无记录时创建），保留 BanCount。
	IncrementFail(ip string) (int, error)
	ResetFail(ip string) error
}

// banStoreGorm implements BanStore using GORM.
type banStoreGorm struct {
	db *gorm.DB
}

// newBanStore creates a new GORM-backed BanStore.
func newBanStore(database *gorm.DB) BanStore {
	return &banStoreGorm{db: database}
}

// Create inserts a new ban entry record.
func (s *banStoreGorm) Create(entry *db.BanEntry) error {
	return s.db.Create(entry).Error
}

// BanIP 手动/直接封禁：事务内删除该 IP 的全部既有记录（含历史 bug
// 产生的重复行与观察计数记录）后插入一条封禁记录，计数清零。
// 与"管理员解封清零"语义一致：手动封禁视为对档位的重新评估。
func (s *banStoreGorm) BanIP(ip, reason string, duration time.Duration) error {
	if ip == "" {
		return fmt.Errorf("empty ip")
	}
	if duration <= 0 {
		return fmt.Errorf("invalid ban duration: %v", duration)
	}
	return s.db.Transaction(func(tx *gorm.DB) error {
		if err := tx.Where("ip_address = ?", ip).Delete(&db.BanEntry{}).Error; err != nil {
			return err
		}
		return tx.Create(&db.BanEntry{
			IPAddress: ip,
			Reason:    reason,
			FailCount: 0,
			BanCount:  0,
			ExpiresAt: time.Now().Add(duration),
		}).Error
	})
}

// GetByIP retrieves the most recent ban entry for a given IP address.
func (s *banStoreGorm) GetByIP(ip string) (*db.BanEntry, error) {
	var entry db.BanEntry
	if err := s.db.Where("ip_address = ?", ip).Order("id DESC").First(&entry).Error; err != nil {
		return nil, err
	}
	return &entry, nil
}

// Update saves changes to an existing ban entry record.
func (s *banStoreGorm) Update(entry *db.BanEntry) error {
	return s.db.Save(entry).Error
}

// Delete removes a ban entry by ID.
func (s *banStoreGorm) Delete(id uint) error {
	return s.db.Delete(&db.BanEntry{}, id).Error
}

// banEpochSentinel 用于区分“未封禁的计数记录”（expires_at 为零值）：
// 所有实际封禁记录的到期时间都晚于 2000 年。
var banEpochSentinel = time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)

// List retrieves a paginated list of ban entries that are or have been
// banned (expires_at set). 仅计数的观察记录（零到期时间、无原因）不返回。
func (s *banStoreGorm) List(page, size int) ([]db.BanEntry, int64, error) {
	var entries []db.BanEntry
	var total int64

	query := s.db.Model(&db.BanEntry{}).Where("expires_at > ?", banEpochSentinel)
	if err := query.Count(&total).Error; err != nil {
		return nil, 0, err
	}

	offset := (page - 1) * size
	if err := query.Order("id DESC").Offset(offset).Limit(size).Find(&entries).Error; err != nil {
		return nil, 0, err
	}
	return entries, total, nil
}

// IsBanned checks whether an IP address is currently banned.
// An IP is considered banned if there is a record with expires_at in the future.
func (s *banStoreGorm) IsBanned(ip string) (bool, *db.BanEntry) {
	var entry db.BanEntry
	if err := s.db.Where("ip_address = ? AND expires_at > ?", ip, time.Now()).First(&entry).Error; err != nil {
		return false, nil
	}
	return true, &entry
}

// IncrementFail increments the fail count for an IP address atomically
// (SQL-side increment, avoiding read-modify-write races). If no record
// exists it creates one with fail_count=1, ban_count=0 and a zero
// expires_at (not yet banned); a concurrent creator wins and the loser's
// insert becomes a no-op via the unique index. Existing BanCount is
// preserved. Returns the current fail count.
func (s *banStoreGorm) IncrementFail(ip string) (int, error) {
	res := s.db.Model(&db.BanEntry{}).
		Where("ip_address = ?", ip).
		Update("fail_count", gorm.Expr("fail_count + 1"))
	if res.Error != nil {
		return 0, res.Error
	}
	if res.RowsAffected == 0 {
		// 无记录：插入首条；ip_address 唯一索引下并发插入用
		// OnConflict DoNothing 兜底，失败方继续走下面的回读。
		err := s.db.Clauses(clause.OnConflict{DoNothing: true}).Create(&db.BanEntry{
			IPAddress: ip,
			FailCount: 1,
			BanCount:  0,
			ExpiresAt: time.Time{}, // Zero time, not yet banned
		}).Error
		if err != nil && !isUniqueConflictErr(err) {
			return 0, err
		}
	}

	// 回读计数（并发下取数据库最终值）
	var count int64
	if err := s.db.Model(&db.BanEntry{}).Where("ip_address = ?", ip).
		Select("fail_count").Scan(&count).Error; err != nil {
		return 0, err
	}
	return int(count), nil
}

// isUniqueConflictErr 判断是否为唯一约束冲突（并发插入竞态的预期结果）。
func isUniqueConflictErr(err error) bool {
	if err == nil {
		return false
	}
	msg := err.Error()
	return strings.Contains(msg, "UNIQUE constraint") || // SQLite
		strings.Contains(msg, "Duplicate entry") || // MySQL
		strings.Contains(msg, "unique constraint") // generic
}

// ResetFail resets the fail count for an IP address by deleting its record.
// 成功登录（或管理员解封）后调用：封禁档位历史随之清零。
func (s *banStoreGorm) ResetFail(ip string) error {
	return s.db.Where("ip_address = ?", ip).Delete(&db.BanEntry{}).Error
}