ophion/internal/api/ratelimit.go

package api

import (
	"fmt"
	"strconv"
	"sync"
	"time"

	"github.com/gofiber/fiber/v2"
)

// ═══════════════════════════════════════════════════════════
// 🚦 RATE LIMITING MIDDLEWARE
// ═══════════════════════════════════════════════════════════

// RateLimitConfig configuração do rate limiter
type RateLimitConfig struct {
	// Requests máximos por janela
	Max int
	
	// Janela de tempo
	Window time.Duration
	
	// Função para extrair identificador (default: IP)
	KeyGenerator func(*fiber.Ctx) string
	
	// Função para resposta quando limitado
	LimitReached func(*fiber.Ctx) error
	
	// Pular rate limit para certos paths
	SkipPaths []string
	
	// Headers customizados
	Headers RateLimitHeaders
}

// RateLimitHeaders headers do rate limit
type RateLimitHeaders struct {
	Limit     string // X-RateLimit-Limit
	Remaining string // X-RateLimit-Remaining
	Reset     string // X-RateLimit-Reset
	RetryAfter string // Retry-After
}

// DefaultRateLimitConfig configuração padrão
func DefaultRateLimitConfig() RateLimitConfig {
	return RateLimitConfig{
		Max:    100,
		Window: time.Minute,
		KeyGenerator: func(c *fiber.Ctx) string {
			return c.IP()
		},
		LimitReached: func(c *fiber.Ctx) error {
			return c.Status(429).JSON(fiber.Map{
				"error":   "Too Many Requests",
				"message": "Rate limit exceeded. Please slow down.",
				"retry_after": 60,
			})
		},
		Headers: RateLimitHeaders{
			Limit:      "X-RateLimit-Limit",
			Remaining:  "X-RateLimit-Remaining",
			Reset:      "X-RateLimit-Reset",
			RetryAfter: "Retry-After",
		},
	}
}

// rateLimitStore armazena contagem de requests
type rateLimitStore struct {
	sync.RWMutex
	entries map[string]*rateLimitEntry
}

type rateLimitEntry struct {
	count     int
	expiresAt time.Time
}

var store = &rateLimitStore{
	entries: make(map[string]*rateLimitEntry),
}

// Limpar entries expirados periodicamente
func init() {
	go func() {
		ticker := time.NewTicker(time.Minute)
		for range ticker.C {
			store.cleanup()
		}
	}()
}

func (s *rateLimitStore) cleanup() {
	s.Lock()
	defer s.Unlock()
	
	now := time.Now()
	for key, entry := range s.entries {
		if now.After(entry.expiresAt) {
			delete(s.entries, key)
		}
	}
}

// RateLimit middleware de rate limiting
func RateLimit(config ...RateLimitConfig) fiber.Handler {
	cfg := DefaultRateLimitConfig()
	if len(config) > 0 {
		cfg = config[0]
	}
	
	return func(c *fiber.Ctx) error {
		// Verificar skip paths
		path := c.Path()
		for _, skip := range cfg.SkipPaths {
			if path == skip {
				return c.Next()
			}
		}
		
		// Gerar key
		key := cfg.KeyGenerator(c)
		
		// Verificar/atualizar contagem
		store.Lock()
		
		now := time.Now()
		entry, exists := store.entries[key]
		
		if !exists || now.After(entry.expiresAt) {
			// Nova janela
			store.entries[key] = &rateLimitEntry{
				count:     1,
				expiresAt: now.Add(cfg.Window),
			}
			store.Unlock()
			
			// Headers
			c.Set(cfg.Headers.Limit, strconv.Itoa(cfg.Max))
			c.Set(cfg.Headers.Remaining, strconv.Itoa(cfg.Max-1))
			c.Set(cfg.Headers.Reset, strconv.FormatInt(now.Add(cfg.Window).Unix(), 10))
			
			return c.Next()
		}
		
		entry.count++
		remaining := cfg.Max - entry.count
		resetTime := entry.expiresAt.Unix()
		
		store.Unlock()
		
		// Headers
		c.Set(cfg.Headers.Limit, strconv.Itoa(cfg.Max))
		c.Set(cfg.Headers.Remaining, strconv.Itoa(max(0, remaining)))
		c.Set(cfg.Headers.Reset, strconv.FormatInt(resetTime, 10))
		
		// Verificar se excedeu
		if remaining < 0 {
			retryAfter := int(time.Until(entry.expiresAt).Seconds())
			c.Set(cfg.Headers.RetryAfter, strconv.Itoa(max(1, retryAfter)))
			return cfg.LimitReached(c)
		}
		
		return c.Next()
	}
}

// ═══════════════════════════════════════════════════════════
// 🎯 RATE LIMIT PRESETS
// ═══════════════════════════════════════════════════════════

// RateLimitAuth rate limit para endpoints de autenticação (mais restritivo)
func RateLimitAuth() fiber.Handler {
	return RateLimit(RateLimitConfig{
		Max:    5,
		Window: time.Minute,
		KeyGenerator: func(c *fiber.Ctx) string {
			// Combinar IP + email para prevenir ataques distribuídos
			email := c.FormValue("email")
			if email == "" {
				email = c.Query("email")
			}
			return fmt.Sprintf("auth:%s:%s", c.IP(), email)
		},
		LimitReached: func(c *fiber.Ctx) error {
			return c.Status(429).JSON(fiber.Map{
				"error":   "Too Many Login Attempts",
				"message": "You have exceeded the maximum number of login attempts. Please wait before trying again.",
				"retry_after": 60,
			})
		},
	})
}

// RateLimitAPI rate limit para API geral
func RateLimitAPI() fiber.Handler {
	return RateLimit(RateLimitConfig{
		Max:    100,
		Window: time.Minute,
		KeyGenerator: func(c *fiber.Ctx) string {
			// Usar API key ou IP
			if apiKey := c.Locals("api_key_id"); apiKey != nil {
				return fmt.Sprintf("api:%v", apiKey)
			}
			if userID := c.Locals("user_id"); userID != nil {
				return fmt.Sprintf("user:%v", userID)
			}
			return fmt.Sprintf("ip:%s", c.IP())
		},
	})
}

// RateLimitIngest rate limit para ingestão de dados (mais permissivo)
func RateLimitIngest() fiber.Handler {
	return RateLimit(RateLimitConfig{
		Max:    1000,
		Window: time.Minute,
		KeyGenerator: func(c *fiber.Ctx) string {
			if orgID := c.Locals("org_id"); orgID != nil {
				return fmt.Sprintf("ingest:%v", orgID)
			}
			return fmt.Sprintf("ingest:ip:%s", c.IP())
		},
		LimitReached: func(c *fiber.Ctx) error {
			return c.Status(429).JSON(fiber.Map{
				"error":   "Ingest Rate Limit Exceeded",
				"message": "Your organization has exceeded the data ingestion rate limit.",
				"retry_after": 10,
			})
		},
	})
}

// RateLimitExport rate limit para exportação de dados
func RateLimitExport() fiber.Handler {
	return RateLimit(RateLimitConfig{
		Max:    10,
		Window: time.Hour,
		KeyGenerator: func(c *fiber.Ctx) string {
			if userID := c.Locals("user_id"); userID != nil {
				return fmt.Sprintf("export:%v", userID)
			}
			return fmt.Sprintf("export:ip:%s", c.IP())
		},
		LimitReached: func(c *fiber.Ctx) error {
			return c.Status(429).JSON(fiber.Map{
				"error":   "Export Rate Limit Exceeded",
				"message": "You have exceeded the maximum number of exports per hour.",
				"retry_after": 3600,
			})
		},
	})
}

// ═══════════════════════════════════════════════════════════
// 🔧 HELPERS
// ═══════════════════════════════════════════════════════════

func max(a, b int) int {
	if a > b {
		return a
	}
	return b
}