handle black and white list, add support to interval

3 files changed, 81 insertions, 19 deletions

diff --git a/lib/Plack/Middleware/Throttle.pm b/lib/Plack/Middleware/Throttle.pm
index 96a66f7..fe1fae1 100644
--- a/lib/Plack/Middleware/Throttle.pm
+++ b/lib/Plack/Middleware/Throttle.pm
@@ -17,6 +17,10 @@ has backend => ( is => 'rw', isa => 'Object', required => 1 );
 has key_prefix =>
     ( is => 'rw', isa => 'Str', lazy => 1, default => 'throttle' );
 has max => ( is => 'rw', isa => 'Int', lazy => 1, default => 100 );
+has white_list =>
+    ( is => 'rw', isa => 'ArrayRef', predicate => 'has_white_list' );
+has black_list =>
+    ( is => 'rw', isa => 'ArrayRef', predicate => 'has_black_list' );
 
 sub prepare_app {
     my $self = shift;
@@ -31,16 +35,17 @@ sub _create_backend {
     }
 
     return $backend if defined $backend && Scalar::Util::blessed $backend;
-    die "backend must be a cache objectn";
+    die "backend must be a cache object";
 }
 
 sub call {
     my ( $self, $env ) = @_;
 
-    my $res          = $self->app->($env);
-    my $request_done = $self->request_done($env);
+    my $res     = $self->app->($env);
+    my $key     = $self->cache_key($env);
+    my $allowed = $self->allowed($key);
 
-    if ( $request_done > $self->max ) {
+    if ( !$allowed ) {
         $self->over_rate_limit();
     }
     else {
@@ -48,16 +53,40 @@ sub call {
             $res,
             sub {
                 my $res = shift;
-                $self->add_headers( $res, $request_done );
+                $self->add_headers($res);
             }
         );
     }
 }
 
+sub allowed {
+    return 1;
+}
+
 sub request_done {
     return 1;
 }
 
+sub is_white_listed {
+    my ( $self, $env ) = @_;
+    return 1 if !$self->has_white_list;
+    my $ip = $env->{REMOTE_ADDR};
+    if ( grep { $_ == $ip } @{ $self->white_list } ) {
+        return 1;
+    }
+    return 0;
+}
+
+sub is_black_listed {
+    my ( $self, $env ) = @_;
+    return 0 if !$self->has_black_list;
+    my $ip = $env->{REMOTE_ADDR};
+    if ( grep { $_ == $ip } @{ $self->black_list } ) {
+        return 1;
+    }
+    return 0;
+}
+
 sub over_rate_limit {
     my $self = shift;
     return [
@@ -71,12 +100,9 @@ sub over_rate_limit {
 }
 
 sub add_headers {
-    my ( $self, $res, $request_done ) = @_;
+    my ( $self, $res ) = @_;
     my $headers = $res->[1];
-    Plack::Util::header_set( $headers, 'X-RateLimit-Limit',
-        $self->max );
-    Plack::Util::header_set( $headers, 'X-RateLimit-Remaining',
-        ( $self->max - $request_done ) );
+    Plack::Util::header_set( $headers, 'X-RateLimit-Limit', $self->max );
     Plack::Util::header_set( $headers, 'X-RateLimit-Reset',
         $self->reset_time );
     return $res;
diff --git a/lib/Plack/Middleware/Throttle/Interval.pm b/lib/Plack/Middleware/Throttle/Interval.pm
index cbe7d59..3e61220 100644
--- a/lib/Plack/Middleware/Throttle/Interval.pm
+++ b/lib/Plack/Middleware/Throttle/Interval.pm
@@ -3,10 +3,29 @@ package Plack::Middleware::Throttle::Interval;
 use Moose;
 extends 'Plack::Middleware::Throttle';
 
+has min => (is => 'rw', isa => 'Int', default => 0, lazy => 1);
+
 sub allowed {
+    my ($self, $key) = @_;
+
+    my $t1 = time();
+    my $t0 = $self->backend->get($key);
+    $self->backend->set($key, $t1);
+
+    if (!$t0 || ($t1 - $t0) > $self->min) {
+        return 1;
+    }else{
+        return 0;
+    }
 }
 
 sub cache_key {
+    my ( $self, $env ) = @_;
+    $self->client_identifier($env);
+}
+
+sub reset_time {
+    time + 1;
 }
 
 1;
diff --git a/lib/Plack/Middleware/Throttle/Limiter.pm b/lib/Plack/Middleware/Throttle/Limiter.pm
index 626732d..4194fe4 100644
--- a/lib/Plack/Middleware/Throttle/Limiter.pm
+++ b/lib/Plack/Middleware/Throttle/Limiter.pm
@@ -3,19 +3,36 @@ package Plack::Middleware::Throttle::Limiter;
 use Moose;
 extends 'Plack::Middleware::Throttle';
 
-sub request_done {
-    my ( $self, $env ) = @_;
-    my $key = $self->cache_key($env);
+has _request_done => (
+    is      => 'rw',
+    isa     => 'Int',
+    default => 0,
+    clearer => '_clear_request_done'
+);
 
-    $self->backend->incr($key);
+sub allowed {
+    my ( $self, $key ) = @_;
 
-    my $request_done = $self->backend->get($key);
+    $self->backend->incr($key);
+    $self->request_done($key);
+    ( $self->_request_done > $self->max ) ? return 0 : return 1;
+}
 
-    if ( !$request_done ) {
-        $self->backend->set( $key, 1 );
-    }
+sub request_done {
+    my ( $self, $key ) = @_;
+    $self->_request_done( $self->backend->get($key) || 0 );
+}
 
-    $request_done;
+sub add_headers {
+    my ( $self, $res ) = @_;
+    my $headers = $res->[1];
+    Plack::Util::header_set( $headers, 'X-RateLimit-Limit', $self->max );
+    Plack::Util::header_set( $headers, 'X-RateLimit-Remaining',
+        ( $self->max - $self->_request_done ) );
+    Plack::Util::header_set( $headers, 'X-RateLimit-Reset',
+        $self->reset_time );
+    $self->_clear_request_done;
+    return $res;
 }
 
 1;