improve: use redis_rate for rate limiter

infrahq · Sep 17, 2022 · 430b9bc · 430b9bc
1 parent f254535
commit 430b9bc
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 12 deletions.
diff --git a/go.mod b/go.mod
@@ -38,6 +38,7 @@ require (
 	github.com/creack/pty v1.1.18
 	github.com/getkin/kin-openapi v0.100.0
 	github.com/go-redis/redis/v8 v8.11.5
+	github.com/go-redis/redis_rate/v9 v9.1.2
 	github.com/google/go-cmp v0.5.9
 	github.com/hinshun/vt10x v0.0.0-20220119200601-820417d04eec
 	github.com/iancoleman/strcase v0.2.0

diff --git a/go.sum b/go.sum
@@ -211,6 +211,8 @@ github.com/go-playground/validator/v10 v10.11.0 h1:0W+xRM511GY47Yy3bZUbJVitCNg2B
 github.com/go-playground/validator/v10 v10.11.0/go.mod h1:i+3WkQ1FvaUjjxh1kSvIA4dMGDBiPU55YFDl0WbKdWU=
 github.com/go-redis/redis/v8 v8.11.5 h1:AcZZR7igkdvfVmQTPnu9WE37LRrO/YrBH5zWyjDC0oI=
 github.com/go-redis/redis/v8 v8.11.5/go.mod h1:gREzHqY1hg6oD9ngVRbLStwAWKhA0FEgq8Jd4h5lpwo=
+github.com/go-redis/redis_rate/v9 v9.1.2 h1:H0l5VzoAtOE6ydd38j8MCq3ABlGLnvvbA1xDSVVCHgQ=
+github.com/go-redis/redis_rate/v9 v9.1.2/go.mod h1:oam2de2apSgRG8aJzwJddXbNu91Iyz1m8IKJE2vpvlQ=
 github.com/go-stack/stack v1.8.0/go.mod h1:v0f6uXyyMGvRgIKkXu+yp6POWl0qKG85gN/melR3HDY=
 github.com/go-test/deep v1.0.2/go.mod h1:wGDj63lr65AM2AQyKZd/NYHGb0R+1RLqB8NKt3aSFNA=
 github.com/go-test/deep v1.0.4 h1:u2CU3YKy9I2pmu9pX0eq50wCgjfGIt539SqR7FbHiho=

diff --git a/internal/server/cache/cache.go b/internal/server/cache/cache.go
@@ -3,9 +3,9 @@ package cache
 import (
 	"context"
 	"fmt"
-	"time"
 
 	"github.com/go-redis/redis/v8"
+	rate "github.com/go-redis/redis_rate/v9"
 
 	"github.com/infrahq/infra/internal/logging"
 )
@@ -44,19 +44,23 @@ func NewCache(options Options) *Cache {
 	}
 }
 
-func (c *Cache) RateOK(key string, limit int64) bool {
+// RateOK checks if the rate per minute is acceptable for the specified key
+func (c *Cache) RateOK(key string, limit int) bool {
 	if c.client != nil {
 		ctx := context.TODO()
-		bucket := time.Now().Round(time.Minute)
-		key := fmt.Sprintf("%s-%v", key, bucket.Unix())
-
-		incr := c.client.Incr(ctx, key)
-		c.client.ExpireAt(ctx, key, bucket.Add(2*time.Minute))
-		if incr.Val() >= limit {
-			c.client.Decr(ctx, key)
-
-			return false
+		limiter := rate.NewLimiter(c.client)
+		result, err := limiter.Allow(ctx, key, rate.PerMinute(limit))
+		if err != nil {
+			panic(err)
 		}
+
+		logging.L.Debug().
+			Int("allowed", result.Allowed).
+			Int("remaining", result.Remaining).
+			Dur("retry_after", result.RetryAfter).
+			Msg("")
+		// TODO: also return result.Remaining and result.RetryAfter so the response headers can be set
+		return result.Allowed > 0
 	}
 
 	return true

diff --git a/internal/server/routes.go b/internal/server/routes.go
@@ -223,7 +223,8 @@ func wrapRoute[Req, Res any](a *API, routeID routeIdentifier, route route[Req, R
 		}
 
 		if org != nil {
-			if !a.server.cache.RateOK(org.ID.String(), 10) {
+			// TODO: limit should be a per-organization setting
+			if !a.server.cache.RateOK(org.ID.String(), 1000) {
 				return internal.ErrTooManyRequests
 			}
 		}