initial commit 2

server/serverconfig/ratelimit.go

@@ -0,0 +1,73 @@
+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())
+}