internal/api/middleware/rate_limit.go

// Package middleware provides HTTP middleware components for the CLI Proxy API server.
// This file contains the rate limiting middleware that integrates with the domain
// rate limiting service to enforce request limits and block abusive clients.
package middleware

import (
	"net/http"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/router-for-me/CLIProxyAPI/v6/internal/domain/ports"
)

// RateLimitMiddleware creates a Gin middleware that enforces rate limiting
// using the provided RateLimitService. It checks if the client IP is blocked
// and records failed authentication attempts.
type RateLimitMiddleware struct {
	service ports.RateLimitService
}

// NewRateLimitMiddleware creates a new rate limiting middleware instance.
func NewRateLimitMiddleware(service ports.RateLimitService) *RateLimitMiddleware {
	return &RateLimitMiddleware{
		service: service,
	}
}

// Middleware returns the Gin middleware function that enforces rate limiting.
// It should be used for management endpoints that require authentication.
func (m *RateLimitMiddleware) Middleware() gin.HandlerFunc {
	return func(c *gin.Context) {
		if m.service == nil {
			c.Next()
			return
		}

		clientIP := c.ClientIP()
		ctx := c.Request.Context()

		// Check if client is blocked
		blocked, blockedUntil, err := m.service.IsBlocked(ctx, clientIP)
		if err != nil {
			c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{
				"error": "rate limit check failed",
			})
			return
		}
		if blocked {
			remaining := time.Until(blockedUntil)
			if remaining < 0 {
				remaining = 0
			}
			c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
				"error":       "IP banned due to too many failed attempts",
				"retry_after": remaining.String(),
			})
			return
		}

		// Check if request is allowed
		allowed, err := m.service.Allow(ctx, clientIP)
		if err != nil {
			c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{
				"error": "rate limit check failed",
			})
			return
		}
		if !allowed {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
				"error": "rate limit exceeded",
			})
			return
		}

		c.Next()
	}
}

// AuthMiddleware wraps the rate limiting middleware with authentication logic.
// It records failed attempts when authentication fails.
type AuthMiddleware struct {
	rateLimitService ports.RateLimitService
	getSecretHash    func() string
	getEnvSecret     func() string
	allowRemote      func() bool
}

// NewAuthMiddleware creates a new authentication middleware with rate limiting.
func NewAuthMiddleware(
	service ports.RateLimitService,
	getSecretHash func() string,
	getEnvSecret func() string,
	allowRemote func() bool,
) *AuthMiddleware {
	return &AuthMiddleware{
		rateLimitService: service,
		getSecretHash:    getSecretHash,
		getEnvSecret:     getEnvSecret,
		allowRemote:      allowRemote,
	}
}

// OnAuthFailure should be called when authentication fails to record the attempt.
func (m *AuthMiddleware) OnAuthFailure(c *gin.Context) {
	if m.rateLimitService == nil {
		return
	}

	clientIP := c.ClientIP()
	ctx := c.Request.Context()

	m.rateLimitService.RecordAttempt(ctx, clientIP, false)
}

// OnAuthSuccess should be called when authentication succeeds to reset attempts.
func (m *AuthMiddleware) OnAuthSuccess(c *gin.Context) {
	if m.rateLimitService == nil {
		return
	}

	clientIP := c.ClientIP()
	ctx := c.Request.Context()

	m.rateLimitService.RecordAttempt(ctx, clientIP, true)
}

// GetRetryAfter returns the duration until the client can retry after being blocked.
func (m *AuthMiddleware) GetRetryAfter(c *gin.Context) time.Duration {
	if m.rateLimitService == nil {
		return 0
	}

	clientIP := c.ClientIP()
	ctx := c.Request.Context()

	blocked, blockedUntil, err := m.rateLimitService.IsBlocked(ctx, clientIP)
	if err != nil {
		return 0
	}
	if blocked {
		remaining := time.Until(blockedUntil)
		if remaining > 0 {
			return remaining
		}
	}

	return 0
}