Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
226 lines
7.2 KiB
Go
226 lines
7.2 KiB
Go
//go:build !dev || test
|
|
|
|
package mw
|
|
|
|
import (
|
|
"crussell/clock"
|
|
"fmt"
|
|
"net/http"
|
|
"time"
|
|
)
|
|
|
|
func NewRateLimiter(limit int, window time.Duration) *RateLimiter {
|
|
rl := &RateLimiter{
|
|
requests: make(map[string][]time.Time),
|
|
limit: limit,
|
|
window: window,
|
|
}
|
|
registerLimiter(rl)
|
|
return rl
|
|
}
|
|
|
|
func (rl *RateLimiter) Allow(key string) bool {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
now := clock.Now()
|
|
windowStart := now.Add(-rl.window)
|
|
|
|
var valid []time.Time
|
|
for _, t := range rl.requests[key] {
|
|
if t.After(windowStart) {
|
|
valid = append(valid, t)
|
|
}
|
|
}
|
|
|
|
if len(valid) >= rl.limit {
|
|
rl.requests[key] = valid
|
|
return false
|
|
}
|
|
|
|
rl.requests[key] = append(valid, now)
|
|
return true
|
|
}
|
|
|
|
func NewProgressiveRateLimiter() *ProgressiveRateLimiter {
|
|
return &ProgressiveRateLimiter{
|
|
requests: make(map[string]*ipProgressiveState),
|
|
}
|
|
}
|
|
|
|
// Check returns the delay in milliseconds. Returns 0 if no delay needed.
|
|
// Strategy:
|
|
// - Count requests in last 5 seconds (burst): allow up to 30
|
|
// - Count requests in last 60 seconds (sustained): allow up to 120
|
|
// - Only delay when BOTH windows are exceeded (high sustained rate with recent bursts)
|
|
// - Progressive: once throttled, delay increases with sustained rate
|
|
func (prl *ProgressiveRateLimiter) Check(ip string) (delayMs int) {
|
|
prl.mu.Lock()
|
|
defer prl.mu.Unlock()
|
|
|
|
now := clock.Now()
|
|
state, exists := prl.requests[ip]
|
|
if !exists {
|
|
prl.requests[ip] = &ipProgressiveState{
|
|
timestamps: []time.Time{now},
|
|
}
|
|
return 0
|
|
}
|
|
|
|
state.timestamps = append(state.timestamps, now)
|
|
|
|
cutoff := now.Add(-60 * time.Second)
|
|
pruned := state.timestamps[:0]
|
|
for _, t := range state.timestamps {
|
|
if t.After(cutoff) {
|
|
pruned = append(pruned, t)
|
|
}
|
|
}
|
|
state.timestamps = pruned
|
|
|
|
burstCutoff := now.Add(-5 * time.Second)
|
|
burstCount := 0
|
|
for _, t := range state.timestamps {
|
|
if t.After(burstCutoff) {
|
|
burstCount++
|
|
}
|
|
}
|
|
|
|
sustainedCutoff := now.Add(-60 * time.Second)
|
|
sustainedCount := 0
|
|
for _, t := range state.timestamps {
|
|
if t.After(sustainedCutoff) {
|
|
sustainedCount++
|
|
}
|
|
}
|
|
|
|
if burstCount <= 30 && sustainedCount <= 120 {
|
|
return 0
|
|
}
|
|
|
|
// Progressive delay based on how far over the sustained limit they are
|
|
// Rate = requests per minute
|
|
switch {
|
|
case sustainedCount <= 140:
|
|
return 500 // 500ms - scraping but not too aggressively
|
|
case sustainedCount <= 200:
|
|
return 2000 // 2s - moderate spam
|
|
case sustainedCount <= 300:
|
|
return 5000 // 5s - heavy spam
|
|
default:
|
|
return 10000 // 10s - abuse
|
|
}
|
|
}
|
|
|
|
// progressiveRejectDelayMs is the delay at which the progressive per-IP limiter
|
|
// stops sleeping and rejects the request 429 immediately (Round 2 Loop B
|
|
// finding 5). ONLY the TOP abuse tier (10s) rejects: the 500ms / 2s / 5s tiers
|
|
// keep sleeping (backoff). Previously every tier past 2s hard-rejected, which
|
|
// under a shared NAT — or any TRUST_PROXY_HEADERS=false deployment where every
|
|
// client collapses onto the proxy's IP — locked out the whole surface behind
|
|
// one abusive client. Now a moderate sustained rate still sleeps (throttling
|
|
// the offender with backoff) instead of hard-rejecting everyone, while the 10s
|
|
// abuse tier keeps the goroutine-parking amplifier bound.
|
|
const progressiveRejectDelayMs = 10000
|
|
|
|
func ProgressiveRateLimit(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
ip := clientIP(r)
|
|
|
|
delay := globalProgressiveLimiter.Check(ip)
|
|
switch {
|
|
case delay >= progressiveRejectDelayMs:
|
|
// Top abuse tier only — far past the sustained budget, reject now
|
|
// instead of parking a goroutine for 10s. The rejection is a clean
|
|
// 429; the client retries after the burst window lapses.
|
|
RespondJSON(w, http.StatusTooManyRequests, map[string]string{"error": "Rate limit exceeded"})
|
|
return
|
|
case delay > 0:
|
|
time.Sleep(time.Duration(delay) * time.Millisecond)
|
|
w.Header().Set("X-RateLimit-Delay", fmt.Sprintf("%d", delay))
|
|
}
|
|
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
|
|
// RateLimit middleware - limits requests per IP
|
|
func RateLimit(limit int, window time.Duration) func(http.Handler) http.Handler {
|
|
limiter := NewRateLimiter(limit, window)
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
ip := clientIP(r)
|
|
|
|
if !limiter.Allow(ip) {
|
|
RespondJSON(w, http.StatusTooManyRequests, map[string]string{"error": "Rate limit exceeded"})
|
|
return
|
|
}
|
|
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
}
|
|
|
|
// RateLimitByUserAndIP limits requests per authenticated user + client IP. The
|
|
// key combines the authenticated userID (mw.UserIDKey, injected by RequireAuth)
|
|
// with the derived client IP, so a per-IP budget can never collapse into a
|
|
// single GLOBAL bucket when the backend sits behind a proxy that does not set
|
|
// TRUST_PROXY_HEADERS=true: without that flag ClientIP keys every request on
|
|
// RemoteAddr = the proxy's IP, so one account holder could otherwise exhaust
|
|
// the shared budget and permanently 429 the whole surface for everyone. With
|
|
// the userID in the key each account gets its own independent budget per IP.
|
|
// When no userID is present (unauthenticated path) the key falls back to
|
|
// ClientIP alone, matching RateLimit's behaviour.
|
|
func RateLimitByUserAndIP(limit int, window time.Duration) func(http.Handler) http.Handler {
|
|
limiter := NewRateLimiter(limit, window)
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
key := ClientIP(r)
|
|
if userID, ok := GetUserID(r.Context()); ok && userID != "" {
|
|
key = userID + "|" + key
|
|
}
|
|
|
|
if !limiter.Allow(key) {
|
|
RespondJSON(w, http.StatusTooManyRequests, map[string]string{"error": "Rate limit exceeded"})
|
|
return
|
|
}
|
|
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
}
|
|
|
|
// RateLimitByUser limits requests per authenticated user ID ALONE, dropping the
|
|
// IP component entirely. This is the B8 safeguard for the 2FA surface (and the
|
|
// B16 gift-card redeem budget): when the IP is part of the key, a client that
|
|
// can rotate its source IP — or that sits behind a proxy which echoes a
|
|
// client-supplied CF-Connecting-IP when TRUST_PROXY_HEADERS is misconfigured
|
|
// true — mints a fresh bucket per IP for the SAME account, collapsing the
|
|
// per-account budget. Keying on the userID alone guarantees exactly one budget
|
|
// per account regardless of IP rotation or proxy configuration. When no userID
|
|
// is present (unauthenticated path) the key falls back to ClientIP so the
|
|
// surface still has a default budget.
|
|
func RateLimitByUser(limit int, window time.Duration) func(http.Handler) http.Handler {
|
|
limiter := NewRateLimiter(limit, window)
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
key, ok := GetUserID(r.Context())
|
|
if !ok || key == "" {
|
|
key = ClientIP(r)
|
|
}
|
|
|
|
if !limiter.Allow(key) {
|
|
RespondJSON(w, http.StatusTooManyRequests, map[string]string{"error": "Rate limit exceeded"})
|
|
return
|
|
}
|
|
|
|
next.ServeHTTP(w, r)
|
|
})
|
|
}
|
|
}
|
|
|
|
// clientIP derives the per-client rate-limit key. It is a thin alias of the
|
|
// exported ClientIP (which lives in ratelimit_shared.go so it is available in
|
|
// every build configuration), kept for backward compatibility with existing
|
|
// callers.
|
|
func clientIP(r *http.Request) string { return ClientIP(r) }
|