pki/vendor/github.com/labstack/echo/v4/middleware/rate_limiter.go

268 lines
8.5 KiB
Go
Raw Normal View History

2024-04-19 16:50:48 +02:00
// SPDX-License-Identifier: MIT
// SPDX-FileCopyrightText: © 2015 LabStack LLC and Echo contributors
2021-04-07 10:58:21 +02:00
package middleware
import (
"net/http"
"sync"
"time"
"github.com/labstack/echo/v4"
"golang.org/x/time/rate"
)
2024-04-19 16:50:48 +02:00
// RateLimiterStore is the interface to be implemented by custom stores.
type RateLimiterStore interface {
// Stores for the rate limiter have to implement the Allow method
Allow(identifier string) (bool, error)
}
2021-04-07 10:58:21 +02:00
2024-04-19 16:50:48 +02:00
// RateLimiterConfig defines the configuration for the rate limiter
type RateLimiterConfig struct {
Skipper Skipper
BeforeFunc BeforeFunc
// IdentifierExtractor uses echo.Context to extract the identifier for a visitor
IdentifierExtractor Extractor
// Store defines a store for the rate limiter
Store RateLimiterStore
// ErrorHandler provides a handler to be called when IdentifierExtractor returns an error
ErrorHandler func(context echo.Context, err error) error
// DenyHandler provides a handler to be called when RateLimiter denies access
DenyHandler func(context echo.Context, identifier string, err error) error
}
2021-04-07 10:58:21 +02:00
2024-04-19 16:50:48 +02:00
// Extractor is used to extract data from echo.Context
type Extractor func(context echo.Context) (string, error)
// ErrRateLimitExceeded denotes an error raised when rate limit is exceeded
var ErrRateLimitExceeded = echo.NewHTTPError(http.StatusTooManyRequests, "rate limit exceeded")
// ErrExtractorError denotes an error raised when extractor function is unsuccessful
var ErrExtractorError = echo.NewHTTPError(http.StatusForbidden, "error while extracting identifier")
2021-04-07 10:58:21 +02:00
// DefaultRateLimiterConfig defines default values for RateLimiterConfig
var DefaultRateLimiterConfig = RateLimiterConfig{
Skipper: DefaultSkipper,
IdentifierExtractor: func(ctx echo.Context) (string, error) {
id := ctx.RealIP()
return id, nil
},
ErrorHandler: func(context echo.Context, err error) error {
return &echo.HTTPError{
Code: ErrExtractorError.Code,
Message: ErrExtractorError.Message,
Internal: err,
}
},
DenyHandler: func(context echo.Context, identifier string, err error) error {
return &echo.HTTPError{
Code: ErrRateLimitExceeded.Code,
Message: ErrRateLimitExceeded.Message,
Internal: err,
}
},
}
/*
RateLimiter returns a rate limiting middleware
e := echo.New()
limiterStore := middleware.NewRateLimiterMemoryStore(20)
e.GET("/rate-limited", func(c echo.Context) error {
return c.String(http.StatusOK, "test")
}, RateLimiter(limiterStore))
*/
func RateLimiter(store RateLimiterStore) echo.MiddlewareFunc {
config := DefaultRateLimiterConfig
config.Store = store
return RateLimiterWithConfig(config)
}
/*
RateLimiterWithConfig returns a rate limiting middleware
e := echo.New()
config := middleware.RateLimiterConfig{
Skipper: DefaultSkipper,
Store: middleware.NewRateLimiterMemoryStore(
middleware.RateLimiterMemoryStoreConfig{Rate: 10, Burst: 30, ExpiresIn: 3 * time.Minute}
)
IdentifierExtractor: func(ctx echo.Context) (string, error) {
id := ctx.RealIP()
return id, nil
},
ErrorHandler: func(context echo.Context, err error) error {
return context.JSON(http.StatusTooManyRequests, nil)
},
DenyHandler: func(context echo.Context, identifier string) error {
return context.JSON(http.StatusForbidden, nil)
},
}
e.GET("/rate-limited", func(c echo.Context) error {
return c.String(http.StatusOK, "test")
}, middleware.RateLimiterWithConfig(config))
*/
func RateLimiterWithConfig(config RateLimiterConfig) echo.MiddlewareFunc {
if config.Skipper == nil {
config.Skipper = DefaultRateLimiterConfig.Skipper
}
if config.IdentifierExtractor == nil {
config.IdentifierExtractor = DefaultRateLimiterConfig.IdentifierExtractor
}
if config.ErrorHandler == nil {
config.ErrorHandler = DefaultRateLimiterConfig.ErrorHandler
}
if config.DenyHandler == nil {
config.DenyHandler = DefaultRateLimiterConfig.DenyHandler
}
if config.Store == nil {
panic("Store configuration must be provided")
}
return func(next echo.HandlerFunc) echo.HandlerFunc {
return func(c echo.Context) error {
if config.Skipper(c) {
return next(c)
}
if config.BeforeFunc != nil {
config.BeforeFunc(c)
}
identifier, err := config.IdentifierExtractor(c)
if err != nil {
c.Error(config.ErrorHandler(c, err))
return nil
}
if allow, err := config.Store.Allow(identifier); !allow {
c.Error(config.DenyHandler(c, identifier, err))
return nil
}
return next(c)
}
}
}
2024-04-19 16:50:48 +02:00
// RateLimiterMemoryStore is the built-in store implementation for RateLimiter
type RateLimiterMemoryStore struct {
visitors map[string]*Visitor
mutex sync.Mutex
rate rate.Limit // for more info check out Limiter docs - https://pkg.go.dev/golang.org/x/time/rate#Limit.
2022-11-02 17:49:03 +01:00
2024-04-19 16:50:48 +02:00
burst int
expiresIn time.Duration
lastCleanup time.Time
2023-10-01 12:07:22 +02:00
2024-04-19 16:50:48 +02:00
timeNow func() time.Time
}
// Visitor signifies a unique user's limiter details
type Visitor struct {
*rate.Limiter
lastSeen time.Time
}
2021-04-07 10:58:21 +02:00
/*
NewRateLimiterMemoryStore returns an instance of RateLimiterMemoryStore with
2023-03-17 13:20:20 +01:00
the provided rate (as req/s).
2022-11-02 17:49:03 +01:00
for more info check out Limiter docs - https://pkg.go.dev/golang.org/x/time/rate#Limit.
2021-08-30 20:04:15 +02:00
Burst and ExpiresIn will be set to default values.
2021-04-07 10:58:21 +02:00
2023-03-17 13:20:20 +01:00
Note that if the provided rate is a float number and Burst is zero, Burst will be treated as the rounded down value of the rate.
2021-04-07 10:58:21 +02:00
Example (with 20 requests/sec):
limiterStore := middleware.NewRateLimiterMemoryStore(20)
*/
func NewRateLimiterMemoryStore(rate rate.Limit) (store *RateLimiterMemoryStore) {
return NewRateLimiterMemoryStoreWithConfig(RateLimiterMemoryStoreConfig{
Rate: rate,
})
}
/*
NewRateLimiterMemoryStoreWithConfig returns an instance of RateLimiterMemoryStore
2023-03-17 13:20:20 +01:00
with the provided configuration. Rate must be provided. Burst will be set to the rounded down value of
2021-04-07 10:58:21 +02:00
the configured rate if not provided or set to 0.
The build-in memory store is usually capable for modest loads. For higher loads other
store implementations should be considered.
Characteristics:
* Concurrency above 100 parallel requests may causes measurable lock contention
* A high number of different IP addresses (above 16000) may be impacted by the internally used Go map
* A high number of requests from a single IP address may cause lock contention
Example:
limiterStore := middleware.NewRateLimiterMemoryStoreWithConfig(
2022-11-02 17:49:03 +01:00
middleware.RateLimiterMemoryStoreConfig{Rate: 50, Burst: 200, ExpiresIn: 5 * time.Minute},
2021-04-07 10:58:21 +02:00
)
*/
func NewRateLimiterMemoryStoreWithConfig(config RateLimiterMemoryStoreConfig) (store *RateLimiterMemoryStore) {
store = &RateLimiterMemoryStore{}
store.rate = config.Rate
store.burst = config.Burst
store.expiresIn = config.ExpiresIn
if config.ExpiresIn == 0 {
store.expiresIn = DefaultRateLimiterMemoryStoreConfig.ExpiresIn
}
if config.Burst == 0 {
store.burst = int(config.Rate)
}
store.visitors = make(map[string]*Visitor)
2023-10-01 12:07:22 +02:00
store.timeNow = time.Now
store.lastCleanup = store.timeNow()
2021-04-07 10:58:21 +02:00
return
}
// RateLimiterMemoryStoreConfig represents configuration for RateLimiterMemoryStore
type RateLimiterMemoryStoreConfig struct {
2022-11-02 17:49:03 +01:00
Rate rate.Limit // Rate of requests allowed to pass as req/s. For more info check out Limiter docs - https://pkg.go.dev/golang.org/x/time/rate#Limit.
2023-03-17 13:20:20 +01:00
Burst int // Burst is maximum number of requests to pass at the same moment. It additionally allows a number of requests to pass when rate limit is reached.
2021-04-07 10:58:21 +02:00
ExpiresIn time.Duration // ExpiresIn is the duration after that a rate limiter is cleaned up
}
// DefaultRateLimiterMemoryStoreConfig provides default configuration values for RateLimiterMemoryStore
var DefaultRateLimiterMemoryStoreConfig = RateLimiterMemoryStoreConfig{
ExpiresIn: 3 * time.Minute,
}
// Allow implements RateLimiterStore.Allow
func (store *RateLimiterMemoryStore) Allow(identifier string) (bool, error) {
store.mutex.Lock()
limiter, exists := store.visitors[identifier]
if !exists {
limiter = new(Visitor)
limiter.Limiter = rate.NewLimiter(store.rate, store.burst)
store.visitors[identifier] = limiter
}
2023-10-01 12:07:22 +02:00
now := store.timeNow()
limiter.lastSeen = now
if now.Sub(store.lastCleanup) > store.expiresIn {
2021-04-07 10:58:21 +02:00
store.cleanupStaleVisitors()
}
store.mutex.Unlock()
2023-10-01 12:07:22 +02:00
return limiter.AllowN(store.timeNow(), 1), nil
2021-04-07 10:58:21 +02:00
}
/*
cleanupStaleVisitors helps manage the size of the visitors map by removing stale records
of users who haven't visited again after the configured expiry time has elapsed
*/
func (store *RateLimiterMemoryStore) cleanupStaleVisitors() {
for id, visitor := range store.visitors {
2023-10-01 12:07:22 +02:00
if store.timeNow().Sub(visitor.lastSeen) > store.expiresIn {
2021-04-07 10:58:21 +02:00
delete(store.visitors, id)
}
}
2023-10-01 12:07:22 +02:00
store.lastCleanup = store.timeNow()
2021-04-07 10:58:21 +02:00
}