74 lines
2.1 KiB
Go
74 lines
2.1 KiB
Go
package serverconfig
|
|
|
|
import (
|
|
"time"
|
|
|
|
"net/http"
|
|
|
|
"github.com/labstack/echo"
|
|
"github.com/labstack/echo/middleware"
|
|
)
|
|
|
|
// RateLimitMiddlewareConfig defines the config for RateLimit middleware.
|
|
type RateLimitMiddlewareConfig struct {
|
|
Skipper middleware.Skipper
|
|
// Max number of allowed requests on a timeframe
|
|
MaxRequests int64
|
|
// The timeframe duration for the requests count
|
|
BucketDuration time.Duration
|
|
// Error response message
|
|
ResponseMessage string
|
|
// Error response
|
|
ResponseStatus int
|
|
}
|
|
|
|
// DefaultRateLimitMiddlewareConfig is the default RateLimit middleware config.
|
|
var DefaultRateLimitMiddlewareConfig = RateLimitMiddlewareConfig{
|
|
Skipper: middleware.DefaultSkipper,
|
|
MaxRequests: 10,
|
|
BucketDuration: time.Second,
|
|
ResponseMessage: "Too Many Requests",
|
|
ResponseStatus: http.StatusServiceUnavailable,
|
|
}
|
|
|
|
// RateLimitMiddleware returns a middleware that protects the API agains massive requests.
|
|
func RateLimitMiddleware() echo.MiddlewareFunc {
|
|
return RateLimitMiddlewareWithConfig(DefaultRateLimitMiddlewareConfig)
|
|
}
|
|
|
|
// RateLimitMiddlewareWithConfig returns a RateLimitMiddleware with configuration parameters.
|
|
// See: `RateLimitMiddleware()`.
|
|
func RateLimitMiddlewareWithConfig(config RateLimitMiddlewareConfig) echo.MiddlewareFunc {
|
|
if config.MaxRequests == 0 {
|
|
config.MaxRequests = DefaultRateLimitMiddlewareConfig.MaxRequests
|
|
}
|
|
|
|
if config.BucketDuration == 0 {
|
|
config.BucketDuration = DefaultRateLimitMiddlewareConfig.BucketDuration
|
|
}
|
|
|
|
// limiter := tollbooth.NewLimiterExpiringBuckets(config.MaxRequests, config.BucketDuration, time.Hour, time.Second*0)
|
|
// limiter.Message = config.ResponseMessage
|
|
// limiter.StatusCode = config.ResponseStatus
|
|
|
|
return func(next echo.HandlerFunc) echo.HandlerFunc {
|
|
return func(c echo.Context) error {
|
|
if config.Skipper(c) {
|
|
return next(c)
|
|
}
|
|
|
|
// err := tollbooth.LimitByRequest(limiter, c.Request())
|
|
// if err != nil {
|
|
// tollbooth.SetResponseHeaders(limiter, c.Response().Writer)
|
|
// return errors.NewHTTPError(limiter.StatusCode, limiter.Message)
|
|
// }
|
|
|
|
return next(c)
|
|
}
|
|
}
|
|
}
|
|
|
|
func setRateLimitMiddleware(e *echo.Echo) {
|
|
e.Pre(RateLimitMiddleware())
|
|
}
|