query-orchestration/internal/serviceconfig/ratelimit/middleware.go

package ratelimit

import (
	"fmt"
	"log/slog"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/labstack/echo/v4"
	"golang.org/x/time/rate"
)

// DualLayerStore implements dual-layer rate limiting with global and per-IP checks
type DualLayerStore struct {
	globalLimiter *rate.Limiter            // Single global limiter
	ipLimiters    map[string]*rate.Limiter // Per-IP limiters
	config        *Config
	mu            sync.RWMutex
	lastCleanup   time.Time
	timeNow       func() time.Time
}

// NewDualLayerStore creates a new dual-layer rate limiter store
func NewDualLayerStore(config *Config) *DualLayerStore {
	return &DualLayerStore{
		globalLimiter: rate.NewLimiter(config.GlobalRate, config.GlobalBurst),
		ipLimiters:    make(map[string]*rate.Limiter),
		config:        config,
		lastCleanup:   time.Now(),
		timeNow:       time.Now,
	}
}

// Allow implements rate limiting with dual-layer protection
func (d *DualLayerStore) Allow(identifier string) (bool, error) {
	// Layer 1: Check global limit first (fail fast for DDoS)
	if !d.globalLimiter.Allow() {
		return false, fmt.Errorf("global rate limit exceeded")
	}

	// Layer 2: Check per-IP limit
	d.mu.Lock()
	ipLimiter, exists := d.ipLimiters[identifier]
	if !exists {
		ipLimiter = rate.NewLimiter(d.config.DefaultRate, d.config.DefaultBurst)
		d.ipLimiters[identifier] = ipLimiter
	}

	// Periodic cleanup of stale IP limiters
	now := d.timeNow()
	if now.Sub(d.lastCleanup) > d.config.ExpiresIn {
		d.cleanupStaleLimiters()
	}
	d.mu.Unlock()

	return ipLimiter.Allow(), nil
}

// cleanupStaleLimiters removes stale IP limiters that haven't been used recently
// Note: This is a simplified cleanup - in production, you'd track last access time per IP
func (d *DualLayerStore) cleanupStaleLimiters() {
	// For now, just clear all limiters periodically
	// In a production implementation, you'd track last access time per IP
	d.ipLimiters = make(map[string]*rate.Limiter)
	d.lastCleanup = d.timeNow()
}

// NewRateLimitMiddleware creates rate limiting middleware from config
func NewRateLimitMiddleware(config *Config, logger *slog.Logger) echo.MiddlewareFunc {
	store := NewDualLayerStore(config)

	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			// Extract identifier (IP address)
			identifier := c.RealIP()

			// Check rate limit
			allowed, err := store.Allow(identifier)
			if err != nil || !allowed {
				// Rate limit exceeded
				route := fmt.Sprintf("%s %s", c.Request().Method, c.Path())

				logger.Warn("Rate limit exceeded",
					"identifier", identifier,
					"route", route,
					"ip", c.RealIP(),
					"error", err,
				)

				// Calculate retry-after based on rate limit
				limit := config.DefaultRate
				if override, ok := config.EndpointOverrides[route]; ok {
					limit = override.Rate
				}

				retryAfter := int(1.0 / float64(limit))
				if retryAfter < 1 {
					retryAfter = 1
				}

				// Set headers
				c.Response().Header().Set("Retry-After", strconv.Itoa(retryAfter))
				setRateLimitHeader(c, config, route)

				return c.JSON(http.StatusTooManyRequests, map[string]interface{}{
					"message": "rate limit exceeded",
				})
			}

			// Execute handler
			err = next(c)

			// Always set RateLimit header on response
			route := fmt.Sprintf("%s %s", c.Request().Method, c.Path())
			setRateLimitHeader(c, config, route)

			return err
		}
	}
}

// setRateLimitHeader sets the RateLimit header on the response
func setRateLimitHeader(c echo.Context, config *Config, route string) {
	burst := config.DefaultBurst
	rateLimit := config.DefaultRate

	if override, ok := config.EndpointOverrides[route]; ok {
		burst = override.Burst
		rateLimit = override.Rate
	}

	// Calculate window: burst / rate = seconds
	window := int(float64(burst) / float64(rateLimit))
	if window < 1 {
		window = 1
	}

	headerValue := fmt.Sprintf("%d;window=%d", burst, window)
	c.Response().Header().Set("RateLimit", headerValue)
}