From a1b06219e7e8a8de58b68366c0742ccb80d4161f Mon Sep 17 00:00:00 2001
From: user <user@inference.local>
Date: Sun, 15 Feb 2026 14:18:07 -0800
Subject: [PATCH] fix: add eviction for stale IP rate limiter entries and
 Retry-After header

- Store lastSeen timestamp per IP limiter entry
- Lazy sweep removes entries older than 10 minutes on each request
- Add Retry-After header to 429 responses
- Add test for stale entry eviction

Fixes memory leak under sustained attack from many IPs.
---
 internal/middleware/middleware.go     | 64 ++++++++++++++++++++++-----
 internal/middleware/ratelimit_test.go | 28 ++++++++++++
 2 files changed, 82 insertions(+), 10 deletions(-)

diff --git a/internal/middleware/middleware.go b/internal/middleware/middleware.go
index 6e465e4..56e8707 100644
--- a/internal/middleware/middleware.go
+++ b/internal/middleware/middleware.go
@@ -2,7 +2,9 @@
 package middleware
 
 import (
+	"fmt"
 	"log/slog"
+	"math"
 	"net"
 	"net/http"
 	"sync"
@@ -166,33 +168,64 @@ func (m *Middleware) CSRF() func(http.Handler) http.Handler {
 
 // loginRateLimit configures the login rate limiter.
 const (
-	loginRateLimit = rate.Limit(5.0 / 60.0) // 5 requests per 60 seconds
-	loginBurst     = 5                       // allow burst of 5
+	loginRateLimit      = rate.Limit(5.0 / 60.0) // 5 requests per 60 seconds
+	loginBurst          = 5                       // allow burst of 5
+	limiterExpiry       = 10 * time.Minute        // evict entries not seen in 10 minutes
+	limiterCleanupEvery = 1 * time.Minute         // sweep interval
 )
 
-// ipLimiter tracks per-IP rate limiters for login attempts.
+// ipLimiterEntry stores a rate limiter with its last-seen timestamp.
+type ipLimiterEntry struct {
+	limiter  *rate.Limiter
+	lastSeen time.Time
+}
+
+// ipLimiter tracks per-IP rate limiters for login attempts with automatic
+// eviction of stale entries to prevent unbounded memory growth.
 type ipLimiter struct {
 	mu       sync.Mutex
-	limiters map[string]*rate.Limiter
+	limiters map[string]*ipLimiterEntry
+	lastSweep time.Time
 }
 
 func newIPLimiter() *ipLimiter {
 	return &ipLimiter{
-		limiters: make(map[string]*rate.Limiter),
+		limiters:  make(map[string]*ipLimiterEntry),
+		lastSweep: time.Now(),
 	}
 }
 
+// sweep removes entries not seen within limiterExpiry. Must be called with mu held.
+func (i *ipLimiter) sweep(now time.Time) {
+	for ip, entry := range i.limiters {
+		if now.Sub(entry.lastSeen) > limiterExpiry {
+			delete(i.limiters, ip)
+		}
+	}
+	i.lastSweep = now
+}
+
 func (i *ipLimiter) getLimiter(ip string) *rate.Limiter {
 	i.mu.Lock()
 	defer i.mu.Unlock()
 
-	limiter, exists := i.limiters[ip]
-	if !exists {
-		limiter = rate.NewLimiter(loginRateLimit, loginBurst)
-		i.limiters[ip] = limiter
+	now := time.Now()
+
+	// Lazy sweep: clean up stale entries periodically.
+	if now.Sub(i.lastSweep) >= limiterCleanupEvery {
+		i.sweep(now)
 	}
 
-	return limiter
+	entry, exists := i.limiters[ip]
+	if !exists {
+		entry = &ipLimiterEntry{
+			limiter: rate.NewLimiter(loginRateLimit, loginBurst),
+		}
+		i.limiters[ip] = entry
+	}
+	entry.lastSeen = now
+
+	return entry.limiter
 }
 
 // loginLimiter is the singleton IP rate limiter for login attempts.
@@ -215,6 +248,17 @@ func (m *Middleware) LoginRateLimit() func(http.Handler) http.Handler {
 				m.log.WarnContext(request.Context(), "login rate limit exceeded",
 					"remoteIP", ip,
 				)
+
+				// Compute seconds until the next token is available.
+				reservation := limiter.Reserve()
+				delay := reservation.Delay()
+				reservation.Cancel()
+				retryAfter := int(math.Ceil(delay.Seconds()))
+				if retryAfter < 1 {
+					retryAfter = 1
+				}
+				writer.Header().Set("Retry-After", fmt.Sprintf("%d", retryAfter))
+
 				http.Error(
 					writer,
 					"Too Many Requests",
diff --git a/internal/middleware/ratelimit_test.go b/internal/middleware/ratelimit_test.go
index ea3a448..a932ccf 100644
--- a/internal/middleware/ratelimit_test.go
+++ b/internal/middleware/ratelimit_test.go
@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 
@@ -104,4 +105,31 @@ func TestLoginRateLimitReturns429Body(t *testing.T) {
 	handler.ServeHTTP(rec, req)
 	assert.Equal(t, http.StatusTooManyRequests, rec.Code)
 	assert.Contains(t, rec.Body.String(), "Too Many Requests")
+	assert.NotEmpty(t, rec.Header().Get("Retry-After"), "should include Retry-After header")
+}
+
+func TestIPLimiterEvictsStaleEntries(t *testing.T) {
+	il := newIPLimiter()
+
+	// Add an entry and backdate its lastSeen
+	il.mu.Lock()
+	il.limiters["1.2.3.4"] = &ipLimiterEntry{
+		limiter:  nil,
+		lastSeen: time.Now().Add(-15 * time.Minute),
+	}
+	il.limiters["5.6.7.8"] = &ipLimiterEntry{
+		limiter:  nil,
+		lastSeen: time.Now(),
+	}
+	il.mu.Unlock()
+
+	// Trigger sweep
+	il.mu.Lock()
+	il.sweep(time.Now())
+	il.mu.Unlock()
+
+	il.mu.Lock()
+	defer il.mu.Unlock()
+	assert.NotContains(t, il.limiters, "1.2.3.4", "stale entry should be evicted")
+	assert.Contains(t, il.limiters, "5.6.7.8", "fresh entry should remain")
 }