✨ feat: added config Rate Limiting and Throttling #15

Author: pnguyen215

config/conf-params.yaml

@@ -101,3 +101,16 @@ curl:
       chat_id:
       - 123456789
       token: <token_here>
+# ################################
+# Rate Limit Seekers Config
+# 2023-11-25 12:02:54
+# ################################
+rate-limit-seekers:
+- key: "psql_rate"
+  usable_default: false 
+  config:
+    enabled: false
+    rate: 2
+    max_burst: 1
+  option:
+    max_retries: 2

go.mod

@@ -15,6 +15,7 @@ require (
 	github.com/sivaosorg/redisconn v1.0.2
 	github.com/sivaosorg/wsconn v1.0.3-beta.1
 	github.com/swaggo/swag v1.16.2
+	golang.org/x/time v0.4.0
 )
 
 require (

go.sum

@@ -208,6 +208,8 @@ golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
 golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
 golang.org/x/text v0.13.0 h1:ablQoSUd0tRdKxZewP80B+BaqeKJuVhuRxj/dkrun3k=
 golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
+golang.org/x/time v0.4.0 h1:Z81tqI5ddIoXDPvVQ7/7CC9TnLM7ubaFG2qXYd5BbYY=
+golang.org/x/time v0.4.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
 golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=

internal/core/routes.go

@@ -16,6 +16,8 @@ func (c *CoreCommand) routes(core *gin.Engine) {
 		v1.GET("/common/psql-status",
 			c.handlers.middlewares.RequestMiddleWare(),
 			c.handlers.middlewares.NoopMiddleWare(),
+			c.handlers.middlewares.RateLimitMiddleWare("psql_rate"),
+			c.handlers.middlewares.NetMiddleware(),
 			c.handlers.commonHandler.OnPsqlStatus)
 		v1.GET("/common/consumer", // endpoint websocket: ws://127.0.0.1:8081/api/v1/common/consumer
 			c.handlers.middlewares.RequestMiddleWare(),

internal/middlewares/throttling_middleware.go

@@ -0,0 +1,94 @@
+package middlewares
+
+import (
+	"fmt"
+	"net"
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	syncconf "github.com/sivaosorg/gocell/internal/syncConf"
+	"github.com/sivaosorg/govm/entity"
+	"github.com/sivaosorg/govm/logger"
+	"github.com/sivaosorg/govm/ratelimitx"
+	"golang.org/x/time/rate"
+)
+
+// https://blog.logrocket.com/rate-limiting-go-application/
+type client struct {
+	limiter *rate.Limiter
+	last    time.Time
+}
+
+var (
+	mu                   sync.Mutex
+	clients              = make(map[string]*client)
+	cleanupTypeWhitelist = time.Minute
+	cleanupWhitelist     = 2 * cleanupTypeWhitelist
+)
+
+func (m *MiddlewareManager) RateLimitMiddleWare(key string) gin.HandlerFunc {
+	rates := syncconf.Params.RateLimits
+	clusters := ratelimitx.NewClusterMultiTenantRateLimitConfig().SetClusters(rates)
+	go m.cleanupWhitelist()
+	return func(c *gin.Context) {
+		conf, err := clusters.FindClusterBy(key)
+		if err != nil || !conf.Config.IsEnabled {
+			return
+		}
+		ip, err := m.decodeNetwork(c)
+		if err != nil {
+			response := entity.NewResponseEntity().BadRequest(http.StatusText(http.StatusBadRequest), nil)
+			response.SetError(err)
+			c.JSON(response.StatusCode, response)
+			c.Abort()
+			return
+		}
+		mu.Lock()
+		if _, found := clients[ip]; !found {
+			clients[ip] = &client{limiter: rate.NewLimiter(rate.Limit(conf.Config.Rate), conf.Config.MaxBurst)}
+		}
+		clients[ip].last = time.Now()
+		if !clients[ip].limiter.Allow() {
+			mu.Unlock()
+			response := entity.NewResponseEntity().TooManyRequest(http.StatusText(http.StatusTooManyRequests), nil)
+			c.JSON(response.StatusCode, response)
+			c.Abort()
+			return
+		}
+		mu.Unlock()
+		c.Next()
+	}
+}
+
+func (m *MiddlewareManager) NetMiddleware() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		ip, err := m.decodeNetwork(c)
+		if err != nil {
+			m.notification(c, err, http.StatusBadRequest)
+		}
+		logger.Debugf(fmt.Sprintf("_endpoint: %v, incoming on IP: %v", c.Request.RequestURI, ip))
+		c.Next()
+	}
+}
+
+func (m *MiddlewareManager) decodeNetwork(c *gin.Context) (ip string, err error) {
+	ip, _, err = net.SplitHostPort(c.Request.RemoteAddr)
+	return ip, err
+}
+
+func (m *MiddlewareManager) cleanupWhitelist() {
+	for {
+		time.Sleep(cleanupTypeWhitelist)
+		mu.Lock()
+		for ip, client := range clients {
+			remain := time.Since(client.last)
+			if remain > cleanupWhitelist {
+				logger.Debugf(fmt.Sprintf("Cleanup whitelist too many requests for IP: %v and remain duration: %v", ip, remain))
+				delete(clients, ip)
+			}
+		}
+		mu.Unlock()
+	}
+}

internal/syncConf/syncConf.go

@@ -6,6 +6,7 @@ import (
 	"github.com/sivaosorg/govm/apix"
 	"github.com/sivaosorg/govm/coltx"
 	"github.com/sivaosorg/govm/configx"
+	"github.com/sivaosorg/govm/ratelimitx"
 	"github.com/sivaosorg/govm/utils"
 )
 
@@ -15,7 +16,8 @@ var Params *keyParams
 var Jobs *jobParams
 
 type keyParams struct {
-	Curl []apix.ApiRequestConfig `json:"curl" yaml:"curl"`
+	Curl       []apix.ApiRequestConfig                 `json:"curl" yaml:"curl"`
+	RateLimits []ratelimitx.MultiTenantRateLimitConfig `json:"rate_limit_seekers" yaml:"rate-limit-seekers"`
 }
 
 type jobParams struct {