rate limiter

conf/xpkg.yaml

@@ -25,7 +25,14 @@ app:
     enable: true
     window: 1m
     maxRequests: 600
-    sync: true
+    type: local # local or sync
+    store:
+      type: redis # 目前分布式仅支持redis
+      redis:
+        dsn: "default:password123@localhost:30079/6"
+        dialTimeout: 10
+        readTimeout: 2
+        writeTimeout: 2
   hideSensitiveFields: "smtpPassword"
   idGeneration:
     type: leaf # snowflake leaf

config/config.go

@@ -274,7 +274,7 @@ type Trace struct {
 
 type IPRateLimiter struct {
 	Enable      bool             `yaml:"enable" json:"enable"`
-	Sync        bool             `yaml:"sync" json:"sync"`
+	Type        string           `yaml:"type" json:"type"` // redis or local
 	Window      string           `yaml:"window" json:"window"`
 	MaxRequests int              `yaml:"maxRequests" json:"maxRequests"`
 	Store       RateLimiterStore `yaml:"store" json:"store"`
@@ -282,9 +282,7 @@ type IPRateLimiter struct {
 
 type RateLimiterStore struct {
 	Type string `yaml:"redis" json:"redis"`
-	Addr     string `yaml:"addr" json:"addr"`
-	Password string `yaml:"password" json:"password"`
-	DB       int    `yaml:"db" json:"db"`
+	Dsn  string `yaml:"dsn" json:"dsn"`
 }
 
 func (x IPRateLimiter) Get() (time.Duration, int) {
@@ -301,14 +299,6 @@ func (x IPRateLimiter) Get() (time.Duration, int) {
 	return window, x.MaxRequests
 }
 
-func (x IPRateLimiter) GetStore() (string, string, int) {
-	addr := "localhost:6379"
-	if x.Store.Addr != "" {
-		addr = x.Store.Addr
-	}
-	return addr, x.Store.Password, x.Store.DB
-}
-
 type GrpcClient struct {
 	ClientSecure          ClientSecure `yaml:"clientSecure" json:"clientSecure"`
 	ClientToken           ClientToken  `yaml:"clientToken" json:"clientToken"`

gin/middleware/rate_limiter.go

 package middleware
 
 import (
+	"context"
 	"net/http"
 	"time"
 
@@ -10,16 +11,22 @@ import (
 	srl "gitlab.wanzhuangkj.com/tush/xpkg/gin/middleware/sync_ratelimiter"
 	"gitlab.wanzhuangkj.com/tush/xpkg/gin/response"
 	"gitlab.wanzhuangkj.com/tush/xpkg/pkg/errcode"
+	"gitlab.wanzhuangkj.com/tush/xpkg/pkg/eventbus"
 	"gitlab.wanzhuangkj.com/tush/xpkg/pkg/ips"
 	"gitlab.wanzhuangkj.com/tush/xpkg/pkg/logger"
 	"gitlab.wanzhuangkj.com/tush/xpkg/utils/ctxUtils"
 )
 
+var (
+	ipRateLimiter      *ratelimiter.IPRateLimiter
+	redisIPRateLimiter *srl.IPRateLimiter
+)
+
 func IPRateLimiter(windowSize time.Duration, maxRequests int) func(c *gin.Context) {
-	rl := ratelimiter.NewIPRateLimiter(windowSize, maxRequests)
+	ipRateLimiter = ratelimiter.NewIPRateLimiter(windowSize, maxRequests)
 	return func(c *gin.Context) {
 		clientIP := ips.GetClientIP(c)
-		if rl.Allow(clientIP) {
+		if ipRateLimiter.Allow(clientIP) {
 			c.Next()
 		} else {
 			Fail(c, errcode.TooManyRequests.Code(), errcode.TooManyRequests.Msg())
@@ -27,17 +34,22 @@ func IPRateLimiter(windowSize time.Duration, maxRequests int) func(c *gin.Contex
 	}
 }
 
-func SyncIPRateLimiter(windowSize time.Duration, maxRequests int64, cli redis.Cmdable) func(c *gin.Context) {
-	rl := srl.NewIPRateLimiter(
+func SyncIPRateLimiter(windowSize time.Duration, maxRequests int, cli *redis.Client) func(c *gin.Context) {
+	redisIPRateLimiter = srl.NewIPRateLimiter(
 		srl.WithWindowSize(windowSize),
-		srl.WithMaxRequests(maxRequests),
+		srl.WithMaxRequests(int64(maxRequests)),
 		srl.WithLogger(logger.Get()),
 		srl.WithRedisCli(cli),
 	)
+	eventbus.Eb.Subscribe(eventbus.TopicApplicationClose, func(ctx context.Context) {
+		if redisIPRateLimiter != nil {
+			redisIPRateLimiter.Close()
+		}
+	})
 	return func(c *gin.Context) {
 		clientIP := ips.GetClientIP(c)
 		ctx := ctxUtils.WrapCtx(c)
-		if rl.Allow(ctx, clientIP) {
+		if redisIPRateLimiter.Allow(ctx, clientIP) {
 			c.Next()
 		} else {
 			Fail(c, errcode.TooManyRequests.Code(), errcode.TooManyRequests.Msg())

gin/middleware/sync_ratelimiter/sync_ratelimiter.go

@@ -3,12 +3,14 @@ package sw
 import (
 	"context"
 	"fmt"
-	"github.com/redis/go-redis/v9"
-	"gitlab.wanzhuangkj.com/tush/xpkg/utils/ctxUtils"
-	"go.uber.org/zap"
 	"strconv"
 	"sync"
 	"time"
+
+	"github.com/redis/go-redis/v9"
+	"gitlab.wanzhuangkj.com/tush/xpkg/pkg/goredis"
+	"gitlab.wanzhuangkj.com/tush/xpkg/utils/ctxUtils"
+	"go.uber.org/zap"
 )
 
 type RedisDatastore struct {
@@ -66,6 +68,13 @@ type IPRateLimiter struct {
 	options *IPRateLimiterOptions
 	store   *RedisDatastore
 	ips     map[string]*SLimiter
+	cli     *redis.Client
+}
+
+func (x *IPRateLimiter) Close() {
+	if x.cli != nil {
+		goredis.Close(x.cli)
+	}
 }
 
 type IPRateLimiterOptions struct {
@@ -75,7 +84,7 @@ type IPRateLimiterOptions struct {
 	redisDB       int
 	redisPassword string
 	logger        *zap.Logger
-	redisCli      redis.Cmdable
+	redisCli      *redis.Client
 }
 
 func defaultOptions() *IPRateLimiterOptions {
@@ -129,7 +138,7 @@ func WithLogger(log *zap.Logger) Option {
 	}
 }
 
-func WithRedisCli(redisCli redis.Cmdable) Option {
+func WithRedisCli(redisCli *redis.Client) Option {
 	return func(o *IPRateLimiterOptions) {
 		o.redisCli = redisCli
 	}
@@ -146,14 +155,16 @@ func NewIPRateLimiter(opts ...Option) *IPRateLimiter {
 		ips:     make(map[string]*SLimiter),
 	}
 	if options.redisCli != nil {
+		rl.cli = options.redisCli
 		rl.store = NewRedisDatastore(options.redisCli, 2*options.windowSize, options.logger)
 	} else {
-		rl.store = NewRedisDatastore(
-			redis.NewClient(&redis.Options{
+		rl.cli = redis.NewClient(&redis.Options{
 			Addr:     options.redisAddr,
 			Password: options.redisPassword,
 			DB:       options.redisDB,
-			}),
+		})
+		rl.store = NewRedisDatastore(
+			rl.cli,
 			2*options.windowSize,
 			options.logger,
 		)

gin/xgin/xgin.go

@@ -20,6 +20,7 @@ import (
 	"gitlab.wanzhuangkj.com/tush/xpkg/gin/validator"
 	"gitlab.wanzhuangkj.com/tush/xpkg/pkg/errcode"
 	"gitlab.wanzhuangkj.com/tush/xpkg/pkg/eventbus"
+	"gitlab.wanzhuangkj.com/tush/xpkg/pkg/goredis"
 	"gitlab.wanzhuangkj.com/tush/xpkg/pkg/ips"
 	"gitlab.wanzhuangkj.com/tush/xpkg/pkg/jwt"
 	"gitlab.wanzhuangkj.com/tush/xpkg/pkg/logger"
@@ -85,7 +86,22 @@ func New() *gin.Engine {
 	if app.Middleware.Trace.Enable {
 		r.Use(middleware.Tracing(app.Name))
 	}
-
+	if app.IPRateLimiter.Enable {
+		if app.IPRateLimiter.Type == "local" {
+			window, max := app.IPRateLimiter.Get()
+			r.Use(middleware.IPRateLimiter(window, max))
+		}
+		if app.IPRateLimiter.Type == "sync" {
+			if app.IPRateLimiter.Store.Type == "redis" {
+				redisCli, err := goredis.Init(app.IPRateLimiter.Store.Dsn, goredis.WithLogger(logger.Get()))
+				if err != nil {
+					logger.Fatal("init redis error: " + err.Error())
+				}
+				window, max := app.IPRateLimiter.Get()
+				r.Use(middleware.SyncIPRateLimiter(window, max, redisCli))
+			}
+		}
+	}
 	if app.Middleware.EnableHTTPProfile {
 		prof.Register(r, prof.WithIOWaitTime())
 	}