Initial commit: GovAI 政务AI平台

server/internal/middleware/ratelimit.go

@@ -0,0 +1,39 @@
+package middleware
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"time"
+
+	"github.com/enterprise-ai-platform/server/internal/response"
+	"github.com/redis/go-redis/v9"
+)
+
+// RateLimit creates a per-user rate limiter using Redis sliding window.
+func RateLimit(rdb *redis.Client, maxRequests int, window time.Duration) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			userID := GetUserID(r.Context())
+			key := fmt.Sprintf("rl:%s:%s", userID.String(), r.URL.Path)
+
+			ctx := context.Background()
+			count, err := rdb.Incr(ctx, key).Result()
+			if err != nil {
+				next.ServeHTTP(w, r)
+				return
+			}
+
+			if count == 1 {
+				rdb.Expire(ctx, key, window)
+			}
+
+			if count > int64(maxRequests) {
+				response.TooManyRequests(w, "请求过于频繁，请稍后再试")
+				return
+			}
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}