improve rate limites

Author: bloom

go.mod

@@ -46,7 +46,7 @@ require (
 	github.com/aws/smithy-go v1.24.0 // indirect
 	github.com/aymerick/douceur v0.2.0 // indirect
 	github.com/clipperhouse/stringish v0.1.1 // indirect
-	github.com/clipperhouse/uax29/v2 v2.3.0 // indirect
+	github.com/clipperhouse/uax29/v2 v2.3.1 // indirect
 	github.com/dlclark/regexp2 v1.11.5 // indirect
 	github.com/fsnotify/fsnotify v1.9.0 // indirect
 	github.com/gorilla/css v1.0.1 // indirect

go.sum

@@ -50,8 +50,8 @@ github.com/aymerick/douceur v0.2.0 h1:Mv+mAeH1Q+n9Fr+oyamOlAkUNPWPlA8PPGR0QAaYuP
 github.com/aymerick/douceur v0.2.0/go.mod h1:wlT5vV2O3h55X9m7iVYN0TBM0NH/MmbLnd30/FjWUq4=
 github.com/clipperhouse/stringish v0.1.1 h1:+NSqMOr3GR6k1FdRhhnXrLfztGzuG+VuFDfatpWHKCs=
 github.com/clipperhouse/stringish v0.1.1/go.mod h1:v/WhFtE1q0ovMta2+m+UbpZ+2/HEXNWYXQgCt4hdOzA=
-github.com/clipperhouse/uax29/v2 v2.3.0 h1:SNdx9DVUqMoBuBoW3iLOj4FQv3dN5mDtuqwuhIGpJy4=
-github.com/clipperhouse/uax29/v2 v2.3.0/go.mod h1:Wn1g7MK6OoeDT0vL+Q0SQLDz/KpfsVRgg6W7ihQeh4g=
+github.com/clipperhouse/uax29/v2 v2.3.1 h1:RjM8gnVbFbgI67SBekIC7ihFpyXwRPYWXn9BZActHbw=
+github.com/clipperhouse/uax29/v2 v2.3.1/go.mod h1:Wn1g7MK6OoeDT0vL+Q0SQLDz/KpfsVRgg6W7ihQeh4g=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=

pkg/errs/errors.go

@@ -72,7 +72,7 @@ func IsInternal(err error) bool {
 
 	switch err.(type) {
 	case *NotFoundError, *InvalidArgumentError,
-		*PermissionDeniedError, *AuthenticationRequiredError, *AlreadyExistsError:
+		*PermissionDeniedError, *AuthenticationRequiredError, *AlreadyExistsError, *TooManyRequestsError:
 		return false
 	default:
 		return true
@@ -136,3 +136,17 @@ func (err *AuthenticationRequiredError) Error() string {
 func AuthenticationRequired(message string) *AuthenticationRequiredError {
 	return &AuthenticationRequiredError{message: message}
 }
+
+// TooManyRequestsError is a wrapper for an error when a client is sending too many requests
+type TooManyRequestsError struct {
+	// message string
+}
+
+func (err *TooManyRequestsError) Error() string {
+	return "Too many requests. Please try again later."
+}
+
+// TooManyRequests returns a new TooManyRequestsError
+func TooManyRequests() *TooManyRequestsError {
+	return &TooManyRequestsError{}
+}

pkg/ratelimit/ratelimit.go

@@ -3,17 +3,19 @@ package ratelimit
 
 import (
 	"encoding/binary"
+	"math/rand/v2"
 	"sync"
 	"time"
 
-	"github.com/skerkour/stdx-go/crypto/blake3"
+	"github.com/skerkour/stdx-go/xxh3"
 )
 
 // Limiter tracks request counts within time buckets.
 type Limiter struct {
-	mutex   sync.Mutex
-	buckets map[[32]byte]*bucket
-	stop    chan struct{}
+	mutex    sync.Mutex
+	buckets  map[uint64]*bucket
+	stop     chan struct{}
+	hashSeed uint64
 }
 
 type bucket struct {
@@ -23,32 +25,35 @@ type bucket struct {
 
 // New creates a new rate limiter with automatic cleanup of expired buckets.
 func New() *Limiter {
+
 	limiter := &Limiter{
-		mutex:   sync.Mutex{},
-		buckets: make(map[[32]byte]*bucket),
-		stop:    make(chan struct{}),
+		mutex:    sync.Mutex{},
+		buckets:  make(map[uint64]*bucket),
+		stop:     make(chan struct{}),
+		hashSeed: rand.Uint64(),
 	}
 	go limiter.cleanupLoop()
 	return limiter
 }
 
-// RateLimit checks if an action by an actor is allowed within the rate limit.
+// IsAllowed checks if an action by an actor is allowed within the rate limit.
 // It returns true if the action is allowed, false if rate limited.
 //
 // Parameters:
+//   - namespace: optional namespace for the action check (e.g. tenant ID). Can be nil.
 //   - action: identifies the type of action being rate limited (e.g., "login", "api-call")
 //   - actor: identifies who is performing the action (e.g., user ID, IP address)
 //   - timeBucket: the duration of each rate limit window
 //   - allowed: maximum number of actions allowed per time bucket
-func (limiter *Limiter) RateLimit(action string, actor []byte, timeBucket time.Duration, allowed uint64) bool {
+func (limiter *Limiter) IsAllowed(action string, namespace []byte, actor []byte, timeBucket time.Duration, allowed uint64) bool {
 	now := time.Now()
 	bucketStart := now.Truncate(timeBucket)
-	key := makeKey(action, actor, uint64(bucketStart.UnixNano()), uint64(timeBucket.Nanoseconds()))
+	key := limiter.makeKey(action, namespace, actor, uint64(bucketStart.UnixNano()), uint64(timeBucket.Nanoseconds()))
 
 	limiter.mutex.Lock()
 	defer limiter.mutex.Unlock()
 
-	b, exists := limiter.buckets[key]
+	existingBucket, exists := limiter.buckets[key]
 	if !exists {
 		limiter.buckets[key] = &bucket{
 			count:   1,
@@ -57,20 +62,20 @@ func (limiter *Limiter) RateLimit(action string, actor []byte, timeBucket time.D
 		return true
 	}
 
-	if b.count >= allowed {
+	if existingBucket.count >= allowed {
 		return false
 	}
 
-	b.count++
+	existingBucket.count++
 	return true
 }
 
 // Count returns the current count for an action/actor in the current time bucket.
 // Useful for showing users how many requests they have remaining.
-func (limiter *Limiter) Count(action string, actor []byte, timeBucket time.Duration) uint64 {
+func (limiter *Limiter) Count(action string, namespace []byte, actor []byte, timeBucket time.Duration) uint64 {
 	now := time.Now()
 	bucketStart := now.Truncate(timeBucket)
-	key := makeKey(action, actor, uint64(bucketStart.UnixNano()), uint64(timeBucket.Nanoseconds()))
+	key := limiter.makeKey(action, namespace, actor, uint64(bucketStart.UnixNano()), uint64(timeBucket.Nanoseconds()))
 
 	limiter.mutex.Lock()
 	defer limiter.mutex.Unlock()
@@ -82,8 +87,8 @@ func (limiter *Limiter) Count(action string, actor []byte, timeBucket time.Durat
 }
 
 // Remaining returns how many requests are remaining for an action/actor.
-func (limiter *Limiter) Remaining(action string, actor []byte, timeBucket time.Duration, allowed uint64) uint64 {
-	count := limiter.Count(action, actor, timeBucket)
+func (limiter *Limiter) Remaining(action string, namespace []byte, actor []byte, timeBucket time.Duration, allowed uint64) uint64 {
+	count := limiter.Count(action, namespace, actor, timeBucket)
 	if count >= allowed {
 		return 0
 	}
@@ -123,15 +128,14 @@ func (limiter *Limiter) cleanup() {
 	}
 }
 
-func makeKey(action string, actor []byte, bucketStartNanos uint64, timeBucketNanos uint64) [32]byte {
-	var hash [32]byte
-
-	hasher := blake3.New(32, nil)
+// makeKey returns a stable key by hashing the inputs. It currently uses xxh3.
+func (limiter *Limiter) makeKey(action string, namespace []byte, actor []byte, bucketStartNanos uint64, timeBucketNanos uint64) uint64 {
+	hasher := xxh3.NewSeed(limiter.hashSeed)
 	hasher.Write([]byte(action))
+	hasher.Write(namespace)
 	hasher.Write(actor)
 	binary.Write(hasher, binary.LittleEndian, bucketStartNanos)
 	binary.Write(hasher, binary.LittleEndian, timeBucketNanos)
 
-	hasher.Sum(hash[:0])
-	return hash
+	return hasher.Sum64()
 }

pkg/ratelimit/ratelimit_test.go

@@ -17,18 +17,18 @@ func TestRateLimit_Basic(t *testing.T) {
 
 	// First 3 should be allowed
 	for i := 0; i < 3; i++ {
-		if !l.RateLimit(action, actor, bucket, allowed) {
+		if !l.IsAllowed(action, nil, actor, bucket, allowed) {
 			t.Errorf("Request %d should have been allowed", i+1)
 		}
 	}
 
 	// 4th should be rate limited
-	if l.RateLimit(action, actor, bucket, allowed) {
+	if l.IsAllowed(action, nil, actor, bucket, allowed) {
 		t.Error("Request 4 should have been rate limited")
 	}
 
 	// 5th should also be rate limited
-	if l.RateLimit(action, actor, bucket, allowed) {
+	if l.IsAllowed(action, nil, actor, bucket, allowed) {
 		t.Error("Request 5 should have been rate limited")
 	}
 }
@@ -37,23 +37,24 @@ func TestRateLimit_DifferentActors(t *testing.T) {
 	l := New()
 	defer l.Stop()
 
+	namespace := []byte("namespace1")
 	actor1 := []byte("user1")
 	actor2 := []byte("user2")
 	action := "test-action"
 	bucket := time.Minute
 	allowed := uint64(2)
 
 	// Actor 1 uses both requests
-	l.RateLimit(action, actor1, bucket, allowed)
-	l.RateLimit(action, actor1, bucket, allowed)
+	l.IsAllowed(action, namespace, actor1, bucket, allowed)
+	l.IsAllowed(action, namespace, actor1, bucket, allowed)
 
 	// Actor 1 should be limited
-	if l.RateLimit(action, actor1, bucket, allowed) {
+	if l.IsAllowed(action, namespace, actor1, bucket, allowed) {
 		t.Error("Actor1 should be rate limited")
 	}
 
 	// Actor 2 should still be allowed
-	if !l.RateLimit(action, actor2, bucket, allowed) {
+	if !l.IsAllowed(action, namespace, actor2, bucket, allowed) {
 		t.Error("Actor2 should be allowed")
 	}
 }
@@ -63,21 +64,22 @@ func TestRateLimit_DifferentActions(t *testing.T) {
 	defer l.Stop()
 
 	actor := []byte("user123")
+	namespace := []byte("namespace")
 	action1 := "action1"
 	action2 := "action2"
 	bucket := time.Minute
 	allowed := uint64(1)
 
 	// Use up action1 limit
-	l.RateLimit(action1, actor, bucket, allowed)
+	l.IsAllowed(action1, namespace, actor, bucket, allowed)
 
 	// Action1 should be limited
-	if l.RateLimit(action1, actor, bucket, allowed) {
+	if l.IsAllowed(action1, namespace, actor, bucket, allowed) {
 		t.Error("Action1 should be rate limited")
 	}
 
 	// Action2 should still be allowed
-	if !l.RateLimit(action2, actor, bucket, allowed) {
+	if !l.IsAllowed(action2, namespace, actor, bucket, allowed) {
 		t.Error("Action2 should be allowed")
 	}
 }
@@ -86,24 +88,25 @@ func TestRateLimit_BucketReset(t *testing.T) {
 	l := New()
 	defer l.Stop()
 
+	namespace := []byte("namespace")
 	actor := []byte("user123")
 	action := "test-action"
 	bucket := 100 * time.Millisecond
 	allowed := uint64(2)
 
 	// Use up the limit
-	l.RateLimit(action, actor, bucket, allowed)
-	l.RateLimit(action, actor, bucket, allowed)
+	l.IsAllowed(action, namespace, actor, bucket, allowed)
+	l.IsAllowed(action, namespace, actor, bucket, allowed)
 
-	if l.RateLimit(action, actor, bucket, allowed) {
+	if l.IsAllowed(action, namespace, actor, bucket, allowed) {
 		t.Error("Should be rate limited")
 	}
 
 	// Wait for next bucket
 	time.Sleep(150 * time.Millisecond)
 
 	// Should be allowed again
-	if !l.RateLimit(action, actor, bucket, allowed) {
+	if !l.IsAllowed(action, namespace, actor, bucket, allowed) {
 		t.Error("Should be allowed after bucket reset")
 	}
 }
@@ -112,18 +115,19 @@ func TestCount(t *testing.T) {
 	l := New()
 	defer l.Stop()
 
+	namespace := []byte("namespace")
 	actor := []byte("user123")
 	action := "test-action"
 	bucket := time.Minute
 
-	if c := l.Count(action, actor, bucket); c != 0 {
+	if c := l.Count(action, namespace, actor, bucket); c != 0 {
 		t.Errorf("Expected count 0, got %d", c)
 	}
 
-	l.RateLimit(action, actor, bucket, 10)
-	l.RateLimit(action, actor, bucket, 10)
+	l.IsAllowed(action, namespace, actor, bucket, 10)
+	l.IsAllowed(action, namespace, actor, bucket, 10)
 
-	if c := l.Count(action, actor, bucket); c != 2 {
+	if c := l.Count(action, namespace, actor, bucket); c != 2 {
 		t.Errorf("Expected count 2, got %d", c)
 	}
 }
@@ -132,19 +136,20 @@ func TestRemaining(t *testing.T) {
 	l := New()
 	defer l.Stop()
 
+	namespace := []byte("namespace")
 	actor := []byte("user123")
 	action := "test-action"
 	bucket := time.Minute
 	allowed := uint64(5)
 
-	if r := l.Remaining(action, actor, bucket, allowed); r != 5 {
+	if r := l.Remaining(action, namespace, actor, bucket, allowed); r != 5 {
 		t.Errorf("Expected 5 remaining, got %d", r)
 	}
 
-	l.RateLimit(action, actor, bucket, allowed)
-	l.RateLimit(action, actor, bucket, allowed)
+	l.IsAllowed(action, namespace, actor, bucket, allowed)
+	l.IsAllowed(action, namespace, actor, bucket, allowed)
 
-	if r := l.Remaining(action, actor, bucket, allowed); r != 3 {
+	if r := l.Remaining(action, namespace, actor, bucket, allowed); r != 3 {
 		t.Errorf("Expected 3 remaining, got %d", r)
 	}
 }
@@ -153,6 +158,7 @@ func TestRateLimit_Concurrent(t *testing.T) {
 	l := New()
 	defer l.Stop()
 
+	namespace := []byte("namespace")
 	actor := []byte("user123")
 	action := "test-action"
 	bucket := time.Minute
@@ -162,12 +168,12 @@ func TestRateLimit_Concurrent(t *testing.T) {
 	allowedCount := 0
 	var mu sync.Mutex
 
-	// Run 200 concurrent requests
-	for i := 0; i < 200; i++ {
+	// Run 300 concurrent requests
+	for i := 0; i < 300; i++ {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			if l.RateLimit(action, actor, bucket, allowed) {
+			if l.IsAllowed(action, namespace, actor, bucket, allowed) {
 				mu.Lock()
 				allowedCount++
 				mu.Unlock()
@@ -192,11 +198,12 @@ func TestRateLimit_BinaryActor(t *testing.T) {
 	action := "test-action"
 	bucket := time.Minute
 	allowed := uint64(1)
+	namespace := []byte("namespace")
 
-	if !l.RateLimit(action, actor, bucket, allowed) {
+	if !l.IsAllowed(action, namespace, actor, bucket, allowed) {
 		t.Error("First request should be allowed")
 	}
-	if l.RateLimit(action, actor, bucket, allowed) {
+	if l.IsAllowed(action, namespace, actor, bucket, allowed) {
 		t.Error("Second request should be rate limited")
 	}
 }

pkg/server/apiutil/apiutil.go

@@ -34,6 +34,7 @@ const (
 	ErrorCodeInternal               apiErrorCode = "INTERNAL"
 	ErrorCodePermissionDenied       apiErrorCode = "PERMISSION_DENIED"
 	ErrorCodeAuthenticationRequired apiErrorCode = "AUTHENTICATION_REQUIRED"
+	ErrorCodeTooManyRequests        apiErrorCode = "TOO_MANY_REQUESTS"
 )
 
 // var jsonEncodingOptions = json.JoinOptions(json.FormatNilMapAsNull(true))
@@ -135,6 +136,9 @@ func SendError(ctx context.Context, w http.ResponseWriter, err error) {
 	case *errs.AuthenticationRequiredError:
 		code = ErrorCodeAuthenticationRequired
 		statusCode = http.StatusUnauthorized
+	case *errs.TooManyRequestsError:
+		code = ErrorCodeTooManyRequests
+		statusCode = http.StatusTooManyRequests
 	default:
 		code = ErrorCodeInternal
 		statusCode = http.StatusInternalServerError

pkg/services/site/service/login.go

@@ -34,8 +34,8 @@ func (service *SiteService) Login(ctx context.Context, input site.LoginInput) (r
 		return
 	}
 
-	if !service.rateLimiter.RateLimit("SiteService.Login", httpCtx.Client.IP.AsSlice(), time.Hour, 20) {
-		err = errs.InvalidArgument("Too many requests. Please try again later.")
+	if !service.rateLimiter.IsAllowed("SiteService.Login", website.ID.Bytes(), httpCtx.Client.IP.AsSlice(), time.Hour, 30) {
+		err = errs.TooManyRequests()
 		return
 	}