Skip to content

Commit

Permalink
✨ feat: added config Rate Limiting and Throttling #15
Browse files Browse the repository at this point in the history
  • Loading branch information
pnguyen215 committed Nov 25, 2023
1 parent 90a6692 commit a3e8fe0
Show file tree
Hide file tree
Showing 6 changed files with 115 additions and 1 deletion.
13 changes: 13 additions & 0 deletions config/conf-params.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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 (
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
2 changes: 2 additions & 0 deletions internal/core/routes.go
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
94 changes: 94 additions & 0 deletions internal/middlewares/throttling_middleware.go
Original file line number Diff line number Diff line change
@@ -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()
}
}
4 changes: 3 additions & 1 deletion internal/syncConf/syncConf.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand All @@ -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 {
Expand Down

0 comments on commit a3e8fe0

Please sign in to comment.