开放原子开发者工作坊

beego 限速

路由之前调用

import "github.com/ulule/limiter/v3"
import "github.com/ulule/limiter/v3/drivers/store/memory"


type rateLimiter struct {
	generalLimiter *limiter.Limiter
	loginLimiter   *limiter.Limiter
}

func main(){
	r := &rateLimiter{}

	rate, err := limiter.NewRateFromFormatted("5-S")
	PanicOnError(err)
	r.generalLimiter = limiter.New(memory.NewStore(), rate)

	loginRate, err := limiter.NewRateFromFormatted("5-M")
	PanicOnError(err)
	r.loginLimiter = limiter.New(memory.NewStore(), loginRate)

	//More on Beego filters here https://beego.me/docs/mvc/controller/filter.md
	beego.InsertFilter("/*", beego.BeforeRouter, func(c *context.Context) {
		rateLimit(r, c)
	}, true)
	
	beego.Run()
}

具体实现

func rateLimit(r *rateLimiter, ctx *context.Context) {
	var (
		limiterCtx limiter.Context
		ip         net.IP
		err        error
		req        = ctx.Request
	)

	if strings.HasPrefix(ctx.Input.URL(), "/v1/user/login") {
		ip = r.loginLimiter.GetIP(req)
		limiterCtx, err = r.loginLimiter.Get(req.Context(), ip.String())
	} else {
		ip = r.generalLimiter.GetIP(req)
		limiterCtx, err = r.generalLimiter.Get(req.Context(), ip.String())
	}
	if err != nil {
		PanicOnError(err)
	}
	h := ctx.ResponseWriter.Header()
	h.Add("X-RateLimit-Limit", strconv.FormatInt(limiterCtx.Limit, 10))
	h.Add("X-RateLimit-Remaining", strconv.FormatInt(limiterCtx.Remaining, 10))
	h.Add("X-RateLimit-Reset", strconv.FormatInt(limiterCtx.Reset, 10))

	if limiterCtx.Reached {
		fmt.Println("Too Many Requests from %s on %s", ip, ctx.Input.URL())
		ctx.Output.SetStatus(429)
		resBytes, err := json.Marshal(controllers.OutResponse(429, nil, ""))
		if err != nil {
			fmt.Println(err)
		}
		ctx.Output.Body(resBytes)
		return
	}

}
``