package utils import ( "github.com/astaxie/beego" "github.com/astaxie/beego/context" "github.com/astaxie/beego/logs" "github.com/ulule/limiter" "github.com/ulule/limiter/v3" "github.com/ulule/limiter/v3/drivers/store/memory" "net/http" "strings" ) // RateLimiter this is a struct type RateLimiter struct { Limiter *limiter.Limiter Username string UserType string UserToken string RemainTimes int MaxTimes int } func RateLimit(rateLimit *RateLimiter, ctx *context.Context) { var ( limiterCtx limiter.Context err error req = ctx.Request ) opt := limiter.Options{ IPv4Mask: limiter.DefaultIPv4Mask, IPv6Mask: limiter.DefaultIPv6Mask, TrustForwardHeader: false, } ip := limiter.GetIP(req, opt) if strings.HasPrefix(ctx.Input.URL(), "/") { limiterCtx, err = rateLimit.Limiter.Get(req.Context(), ip.String()) } else { logs.Info("The api request is not track ") } if err != nil { ctx.Abort(http.StatusInternalServerError, err.Error()) return } if limiterCtx.Reached { logs.Debug("Too Many Requests from %s on %s", ip, ctx.Input.URL()) // refer to https://beego.me/docs/mvc/controller/errors.md for error handling ctx.Abort(http.StatusTooManyRequests, "429") return } } func PanicError(e error) { if e != nil { panic(e) } } func RunRate() { // 限制每秒登錄的請求次數 theRateLimit := &RateLimiter{} // 100 reqs/second: "100-S" "100-s" loginMaxRate := beego.AppConfig.String("total_rule::reqrate") loginRate, err := limiter.NewRateFromFormatted(loginMaxRate + "-s") PanicError(err) theRateLimit.Limiter = limiter.New(memory.NewStore(), loginRate) beego.InsertFilter("/*", beego.BeforeRouter, func(ctx *context.Context) { RateLimit(theRateLimit, ctx) }, true) }
在main.go 里面調用方法即可