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

// clientIP derives the per-client rate-limit key. Priority:
//  1. CF-Connecting-IP header (set only when Cloudflare is the edge;
//     nginx never sets it, so it cannot be spoofed through our proxy)
//  2. middleware.GetClientIP(r.Context()) — the X-Real-IP value nginx sets,
//     captured by middleware.ClientIPFromHeader("X-Real-IP") in main.go
//  3. net.SplitHostPort(r.RemoteAddr) / r.RemoteAddr fallback
func clientIP(r *http.Request) string {
	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
}