adds middleware for rate limiting

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,97 @@
+package middleware
+
+import (
+	"github.com/labstack/echo/v4"
+	"golang.org/x/time/rate"
+	"net/http"
+	"sync"
+)
+
+// TokenStore is the interface to be implemented by custom stores.
+type TokenStore interface {
+	// Stores for the rate limiter have to implement the ShouldAllow method
+	ShouldAllow(identifier string) bool
+}
+
+type (
+	// RateLimiterConfig defines the configuration for the rate limiter
+	RateLimiterConfig struct {
+		Skipper    Skipper
+		BeforeFunc BeforeFunc
+		// SourceFunc uses echo.Context to extract the identifier for a visitor
+		SourceFunc func(context echo.Context) string
+		// Store defines a store for the rate limiter
+		Store TokenStore
+	}
+)
+
+// DefaultRateLimiterConfig defines default values for RateLimiterConfig
+var DefaultRateLimiterConfig = RateLimiterConfig{
+	Skipper: DefaultSkipper,
+}
+
+// RateLimiter returns a rate limiting middleware
+func RateLimiter(source func(context echo.Context) string, store TokenStore) echo.MiddlewareFunc {
+	config := DefaultRateLimiterConfig
+	config.SourceFunc = source
+	config.Store = store
+
+	return RateLimiterWithConfig(config)
+}
+
+// RateLimiterWithConfig returns a rate limiting middleware
+func RateLimiterWithConfig(config RateLimiterConfig) echo.MiddlewareFunc {
+	if config.Skipper == nil {
+		config.Skipper = DefaultRateLimiterConfig.Skipper
+	}
+	if config.SourceFunc == nil {
+		panic("Source function must be provided")
+	}
+	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 := config.SourceFunc(c)
+
+			allowed := config.Store.ShouldAllow(identifier)
+			if !allowed {
+				return c.JSON(http.StatusTooManyRequests, nil)
+			}
+			return next(c)
+		}
+	}
+}
+
+// InMemoryStore is the built-in store implementation for RateLimiter
+type InMemoryStore struct {
+	visitors map[string]*rate.Limiter
+	mutex    sync.Mutex
+	rate     rate.Limit
+	burst    int
+}
+
+// ShouldAllow implements TokenStore.ShouldAllow
+func (store *InMemoryStore) ShouldAllow(identifier string) bool {
+	store.mutex.Lock()
+	defer store.mutex.Unlock()
+
+	if store.visitors == nil {
+		store.visitors = make(map[string]*rate.Limiter)
+	}
+
+	limiter, exists := store.visitors[identifier]
+	if !exists {
+		limiter = rate.NewLimiter(store.rate, store.burst)
+		store.visitors[identifier] = limiter
+	}
+
+	return limiter.Allow()
+}

middleware/rate_limiter_test.go

@@ -0,0 +1,192 @@
+package middleware
+
+import (
+	"github.com/labstack/echo/v4"
+	"github.com/stretchr/testify/assert"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+
+func TestRateLimiter(t *testing.T) {
+	var inMemoryStore = new(InMemoryStore)
+	inMemoryStore.rate = 1
+	inMemoryStore.burst = 3
+
+	e := echo.New()
+
+	handler := func(c echo.Context) error {
+		return c.String(http.StatusOK, "test")
+	}
+
+	assert.Panics(t, func() {
+		RateLimiter(nil, nil)
+	})
+
+	assert.Panics(t, func() {
+		RateLimiter(func(ctx echo.Context) string {
+			return "127.0.0.1"
+		}, nil)
+	})
+
+	assert.NotPanics(t, func() {
+		RateLimiter(func(ctx echo.Context) string {
+			return "127.0.0.1"
+		}, inMemoryStore)
+	})
+
+	{
+		var skipped bool
+		var inMemoryStore = new(InMemoryStore)
+		inMemoryStore.rate = 1
+		inMemoryStore.burst = 3
+
+		req := httptest.NewRequest(http.MethodGet, "/", nil)
+		req.Header.Add(echo.HeaderXRealIP, "127.0.0.1")
+
+		rec := httptest.NewRecorder()
+
+		c := e.NewContext(req, rec)
+
+		mw := RateLimiterWithConfig(RateLimiterConfig{
+			Skipper: func(c echo.Context) bool {
+				skipped = true
+				return true
+			},
+			Store: inMemoryStore,
+			SourceFunc: func(ctx echo.Context) string {
+				return "127.0.0.1"
+			},
+		})
+
+		_ = mw(handler)(c)
+
+		assert.Equal(t, true, skipped)
+	}
+
+	{
+		var beforeRan bool
+		var inMemoryStore = new(InMemoryStore)
+		inMemoryStore.rate = 1
+		inMemoryStore.burst = 3
+
+		req := httptest.NewRequest(http.MethodGet, "/", nil)
+		req.Header.Add(echo.HeaderXRealIP, "127.0.0.1")
+
+		rec := httptest.NewRecorder()
+
+		c := e.NewContext(req, rec)
+
+		mw := RateLimiterWithConfig(RateLimiterConfig{
+			BeforeFunc: func(c echo.Context) {
+				beforeRan = true
+			},
+			Store: inMemoryStore,
+			SourceFunc: func(ctx echo.Context) string {
+				return "127.0.0.1"
+			},
+		})
+
+		_ = mw(handler)(c)
+
+		assert.Equal(t, true, beforeRan)
+	}
+
+	testCases := []struct {
+		id   string
+		code int
+	}{
+		{"127.0.0.1", 200},
+		{"127.0.0.1", 200},
+		{"127.0.0.1", 200},
+		{"127.0.0.1", 429},
+		{"127.0.0.1", 429},
+		{"127.0.0.1", 429},
+		{"127.0.0.1", 429},
+	}
+
+	for _, tc := range testCases {
+		req := httptest.NewRequest(http.MethodGet, "/", nil)
+		req.Header.Add(echo.HeaderXRealIP, tc.id)
+
+		rec := httptest.NewRecorder()
+
+		c := e.NewContext(req, rec)
+		mw := RateLimiter(func(c echo.Context) string {
+			return c.Request().Header.Get(echo.HeaderXRealIP)
+		}, inMemoryStore)
+
+		_ = mw(handler)(c)
+
+		assert.Equal(t, tc.code, rec.Code)
+	}
+}
+
+func TestRateLimiterWithConfig(t *testing.T) {
+	var inMemoryStore = new(InMemoryStore)
+	inMemoryStore.rate = 1
+	inMemoryStore.burst = 3
+
+	e := echo.New()
+
+	handler := func(c echo.Context) error {
+		return c.String(http.StatusOK, "test")
+	}
+
+	testCases := []struct {
+		id   string
+		code int
+	}{
+		{"127.0.0.1", 200},
+		{"127.0.0.1", 200},
+		{"127.0.0.1", 200},
+		{"127.0.0.1", 429},
+		{"127.0.0.1", 429},
+		{"127.0.0.1", 429},
+		{"127.0.0.1", 429},
+	}
+
+	for _, tc := range testCases {
+		req := httptest.NewRequest(http.MethodGet, "/", nil)
+		req.Header.Add(echo.HeaderXRealIP, tc.id)
+
+		rec := httptest.NewRecorder()
+
+		c := e.NewContext(req, rec)
+		mw := RateLimiterWithConfig(RateLimiterConfig{
+			SourceFunc: func(c echo.Context) string {
+				return c.Request().Header.Get(echo.HeaderXRealIP)
+			},
+			Store: inMemoryStore,
+		})
+
+		_ = mw(handler)(c)
+
+		assert.Equal(t, tc.code, rec.Code)
+	}
+}
+
+func TestInMemoryStore_ShouldAllow(t *testing.T) {
+	var inMemoryStore = new(InMemoryStore)
+	inMemoryStore.rate = 1
+	inMemoryStore.burst = 3
+
+	testCases := []struct {
+		id      string
+		allowed bool
+	}{
+		{"127.0.0.1", true},
+		{"127.0.0.1", true},
+		{"127.0.0.1", true},
+		{"127.0.0.1", false},
+		{"127.0.0.1", false},
+		{"127.0.0.1", false},
+		{"127.0.0.1", false},
+	}
+
+	for _, tc := range testCases {
+		allowed := inMemoryStore.ShouldAllow(tc.id)
+
+		assert.Equal(t, tc.allowed, allowed)
+	}
+}