42 lines
1.2 KiB
Go
42 lines
1.2 KiB
Go
package middleware
|
||
|
||
import (
|
||
"time"
|
||
|
||
"github.com/gofiber/fiber/v2"
|
||
"github.com/gofiber/fiber/v2/middleware/limiter"
|
||
"github.com/gofiber/storage/redis/v3"
|
||
|
||
"github.com/break/junhong_cmp_fiber/pkg/constants"
|
||
"github.com/break/junhong_cmp_fiber/pkg/errors"
|
||
"github.com/break/junhong_cmp_fiber/pkg/response"
|
||
)
|
||
|
||
// RateLimiter 创建基于 IP 的限流中间件
|
||
// storage 参数：nil = 内存存储，传入 Redis storage = 分布式限流
|
||
func RateLimiter(max int, expiration time.Duration, storage fiber.Storage) fiber.Handler {
|
||
return limiter.New(limiter.Config{
|
||
Max: max,
|
||
Expiration: expiration,
|
||
KeyGenerator: func(c *fiber.Ctx) string {
|
||
// 使用统一的 Redis 键生成函数
|
||
return constants.RedisRateLimitKey(c.IP())
|
||
},
|
||
LimitReached: func(c *fiber.Ctx) error {
|
||
return response.Error(c, 400, errors.CodeTooManyRequests, errors.GetMessage(errors.CodeTooManyRequests, "zh"))
|
||
},
|
||
Storage: storage, // 支持内存或 Redis 存储
|
||
})
|
||
}
|
||
|
||
// NewRedisStorage 创建 Redis 存储用于限流
|
||
func NewRedisStorage(addr, password string, db, prot int) fiber.Storage {
|
||
return redis.New(redis.Config{
|
||
Host: addr,
|
||
Port: prot,
|
||
Password: password,
|
||
Database: db,
|
||
Reset: false,
|
||
})
|
||
}
|