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 之后仍处于锁定") } }