feat: Rate Limiting

go.mod

@@ -20,4 +20,5 @@ require (
 	go.uber.org/multierr v1.11.0 // indirect
 	go.yaml.in/yaml/v3 v3.0.4 // indirect
 	golang.org/x/sys v0.39.0 // indirect
+	golang.org/x/time v0.14.0
 )

go.sum

@@ -40,6 +40,8 @@ go.yaml.in/yaml/v3 v3.0.4 h1:tfq32ie2Jv2UxXFdLJdh3jXuOzWiL1fo0bu/FbuKpbc=
 go.yaml.in/yaml/v3 v3.0.4/go.mod h1:DhzuOOF2ATzADvBadXxruRBLzYTpT36CKvDb3+aBEFg=
 golang.org/x/sys v0.39.0 h1:CvCKL8MeisomCi6qNZ+wbb0DN9E5AATixKsvNtMoMFk=
 golang.org/x/sys v0.39.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks=
+golang.org/x/time v0.14.0 h1:MRx4UaLrDotUKUdCIqzPC48t1Y9hANFKIRpNx+Te8PI=
+golang.org/x/time v0.14.0/go.mod h1:eL/Oa2bBBK0TkX57Fyni+NgnyQQN4LitPmob2Hjnqw4=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

internals/config/structure/structure.go

@@ -30,11 +30,11 @@ type API struct {
 }
 
 type SETTINGS struct {
-	ACCESS 				ACCESS_SETTINGS 			`koanf:"access"`
-	MESSAGE				MESSAGE_SETTINGS			`koanf:"message"`
+	ACCESS 				ACCESS 			`koanf:"access"`
+	MESSAGE				MESSAGE			`koanf:"message"`
 }
 
-type MESSAGE_SETTINGS struct {
+type MESSAGE struct {
 	VARIABLES         	map[string]any              `koanf:"variables"       childtransform:"upper"`
 	FIELD_MAPPINGS      map[string][]FieldMapping	`koanf:"fieldmappings"   childtransform:"default"`
 	TEMPLATE  			string                      `koanf:"template"`
@@ -45,12 +45,18 @@ type FieldMapping struct {
 	Score 				int    						`koanf:"score"`
 }
 
-type ACCESS_SETTINGS struct {
+type ACCESS struct {
 	ENDPOINTS			[]string					`koanf:"endpoints"`
 	FIELD_POLICIES		map[string]FieldPolicy		`koanf:"fieldpolicies"   childtransform:"default"`
+	RATE_LIMITING		RateLimiting				`koanf:"ratelimiting"`
 }
 
 type FieldPolicy struct {
 	Value				any						    `koanf:"value"`
 	Action				string						`koanf:"action"`
+}
+
+type RateLimiting struct {
+	Limit				int							`koanf:"limit"`
+	Period				string						`koanf:"period"`
 }

internals/proxy/middlewares/ratelimit.go

@@ -0,0 +1,84 @@
+package middlewares
+
+import (
+	"net/http"
+	"strings"
+	"time"
+
+	"golang.org/x/time/rate"
+)
+
+var RateLimit Middleware = Middleware{
+	Name: "Rate Limiting",
+	Use: ratelimitHandler,
+}
+
+type TokenLimiter struct {
+	limiter *rate.Limiter
+}
+
+func NewTokenLimiter(limit int, period time.Duration) *TokenLimiter {
+	r := rate.Every(period / time.Duration(limit))
+
+	return &TokenLimiter{
+		limiter: rate.NewLimiter(r, limit),
+	}
+}
+
+func (t *TokenLimiter) Allow() bool {
+	return t.limiter.Allow()
+}
+
+var tokenLimiters = map[string]*TokenLimiter{}
+
+func ratelimitHandler(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
+		logger := getLogger(req)
+
+		conf := getConfigByReq(req)
+
+		rateLimiting := conf.SETTINGS.ACCESS.RATE_LIMITING
+
+		limit := rateLimiting.Limit
+
+		if limit == 0 {
+			limit = getConfig("").SETTINGS.ACCESS.RATE_LIMITING.Limit
+		}
+
+		periodStr := rateLimiting.Period
+
+		if strings.TrimSpace(periodStr) == "" {
+			periodStr = conf.SETTINGS.ACCESS.RATE_LIMITING.Period
+		}
+
+		if strings.TrimSpace(periodStr) != "" && limit != 0 {
+			period, err := time.ParseDuration(periodStr)
+
+			if err != nil {
+				logger.Error("Could not parse Duration: ", err.Error())
+				http.Error(w, "Internal Server Error", http.StatusInternalServerError)
+				return
+			}
+
+			token := getToken(req)
+
+			tokenLimiter, exists := tokenLimiters[token]
+
+			if !exists {
+				tokenLimiter = NewTokenLimiter(limit, period)
+				tokenLimiters[token] = tokenLimiter
+			}
+
+			if !tokenLimiter.Allow() {
+				logger.Warn("Token exceeded Rate Limit")
+
+				http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
+				w.Header().Set("Retry-After", "60")
+
+				return
+			}
+		}
+
+		next.ServeHTTP(w, req)
+	})
+}

internals/proxy/proxy.go

@@ -36,6 +36,7 @@ func (proxy Proxy) Init() http.Handler {
 		Use(m.RequestLogger).
 		Use(m.InternalAuthRequirement).
 		Use(m.Port).
+		Use(m.RateLimit).
 		Use(m.Template).
 		Use(m.Endpoints).
 		Use(m.Mapping).