vikunja-api/pkg/routes/rate_limit.go

108 lines
3.6 KiB
Go
Raw Permalink Normal View History

2020-12-29 02:04:20 +01:00
// Vikunja is a to-do list application to facilitate your life.
2021-02-02 20:19:13 +01:00
// Copyright 2018-2021 Vikunja and contributors. All rights reserved.
2019-07-21 23:27:30 +02:00
//
2020-12-29 02:04:20 +01:00
// This program is free software: you can redistribute it and/or modify
2020-12-23 16:41:52 +01:00
// it under the terms of the GNU Affero General Public Licensee as published by
2019-07-21 23:27:30 +02:00
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
2020-12-29 02:04:20 +01:00
// This program is distributed in the hope that it will be useful,
2019-07-21 23:27:30 +02:00
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
2020-12-23 16:41:52 +01:00
// GNU Affero General Public Licensee for more details.
2019-07-21 23:27:30 +02:00
//
2020-12-23 16:41:52 +01:00
// You should have received a copy of the GNU Affero General Public Licensee
2020-12-29 02:04:20 +01:00
// along with this program. If not, see <https://www.gnu.org/licenses/>.
2019-07-21 23:27:30 +02:00
package routes
import (
"net/http"
"strconv"
"time"
2019-07-21 23:27:30 +02:00
"code.vikunja.io/api/pkg/config"
"code.vikunja.io/api/pkg/log"
auth2 "code.vikunja.io/api/pkg/modules/auth"
2019-07-21 23:27:30 +02:00
"code.vikunja.io/api/pkg/red"
"github.com/labstack/echo/v4"
"github.com/ulule/limiter/v3"
"github.com/ulule/limiter/v3/drivers/store/memory"
"github.com/ulule/limiter/v3/drivers/store/redis"
)
// RateLimit is the rate limit middleware
func RateLimit(rateLimiter *limiter.Limiter, rateLimitKind string) echo.MiddlewareFunc {
2019-07-21 23:27:30 +02:00
return func(next echo.HandlerFunc) echo.HandlerFunc {
return func(c echo.Context) (err error) {
var rateLimitKey string
switch rateLimitKind {
2019-07-21 23:27:30 +02:00
case "ip":
rateLimitKey = c.RealIP()
case "user":
auth, err := auth2.GetAuthFromClaims(c)
2019-07-21 23:27:30 +02:00
if err != nil {
log.Errorf("Error getting auth from jwt claims: %v", err)
2019-07-21 23:27:30 +02:00
}
rateLimitKey = "user_" + strconv.FormatInt(auth.GetID(), 10)
2019-07-21 23:27:30 +02:00
default:
log.Errorf("Unknown rate limit kind configured: %s", rateLimitKind)
2019-07-21 23:27:30 +02:00
}
limiterCtx, err := rateLimiter.Get(c.Request().Context(), rateLimitKey)
if err != nil {
log.Errorf("IPRateLimit - rateLimiter.Get - err: %v, %s on %s", err, rateLimitKey, c.Request().URL)
return c.JSON(http.StatusInternalServerError, echo.Map{
"message": err,
})
}
h := c.Response().Header()
h.Set("X-RateLimit-Limit", strconv.FormatInt(limiterCtx.Limit, 10))
h.Set("X-RateLimit-Remaining", strconv.FormatInt(limiterCtx.Remaining, 10))
h.Set("X-RateLimit-Reset", strconv.FormatInt(limiterCtx.Reset, 10))
if limiterCtx.Reached {
log.Infof("Too Many Requests from %s on %s", rateLimitKey, c.Request().URL)
return c.JSON(http.StatusTooManyRequests, echo.Map{
"message": "Too Many Requests on " + c.Request().URL.String(),
})
}
// log.Printf("%s request continue", c.RealIP())
return next(c)
}
}
}
func createRateLimiter(rate limiter.Rate) *limiter.Limiter {
var store limiter.Store
var err error
switch config.RateLimitStore.GetString() {
case "memory":
store = memory.NewStore()
case "redis":
if !config.RedisEnabled.GetBool() {
log.Fatal("Redis is configured for rate limiting, but not enabled!")
}
store, err = redis.NewStore(red.GetRedis())
if err != nil {
log.Fatalf("Error while creating rate limit redis store: %s", err)
}
default:
log.Fatalf("Unknown Rate limit store \"%s\"", config.RateLimitStore.GetString())
}
return limiter.New(store, rate)
}
func setupRateLimit(a *echo.Group, rateLimitKind string) {
2019-07-21 23:27:30 +02:00
if config.RateLimitEnabled.GetBool() {
rate := limiter.Rate{
Period: config.RateLimitPeriod.GetDuration() * time.Second,
Limit: config.RateLimitLimit.GetInt64(),
}
rateLimiter := createRateLimiter(rate)
2019-07-21 23:27:30 +02:00
log.Debugf("Rate limit configured with %s and %v requests per %v", config.RateLimitStore.GetString(), rate.Limit, rate.Period)
a.Use(RateLimit(rateLimiter, rateLimitKind))
2019-07-21 23:27:30 +02:00
}
}