ratelimiter

README

ratelimiter

Rate limiter for any resource type (not just http requests), inspired by Cloudflare's approach: How we built rate limiting capable of scaling to millions of domains.

Usage

package main

import (
	"fmt"
	"log"
	"time"

	"github.com/theopenlane/core/pkg/middleware/ratelimiter"
)

func main() {
	limitedKey := "key"
	windowSize := 1 * time.Minute
	// create map data store for rate limiter and set each element's expiration time to 2*windowSize and old data flush interval to 10*time.Second
	dataStore := ratelimiter.NewMapLimitStore(2*windowSize, 10*time.Second)

	var maxLimit int64 = 5
	// allow 5 requests per windowSize (1 minute)
	rateLimiter := ratelimiter.New(dataStore, maxLimit, windowSize)

	for i := 0; i < 10; i++ {
		limitStatus, err := rateLimiter.Check(limitedKey)
		if err != nil {
			log.Fatal(err)
		}

		if limitStatus.IsLimited {
			fmt.Printf("too high rate for key: %s: rate: %f, limit: %d\nsleep: %s", limitedKey, limitStatus.CurrentRate, maxLimit, *limitStatus.LimitDuration)
			time.Sleep(*limitStatus.LimitDuration)
		} else {
			err := rateLimiter.Inc(limitedKey)
			if err != nil {
				log.Fatal(err)
			}
		}
	}
}

Rate-limit IP requests in http middleware

func rateLimitMiddleware(rateLimiter *ratelimiter.RateLimiter) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			remoteIP := GetRemoteIP([]string{"X-Forwarded-For", "RemoteAddr", "X-Real-IP"}, 0, r)
			key := fmt.Sprintf("%s_%s_%s", remoteIP, r.URL.String(), r.Method)

			limitStatus, err := rateLimiter.Check(key)
			if err != nil {
				// if rate limit error then pass the request
				next.ServeHTTP(w, r)
			}

			if limitStatus.IsLimited {
				w.WriteHeader(http.StatusTooManyRequests)
				return
			}

			rateLimiter.Inc(key)
			next.ServeHTTP(w, r)
		})
	}
}

func hello(w http.ResponseWriter, r *http.Request) {
	fmt.Fprintf(w, "Hello, %q", html.EscapeString(r.URL.Path))
}

func main() {
	windowSize := 1 * time.Minute
	// create map data store for rate limiter and set each element's expiration time to 2*windowSize and old data flush interval to 10*time.Second
	dataStore := ratelimiter.NewMapLimitStore(2*windowSize, 10*time.Second)
	// allow 5 requests per windowSize (1 minute)
	rateLimiter := ratelimiter.New(dataStore, 5, windowSize)

	rateLimiterHandler := rateLimitMiddleware(rateLimiter)
	helloHandler := http.HandlerFunc(hello)
	http.Handle("/", rateLimiterHandler(helloHandler))

	log.Fatal(http.ListenAndServe(":8080", nil))

}

See full example

Implement your own limit data store (or external persistence method)

To use custom data store (memcached, Redis, MySQL etc.) you just need to implement the LimitStore interface, for example:

type FakeDataStore struct{}

func (f FakeDataStore) Inc(key string, window time.Time) error {
	return nil
}

func (f FakeDataStore) Get(key string, previousWindow, currentWindow time.Time) (prevValue int64, currValue int64, err error) {
	return 0, 0, nil
}

rateLimiter := ratelimiter.New(FakeDataStore{}, maxLimit, windowSize)

Documentation

Overview

Package ratelimiter is a ratelimiter based on cloudflare's approach

Index

• type LimitStatus
• type LimitStore
• type MapLimitStore
  • func NewMapLimitStore(expirationTime time.Duration, flushInterval time.Duration) (m *MapLimitStore)
  • func (m *MapLimitStore) Get(key string, previousWindow, currentWindow time.Time) (prevValue int64, currValue int64, err error)
  • func (m *MapLimitStore) Inc(key string, window time.Time) error
  • func (m *MapLimitStore) Size() int
• type RateLimiter
  • func New(dataStore LimitStore, requestsLimit int64, windowSize time.Duration) *RateLimiter
  • func (r *RateLimiter) Check(key string) (limitStatus *LimitStatus, err error)
  • func (r *RateLimiter) Inc(key string) error

Types

type LimitStatus

type LimitStatus struct {
	// IsLimited is true when a given key should be rate-limited
	IsLimited bool
	// LimitDuration is the time for which a given key should be blocked
	LimitDuration *time.Duration
	// CurrentRate is approximated current requests rate per window size
	CurrentRate float64
}

LimitStatus represents current status of limitation for a given key

type LimitStore

type LimitStore interface {
	// Inc increments current window limit counter for key
	Inc(key string, window time.Time) error
	// Get gets value of previous window counter and current window counter for key
	Get(key string, previousWindow, currentWindow time.Time) (prevValue int64, currValue int64, err error)
}

LimitStore is the interface that represents limiter internal data store

type MapLimitStore

type MapLimitStore struct {
	// contains filtered or unexported fields
}

MapLimitStore represents a data structure for in-memory storage of ratelimiter information

func NewMapLimitStore

func NewMapLimitStore(expirationTime time.Duration, flushInterval time.Duration) (m *MapLimitStore)

NewMapLimitStore creates new in-memory data store for internal limiter data

func (*MapLimitStore) Get

func (m *MapLimitStore) Get(key string, previousWindow, currentWindow time.Time) (prevValue int64, currValue int64, err error)

Get gets value of previous window counter and current window counter

func (*MapLimitStore) Inc

func (m *MapLimitStore) Inc(key string, window time.Time) error

Inc increments current window limit counter

func (*MapLimitStore) Size

func (m *MapLimitStore) Size() int

Size returns current length of data map

type RateLimiter

type RateLimiter struct {
	// contains filtered or unexported fields
}

RateLimiter is defining the structure of a rate limiter object

func New

func New(dataStore LimitStore, requestsLimit int64, windowSize time.Duration) *RateLimiter

New creates new rate limiter

func (*RateLimiter) Check

func (r *RateLimiter) Check(key string) (limitStatus *LimitStatus, err error)

Check checks status of the rate limit key

func (*RateLimiter) Inc

func (r *RateLimiter) Inc(key string) error

Inc increments the limit counter for a given key