diff --git a/internal/mcp/ratelimit.go b/internal/mcp/ratelimit.go new file mode 100644 index 0000000..99fe232 --- /dev/null +++ b/internal/mcp/ratelimit.go @@ -0,0 +1,117 @@ +package mcp + +import ( + "net/http" + "sync" + "time" + + "github.com/yan1h/agent-coding-workflow/internal/infra/errs" + httpx "github.com/yan1h/agent-coding-workflow/internal/transport/http" + "github.com/yan1h/agent-coding-workflow/internal/transport/http/middleware" +) + +// RateLimiter 是一对 token bucket:per-token + 全局。 +// +// 实现选择:固定时间窗口 + 计数器(每分钟重置)。简单、无第三方依赖、 +// 内存占用极低。10–50 人规模 + 默认 100 req/min/token 不需要更精细的 +// sliding window / leaky bucket。 +// +// 并发:sync.Mutex + map;rate-limit 调用是热路径(每请求一次),但 mutex +// 持有时间极短(map lookup + 1 个 if)。 +type RateLimiter struct { + cfg RateLimitConfig + mu sync.Mutex + bucket map[string]*tokenBucket + global *tokenBucket + now func() time.Time // 测试注入 +} + +// RateLimitConfig 是 limiter 的配置。 +type RateLimitConfig struct { + PerTokenPerMinute int + GlobalPerMinute int +} + +type tokenBucket struct { + count int + windowAt time.Time +} + +// NewRateLimiter 用配置构造 limiter;now=nil 时用 time.Now。 +func NewRateLimiter(cfg RateLimitConfig, now func() time.Time) *RateLimiter { + if now == nil { + now = time.Now + } + return &RateLimiter{ + cfg: cfg, + bucket: map[string]*tokenBucket{}, + global: &tokenBucket{}, + now: now, + } +} + +// Allow 返回 (allowed, reason)。reason 仅在 !allowed 时有值, +// 取值 ∈ {"per_token", "global"} 用于错误消息区分。 +func (l *RateLimiter) Allow(tokenID string) (bool, string) { + l.mu.Lock() + defer l.mu.Unlock() + now := l.now() + + // per-token + b := l.bucket[tokenID] + if b == nil { + b = &tokenBucket{} + l.bucket[tokenID] = b + } + if !sameMinute(b.windowAt, now) { + b.windowAt = now + b.count = 0 + } + if b.count >= l.cfg.PerTokenPerMinute { + return false, "per_token" + } + + // global + if !sameMinute(l.global.windowAt, now) { + l.global.windowAt = now + l.global.count = 0 + } + if l.global.count >= l.cfg.GlobalPerMinute { + return false, "global" + } + + // 两层桶都通过:增量 + b.count++ + l.global.count++ + return true, "" +} + +// sameMinute 判定两时间是否在同一分钟窗口内。 +func sameMinute(a, b time.Time) bool { + return a.Truncate(time.Minute).Equal(b.Truncate(time.Minute)) +} + +// NewRateLimitMiddleware 必须挂载在 NewAuthMiddleware 之后。 +// 从 ctx 取 AuthSession.TokenID 作为 bucket key;找不到 = middleware 顺序错, +// 返 500(避免 fail-open 让无鉴权请求通过 rate limit)。 +func NewRateLimitMiddleware(limiter *RateLimiter) func(http.Handler) http.Handler { + return func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + sess, ok := FromContext(r.Context()) + if !ok { + httpx.WriteError(w, middleware.RequestIDFromContext(r.Context()), + errs.New(errs.CodeInternal, + "rate limit middleware: missing auth session (middleware order bug)")) + return + } + allowed, reason := limiter.Allow(sess.TokenID) + if !allowed { + httpx.WriteError(w, middleware.RequestIDFromContext(r.Context()), + errs.New(errs.CodeMcpRateLimited, + "rate limit exceeded ("+reason+")")) + return + } + next.ServeHTTP(w, r) + }) + } +}