bl/login/internal/cmd/cmd.go

package cmd

import (
	"context"
	"runtime"

	"blazing/cool"

	i18n "blazing/modules/base/middleware"

	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/net/ghttp"
	"github.com/gogf/gf/v2/os/gcmd"
	"github.com/gogf/gf/v2/os/gfile"
	"github.com/xiaoqidun/limit"
	"golang.org/x/time/rate"
)

var (
	limiter = limit.New()
	Main    = gcmd.Command{
		Name:  "main",
		Usage: "main",
		Brief: "start http server",
		Func: func(ctx context.Context, parser *gcmd.Parser) (err error) {
			// g.Dump(g.DB("test").GetConfig())
			if cool.IsRedisMode {
				go cool.ListenFunc(ctx)
			}
			//go robot()
			s := g.Server()
			s.Use(Limiter, ghttp.MiddlewareHandlerResponse)
			s.EnableAdmin()
			s.SetServerAgent(cool.Config.Name)
			s.BindHookHandler("/*", ghttp.HookBeforeServe, beforeServeHook)
			runtime.SetMutexProfileFraction(1) // (非必需)开启对锁调用的跟踪
			runtime.SetBlockProfileRate(1)     // (非必需)开启对阻塞操作的跟踪
			s.EnablePProf()
			// 如果存在 data/cool-admin-vue/dist 目录，则设置为主目录
			if gfile.IsDir("public") {
				s.SetServerRoot("public")
			}
			// i18n 信息
			s.BindHandler("/i18n", i18n.I18nInfo)
			//	g.Server().BindMiddleware("/*", MiddlewareCORS)
			s.Run()
			return nil
		},
	}
)

func beforeServeHook(r *ghttp.Request) {
	//glog.Debugf(r.GetCtx(), "beforeServeHook [is file:%v] URI:%s", r.IsFileRequest(), r.RequestURI)
	r.Response.CORSDefault()
}

// Limiter is a middleware that implements rate limiting for all HTTP requests.
// It returns HTTP 429 (Too Many Requests) when the rate limit is exceeded.
func Limiter(r *ghttp.Request) {
	// 3. 为任意键 "some-key" 获取一个速率限制器
	//    - rate.Limit(2): 表示速率为 "每秒2个请求"
	//    - 2: 表示桶的容量 (Burst)，允许瞬时处理2个请求
	rateLimiter := limiter.Get(r.GetClientIp(), rate.Limit(10), 5)

	if !rateLimiter.Allow() {
		r.Response.WriteStatusExit(429) // Return 429 Too Many Requests
		r.ExitAll()
	}
	r.Middleware.Next()
}