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.
This commit is contained in:
parent
66661d1b1d
commit
a1b06219e7
@ -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",
|
||||
|
||||
@ -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")
|
||||
}
|
||||
|
||||
Loading…
Reference in New Issue
Block a user