52 lines
1.1 KiB
Go
52 lines
1.1 KiB
Go
package ratelimit
|
||
|
||
import (
|
||
"fmt"
|
||
"net/http"
|
||
|
||
"github.com/gin-gonic/gin"
|
||
|
||
"github.com/hhs/camtalk/internal/trace"
|
||
)
|
||
|
||
// Middleware 返回 Gin 中间件，按 key 维度限流。
|
||
// keyFunc 从请求中提取限流 key（如 IP、用户 ID）。
|
||
func Middleware(limiter Limiter, keyFunc func(*gin.Context) string) gin.HandlerFunc {
|
||
return func(c *gin.Context) {
|
||
if limiter == nil {
|
||
c.Next()
|
||
return
|
||
}
|
||
|
||
key := keyFunc(c)
|
||
if key == "" {
|
||
// key 为空时跳过限流
|
||
c.Next()
|
||
return
|
||
}
|
||
|
||
allowed, retryAfter := limiter.Allow(c.Request.Context(), key)
|
||
|
||
if !allowed {
|
||
log := trace.FromContext(c.Request.Context())
|
||
log.Warnw("rate limited",
|
||
"client_ip", c.ClientIP(),
|
||
"path", c.Request.URL.Path,
|
||
"limit_key", key,
|
||
"retry_after_sec", int(retryAfter.Seconds()+0.5))
|
||
|
||
// 设置 Retry-After header（秒）
|
||
c.Header("Retry-After", fmt.Sprintf("%d", int(retryAfter.Seconds()+0.5)))
|
||
|
||
c.JSON(http.StatusTooManyRequests, gin.H{
|
||
"code": "RATE_LIMITED",
|
||
"message": fmt.Sprintf("too many requests, retry after %s", retryAfter.Round(1)),
|
||
})
|
||
c.Abort()
|
||
return
|
||
}
|
||
|
||
c.Next()
|
||
}
|
||
}
|