beego-接口频率限制

beego-接口频率限制

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 里面调用方法即可



posted @ Jeff的技术栈 阅读 (9999+) 评论 (99) 点赞(99) 编辑 MD 收藏
编辑推荐:
· DeepSeek 解答了困扰我五年的技术问题
· 为什么说在企业级应用开发中,后端往往是效率杀手?
· 用 C# 插值字符串处理器写一个 sscanf
· Java 中堆内存和栈内存上的数据分布和特点
· 开发中对象命名的一点思考
阅读排行:
· 为什么说在企业级应用开发中,后端往往是效率杀手?
· DeepSeek 解答了困扰我五年的技术问题。时代确实变了!
· 本地部署DeepSeek后,没有好看的交互界面怎么行!
· 趁着过年的时候手搓了一个低代码框架
· 推荐一个DeepSeek 大模型的免费 API 项目!兼容OpenAI接口!
回顶部
点击右上角即可分享
微信分享提示