fix Time-of-Check-Time-of-Use bug in rate limiter

Author: aldas

middleware/rate_limiter.go

@@ -249,8 +249,9 @@ func (store *RateLimiterMemoryStore) Allow(identifier string) (bool, error) {
 	if now.Sub(store.lastCleanup) > store.expiresIn {
 		store.cleanupStaleVisitors()
 	}
+	allowed := limiter.AllowN(now, 1)
 	store.mutex.Unlock()
-	return limiter.AllowN(store.timeNow(), 1), nil
+	return allowed, nil
 }
 
 /*

middleware/rate_limiter_test.go

@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -457,3 +458,142 @@ func BenchmarkRateLimiterMemoryStore_conc100_10000(b *testing.B) {
 	var store = NewRateLimiterMemoryStoreWithConfig(RateLimiterMemoryStoreConfig{Rate: 100, Burst: 200, ExpiresIn: testExpiresIn})
 	benchmarkStore(store, 100, 10000, b)
 }
+
+// TestRateLimiterMemoryStore_TOCTOUFix verifies that the TOCTOU race condition is fixed
+// by ensuring timeNow() is only called once per Allow() call
+func TestRateLimiterMemoryStore_TOCTOUFix(t *testing.T) {
+	t.Parallel()
+
+	store := NewRateLimiterMemoryStoreWithConfig(RateLimiterMemoryStoreConfig{
+		Rate:      1,
+		Burst:     1,
+		ExpiresIn: 2 * time.Second,
+	})
+
+	// Track time calls to verify we use the same time value
+	timeCallCount := 0
+	baseTime := time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC)
+
+	store.timeNow = func() time.Time {
+		timeCallCount++
+		return baseTime
+	}
+
+	// First request - should succeed
+	allowed, err := store.Allow("127.0.0.1")
+	assert.NoError(t, err)
+	assert.True(t, allowed, "First request should be allowed")
+
+	// Verify timeNow() was only called once
+	assert.Equal(t, 1, timeCallCount, "timeNow() should only be called once per Allow()")
+}
+
+// TestRateLimiterMemoryStore_ConcurrentAccess verifies rate limiting correctness under concurrent load
+func TestRateLimiterMemoryStore_ConcurrentAccess(t *testing.T) {
+	t.Parallel()
+
+	store := NewRateLimiterMemoryStoreWithConfig(RateLimiterMemoryStoreConfig{
+		Rate:      10,
+		Burst:     5,
+		ExpiresIn: 5 * time.Second,
+	})
+
+	const goroutines = 50
+	const requestsPerGoroutine = 20
+
+	var wg sync.WaitGroup
+	var allowedCount, deniedCount int32
+
+	for i := 0; i < goroutines; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for j := 0; j < requestsPerGoroutine; j++ {
+				allowed, err := store.Allow("test-user")
+				assert.NoError(t, err)
+				if allowed {
+					atomic.AddInt32(&allowedCount, 1)
+				} else {
+					atomic.AddInt32(&deniedCount, 1)
+				}
+				time.Sleep(time.Millisecond)
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	totalRequests := goroutines * requestsPerGoroutine
+	allowed := int(allowedCount)
+	denied := int(deniedCount)
+
+	assert.Equal(t, totalRequests, allowed+denied, "All requests should be processed")
+	assert.Greater(t, denied, 0, "Some requests should be denied due to rate limiting")
+	assert.Greater(t, allowed, 0, "Some requests should be allowed")
+}
+
+// TestRateLimiterMemoryStore_RaceDetection verifies no data races with high concurrency
+// Run with: go test -race ./middleware -run TestRateLimiterMemoryStore_RaceDetection
+func TestRateLimiterMemoryStore_RaceDetection(t *testing.T) {
+	t.Parallel()
+
+	store := NewRateLimiterMemoryStoreWithConfig(RateLimiterMemoryStoreConfig{
+		Rate:      100,
+		Burst:     200,
+		ExpiresIn: 1 * time.Second,
+	})
+
+	const goroutines = 100
+	const requestsPerGoroutine = 100
+
+	var wg sync.WaitGroup
+	identifiers := []string{"user1", "user2", "user3", "user4", "user5"}
+
+	for i := 0; i < goroutines; i++ {
+		wg.Add(1)
+		go func(routineID int) {
+			defer wg.Done()
+			for j := 0; j < requestsPerGoroutine; j++ {
+				identifier := identifiers[routineID%len(identifiers)]
+				_, err := store.Allow(identifier)
+				assert.NoError(t, err)
+			}
+		}(i)
+	}
+
+	wg.Wait()
+}
+
+// TestRateLimiterMemoryStore_TimeOrdering verifies time ordering consistency in rate limiting decisions
+func TestRateLimiterMemoryStore_TimeOrdering(t *testing.T) {
+	t.Parallel()
+
+	store := NewRateLimiterMemoryStoreWithConfig(RateLimiterMemoryStoreConfig{
+		Rate:      1,
+		Burst:     2,
+		ExpiresIn: 5 * time.Second,
+	})
+
+	currentTime := time.Date(2009, time.November, 10, 23, 0, 0, 0, time.UTC)
+	store.timeNow = func() time.Time {
+		return currentTime
+	}
+
+	// First two requests should succeed (burst=2)
+	allowed1, _ := store.Allow("user1")
+	assert.True(t, allowed1, "Request 1 should be allowed (burst)")
+
+	allowed2, _ := store.Allow("user1")
+	assert.True(t, allowed2, "Request 2 should be allowed (burst)")
+
+	// Third request should be denied
+	allowed3, _ := store.Allow("user1")
+	assert.False(t, allowed3, "Request 3 should be denied (burst exhausted)")
+
+	// Advance time by 1 second
+	currentTime = currentTime.Add(1 * time.Second)
+
+	// Fourth request should succeed
+	allowed4, _ := store.Allow("user1")
+	assert.True(t, allowed4, "Request 4 should be allowed (1 token available)")
+}