Crussell/backend/mw/ratelimit.go

//go:build !dev || test

package mw

import (
	"crussell/clock"
	"fmt"
	"net"
	"net/http"
	"time"

	"github.com/go-chi/chi/v5/middleware"
)

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 60
//   - 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)

	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
	}
}

func ProgressiveRateLimit(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ip := clientIP(r)

		delay := globalProgressiveLimiter.Check(ip)
		if 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)
		})
	}
}

// clientIP derives the per-client rate-limit key. Priority:
//  1. CF-Connecting-IP header — honored ONLY when TRUST_PROXY_HEADERS=true
//     (see trustProxyHeaders). A trusted edge (Cloudflare, or nginx whose
//     real_ip module validated it against the set_real_ip_from ranges) has
//     already overwritten it with the real client IP, so it is unspoofable
//     there. Ignored by default because an origin-exposed backend must never
//     trust a client-controlled value.
//  2. middleware.GetClientIP(r.Context()) — the X-Real-IP value nginx sets
//     from $remote_addr, captured by middleware.ClientIPFromHeader("X-Real-IP")
//     in main.go. That middleware is registered only when
//     TRUST_PROXY_HEADERS=true, so it too is trusted solely behind a proxy.
//  3. net.SplitHostPort(r.RemoteAddr) / r.RemoteAddr fallback — the actual
//     TCP peer; the only key source usable when the backend is origin-exposed.
func clientIP(r *http.Request) string {
	if trustProxyHeaders {
		if ip := r.Header.Get("CF-Connecting-IP"); ip != "" {
			return ip
		}
	}
	if ip := middleware.GetClientIP(r.Context()); ip != "" {
		return ip
	}
	if ip, _, err := net.SplitHostPort(r.RemoteAddr); err == nil && ip != "" {
		return ip
	}
	return r.RemoteAddr
}