Lithium/middlewares/ratelimiter.go

package middlewares

import (
	"net/http"

	"github.com/throttled/throttled/store/memstore"
	"github.com/throttled/throttled/v2"
)

type RateLimiterMiddleware struct {
	rateLimiter throttled.HTTPRateLimiter
}

func (middleware RateLimiterMiddleware) Middleware(next http.Handler) http.Handler {
	return middleware.rateLimiter.RateLimit(next)
}

func CreateRateLimiterMiddleware(requestsPerMinute int, allowedBurst int) (*RateLimiterMiddleware, error) {
	store, err := memstore.New(65536)

	if err != nil {
		return nil, err
	}

	quota := throttled.RateQuota{
		MaxRate:  throttled.PerMin(requestsPerMinute),
		MaxBurst: allowedBurst,
	}

	rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
	if err != nil {
		return nil, err
	}

	httpRateLimiter := throttled.HTTPRateLimiter{
		RateLimiter: rateLimiter,
		VaryBy:      &throttled.VaryBy{Path: true},
	}

	return &RateLimiterMiddleware{
		rateLimiter: httpRateLimiter,
	}, nil
}