File Coverage

blib/lib/Trickster/Middleware/RateLimiter.pm

Criterion	Covered	Total	%
statement	38	40	95.0
branch	8	14	57.1
condition	3	4	75.0
subroutine	10	10	100.0
pod	2	2	100.0
total	61	70	87.1

package Trickster::Middleware::RateLimiter;

use strict;
use warnings;
use v5.14;

use parent 'Plack::Middleware';
use Plack::Util::Accessor qw(
    requests
    window
    storage
    key_generator
    error_handler
);

sub prepare_app {
    my ($self) = @_;

    $self->requests(60) unless $self->requests;
    $self->window(60) unless $self->window;
    $self->storage({}) unless $self->storage;

    $self->key_generator(sub {
        my $env = shift;
        return $env->{REMOTE_ADDR} || 'unknown';
    }) unless $self->key_generator;

    $self->error_handler(sub {
        my ($env, $remaining, $reset) = @_;
        return [
            429,
            [
                'Content-Type' => 'application/json',
                'X-RateLimit-Limit' => $self->requests,
                'X-RateLimit-Remaining' => 0,
                'X-RateLimit-Reset' => $reset,
                'Retry-After' => $reset - time,
            ],
            ['{"error":"Rate limit exceeded"}'],
        ];
    }) unless $self->error_handler;
}

sub call {
    my ($self, $env) = @_;

    my $key = $self->key_generator->($env);
    my $now = time;

    # Get or initialize bucket
    my $bucket = $self->storage->{$key} ||= {
        count => 0,
        reset => $now + $self->window,
    };

    # Reset bucket if window expired
    if ($now >= $bucket->{reset}) {
        $bucket->{count} = 0;
        $bucket->{reset} = $now + $self->window;
    }

    # Check rate limit
    if ($bucket->{count} >= $self->requests) {
        return $self->error_handler->($env, 0, $bucket->{reset});
    }

    # Increment counter
    $bucket->{count}++;

    my $remaining = $self->requests - $bucket->{count};

    # Call app
    my $res = $self->app->($env);

    # Add rate limit headers
    return $self->response_cb($res, sub {
        my $res = shift;
        push @{$res->[1]},
            'X-RateLimit-Limit' => $self->requests,
            'X-RateLimit-Remaining' => $remaining,
            'X-RateLimit-Reset' => $bucket->{reset};
    });
}

1;

__END__