adds middleware for rate limiting

.gitignore

@@ -5,3 +5,4 @@ vendor
 .idea
 *.iml
 *.out
+.vscode

go.mod

@@ -12,4 +12,5 @@ require (
 	golang.org/x/net v0.0.0-20200822124328-c89045814202
 	golang.org/x/sys v0.0.0-20200826173525-f9321e4c35a6 // indirect
 	golang.org/x/text v0.3.3 // indirect
+	golang.org/x/time v0.0.0-20201208040808-7e3f01d25324
 )

go.sum

@@ -46,6 +46,8 @@ golang.org/x/text v0.3.0 h1:g61tztE5qeGQ89tm6NTjjM9VPIm088od1l6aSorWRWg=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3 h1:cokOdA+Jmi5PJGXLlLllQSgYigAEfHXJAERHVMaCc2k=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
+golang.org/x/time v0.0.0-20201208040808-7e3f01d25324 h1:Hir2P/De0WpUhtrKGGjvSb2YxUgyZ7EFOSLIcSSpiwE=
+golang.org/x/time v0.0.0-20201208040808-7e3f01d25324/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=

middleware/rate_limiter.go

@@ -0,0 +1,219 @@
+package middleware
+
+import (
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/labstack/echo/v4"
+	"golang.org/x/time/rate"
+)
+
+type (
+	// RateLimiterStore is the interface to be implemented by custom stores.
+	RateLimiterStore interface {
+		// Stores for the rate limiter have to implement the Allow method
+		Allow(identifier string) bool
+	}
+)
+
+type (
+	// RateLimiterConfig defines the configuration for the rate limiter
+	RateLimiterConfig struct {
+		Skipper    Skipper
+		BeforeFunc BeforeFunc
+		// IdentifierExtractor uses echo.Context to extract the identifier for a visitor
+		IdentifierExtractor Extractor
+		// Store defines a store for the rate limiter
+		Store        RateLimiterStore
+		ErrorHandler ErrorHandler
+		DenyHandler  ErrorHandler
+	}
+	// ErrorHandler provides a handler for returning errors from the middleware
+	ErrorHandler func(context echo.Context) error
+	// Extractor is used to extract data from echo.Context
+	Extractor func(context echo.Context) (string, error)
+)
+
+// DefaultRateLimiterConfig defines default values for RateLimiterConfig
+var DefaultRateLimiterConfig = RateLimiterConfig{
+	Skipper: DefaultSkipper,
+	IdentifierExtractor: func(ctx echo.Context) (string, error) {
+		id := ctx.RealIP()
+		return id, nil
+	},
+	ErrorHandler: func(context echo.Context) error {
+		return context.JSON(http.StatusTooManyRequests, nil)
+	},
+	DenyHandler: func(context echo.Context) error {
+		return context.JSON(http.StatusForbidden, nil)
+	},
+}
+
+/*
+RateLimiter returns a rate limiting middleware
+
+	e := echo.New()
+
+	var inMemoryStore = NewRateLimiterMemoryStore(
+		RateLimiterMemoryStoreConfig{rate: 1, burst: 3, expiresIn: 3 * time.Minute}
+	)
+
+	e.GET("/rate-limited", func(c echo.Context) error {
+		return c.String(http.StatusOK, "test")
+	}, RateLimiter(inMemoryStore))
+*/
+func RateLimiter(store RateLimiterStore) echo.MiddlewareFunc {
+	config := DefaultRateLimiterConfig
+	config.Store = store
+
+	return RateLimiterWithConfig(config)
+}
+
+/*
+RateLimiterWithConfig returns a rate limiting middleware
+
+	e := echo.New()
+
+	var inMemoryStore = NewRateLimiterMemoryStore(
+		RateLimiterMemoryStoreConfig{rate: 1, burst: 3, expiresIn: 3 * time.Minute}
+	)
+
+	config := RateLimiterConfig{
+		Skipper: DefaultSkipper,
+		IdentifierExtractor: func(ctx echo.Context) (string, error) {
+			id := ctx.RealIP()
+			return id, nil
+		},
+		ErrorHandler: func(context echo.Context) error {
+			return context.JSON(http.StatusTooManyRequests, nil)
+		},
+		DenyHandler: func(context echo.Context) error {
+			return context.JSON(http.StatusForbidden, nil)
+		},
+	}
+
+	e.GET("/rate-limited", func(c echo.Context) error {
+		return c.String(http.StatusOK, "test")
+	}, RateLimiterWithConfig(config))
+*/
+func RateLimiterWithConfig(config RateLimiterConfig) echo.MiddlewareFunc {
+	if config.Skipper == nil {
+		config.Skipper = DefaultRateLimiterConfig.Skipper
+	}
+	if config.IdentifierExtractor == nil {
+		config.IdentifierExtractor = DefaultRateLimiterConfig.IdentifierExtractor
+	}
+	if config.ErrorHandler == nil {
+		config.ErrorHandler = DefaultRateLimiterConfig.ErrorHandler
+	}
+	if config.DenyHandler == nil {
+		config.DenyHandler = DefaultRateLimiterConfig.DenyHandler
+	}
+	if config.Store == nil {
+		panic("Store configuration must be provided")
+	}
+	return func(next echo.HandlerFunc) echo.HandlerFunc {
+		return func(c echo.Context) error {
+			if config.Skipper(c) {
+				return next(c)
+			}
+			if config.BeforeFunc != nil {
+				config.BeforeFunc(c)
+			}
+
+			identifier, err := config.IdentifierExtractor(c)
+			if err != nil {
+				return config.DenyHandler(c)
+			}
+
+			if !config.Store.Allow(identifier) {
+				return config.ErrorHandler(c)
+			}
+			return next(c)
+		}
+	}
+}
+
+type (
+	// RateLimiterMemoryStore is the built-in store implementation for RateLimiter
+	RateLimiterMemoryStore struct {
+		visitors    map[string]*Visitor
+		mutex       sync.Mutex
+		rate        rate.Limit
+		burst       int
+		expiresIn   time.Duration
+		lastCleanup time.Time
+	}
+	// Visitor signifies a unique user's limiter details
+	Visitor struct {
+		*rate.Limiter
+		lastSeen time.Time
+	}
+)
+
+// NewRateLimiterMemoryStore returns an instance of RateLimiterMemoryStore
+func NewRateLimiterMemoryStore(config RateLimiterMemoryStoreConfig) (store *RateLimiterMemoryStore) {
+	store = &RateLimiterMemoryStore{}
+
+	store.rate = config.rate
+	store.burst = config.burst
+	if config.expiresIn == 0 {
+		store.expiresIn = DefaultRateLimiterMemoryStoreConfig.expiresIn
+	} else {
+		store.expiresIn = config.expiresIn
+	}
+	store.visitors = make(map[string]*Visitor)
+	store.lastCleanup = now()
+	return
+}
+
+// RateLimiterMemoryStoreConfig represents configuration for RateLimiterMemoryStore
+type RateLimiterMemoryStoreConfig struct {
+	rate      rate.Limit
+	burst     int
+	expiresIn time.Duration
+}
+
+// DefaultRateLimiterMemoryStoreConfig provides default configuration values for RateLimiterMemoryStore
+var DefaultRateLimiterMemoryStoreConfig = RateLimiterMemoryStoreConfig{
+	expiresIn: 3 * time.Minute,
+}
+
+// Allow implements RateLimiterStore.Allow
+func (store *RateLimiterMemoryStore) Allow(identifier string) bool {
+	store.mutex.Lock()
+	limiter, exists := store.visitors[identifier]
+	if !exists {
+		limiter = new(Visitor)
+		limiter.Limiter = rate.NewLimiter(store.rate, store.burst)
+		limiter.lastSeen = now()
+		store.visitors[identifier] = limiter
+	}
+	limiter.lastSeen = now()
+	if now().Sub(store.lastCleanup) > store.expiresIn {
+		store.cleanupStaleVisitors()
+	}
+	store.mutex.Unlock()
+	return limiter.AllowN(now(), 1)
+}
+
+/*
+cleanupStaleVisitors helps manage the size of the visitors map by removing stale records
+of users who haven't visited again after the configured expiry time has elapsed
+*/
+func (store *RateLimiterMemoryStore) cleanupStaleVisitors() {
+	for id, visitor := range store.visitors {
+		if now().Sub(visitor.lastSeen) > store.expiresIn {
+			delete(store.visitors, id)
+		}
+	}
+	store.lastCleanup = now()
+}
+
+/*
+actual time method which is mocked in test file
+*/
+var now = func() time.Time {
+	return time.Now()
+}

middleware/rate_limiter_shard_store.go

@@ -0,0 +1,118 @@
+package middleware
+
+import (
+	"sync"
+	"time"
+
+	"golang.org/x/time/rate"
+)
+
+const (
+	numShards = 15
+)
+
+type (
+	// RateLimiterShardedMemoryStore is the built-in store implementation for RateLimiter
+	RateLimiterShardedMemoryStore struct {
+		shards    []*VisitorShard
+		rate      rate.Limit
+		burst     int
+		expiresIn time.Duration
+	}
+	// VisitorShard is a lockable memory map
+	VisitorShard struct {
+		visitors  map[string]*Visitor
+		mutex     sync.Mutex
+		expiresIn time.Duration
+		expires   time.Time
+	}
+)
+
+// NewRateLimiterShardedMemoryStore returns an instance of RateLimiterMemoryStore
+func NewRateLimiterShardedMemoryStore(config RateLimiterMemoryStoreConfig) (store *RateLimiterShardedMemoryStore) {
+	store = &RateLimiterShardedMemoryStore{}
+
+	store.rate = config.rate
+	store.burst = config.burst
+	if config.expiresIn == 0 {
+		store.expiresIn = DefaultRateLimiterShardedMemoryStoreConfig.expiresIn
+	} else {
+		store.expiresIn = config.expiresIn
+	}
+	store.shards = make([]*VisitorShard, numShards)
+	for i := 0; i < numShards; i++ {
+		store.shards[i] = NewVisitorShard(store.expiresIn)
+	}
+
+	return store
+}
+
+// RateLimiterShardedMemoryStoreConfig represents configuration for RateLimiterShardedMemoryStore
+type RateLimiterShardedMemoryStoreConfig struct {
+	rate      rate.Limit
+	burst     int
+	expiresIn time.Duration
+}
+
+// DefaultRateLimiterShardedMemoryStoreConfig provides default configuration values for RateLimiterShardedMemoryStore
+var DefaultRateLimiterShardedMemoryStoreConfig = RateLimiterMemoryStoreConfig{
+	expiresIn: 3 * time.Minute,
+}
+
+// Allow implements RateLimiterStore.Allow
+func (store *RateLimiterShardedMemoryStore) Allow(identifier string) bool {
+
+	shardIdx := byte(0)
+	for _, b := range []byte(identifier) {
+		shardIdx += b
+	}
+	shardIdx = shardIdx % numShards
+	// println(shardIdx)
+	shard := store.shards[shardIdx]
+
+	shard.mutex.Lock()
+	limiter, exists := shard.visitors[identifier]
+	if !exists {
+		// Not found in current shard, attempt other shard too
+		limiter, exists = shard.visitors[identifier]
+		if !exists || now().Sub(limiter.lastSeen) < store.expiresIn {
+			// Create a new limiter unless exists and not expired
+			limiter = new(Visitor)
+			limiter.Limiter = rate.NewLimiter(store.rate, store.burst)
+			limiter.lastSeen = now()
+		}
+		shard.visitors[identifier] = limiter
+	}
+	limiter.lastSeen = now()
+
+	if now().After(shard.expires) {
+		shard.expires = now().Add(shard.expiresIn)
+		go shard.Cleanup()
+	}
+
+	shard.mutex.Unlock()
+
+	return limiter.AllowN(now(), 1)
+}
+
+// NewVisitorShard returns a visitor shard
+func NewVisitorShard(expiresIn time.Duration) *VisitorShard {
+	v := &VisitorShard{
+		visitors:  make(map[string]*Visitor),
+		expiresIn: expiresIn,
+		expires:   now().Add(expiresIn),
+	}
+	return v
+}
+
+// Cleanup a shards completely
+func (shard *VisitorShard) Cleanup() {
+	shard.mutex.Lock()
+	for id, visitor := range shard.visitors {
+		if now().Sub(visitor.lastSeen) > shard.expiresIn {
+			delete(shard.visitors, id)
+		}
+	}
+	shard.expires = now().Add(shard.expiresIn)
+	shard.mutex.Unlock()
+}