Files
MailUI4Agents/gateway/internal/handler/ratelimit.go
JianFeeeee 9d4718a412 feat: SSE Last-Event-ID 补投 + 连接状态指示 + 限速器 DB 化
## SSE Last-Event-ID 补投

EventSource 断线重连时自带 Last-Event-ID 头,但服务端直接忽略了——
所有断线期间的邮件通知都丢失。用户刷新页面也会错过已推的事件。

改为 per-user 事件环形缓冲区(500 条,~100KB/用户,20 在线 ≈ 2MB):
每次 Broadcast/SendToUser/SendToAgent 同时写入对应用户的缓冲区;
AddClient 时取 Last-Event-ID 头,找到该 ID 的位置后从下一条回放。
找不到 ID 说明事件已被覆盖(缓冲区溢出),从头回放全部。

事件 ID 用全局递增序列号(非 UUID),EventSource 的 Last-Event-ID
就是靠这个 ID 记住断点的。

新增测试:缓冲区回放、溢出行为、并发安全(10 goroutine × 200 次 push)、
端到端重连验证(SendToUser → 带 Last-Event-ID 的 AddClient → 补投)。

## 连接状态指示器

Sidebar 用户头像右下角的小圆点:绿=已连接,黄=连接中,橙=重连中,红=断开。
NarrowNav 底栏也有(移动端)。

SSE 模块新增 onSSEStatus/getSSEStatus 接口,onerror/onopen 驱动状态变化。
状态点用 absolute 定位在头像边缘,不遮挡文字。

## 限速器 DB 化(解决多实例部署时的计数漂移)

原实现:LoginLimiter 与 sessionRateLimiter 都是进程内内存计数器。
多实例部署时各自独立计数,等效上限变成 N 倍。

改为 rate_limits 表(bucket + ts),两个限速器共享同一套基础设施:
- LoginLimiter:bucket="login:<username>",COUNT(*) >= 5 → 锁定 5 分钟
- sessionRateLimiter:bucket="session:<agent_name>",COUNT(*) >= 20/h → 拒绝

判断与写入在同一个 BEGIN IMMEDIATE 事务里——SQLite 的 IMMEDIATE
在事务开始时获取 RESERVED 锁,防并发写事务同时进入 COMMIT 阶段。
实测 80 并发下恰好放行 20 次(旧内存版同样通过,但 DB 版才能多实例共享)。

DB 不可用时放行(宁可放开限速也不能让用户完全无法使用)。
新建 rate_limits 表迁移(SQLite + PG 两版)。
2026-09-02 14:33:41 +08:00

76 lines
2.1 KiB
Go

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)
}