1 files changed, 26 insertions, 9 deletions
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;