码迷,mamicode.com
首页 > 其他好文 > 详细

beego 限速

时间:2020-05-19 10:53:40      阅读:79      评论:0      收藏:0      [点我收藏+]

标签:from   error   limit   print   body   ini   control   log   new   

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
	}

}
``

beego 限速

标签:from   error   limit   print   body   ini   control   log   new   

原文地址:https://www.cnblogs.com/tomtellyou/p/12915313.html

(0)
(0)
   
举报
评论 一句话评论(0
登录后才能评论!
© 2014 mamicode.com 版权所有  联系我们:gaon5@hotmail.com
迷上了代码!