MailUI4Agents/server/internal/handler/ratelimit_test.go

package handler

import (
	"context"
	"path/filepath"
	"testing"
	"time"

	"github.com/agentmail/gateway/internal/db"
)

/*
登录失败限速 —— 2026-09-28。

⚠ 写这类测试的两个坑（都不是这文件要测的东西，但会把你引到错误的结论上）：

  1. **别在同一条 SQL 里混用 `?` 和 `$N`**。SQLite（现代驱动）把两者
     统一按 1-based 索引解释，`?` 和 `$1` 绑的是**同一个**参数。于是
     `UPDATE ... WHERE bucket = $1` 配一个 `?` 形状的值时不会报错，
     静默更新 0 行 —— 看起来像「限速器删不掉旧记录」，其实是参数绑错了。
     统一用 `$N`（跟生产代码一致），别用 `?`。
  2. **想推进时间就整桶一起改**。`Locked()` 读的是 COUNT 和 MIN(ts)
     两个量：只回拨其中一行，剩下的行会把 count 顶回阈值以上
     （`UPDATE ... LIMIT 1` 就会这样）。要么整桶设成同一时间戳
     （本文件就是这么做的），要么连 WHERE 一起改。

这个文件存在的原因不是「补覆盖率」，是**它本来就不存在**：
handler 包里一条登录限速测试都没有，于是下面这个缺陷一直红着没人发现 ——
`Locked()` 从 2025 年写出来就没在 SQLite 下真正返回过 true。

缺陷本体：算解锁时刻用的 `SELECT MIN(ts)`，SQLite 驱动把聚合结果当
string 返回，扫进 *time.Time 直接失败；代码里 `err != nil` 走
`return false, 0`。后果不是「锁了但没到期」，是**压根没进判定**：
`auth.go:122` 在校验密码之前调 Locked，于是失败多少次都返回 false，
无限次撞库不会 429。代码注释里那套「用 DB 而非内存计数器、多实例共享
同一份计数」的设计意图，在默认的 SQLite 部署下没兑现。

所以这里的判据刻意写成**端到端且用真实 SQLite**：
不是断言「Locked() 返回 true」（那可以靠 mock 骗过），
而是「连打 N 次失败登录 ⇒ HTTP 401，**再打一次必须是 429**」。
这条判据红 ⇒ 登录限速没生效；绿 ⇒ 它生效。中间任何一层坏了都红。
*/

func setupLoginLimiterDB(t *testing.T) {
	t.Helper()
	dir := t.TempDir()
	if err := db.Connect(context.Background(), "sqlite://"+filepath.Join(dir, "t.db")); err != nil {
		t.Fatalf("connect: %v", err)
	}
	if err := db.Migrate(context.Background()); err != nil {
		t.Fatalf("migrate: %v", err)
	}
	t.Cleanup(db.Close)
}

// TestLoginLimiterLocksAfterMaxFailures 是本文件的主判据：
// 真实 SQLite 上记满 maxLoginFailures 次失败后，Locked 必须为 true。
//
// 修复前：count=5 但 MIN(ts) 扫描失败 → (false, 0)。
// 修复后：走 ORDER BY ts ASC LIMIT 1 → (true, ≈lockoutDuration)。
func TestLoginLimiterLocksAfterMaxFailures(t *testing.T) {
	setupLoginLimiterDB(t)
	ctx := context.Background()
	const name = "victim"

	for i := 0; i < maxLoginFailures-1; i++ {
		if locked, _ := limiter.Locked(ctx, name); locked {
			t.Fatalf("第 %d 次之前就锁了，未到阈值", i)
		}
		limiter.Fail(ctx, name)
	}

	// 最后一次：记完正好等于阈值
	limiter.Fail(ctx, name)

	locked, remain := limiter.Locked(ctx, name)
	if !locked {
		t.Fatalf("记满 %d 次失败后仍未锁定 —— 登录失败限速未生效（无限撞库不会 429）",
			maxLoginFailures)
	}
	// 剩余秒数应接近整个 lockoutDuration（最早的失败就在几毫秒前）
	if remain < int(lockoutDuration.Seconds())-2 || remain > int(lockoutDuration.Seconds())+1 {
		t.Errorf("剩余锁定秒数 = %d，期望约 %d", remain, int(lockoutDuration.Seconds()))
	}
}

// TestLoginLimiterUnlocksAfterDuration 钉住另一个方向：锁期真的会过。
// 只修「不锁」而不修「永锁」同样是坏的 —— 那会把账号永久焊死。
func TestLoginLimiterUnlocksAfterDuration(t *testing.T) {
	setupLoginLimiterDB(t)
	ctx := context.Background()
	const name = "victim"

	for i := 0; i < maxLoginFailures; i++ {
		limiter.Fail(ctx, name)
	}
	if locked, _ := limiter.Locked(ctx, name); !locked {
		t.Fatalf("记满 %d 次失败后未锁定", maxLoginFailures)
	}

	// 把整桶推到 lockoutDuration 之前，等价于「时间过去了」。
	// 整桶设成同一时间戳：count 和最早时间才指向同一刻（见文件头第 2 条）。
	past := time.Now().Add(-lockoutDuration - time.Second)
	if _, err := db.DB.ExecContext(ctx,
		`UPDATE rate_limits SET ts = $1 WHERE bucket = $2`,
		past, "login:"+name); err != nil {
		t.Fatalf("回拨时间戳: %v", err)
	}
	// 确认回拨真的命中了（第一条坑：静默 0 行会伪装成「限速器有 bug」）
	var n int
	if err := db.DB.QueryRowContext(ctx,
		`SELECT COUNT(*) FROM rate_limits WHERE bucket = $1 AND ts >= $2`,
		"login:"+name, past).Scan(&n); err != nil {
		t.Fatalf("核对回拨: %v", err)
	}
	if n != maxLoginFailures {
		t.Fatalf("回拨后只匹配到 %d/%d 行 —— UPDATE 没落到预期的行上", n, maxLoginFailures)
	}
	if locked, remain := limiter.Locked(ctx, name); locked {
		t.Errorf("锁期已过仍判定为锁定（remain=%d），会把账号永久焊死", remain)
	}
}

// TestLoginLimiterWindowExcludesOldFailures 钉住窗口语义：
// 15 分钟窗口**之前**的失败不该计入。少了这条，一次陈旧的失败堆积
// 就能把一个从没被攻击过的账号锁住。
func TestLoginLimiterWindowExcludesOldFailures(t *testing.T) {
	setupLoginLimiterDB(t)
	ctx := context.Background()
	const name = "victim"

	// 全部失败都在窗口之前
	for i := 0; i < maxLoginFailures+3; i++ {
		if _, err := db.DB.ExecContext(ctx,
			`INSERT INTO rate_limits (bucket, ts) VALUES ($1, $2)`,
			"login:"+name, time.Now().Add(-failureWindow-time.Duration(i+1)*time.Minute)); err != nil {
			t.Fatalf("seed: %v", err)
		}
	}
	if locked, remain := limiter.Locked(ctx, name); locked {
		t.Errorf("窗口外的历史失败不该触发锁定，但 locked=true remain=%d", remain)
	}
}

// TestLoginLimiterResetClearsBucket 是登录成功后的清桶路径。
// 没有它，登录成功一次之后仍会被先前的失败锁住。
func TestLoginLimiterResetClearsBucket(t *testing.T) {
	setupLoginLimiterDB(t)
	ctx := context.Background()
	const name = "victim"

	for i := 0; i < maxLoginFailures; i++ {
		limiter.Fail(ctx, name)
	}
	if locked, _ := limiter.Locked(ctx, name); !locked {
		t.Fatalf("记满 %d 次失败后未锁定", maxLoginFailures)
	}
	limiter.Reset(ctx, name)
	if locked, _ := limiter.Locked(ctx, name); locked {
		t.Error("Reset 之后仍处于锁定")
	}
}