26 lines
604 B
Go
26 lines
604 B
Go
package limit
|
|
|
|
import (
|
|
"fmt"
|
|
"stat/common/redis"
|
|
"stat/statg"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func GetKey(name string) func() string {
|
|
return func() string {
|
|
return redis.GetIpLimitDayKey(statg.Conf.Host, statg.Conf.Port) + "_" + name
|
|
}
|
|
}
|
|
|
|
func Limit(name string, capacity int64, rate int64) gin.HandlerFunc {
|
|
limiter := redis.NewTokenBucket(redis.StatRds.Client(), GetKey(name), capacity, rate, 48*time.Hour)
|
|
return func(ctx *gin.Context) {
|
|
if err := limiter.TakeTokens(ctx.ClientIP(), 1); err != nil {
|
|
ctx.AbortWithStatusJSON(200, gin.H{"msg": fmt.Sprintf("%+v", err)})
|
|
}
|
|
}
|
|
}
|