// Copyright (C) 2024 veypi // 2025-03-04 16:08:06 // Distributed under terms of the MIT license. package middleware import ( "fmt" "time" "github.com/veypi/vbase/libs/cache" "github.com/veypi/vigo" ) // RateLimiter 限流中间件 func RateLimiter(requests int, window time.Duration) func(*vigo.X) error { return func(x *vigo.X) error { // 获取标识(优先用户ID,其次IP) identifier := "" if uid, ok := x.Get("user_id").(string); ok && uid != "" { identifier = uid } else { identifier = x.GetRemoteIP() } key := fmt.Sprintf("ratelimit:%s:%s", identifier, x.Request.URL.Path) count, err := cache.Incr(key) if err != nil { // Redis未启用,跳过限流 return nil } // 第一次请求设置过期时间 if count == 1 { cache.Expire(key, window) } if count > int64(requests) { return vigo.ErrForbidden.WithString("rate limit exceeded") } return nil } }