pulse/internal/api/rate_limit_config.go

package api

import (
	"net/http"
	"os"
	"strconv"
	"strings"
	"time"
)

const (
	defaultGeneralAPIRateLimit = 500
	maxDevGeneralAPIRateLimit  = 100000
)

// EndpointRateLimitConfig defines rate limiting configuration for different endpoint categories
type EndpointRateLimitConfig struct {
	AuthEndpoints      *RateLimiter // Login, logout, password change
	ConfigEndpoints    *RateLimiter // Node configuration changes
	ExportEndpoints    *RateLimiter // Export/import operations
	RecoveryEndpoints  *RateLimiter // Recovery operations
	UpdateEndpoints    *RateLimiter // Update checks and operations
	WebSocketEndpoints *RateLimiter // WebSocket connections
	GeneralAPI         *RateLimiter // General API calls
	PublicEndpoints    *RateLimiter // Public endpoints (health, version)
}

var globalRateLimitConfig *EndpointRateLimitConfig

func newEndpointRateLimitConfig() *EndpointRateLimitConfig {
	return &EndpointRateLimitConfig{
		// Authentication endpoints: strict limits to prevent brute force
		AuthEndpoints: NewRateLimiter(10, 1*time.Minute), // 10 attempts per minute

		// Configuration changes: moderate limits
		ConfigEndpoints: NewRateLimiter(30, 1*time.Minute), // 30 changes per minute

		// Export/import: very strict limits
		ExportEndpoints: NewRateLimiter(5, 5*time.Minute), // 5 exports per 5 minutes

		// Recovery operations: extremely strict
		RecoveryEndpoints: NewRateLimiter(3, 10*time.Minute), // 3 attempts per 10 minutes

		// Update operations: allow frequent polling during updates
		UpdateEndpoints: NewRateLimiter(60, 1*time.Minute), // 60 checks per minute (modal polls every 2s)

		// WebSocket connections: per-connection limits
		WebSocketEndpoints: NewRateLimiter(30, 1*time.Minute), // 30 new connections per minute

		// General API: higher limits for normal operations. Browser E2E runs can
		// explicitly raise this in PULSE_DEV because many isolated test sessions
		// share one Docker bridge address; production always keeps the default.
		GeneralAPI: NewRateLimiter(generalAPIRateLimit(), 1*time.Minute),

		// Public endpoints: very high limits (health checks, etc.)
		PublicEndpoints: NewRateLimiter(1000, 1*time.Minute), // 1000 requests per minute
	}
}

// Stop joins only the workers owned by this endpoint configuration. A Router
// must not stop the package-global configuration or another Router's budgets.
func (cfg *EndpointRateLimitConfig) Stop() {
	if cfg == nil {
		return
	}
	for _, limiter := range []*RateLimiter{
		cfg.AuthEndpoints, cfg.ConfigEndpoints, cfg.ExportEndpoints,
		cfg.RecoveryEndpoints, cfg.UpdateEndpoints, cfg.WebSocketEndpoints,
		cfg.GeneralAPI, cfg.PublicEndpoints,
	} {
		limiter.Stop()
	}
}

func generalAPIRateLimit() int {
	if os.Getenv("PULSE_DEV") != "true" {
		return defaultGeneralAPIRateLimit
	}

	configured, err := strconv.Atoi(strings.TrimSpace(os.Getenv("PULSE_DEV_GENERAL_API_RATE_LIMIT")))
	if err != nil || configured < defaultGeneralAPIRateLimit || configured > maxDevGeneralAPIRateLimit {
		return defaultGeneralAPIRateLimit
	}
	return configured
}

func ensureGlobalRateLimitConfig() *EndpointRateLimitConfig {
	if globalRateLimitConfig == nil {
		InitializeRateLimiters()
	}
	return globalRateLimitConfig
}

// InitializeRateLimiters sets up rate limiters for all endpoint categories
func InitializeRateLimiters() {
	globalRateLimitConfig = newEndpointRateLimitConfig()
}

// GetRateLimiterForEndpoint returns the appropriate rate limiter for a given endpoint
func GetRateLimiterForEndpoint(path string, method string) *RateLimiter {
	return getRateLimiterForEndpoint(ensureGlobalRateLimitConfig(), path, method)
}

func getRateLimiterForEndpoint(cfg *EndpointRateLimitConfig, path string, method string) *RateLimiter {
	// Normalize path
	path = strings.ToLower(path)

	// Authentication endpoints (including bootstrap token validation — critical security credential)
	if strings.Contains(path, "/api/login") ||
		strings.Contains(path, "/api/logout") ||
		strings.Contains(path, "/api/security/change-password") ||
		strings.Contains(path, "/api/security/validate-bootstrap-token") ||
		strings.Contains(path, "/api/security/regenerate-token") ||
		strings.Contains(path, "/api/security/validate-token") ||
		strings.Contains(path, "/api/auth") {
		return cfg.AuthEndpoints
	}

	// Recovery endpoints
	if strings.Contains(path, "/api/security/recovery") {
		return cfg.RecoveryEndpoints
	}

	// Export/Import endpoints
	if strings.Contains(path, "/api/config/export") ||
		strings.Contains(path, "/api/config/import") {
		return cfg.ExportEndpoints
	}

	// Configuration endpoints (write operations only)
	if method != "GET" && (strings.Contains(path, "/api/config/nodes") ||
		strings.Contains(path, "/api/config/system") ||
		strings.Contains(path, "/api/config/webhooks") ||
		strings.Contains(path, "/api/config/alerts")) {
		return cfg.ConfigEndpoints
	}

	// Configuration read endpoints get higher limits to prevent UI issues
	if method == "GET" && (strings.Contains(path, "/api/config/") ||
		strings.Contains(path, "/api/discover") ||
		strings.Contains(path, "/api/security/status")) {
		return cfg.PublicEndpoints // Use higher limit for reads
	}

	// Update endpoints
	if strings.Contains(path, "/api/updates") {
		return cfg.UpdateEndpoints
	}

	// WebSocket endpoints
	if strings.Contains(path, "/ws") {
		return cfg.WebSocketEndpoints
	}

	// Public endpoints (no auth required)
	if strings.Contains(path, "/api/health") ||
		strings.Contains(path, "/api/version") ||
		strings.Contains(path, "/api/security/status") ||
		strings.Contains(path, "/metrics") {
		return cfg.PublicEndpoints
	}

	// Default to general API rate limiter
	return cfg.GeneralAPI
}

// UniversalRateLimitMiddleware applies appropriate rate limiting to all endpoints
func UniversalRateLimitMiddleware(next http.Handler) http.Handler {
	return UniversalRateLimitMiddlewareWithConfig(ensureGlobalRateLimitConfig(), next)
}

func UniversalRateLimitMiddlewareWithConfig(cfg *EndpointRateLimitConfig, next http.Handler) http.Handler {
	if cfg == nil {
		cfg = newEndpointRateLimitConfig()
	}
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		// Skip rate limiting for static assets
		if !strings.HasPrefix(r.URL.Path, "/api") && !strings.HasPrefix(r.URL.Path, "/ws") {
			next.ServeHTTP(w, r)
			return
		}

		// Skip rate limiting for real-time data endpoints that are polled frequently
		// These endpoints are essential for UI functionality and should not be rate limited
		skipPaths := []string{
			"/api/state",           // Real-time state updates
			"/api/guests/metadata", // Guest metadata (polled frequently)
			"/ws",                  // WebSocket connections (authenticated, reconnections are normal)
		}
		for _, path := range skipPaths {
			if strings.Contains(r.URL.Path, path) {
				next.ServeHTTP(w, r)
				return
			}
		}

		// Extract client IP
		ip := GetClientIP(r)

		// Skip rate limiting only for direct loopback requests (no proxy headers)
		if (ip == "127.0.0.1" || ip == "::1" || ip == "localhost") && isDirectLoopbackRequest(r) {
			next.ServeHTTP(w, r)
			return
		}

		// Get appropriate rate limiter for this endpoint
		limiter := getRateLimiterForEndpoint(cfg, r.URL.Path, r.Method)

		// Check rate limit
		if !limiter.Allow(ip) {
			// Add retry-after header matching the limiter's actual window
			w.Header().Set("Retry-After", strconv.Itoa(int(limiter.window.Seconds())))
			w.Header().Set("X-RateLimit-Limit", strconv.Itoa(limiter.limit))
			w.Header().Set("X-RateLimit-Remaining", "0")
			w.Header().Set("X-RateLimit-Reset", time.Now().Add(limiter.window).Format(time.RFC3339))

			http.Error(w, "Rate limit exceeded. Please try again later.", http.StatusTooManyRequests)
			return
		}

		// Continue to next handler
		next.ServeHTTP(w, r)
	})
}

// ResetRateLimitForIP resets rate limit counters for a specific IP (use carefully)
func ResetRateLimitForIP(ip string) {
	if globalRateLimitConfig == nil {
		return
	}

	// Reset for all rate limiters
	limiters := []*RateLimiter{
		globalRateLimitConfig.AuthEndpoints,
		globalRateLimitConfig.ConfigEndpoints,
		globalRateLimitConfig.ExportEndpoints,
		globalRateLimitConfig.RecoveryEndpoints,
		globalRateLimitConfig.UpdateEndpoints,
		globalRateLimitConfig.WebSocketEndpoints,
		globalRateLimitConfig.GeneralAPI,
		globalRateLimitConfig.PublicEndpoints,
	}

	for _, limiter := range limiters {
		limiter.Reset(ip)
	}
}

// Reset clears rate limit history for a specific IP
func (rl *RateLimiter) Reset(ip string) {
	rl.mu.Lock()
	defer rl.mu.Unlock()
	delete(rl.attempts, ip)
}