wp-agentic-writer/includes/class-rate-limiter.php

<?php
/**
 * Rate Limiter for WP Agentic Writer REST Endpoints
 *
 * @package WP_Agentic_Writer
 */

/**
 * Class WPAW_Rate_Limiter
 *
 * Simple WordPress transient-based rate limiter for REST endpoints.
 *
 * @since 0.3.0
 */
class WPAW_Rate_Limiter {

    /**
     * Default rate limits per endpoint type.
     *
     * @var array
     */
    private static $default_limits = [
        'chat'             => [ 'limit' => 60, 'window' => 60 ],   // 60/min
        'chat_stream'      => [ 'limit' => 60, 'window' => 60 ],   // 60/min
        'generate_plan'    => [ 'limit' => 10, 'window' => 60 ],   // 10/min
        'stream_plan'     => [ 'limit' => 10, 'window' => 60 ],   // 10/min
        'execute_article' => [ 'limit' => 10, 'window' => 60 ],   // 10/min
        'stream_article'  => [ 'limit' => 10, 'window' => 60 ],   // 10/min
        'generate_image'  => [ 'limit' => 10, 'window' => 60 ],   // 10/min
        'seo_audit'       => [ 'limit' => 20, 'window' => 60 ],   // 20/min
        'refine'          => [ 'limit' => 30, 'window' => 60 ],   // 30/min
        'default'         => [ 'limit' => 30, 'window' => 60 ],   // 30/min fallback
    ];

    /**
     * Generate cache key for rate limit tracking.
     *
     * @since 0.3.0
     * @param int    $user_id  User ID.
     * @param string $endpoint Endpoint name.
     * @return string Cache key.
     */
    private static function get_key( $user_id, $endpoint ) {
        return "wpaw_rl_{$endpoint}_{$user_id}";
    }

    /**
     * Check if request is within rate limit.
     *
     * @since 0.3.0
     * @param string $endpoint Endpoint name.
     * @param int    $limit    Max requests per window.
     * @param int    $window   Time window in seconds.
     * @return true|WP_Error True if allowed, WP_Error if rate limited.
     */
    public static function check( $endpoint, $limit = null, $window = null ) {
        // Use default limits if not specified.
        if ( null === $limit || null === $window ) {
            $defaults = self::$default_limits[ $endpoint ] ?? self::$default_limits['default'];
            $limit    = $defaults['limit'];
            $window   = $defaults['window'];
        }

        $user_id = get_current_user_id();

        // For non-authenticated users, use IP address as fallback.
        if ( 0 === $user_id ) {
            $user_id = self::get_client_ip();
        }

        $key   = self::get_key( $user_id, $endpoint );
        $count = (int) get_transient( $key );

        if ( $count >= $limit ) {
            $retry_after = (int) get_transient( $key . '_expires' );
            if ( false === $retry_after ) {
                $retry_after = $window;
            }

            return new WP_Error(
                'rate_limited',
                sprintf(
                    // translators: %d is the number of seconds to wait.
                    __( 'Too many requests. Please wait %d seconds.', 'wp-agentic-writer' ),
                    $retry_after
                ),
                [
                    'status'      => 429,
                    'retry_after' => $retry_after,
                ]
            );
        }

        // Increment counter.
        if ( false === get_transient( $key ) ) {
            // First request in window - set expiration tracking.
            set_transient( $key, 1, $window );
            set_transient( $key . '_expires', $window, $window );
        } else {
            // Increment existing counter.
            set_transient( $key, $count + 1, $window );
        }

        return true;
    }

    /**
     * Get current request count for an endpoint.
     *
     * @since 0.3.0
     * @param string $endpoint Endpoint name.
     * @return int Current request count.
     */
    public static function get_current_count( $endpoint ) {
        $user_id = get_current_user_id();

        if ( 0 === $user_id ) {
            $user_id = self::get_client_ip();
        }

        $key   = self::get_key( $user_id, $endpoint );
        $count = (int) get_transient( $key );

        return false === $count ? 0 : $count;
    }

    /**
     * Get remaining requests for an endpoint.
     *
     * @since 0.3.0
     * @param string $endpoint Endpoint name.
     * @param int    $limit    Max requests per window.
     * @return int Remaining requests.
     */
    public static function get_remaining( $endpoint, $limit = null ) {
        if ( null === $limit ) {
            $defaults = self::$default_limits[ $endpoint ] ?? self::$default_limits['default'];
            $limit    = $defaults['limit'];
        }

        $current = self::get_current_count( $endpoint );
        return max( 0, $limit - $current );
    }

    /**
     * Get client IP address for unauthenticated rate limiting.
     *
     * @since 0.3.0
     * @return string Client IP address.
     */
    private static function get_client_ip() {
        $ip = '';

        if ( ! empty( $_SERVER['HTTP_CF_CONNECTING_IP'] ) ) {
            // Cloudflare.
            $ip = sanitize_text_field( wp_unslash( $_SERVER['HTTP_CF_CONNECTING_IP'] ) );
        } elseif ( ! empty( $_SERVER['HTTP_X_FORWARDED_FOR'] ) ) {
            // X-Forwarded-For header (may contain multiple IPs).
            $ips = explode( ',', sanitize_text_field( wp_unslash( $_SERVER['HTTP_X_FORWARDED_FOR'] ) ) );
            $ip  = trim( $ips[0] );
        } elseif ( ! empty( $_SERVER['HTTP_X_REAL_IP'] ) ) {
            // X-Real-IP header.
            $ip = sanitize_text_field( wp_unslash( $_SERVER['HTTP_X_REAL_IP'] ) );
        } elseif ( ! empty( $_SERVER['REMOTE_ADDR'] ) ) {
            // Direct IP.
            $ip = sanitize_text_field( wp_unslash( $_SERVER['REMOTE_ADDR'] ) );
        }

        // Validate IP format.
        if ( ! filter_var( $ip, FILTER_VALIDATE_IP ) ) {
            $ip = '0.0.0.0';
        }

        return $ip;
    }

    /**
     * Reset rate limit for a user/endpoint (for testing).
     *
     * @since 0.3.0
     * @param int    $user_id  User ID.
     * @param string $endpoint Endpoint name.
     * @return bool Success.
     */
    public static function reset( $user_id, $endpoint ) {
        $key = self::get_key( $user_id, $endpoint );
        delete_transient( $key );
        delete_transient( $key . '_expires' );
        return true;
    }

    /**
     * Add rate limit headers to response.
     *
     * @since 0.3.0
     * @param WP_REST_Response $response Response object.
     * @param string           $endpoint Endpoint name.
     * @param int              $limit    Max requests per window.
     * @return WP_REST_Response Modified response with headers.
     */
    public static function add_headers( $response, $endpoint, $limit = null ) {
        if ( null === $limit ) {
            $defaults = self::$default_limits[ $endpoint ] ?? self::$default_limits['default'];
            $limit    = $defaults['limit'];
        }

        $remaining = self::get_remaining( $endpoint, $limit );
        $current   = self::get_current_count( $endpoint );

        $response->header( 'X-RateLimit-Limit', $limit );
        $response->header( 'X-RateLimit-Remaining', $remaining );
        $response->header( 'X-RateLimit-Used', $current );

        return $response;
    }
}