Close #3: Add support for rate limiting

Author: TwiN

README.md

@@ -166,3 +166,9 @@ have the `backup` permission:
 ```go
 router.Handle("/backup", gate.ProtectWithPermissions(&testHandler{}, []string{"read", "backup"}))
 ```
+
+## Rate limiting
+To add a rate limit of 100 requests per second:
+```
+gate := g8.NewGate(nil).WithRateLimit(100)
+```

gate.go

@@ -9,22 +9,28 @@ const (
 	// AuthorizationHeader is the header in which g8 looks for the authorization bearer token
 	AuthorizationHeader = "Authorization"
 
-	// DefaultUnauthorizedResponseBody is the default response body returned if a request was sent with a missing or
-	// invalid token
+	// DefaultUnauthorizedResponseBody is the default response body returned if a request was sent with a missing or invalid token
 	DefaultUnauthorizedResponseBody = "Authorization Bearer token is missing or invalid"
+
+	// DefaultTooManyRequestsResponseBody is the default response body returned if a request exceeded the allowed rate limit
+	DefaultTooManyRequestsResponseBody = "Too Many Requests"
 )
 
 // Gate is lock to the front door of your API, letting only those you allow through.
 type Gate struct {
 	authorizationService     *AuthorizationService
 	unauthorizedResponseBody []byte
+
+	rateLimiter                 *RateLimiter
+	tooManyRequestsResponseBody []byte
 }
 
 // NewGate creates a new Gate.
 func NewGate(authorizationService *AuthorizationService) *Gate {
 	return &Gate{
-		unauthorizedResponseBody: []byte(DefaultUnauthorizedResponseBody),
-		authorizationService:     authorizationService,
+		authorizationService:        authorizationService,
+		unauthorizedResponseBody:    []byte(DefaultUnauthorizedResponseBody),
+		tooManyRequestsResponseBody: []byte(DefaultTooManyRequestsResponseBody),
 	}
 }
 
@@ -34,6 +40,12 @@ func (gate *Gate) WithCustomUnauthorizedResponseBody(unauthorizedResponseBody []
 	return gate
 }
 
+// WithRateLimit adds rate limiting to the Gate
+func (gate *Gate) WithRateLimit(maximumRequestsPerSecond int) *Gate {
+	gate.rateLimiter = NewRateLimiter(maximumRequestsPerSecond)
+	return gate
+}
+
 // Protect secures a handler, requiring requests going through to have a valid Authorization Bearer token.
 // Unlike ProtectWithPermissions, Protect will allow access to any registered tokens, regardless of their permissions
 // or lack thereof.
@@ -106,6 +118,13 @@ func (gate *Gate) ProtectFunc(handlerFunc http.HandlerFunc) http.HandlerFunc {
 //
 func (gate *Gate) ProtectFuncWithPermissions(handlerFunc http.HandlerFunc, permissions []string) http.HandlerFunc {
 	return func(writer http.ResponseWriter, request *http.Request) {
+		if gate.rateLimiter != nil {
+			if !gate.rateLimiter.Try() {
+				writer.WriteHeader(http.StatusTooManyRequests)
+				_, _ = writer.Write(gate.tooManyRequestsResponseBody)
+				return
+			}
+		}
 		if gate.authorizationService != nil {
 			token := extractTokenFromRequest(request)
 			if !gate.authorizationService.IsAuthorized(token, permissions) {

gate_test.go

@@ -399,3 +399,37 @@ func TestGate_ProtectWithNilAuthorizationService(t *testing.T) {
 		t.Errorf("%s %s should have returned %d, but returned %d instead", request.Method, request.URL, http.StatusOK, responseRecorder.Code)
 	}
 }
+
+func TestGate_ProtectWithNilAuthorizationServiceAndRateLimit(t *testing.T) {
+	gate := NewGate(nil).WithRateLimit(2)
+	request, _ := http.NewRequest("GET", "/handle", nil)
+	router := http.NewServeMux()
+	router.Handle("/handle", gate.Protect(&testHandler{}))
+
+	responseRecorder := httptest.NewRecorder()
+	router.ServeHTTP(responseRecorder, request)
+	if responseRecorder.Code != http.StatusOK {
+		t.Errorf("%s %s should have returned %d, but returned %d instead", request.Method, request.URL, http.StatusOK, responseRecorder.Code)
+	}
+
+	responseRecorder = httptest.NewRecorder()
+	router.ServeHTTP(responseRecorder, request)
+	if responseRecorder.Code != http.StatusOK {
+		t.Errorf("%s %s should have returned %d, but returned %d instead", request.Method, request.URL, http.StatusOK, responseRecorder.Code)
+	}
+
+	responseRecorder = httptest.NewRecorder()
+	router.ServeHTTP(responseRecorder, request)
+	if responseRecorder.Code != http.StatusTooManyRequests {
+		t.Errorf("%s %s should have returned %d, but returned %d instead", request.Method, request.URL, http.StatusTooManyRequests, responseRecorder.Code)
+	}
+
+	// Wait for rate limit time window to pass
+	time.Sleep(time.Second)
+
+	responseRecorder = httptest.NewRecorder()
+	router.ServeHTTP(responseRecorder, request)
+	if responseRecorder.Code != http.StatusOK {
+		t.Errorf("%s %s should have returned %d, but returned %d instead", request.Method, request.URL, http.StatusOK, responseRecorder.Code)
+	}
+}

ratelimiter.go

@@ -0,0 +1,42 @@
+package g8
+
+import (
+	"sync"
+	"time"
+)
+
+// RateLimiter is a fixed rate limiter
+type RateLimiter struct {
+	maximumExecutionsPerSecond int
+	executionsLeftInWindow     int
+	windowStartTime            time.Time
+	mutex                      sync.Mutex
+}
+
+// NewRateLimiter creates a RateLimiter
+func NewRateLimiter(maximumExecutionsPerSecond int) *RateLimiter {
+	return &RateLimiter{
+		windowStartTime:            time.Now(),
+		executionsLeftInWindow:     maximumExecutionsPerSecond,
+		maximumExecutionsPerSecond: maximumExecutionsPerSecond,
+	}
+}
+
+// Try updates the number of executions if the rate limit quota hasn't been reached and returns whether the
+// attempt was successful or not.
+//
+// Returns false if the execution was not successful (rate limit quota has been reached)
+// Returns true if the execution was successful (rate limit quota has not been reached)
+func (r *RateLimiter) Try() bool {
+	r.mutex.Lock()
+	defer r.mutex.Unlock()
+	if time.Now().Add(-time.Second).After(r.windowStartTime) {
+		r.windowStartTime = time.Now()
+		r.executionsLeftInWindow = r.maximumExecutionsPerSecond
+	}
+	if r.executionsLeftInWindow == 0 {
+		return false
+	}
+	r.executionsLeftInWindow--
+	return true
+}

ratelimiter_test.go

@@ -0,0 +1,73 @@
+package g8
+
+import (
+	"testing"
+	"time"
+)
+
+func TestNewRateLimiter(t *testing.T) {
+	rl := NewRateLimiter(2)
+	if rl.maximumExecutionsPerSecond != 2 {
+		t.Errorf("expected maximumExecutionsPerSecond to be %d, got %d", 2, rl.maximumExecutionsPerSecond)
+	}
+	if rl.executionsLeftInWindow != 2 {
+		t.Errorf("expected executionsLeftInWindow to be %d, got %d", 2, rl.executionsLeftInWindow)
+	}
+	// First execution: should not be rate limited
+	if notRateLimited := rl.Try(); !notRateLimited {
+		t.Error("expected Try to return true")
+	}
+	if rl.maximumExecutionsPerSecond != 2 {
+		t.Errorf("expected maximumExecutionsPerSecond to be %d, got %d", 2, rl.maximumExecutionsPerSecond)
+	}
+	if rl.executionsLeftInWindow != 1 {
+		t.Errorf("expected executionsLeftInWindow to be %d, got %d", 1, rl.executionsLeftInWindow)
+	}
+	// Second execution: should not be rate limited
+	if notRateLimited := rl.Try(); !notRateLimited {
+		t.Error("expected Try to return true")
+	}
+	if rl.maximumExecutionsPerSecond != 2 {
+		t.Errorf("expected maximumExecutionsPerSecond to be %d, got %d", 2, rl.maximumExecutionsPerSecond)
+	}
+	if rl.executionsLeftInWindow != 0 {
+		t.Errorf("expected executionsLeftInWindow to be %d, got %d", 0, rl.executionsLeftInWindow)
+	}
+	// Third execution: should be rate limited
+	if notRateLimited := rl.Try(); notRateLimited {
+		t.Error("expected Try to return false")
+	}
+	if rl.maximumExecutionsPerSecond != 2 {
+		t.Errorf("expected maximumExecutionsPerSecond to be %d, got %d", 2, rl.maximumExecutionsPerSecond)
+	}
+	if rl.executionsLeftInWindow != 0 {
+		t.Errorf("expected executionsLeftInWindow to be %d, got %d", 0, rl.executionsLeftInWindow)
+	}
+}
+
+func TestRateLimiter_Try(t *testing.T) {
+	rl := NewRateLimiter(5)
+	for i := 0; i < 20; i++ {
+		notRateLimited := rl.Try()
+		if i < 5 {
+			if !notRateLimited {
+				t.Fatal("expected to not be rate limited")
+			}
+		} else {
+			if notRateLimited {
+				t.Fatal("expected to be rate limited")
+			}
+		}
+	}
+}
+
+func TestRateLimiter_TryAlwaysUnderRateLimit(t *testing.T) {
+	rl := NewRateLimiter(20)
+	for i := 0; i < 45; i++ {
+		notRateLimited := rl.Try()
+		if !notRateLimited {
+			t.Fatal("expected to not be rate limited")
+		}
+		time.Sleep(51 * time.Millisecond)
+	}
+}