File Coverage

blib/lib/Plack/Middleware/Greylist.pm

Criterion	Covered	Total	%
statement	121	122	99.1
branch	38	50	76.0
condition	25	37	67.5
subroutine	17	17	100.0
pod	2	2	100.0
total	203	228	89.0

package Plack::Middleware::Greylist;

# ABSTRACT: throttle requests with different rates based on net blocks

# RECOMMEND PREREQ: Cache::FastMmap 1.52
# RECOMMEND PREREQ: Ref::Util::XS

use v5.20;
use warnings;

use parent qw( Plack::Middleware );

use HTTP::Status    qw/ HTTP_FORBIDDEN HTTP_TOO_MANY_REQUESTS /;
use List::Util 1.29 qw/ pairs /;
use Module::Load    qw/ load /;
use Net::IP::LPM;
use Plack::Util;
use Plack::Util::Accessor qw/ default_rate rules cache file _match greylist retry_after cache_config callback /;
use Ref::Util             qw/ is_plain_arrayref is_coderef /;
use Time::Seconds         qw/ ONE_MINUTE /;

use experimental qw/ postderef signatures /;

our $VERSION = 'v0.8.1';


sub prepare_app($self) {

    $self->default_rate(-1) unless defined $self->default_rate;

    die "default_rate must be a positive integer" unless $self->default_rate =~ /^[1-9][0-9]*$/;

    my $config = $self->cache_config;
    $self->cache_config( $config //= {} ) unless defined $config;

    $config->{init_file}                //= 0;
    $config->{unlink_on_exit}           //= !$config->{init_file};
    $config->{serializer}               //= '';
    my $expiry = $config->{expire_time} //= ONE_MINUTE;

    $self->retry_after( $config->{expire_time} + 1 ) unless defined $self->retry_after;
    die "retry_after must be a positive integer greater than $expiry seconds"
      unless $self->retry_after =~ /^[1-9][0-9]*$/ && $self->retry_after > $expiry;

    unless ( $self->cache ) {

        my $file = $self->file // $config->{share_file};
        die "No cache was set" unless defined $file;
        $config->{share_file} = "$file";

        load Cache::FastMmap;
        die "Cache::FastMmap version 1.52 or newer is required" if Cache::FastMmap->VERSION < 1.52;

        my $cache = Cache::FastMmap->new(%$config);

        $self->cache(
            sub($ip) {
                return $cache->get_and_set(
                    $ip,
                    sub( $, $count, $opts ) {
                        $count //= 0;
                        return ( $count + 1, { expire_on => $opts->{expire_on} } );
                    }
                );
            }
        );

    }

    my $match = Net::IP::LPM->new;

    $self->_match( sub($ip) { $match->lookup($ip) } );

    my @blocks;

    if ( my $greylist = $self->greylist ) {
        push @blocks, ( $greylist->%* );
    }

    $self->rules( my $rules = {} );

    my %codes = ( whitelist => -1, allowed => -1, blacklist => 0, rejected => 0, norobots => 0 );
    my %types = ( ip => '', netblock => 1 );

    for my $line ( pairs @blocks ) {

        my ( $block, $rule ) = $line->@*;
        $rule = [ split /\s+/, $rule ] unless is_plain_arrayref($rule);

        my ( $rate, $type ) = $rule->@*;

        $type //= "ip";
        my $mask = $types{$type} // $type;
        $mask = $block if $mask eq "1";

        $rate //= "rejected";
        if ( exists $codes{$rate} ) {
            $mask = $rate if $mask eq "";
            $rate = $codes{$rate};
        }

        $rules->{$block} = [ $rate, $mask ];
        $match->add( $block => $block );
    }

    if ( my $fn = $self->callback ) {
        die "callback must be a code reference" unless is_coderef($fn);
    }
    else {

        $self->callback(
            sub($info) {
                my $env    = $info->{env};
                my $msg    = $info->{message};
                if ( my $log = $env->{'psgix.logger'} ) {
                    $log->( { message => $msg, level => 'warn' } );
                }
                else {
                    $env->{'psgi.errors'}->print($msg);
                }
                return 1;
            }
        );
    }

}

sub call( $self, $env ) {

    my $ip   = $env->{REMOTE_ADDR};
    my $name = $self->_match->($ip);
    my $rule = $name ? $self->rules->{$name} : [ $self->default_rate ];

    my $rate = $rule->[0];

    if ( $rate == 0 && $rule->[1] && $rule->[1] eq "norobots" ) {
        if ( $env->{PATH_INFO} eq "/robots.txt" ) {
            $rate = ONE_MINUTE;    # one request/second
        }
    }

    if ( $rate >= 0 ) {

        my $limit = $rate == 0;

        my ($hits) = $self->cache->( $rule->[1] || $ip );
        $limit = $hits > $rate ? $hits : 0;

        if ($limit) {

            my $block = $name || "default";

            if ( my $fn = $self->callback ) {
                $fn->(
                    {
                        env     => $env,
                        ip      => $ip,
                        hits    => $limit,
                        rate    => $rate,
                        block   => $block,
                        message => "Rate limiting ${ip} after ${limit}/${rate} for ${block}",
                    }
                ) or return $self->app->($env);
            }

            if ( $rate == 0 ) {

                return [ HTTP_FORBIDDEN, [], ["Forbbidden"] ];

            }
            else {

                return [
                    HTTP_TOO_MANY_REQUESTS,
                    [
                        "Retry-After" => $self->retry_after,
                    ],
                    ["Too Many Requests"]
                ];

            }
        }

    }

    return $self->app->($env);
}


1;

__END__