Crussell/backend/mw/ratelimit.go

//go:build !dev
// +build !dev

package mw

import (
	"fmt"
	"net/http"
	"sync"
	"net"
	"time"
)

// RateLimiter implements a simple in-memory rate limiter
type RateLimiter struct {
	requests map[string][]time.Time
	mu       sync.RWMutex
	limit    int
	window   time.Duration
}

func NewRateLimiter(limit int, window time.Duration) *RateLimiter {
	rl := &RateLimiter{
		requests: make(map[string][]time.Time),
		limit:    limit,
		window:   window,
	}
	// Cleanup old entries periodically
	go func() {
		for {
			time.Sleep(window)
			rl.cleanup()
		}
	}()
	return rl
}

func (rl *RateLimiter) cleanup() {
	rl.mu.Lock()
	defer rl.mu.Unlock()
	now := time.Now()
	for key, times := range rl.requests {
		var valid []time.Time
		for _, t := range times {
			if now.Sub(t) < rl.window {
				valid = append(valid, t)
			}
		}
		if len(valid) == 0 {
			delete(rl.requests, key)
		} else {
			rl.requests[key] = valid
		}
	}
}

func (rl *RateLimiter) Allow(key string) bool {
	rl.mu.Lock()
	defer rl.mu.Unlock()
	now := time.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
}

// ProgressiveRateLimiter implements per-IP rate limiting with increasing backoff
// Designed for bot-spam prevention across accounts (not account-specific lockout)
type ProgressiveRateLimiter struct {
	requests map[string]*ipProgressiveState
	mu       sync.RWMutex
}

type ipProgressiveState struct {
	// Timestamps of all requests within the tracking window
	timestamps []time.Time
}

func NewProgressiveRateLimiter() *ProgressiveRateLimiter {
	prl := &ProgressiveRateLimiter{
		requests: make(map[string]*ipProgressiveState),
	}
	go func() {
		for {
			time.Sleep(30 * time.Second)
			prl.cleanup()
		}
	}()
	return prl
}

func (prl *ProgressiveRateLimiter) cleanup() {
	prl.mu.Lock()
	defer prl.mu.Unlock()
	cutoff := time.Now().Add(-60 * time.Second)
	for ip, state := range prl.requests {
		var valid []time.Time
		for _, t := range state.timestamps {
			if t.After(cutoff) {
				valid = append(valid, t)
			}
		}
		if len(valid) == 0 {
			delete(prl.requests, ip)
		} else {
			state.timestamps = valid
		}
	}
}

// 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 := time.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
	if sustainedCount <= 140 {
		return 500 // 500ms - scraping but not too aggressively
	} else if sustainedCount <= 200 {
		return 2000 // 2s - moderate spam
	} else if sustainedCount <= 300 {
		return 5000 // 5s - heavy spam
	} else {
		return 10000 // 10s - abuse
	}
}

var globalProgressiveLimiter = NewProgressiveRateLimiter()

func ProgressiveRateLimit(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		ip := r.Header.Get("CF-Connecting-IP")
		if ip == "" {
			ip, _, _ = net.SplitHostPort(r.RemoteAddr)
			if ip == "" {
				ip = r.RemoteAddr
			}
		}

		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) {
			// Get client IP
			ip := r.Header.Get("CF-Connecting-IP")
			if ip == "" {
				ip, _, _ = net.SplitHostPort(r.RemoteAddr)
				if ip == "" {
					ip = r.RemoteAddr
				}
			}

			if !limiter.Allow(ip) {
				http.Error(w, "Rate limit exceeded", http.StatusTooManyRequests)
				return
			}

			next.ServeHTTP(w, r)
		})
	}
}