File Coverage

blib/lib/CGI/Application/Plugin/RateLimit.pm

Criterion	Covered	Total	%
statement	24	113	21.2
branch	0	44	0.0
condition	0	19	0.0
subroutine	7	21	33.3
pod	5	10	50.0
total	36	207	17.3

package CGI::Application::Plugin::RateLimit;

use 5.006;
use strict;
use warnings;

use Carp qw(croak);

our $VERSION = '1.0';

# export the rate_limit method into the using CGI::App and setup the
# prerun callback
sub import {
    my $pkg     = shift;
    my $callpkg = caller;

    {
        no strict qw(refs);
        *{$callpkg . '::rate_limit'} = \&rate_limit;
    }

    $callpkg->add_callback(prerun => \&prerun_callback);
}

# setup accessor/mutators for simple stuff
use base qw(Class::Accessor::Fast);
__PACKAGE__->mk_accessors(
    qw(dbh table violation_mode violation_callback identity_callback
      violated_mode violated_action violated_limits));

# setup a new object the first time it's called
sub rate_limit {
    my $cgi_app = shift;
    return $cgi_app->{__rate_limit_obj} if $cgi_app->{__rate_limit_obj};

    my $rate_limit = $cgi_app->{__rate_limit_obj} = __PACKAGE__->new();

    # setup defaults
    $rate_limit->table('rate_limit_hits');
    $rate_limit->identity_callback(
        sub {
            return $ENV{REMOTE_USER} || $ENV{REMOTE_IP};
        });

    return $rate_limit;
}

# intercept the run-mode call
sub prerun_callback {
    my $cgi_app = shift;
    my $self    = $cgi_app->rate_limit;
    my $query   = $cgi_app->query;

    # see if this mode is protected
    my $mode = $query->param($cgi_app->mode_param)
      || $cgi_app->start_mode;
    my $protected = $self->protected_modes || {};
    my $limits = $protected->{$mode};
    return unless $limits;

    $self->_verify_attributes();

    # record the hit
    my $action = ref($cgi_app) . "::$mode";
    $self->record_hit(action => $action);

    # check for a violation
    if ($self->check_violation(action => $action, limits => $limits)) {

        # deal with it by jumping to violation_mode or calling the
        # violation callback
        if ($self->violation_mode) {
            $cgi_app->prerun_mode($self->violation_mode);
        } else {
            my $violation_callback = $self->violation_callback();
            $cgi_app->prerun_mode($violation_callback->($cgi_app));
        }
    }
}

# make sure we're ready to rumble
sub _verify_attributes {
    my $self = shift;

    for my $name (qw(dbh table identity_callback)) {
        croak(  "You forgot to set the required '$name' attribute on your "
              . __PACKAGE__
              . " object.")
          unless $self->{$name};
    }
    croak(  "You forgot to set the required 'violation_mode' or "
          . "'violation_callback' attribute on your "
          . __PACKAGE__
          . " object.")
      unless $self->{violation_mode}
      or $self->{violation_callback};
}

# translate a timeframe like 10s, 5m or 1h into seconds
sub _timeframe_to_seconds {
    my $time = shift;
    my ($digits, $modifier) = $time =~ /^(\d+)([smh])$/;
    croak(  "Invalid timeframe found: '$time'.  "
          . "Should be a number followed by s, m or h.")
      unless $digits and $modifier;

    return $digits           if $modifier eq 's';
    return $digits * 60      if $modifier eq 'm';
    return $digits * 60 * 60 if $modifier eq 'h';
}

sub protected_modes {
    my ($self, %args) = @_;
    return $self->{protected_modes} unless @_ > 1;
    $self->_check_limits(\%args);
    $self->{protected_modes} = \%args;
}

sub protected_actions {
    my ($self, %args) = @_;
    return $self->{protected_actions} unless @_ > 1;
    $self->_check_limits(\%args);
    $self->{protected_actions} = \%args;
}

sub _check_limits {
    my ($self, $args) = @_;
    foreach my $limits (values %$args) {
        defined $limits->{$_}
          or croak("Missing required value in protected limits hash: '$_'.")
          for (qw(timeframe max_hits));
        croak("Unknown keys found in protected limits hash.")
          unless keys(%$limits) == 2;
    }
}

sub record_hit {
    my ($self, %args) = @_;
    $self->_verify_attributes();

    my $dbh       = $self->dbh;
    my $timestamp = time;

    my $id_callback = $self->identity_callback();
    my $user_id     = $id_callback->();
    croak(  "Identity callback failed to return a value to "
          . __PACKAGE__
          . "::record_hit.")
      unless $user_id;

    $self->record_hit_sth($dbh)->execute($user_id, $args{action}, $timestamp)
      or croak(  "Failed to insert hit into table '"
               . $self->table . "': "
               . $dbh->errstr);

    # record particulars of last hit for revoke
    $self->{last_hit} = {user_id   => $user_id,
                         action    => $args{action},
                         timestamp => $timestamp};
}

sub revoke_hit {
    my $self = shift;
    $self->_verify_attributes();

    my $dbh      = $self->dbh;
    my $last_hit = $self->{last_hit}
      or croak("revoke_hit called without previous hit!");

    my $sth = $self->revoke_hit_sth($dbh);
    $sth->execute($last_hit->{user_id}, $last_hit->{action},
                  $last_hit->{timestamp})
      or croak(  "Failed to delete hit from table '"
               . $self->table . "': "
               . $dbh->errstr);

}

sub check_violation {
    my ($self, %args) = @_;

    my $dbh = $self->dbh;

    my $id_callback = $self->identity_callback();
    my $user_id     = $id_callback->();
    croak(  "Identity callback failed to return a value to "
          . __PACKAGE__
          . "::check_violation.")
      unless $user_id;

    # get limits passed-in for protected modes, else lookup for actions
    my ($limits, $is_mode);
    if ($args{limits}) {
        $limits  = $args{limits};
        $is_mode = 1;
    } else {
        my $protected_actions = $self->protected_actions() || {};
        $limits = $protected_actions->{$args{action}};
        croak(  "Called check_violation() for unknown protected action "
              . "'$args{action}'.")
          unless $limits;
    }

    my $seconds = _timeframe_to_seconds($limits->{timeframe});

    my $sth = $self->check_violation_sth($dbh);
    $sth->execute($user_id, $args{action}, time - $seconds);
    my ($count) = $sth->fetchrow_array();
    $sth->finish;

    if ($count > $limits->{max_hits}) {

        # setup violation details
        if ($is_mode) {
            $self->violated_mode($args{action});
        } else {
            $self->violated_action($args{action});
        }
        $self->violated_limits($limits);

        return 1;
    }

    return 0;
}

#
# SQL code.  If you want to port this module to a new DB, add some
# magic here.  With any luck you won't have to - this SQL is pretty
# bland.
#

sub record_hit_sth {
    my ($self, $dbh) = @_;

    return $dbh->prepare_cached('INSERT INTO '
                            . $dbh->quote_identifier($self->table)
                            . ' (user_id, action, timestamp) VALUES (?,?,?)');
}

sub check_violation_sth {
    my ($self, $dbh) = @_;

    return $dbh->prepare_cached('SELECT COUNT(*) FROM '
                     . $dbh->quote_identifier($self->table)
                     . ' WHERE user_id = ? AND action = ? AND timestamp > ?');
}

sub revoke_hit_sth {
    my ($self, $dbh) = @_;

    return $dbh->prepare_cached('DELETE FROM '
                     . $dbh->quote_identifier($self->table)
                     . ' WHERE user_id = ? AND action = ? AND timestamp = ?');
}

1;
__END__