File Coverage

blib/lib/Plack/Middleware/Throttle/Lite.pm

Criterion	Covered	Total	%
statement	124	124	100.0
branch	36	40	90.0
condition	16	22	72.7
subroutine	25	25	100.0
pod	9	9	100.0
total	210	220	95.4

package Plack::Middleware::Throttle::Lite;

# ABSTRACT: Requests throttling for Plack

use strict;
use warnings;
use parent 'Plack::Middleware';
use Plack::Util::Accessor qw(limits maxreq units backend routes blacklist whitelist defaults privileged header_prefix);
use List::MoreUtils qw(any);
use Plack::Util;
use Carp ();
use Net::CIDR::Lite;

our $VERSION = '0.05'; # VERSION
our $AUTHORITY = 'cpan:CHIM'; # AUTHORITY

#
# Some important routines
sub prepare_app {
    my ($self) = @_;

    # setting up defaults
    $self->defaults({
        requests        => 199,
        units           => 'req/hour',
        backend         => 'Simple',
        header_prefix   => 'Throttle-Lite',
        username        => 'nobody',
    });

    $self->_normalize_header_prefix;
    $self->_normalize_limits;
    $self->_initialize_backend;
    $self->_normalize_routes;
    $self->blacklist($self->_initialize_accesslist($self->blacklist));
    $self->whitelist($self->_initialize_accesslist($self->whitelist));

    $self->backend->reqs_max($self->maxreq);
    $self->backend->units($self->units);
}

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

    my $response;

    if ($self->have_to_throttle($env)) {

        return $self->reject_request(blacklist => 403) if $self->is_remote_blacklisted($env);

        # update client id
        $self->backend->requester_id($self->requester_id($env));

        # update
        $self->privileged($self->is_remote_whitelisted($env));

        $response = $self->is_allowed
            ? $self->app->($env)
            : $self->reject_request(ratelimit => 429);

        $self->response_cb($response, sub {
            $self->modify_headers(@_);
        });
    }
    else {
        $response = $self->app->($env);
    }

    $response;
}

#
# Rejects incoming request with some reason
sub reject_request {
    my ($self, $reason, $code) = @_;

    my $reasons = {
        blacklist => 'IP Address Blacklisted',
        ratelimit => 'Rate Limit Exceeded',
    };

    [ $code, [ 'Content-Type' => 'text/plain', ], [ $reasons->{$reason} ] ];
}

#
# Set prefix for headers
sub _normalize_header_prefix {
    my ($self) = @_;

    my $prefix = $self->defaults->{header_prefix};

    if ($self->header_prefix) {
        $prefix = $self->header_prefix;

        # remove invalid chars
        $prefix =~ s/[^0-9a-zA-Z\s]//g;

        # trim spaces
        $prefix =~ s/^\s+//g;
        $prefix =~ s/\s+$//g;

        # camelize
        $prefix = join '-' => map { ucfirst } split /\s+/, $prefix;

        # set default value in case of empty prefix
        $prefix = $prefix || $self->defaults->{header_prefix};
    }

    $self->header_prefix($prefix);
}

#
# Rate limit normalization
sub _normalize_limits {
    my ($self) = @_;

    my $units = {
        'm' => 'req/min',
        'h' => 'req/hour',
        'd' => 'req/day',
    };

    my $limits_re = qr{^(\d*)(\s*)(r|req)(\/|\sper\s)(h|hour|d|day|m|min).*};

    if ($self->limits) {
        my $t_limits = lc($self->limits);
        $t_limits =~ s/\s+/ /g;
        $t_limits =~ /$limits_re/;
        $self->maxreq($1 || $self->defaults->{requests});
        $self->units($units->{$5} || $self->defaults->{units})
    }
    else {
        $self->maxreq($self->defaults->{requests});
        $self->units($self->defaults->{units})
    }
}

#
# Storage backend
sub _initialize_backend {
    my ($self) = @_;

    my ($class, $args) = ($self->defaults->{backend}, {});

    if ($self->backend) {
        my $reft = uc(ref($self->backend) || 'NA');
        if ($reft eq 'NA') { # SCALAR
            ($class, $args) = ($self->backend, {});
        }
        elsif ($reft eq 'ARRAY') {
            ($class, $args) = @{ $self->backend };
        }
        else {
            Carp::croak 'Expected scalar or array reference!';
        }
    }

    my $backend = Plack::Util::load_class($class, 'Plack::Middleware::Throttle::Lite::Backend');

    $self->backend($backend->new($args));
}

#
# Routes' normalization
sub _normalize_routes {
    my ($self) = @_;

    my $routes = [];

    if ($self->routes) {
        my $reft = uc(ref($self->routes) || 'NA');
        if ($reft eq 'NA') { # SCALAR
            $routes = [ $self->routes ];
        }
        elsif ($reft eq 'REGEXP') {
            $routes = [ $self->routes ];
        }
        elsif ($reft eq 'ARRAY') {
            $routes = $self->routes;
        }
        else {
            Carp::croak 'Expected scalar, regex or array reference!';
        }
    }

    $self->routes($routes);
}

#
# Adds extra headers to response
sub modify_headers {
    my ($self, $response) = @_;
    my $headers = $response->[1];

    my $prefix = $self->header_prefix;

    my %inject = (
        "X-${prefix}-Limit" => $self->privileged ? 'unlimited' : $self->maxreq,
        "X-${prefix}-Units" => $self->units,
        "X-${prefix}-Used"  => $self->backend->reqs_done,
    );

    if (($self->backend->reqs_done >= $self->maxreq) && !$self->privileged) {
        $inject{"X-${prefix}-Expire"} = $inject{"Retry-After"} = $self->backend->expire_in;
    }

    map { Plack::Util::header_set($headers, $_, $inject{$_}) } sort keys %inject;

    $response;
}

#
# Checks if requested path should be throttled
sub have_to_throttle {
    my ($self, $env) = @_;

    any { $env->{PATH_INFO} =~ /$_/ } @{ $self->routes };
}

#
# Checks if the requester's IP in the blacklist
sub is_remote_blacklisted {
    my ($self, $env) = @_;

    $self->_is_listed_in(blacklist => $env);
}

#
# Checks if the requester's IP in the whitelist
sub is_remote_whitelisted {
    my ($self, $env) = @_;

    $self->_is_listed_in(whitelist => $env);
}

#
# Checks if remote IP address in accesslist
sub _is_listed_in {
    my ($self, $list, $env) = @_;

    return unless $self->$list;
    return $self->$list->find($env->{REMOTE_ADDR});
}

#
# Populates the blacklist/whitelist
sub _initialize_accesslist {
    my ($self, $items) = @_;

    my $list = Net::CIDR::Lite->new;

    if ($items) {
        map { $list->add_any($_) } ref($items) eq 'ARRAY' ? @$items : ( $items );
    }

    $list;
}

#
# Check if limits is not exceeded
sub is_allowed {
    my ($self) = @_;

    if (($self->backend->reqs_done < $self->backend->reqs_max) || $self->privileged) {
        $self->backend->increment;
        return $self->privileged
            ? 1 : $self->backend->reqs_done <= $self->backend->reqs_max ? 1 : 0;
    }
    else {
        return 0;
    }
}

#
# Requester's ID
sub requester_id {
    my ($self, $env) = @_;
    join ':' => 'throttle', $env->{REMOTE_ADDR}, ($env->{REMOTE_USER} || $self->defaults->{username});
}

1; # End of Plack::Middleware::Throttle::Lite

__END__