diff --git a/go.mod b/go.mod index 8fc41b9c..495eafa0 100644 --- a/go.mod +++ b/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 ) diff --git a/go.sum b/go.sum index 72f06e3e..1e267259 100644 --- a/go.sum +++ b/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= diff --git a/internals/config/structure/structure.go b/internals/config/structure/structure.go index b276d57d..3b39cab0 100644 --- a/internals/config/structure/structure.go +++ b/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"` } \ No newline at end of file diff --git a/internals/proxy/middlewares/ratelimit.go b/internals/proxy/middlewares/ratelimit.go new file mode 100644 index 00000000..775b3c00 --- /dev/null +++ b/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) + }) +} \ No newline at end of file diff --git a/internals/proxy/proxy.go b/internals/proxy/proxy.go index 1ee5def7..aae6345f 100644 --- a/internals/proxy/proxy.go +++ b/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).