feat(mcp): rate limit middleware (per-token + global token bucket, 1-min window)

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