ratelimiter

package module
v0.2.0 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: May 6, 2022 License: Apache-2.0 Imports: 10 Imported by: 0

README

go-ratelimiter

Go Report codebeat badge codecov Build status GitHub license

A super easy rate limiting package for Go. Package provide Store interface, for which you can use your own implementations

Install

go get github.com/robotomize/go-ratelimiter

Usage

Example of using redis datastore


package main

import (
	"context"
	"fmt"
	"log"
	"time"

	"github.com/robotomize/go-ratelimiter"
)

func main() {
	// set a limit of 10 request per 1 seconds per api key
	redisStore, err := ratelimiter.DefaultRedisStore(context.Background(), ":6379", 1*time.Second, 10)
	if err != nil {
		log.Fatal(err)
	}

	// Retrieve data by the api key in the datastore
	limit, remaining, resetTime, ok, err := redisStore.Take(context.Background(), "apikey-1")
	if err != nil {
		log.Fatal(err)
	}

	if !ok {
		fmt.Println("limit exceeded")
	}

	// Print the constraints from the datastore
	fmt.Printf(
		"resource: maximum %d, remaining %d, reset time %s",
		limit, remaining, time.Unix(0, int64(resetTime)).UTC().Format(time.RFC1123),
	)
}

To limit access at the http level, you can use middleware, which can block the request by providing the http code 429 Too Many Requests

Example of using http middleware with redis datastore

package main

import (
	"context"
	"errors"
	"fmt"
	"log"
	"net/http"
	"os"
	"time"

	"github.com/robotomize/go-ratelimiter"
)

func main() {
	// set a limit of 5 request per 1 seconds per api key
	redisStore, err := ratelimiter.DefaultRedisStore(context.Background(), ":6379", 1*time.Second, 5)
	if err != nil {
		log.Fatal(err)
	}

	mx := http.NewServeMux()
	// Let's create a test handler
	healthHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(`{"status": "ok"}`))
	})
    
	mx.Handle(
		"/health", ratelimiter.LimiterMiddleware(
			redisStore, func(r *http.Request) (string, error) {
				// Example key func
				ctx := r.Context()
				// Get key value out of context
				ctxValue := ctx.Value("apikey")
				if key, ok := ctxValue.(string); ok {
					return key, nil
				}

				return "", errors.New("get api key from ctx")
			}, ratelimiter.WithSkipper(func() bool {
				// set a skipper, skip ratelimiter if DEBUG == 1
				return os.Getenv("DEBUG") == "1" 
			}, ),
		)(healthHandler),
	)

	// start listener
	if err = http.ListenAndServe(":8888", mx); err != nil {
		log.Fatal(err)
	}
}

TODO

  • add http middleware
  • add redis datastore
  • improve unit coverage
  • improve redis datastore
  • add tarantool datastore
  • add aerospike datastore

Contributing

License

go-ratelimiter is under the Apache 2.0 license. See the LICENSE file for details.

Documentation

Overview

Package ratelimiter is a generated GoMock package.

Index

Constants

View Source
const (
	// HeaderRateLimitLimit - maximum number of calls
	HeaderRateLimitLimit = "X-RateLimit-Limit"
	// HeaderRateLimitRemaining - Number of calls before restrictions apply
	HeaderRateLimitRemaining = "X-RateLimit-Remaining"
	// HeaderRateLimitReset - Limit reset time
	HeaderRateLimitReset = "X-RateLimit-Reset"

	// HeaderRetryAfter is the header used to indicate when a client should retry
	// requests (when the rate limit expires), in UTC time.
	HeaderRetryAfter = "Retry-After"
)

Variables

This section is empty.

Functions

func LimiterMiddleware

func LimiterMiddleware(s Store, keyFunc KeyFunc, opts ...Option) func(next http.Handler) http.Handler

LimiterMiddleware returns a handler, which is a rate limiter with data storage in store

Types

type ExclFunc

type ExclFunc func(key string) (ok bool, limit uint64, interval time.Duration)

ExclFunc defines exceptional parameters for a given key

type KeyFunc

type KeyFunc func(r *http.Request) (string, error)

KeyFunc is a function that accepts an http request and returns a string key that uniquely identifies this request for the purpose of rate limiting.

KeyFuncs are called on each request, so be mindful of performance and implement caching where possible. If a KeyFunc returns an error, the HTTP handler will return Internal Server Error and will NOT take from the limiter store.

func IPKeyFunc

func IPKeyFunc(headers ...string) KeyFunc

IPKeyFunc rate limit by ip

type MockStore

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

MockStore is a mock of Store interface.

func NewMockStore

func NewMockStore(ctrl *gomock.Controller) *MockStore

NewMockStore creates a new mock instance.

func (*MockStore) EXPECT

func (m *MockStore) EXPECT() *MockStoreMockRecorder

EXPECT returns an object that allows the caller to indicate expected use.

func (*MockStore) Reset

func (m *MockStore) Reset(ctx context.Context) error

Reset mocks base method.

func (*MockStore) Take

func (m *MockStore) Take(ctx context.Context, key string) (uint64, uint64, uint64, bool, error)

Take mocks base method.

func (*MockStore) TakeExcl

func (m *MockStore) TakeExcl(ctx context.Context, key string, f ExclFunc) (uint64, uint64, uint64, bool, error)

TakeExcl mocks base method.

type MockStoreMockRecorder

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

MockStoreMockRecorder is the mock recorder for MockStore.

func (*MockStoreMockRecorder) Reset

func (mr *MockStoreMockRecorder) Reset(ctx interface{}) *gomock.Call

Reset indicates an expected call of Reset.

func (*MockStoreMockRecorder) Take

func (mr *MockStoreMockRecorder) Take(ctx, key interface{}) *gomock.Call

Take indicates an expected call of Take.

func (*MockStoreMockRecorder) TakeExcl

func (mr *MockStoreMockRecorder) TakeExcl(ctx, key, f interface{}) *gomock.Call

TakeExcl indicates an expected call of TakeExcl.

type Option

type Option func(*Options)

func WithDateFormat

func WithDateFormat(format string) Option

WithDateFormat set custom date format into HeaderRetryAfter/HeaderRateLimitReset

func WithSkipper

func WithSkipper(skipper func() bool) Option

WithSkipper set skipper function for skipping

type Options

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

Options for middleware

type RedisClient

type RedisClient interface {
	HMGet(ctx context.Context, key string, fields ...string) *redis.SliceCmd
	HSet(ctx context.Context, key string, values ...interface{}) *redis.IntCmd
	Del(ctx context.Context, keys ...string) *redis.IntCmd
	Expire(ctx context.Context, key string, expiration time.Duration) *redis.BoolCmd
	TxPipeline() redis.Pipeliner
	SMembers(ctx context.Context, key string) *redis.StringSliceCmd
}

type RedisConfig

type RedisConfig struct {
	// redis tags for invalidating keys
	Tags []string
	// tagging keys
	TagsPrefix string
	// redis key prefix for redis store
	Prefix string
	// limiter interval
	Interval time.Duration
	// limiter max points
	Points uint64
}

RedisConfig - struct for configure redis store

For example if you want to set the limit of requests per second to 10 req per sec RedisConfig{ Interval: time.Second * 1, Points: 10} or 20 req per 2 minutes RedisConfig{ Interval: time.Minute * 2, Points: 20}

type Store

type Store interface {
	// Take actual rate limit constraint by key or create new
	Take(ctx context.Context, key string) (limit, remaining, resetTime uint64, ok bool, err error)
	// Reset completely clears the store and resets all tokens
	Reset(ctx context.Context) error
	// TakeExcl actual rate limit constraint by key or create new with exclusive function
	TakeExcl(ctx context.Context, key string, f ExclFunc) (limit, remaining, resetTime uint64, ok bool, err error)
}

func DefaultRedisStore added in v0.2.0

func DefaultRedisStore(ctx context.Context, addr string, interval time.Duration, points uint64) (Store, error)

DefaultRedisStore return Store instance of default redis options

func NewNoop

func NewNoop() Store

func NewRedisStore

func NewRedisStore(instance RedisClient, cfg RedisConfig) Store

NewRedisStore make redis store

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL