handle black and white list, add support to interval

2 files changed, 45 insertions, 9 deletions

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;