package handler import ( "context" "time" "github.com/agentmail/gateway/internal/db" ) // 登录失败限速:同一用户名连续 N 次失败后锁定一段时间。 // 用 DB 而非进程内内存计数器,多实例部署时各实例共享同一份计数。 const ( maxLoginFailures = 5 lockoutDuration = 5 * time.Minute failureWindow = 15 * time.Minute ) // LoginLimiter 通过 DB 实现的登录失败限速器。 type LoginLimiter struct{} var limiter = &LoginLimiter{} // Locked 返回该用户名是否处于锁定期,以及剩余秒数。 // 不记账,只读。 func (l *LoginLimiter) Locked(ctx context.Context, name string) (bool, int) { bucket := "login:" + name cutoff := time.Now().Add(-failureWindow) var count int err := db.DB.QueryRowContext(ctx, `SELECT COUNT(*) FROM rate_limits WHERE bucket = $1 AND ts >= $2`, bucket, cutoff).Scan(&count) if err != nil || count < maxLoginFailures { return false, 0 } // 找到最早那条记录 + lockoutDuration = 解锁时间 var earliest time.Time err = db.DB.QueryRowContext(ctx, `SELECT MIN(ts) FROM rate_limits WHERE bucket = $1 AND ts >= $2`, bucket, cutoff).Scan(&earliest) if err != nil || earliest.IsZero() { return false, 0 } unlockAt := earliest.Add(lockoutDuration) now := time.Now() if now.Before(unlockAt) { remain := int(unlockAt.Sub(now).Seconds()) + 1 if remain < 1 { remain = 1 } return true, remain } // 锁定期已过,清理旧记录 db.DB.ExecContext(ctx, `DELETE FROM rate_limits WHERE bucket = $1 AND ts < $2`, bucket, unlockAt) return false, 0 } // Fail 记录一次登录失败。达到阈值时不额外标记 —— // Locked() 用 COUNT >= maxLoginFailures 自然判定锁定。 func (l *LoginLimiter) Fail(ctx context.Context, name string) { bucket := "login:" + name db.DB.ExecContext(ctx, `INSERT INTO rate_limits (bucket, ts) VALUES ($1, $2)`, bucket, time.Now()) } // Reset 登录成功后清除失败计数。 func (l *LoginLimiter) Reset(ctx context.Context, name string) { bucket := "login:" + name db.DB.ExecContext(ctx, `DELETE FROM rate_limits WHERE bucket = $1`, bucket) }