Changeset 0.32.0 (#206)

server/middleware/ratelimit.go

@@ -1,70 +1,46 @@
 package middleware
 
 import (
+	"log"
 	"net/http"
-	"sync"
-	"time"
 
 	"github.com/gin-gonic/gin"
+
+	"chat-switchboard/store"
 )
 
-type visitor struct {
-	tokens   float64
-	lastSeen time.Time
-}
-
-// RateLimiter implements a per-IP token bucket rate limiter.
+// RateLimiter implements per-IP rate limiting backed by a shared store.
+// v0.32.0: replaces in-memory token bucket for cross-pod consistency.
+//
+// Fail-open: if the store is unavailable (PG down), requests are allowed.
+// Auth endpoints have their own protections (bcrypt, lockout); blocking
+// legitimate logins due to a rate-limit DB failure is worse than a brief burst.
 type RateLimiter struct {
-	mu       sync.Mutex
-	visitors map[string]*visitor
-	rate     float64 // tokens per second
-	burst    int     // max tokens
+	store store.RateLimitStore
+	rate  float64
+	burst int
 }
 
-// NewRateLimiter creates a rate limiter.
-// rate is requests per second, burst is max burst size.
-func NewRateLimiter(rate float64, burst int) *RateLimiter {
-	rl := &RateLimiter{
-		visitors: make(map[string]*visitor),
-		rate:     rate,
-		burst:    burst,
-	}
-
-	// Cleanup stale entries every 5 minutes
-	go func() {
-		for {
-			time.Sleep(5 * time.Minute)
-			rl.cleanup()
-		}
-	}()
-
-	return rl
+// NewRateLimiter creates a rate limiter backed by the given store.
+// rate is requests per second, burst is max requests per window.
+func NewRateLimiter(s store.RateLimitStore, rate float64, burst int) *RateLimiter {
+	return &RateLimiter{store: s, rate: rate, burst: burst}
 }
 
 // Limit returns a Gin middleware that enforces the rate limit.
 func (rl *RateLimiter) Limit() gin.HandlerFunc {
 	return func(c *gin.Context) {
-		ip := c.ClientIP()
+		key := "auth:" + c.ClientIP()
 
-		rl.mu.Lock()
-		v, exists := rl.visitors[ip]
-		now := time.Now()
-
-		if !exists {
-			v = &visitor{tokens: float64(rl.burst), lastSeen: now}
-			rl.visitors[ip] = v
+		allowed, err := rl.store.Allow(c.Request.Context(), key, rl.rate, rl.burst)
+		if err != nil {
+			// Fail open — don't block auth on rate limit DB failure.
+			log.Printf("[ratelimit] store error (allowing request): %v", err)
+			c.Next()
+			return
 		}
 
-		// Refill tokens based on elapsed time
-		elapsed := now.Sub(v.lastSeen).Seconds()
-		v.tokens += elapsed * rl.rate
-		if v.tokens > float64(rl.burst) {
-			v.tokens = float64(rl.burst)
-		}
-		v.lastSeen = now
-
-		if v.tokens < 1 {
-			rl.mu.Unlock()
+		if !allowed {
 			c.Header("Retry-After", "1")
 			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
 				"error": "rate limit exceeded",
@@ -72,21 +48,6 @@ func (rl *RateLimiter) Limit() gin.HandlerFunc {
 			return
 		}
 
-		v.tokens--
-		rl.mu.Unlock()
-
 		c.Next()
 	}
 }
-
-func (rl *RateLimiter) cleanup() {
-	rl.mu.Lock()
-	defer rl.mu.Unlock()
-
-	cutoff := time.Now().Add(-10 * time.Minute)
-	for ip, v := range rl.visitors {
-		if v.lastSeen.Before(cutoff) {
-			delete(rl.visitors, ip)
-		}
-	}
-}