freedakgmail
40 lines
918 B
Go
40 lines
918 B
Go
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)
|
|
})
|
|
}
|
|
}
|