a middleware to throttle request on an API

5 files changed, 261 insertions, 6 deletions

diff --git a/lib/Plack/Middleware/APIRateLimit.pm b/lib/Plack/Middleware/APIRateLimit.pm
index 905e944..e94407c 100644
--- a/lib/Plack/Middleware/APIRateLimit.pm
+++ b/lib/Plack/Middleware/APIRateLimit.pm
@@ -1,27 +1,161 @@
 package Plack::Middleware::APIRateLimit;
 
-use strict;
-use warnings;
+use Moose;
+use Carp;
+use Scalar::Util;
+use Plack::Util;
+use DateTime;
+
 our $VERSION = '0.01';
 
+extends 'Plack::Middleware';
+
+has backend =>
+    ( is => 'rw', isa => 'Plack::Middleware::APIRateLimit::Backend', );
+has requests_per_hour =>
+    ( is => 'rw', isa => 'Int', lazy => 1, default => 60 );
+has key => ( is => 'rw', isa => 'Str', predicate => 'has_key' );
+
+sub prepare_app {
+    my $self = shift;
+    $self->backend( $self->_create_backend( $self->backend ) );
+}
+
+sub _create_backend {
+    my ( $self, $backend ) = @_;
+
+    return $backend if defined $backend && Scalar::Util::blessed $backend;
+
+    my ( $backend_name, $backend_options ) = ( undef, {} );
+
+    if ( !defined $backend ) {
+        $backend_name = "Hash";
+    }
+    elsif ( ref $backend eq 'ARRAY' ) {
+        $backend_name    = shift @$backend;
+        $backend_options = shift @$backend;
+    }
+    else {
+        $backend_name = $backend;
+    }
+
+    Plack::Util::load_class(
+        "Plack::Middleware::APIRateLimit::Backend::" . $backend_name )
+        ->new($backend_options);
+}
+
+sub call {
+    my ( $self, $env ) = @_;
+
+    my $res = $self->app->($env);
+
+    my $key = $self->_generate_key($env);
+
+    $self->backend->incr($key);
+    my $request_done = $self->backend->get($key);
+
+    return $self->over_rate_limit()
+        if $request_done > $self->requests_per_hour;
+
+    my $headers = $res->[1];
+    Plack::Util::header_set( $headers, 'X-RateLimit-Limit',
+        $self->requests_per_hour );
+    Plack::Util::header_set( $headers, 'X-RateLimit-Remaining',
+        ( $self->requests_per_hour - $request_done ) );
+    Plack::Util::header_set( $headers, 'X-RateLimit-Reset',
+        $self->_reset_time );
+    return $res;
+}
+
+sub _generate_key {
+    my ( $self, $env ) = @_;
+    return $self->key if $self->has_key;
+    if ( $env->{REMOTE_USER} ) {
+        return $env->{REMOTE_USER} . "_"
+            . DateTime->now->strftime("%Y-%m-%d-%H");
+    }
+    else {
+        return $env->{REMOTE_ADDR} . "_"
+            . DateTime->now->strftime("%Y-%m-%d-%H");
+    }
+}
+
+sub _reset_time {
+    my $reset = time + ( 60 - DateTime->now->minute ) * 60;
+}
+
+sub over_rate_limit {
+    my ($self) = @_;
+    return [
+        503,
+        [
+            'Content-Type'      => 'text/plain',
+            'X-RateLimit-Reset' => $self->_reset_time
+        ],
+        ['Over Rate Limit']
+    ];
+}
+
 1;
 __END__
 
 =head1 NAME
 
-Plack::Middleware::APIRateLimit -
+Plack::Middleware::APIRateLimit - A Plack Middleware for API Throttling
 
 =head1 SYNOPSIS
 
-  use Plack::Middleware::APIRateLimit;
+  my $handler = builder {
+    enable "APIRateLimit";
+    # or
+    enable "APIRateLimit", requests_per_hour => 2, backend => "Hash";
+    # or
+    enable "APIRateLimit", requests_per_hour => 2, backend => ["Redis", {port => 6379, server => '127.0.0.1'}];
+    sub { [ '200', [ 'Content-Type' => 'text/html' ], ['hello world'] ] };
+  };
 
 =head1 DESCRIPTION
 
-Plack::Middleware::APIRateLimit is
+Plack::Middleware::APIRateLimit is a Plack middleware for controlling API
+access.
+
+Set a limit on how many requests per hour is allowed on your API. In the case
+of a authorized request, 3 headers are added:
+
+=over 2
+
+=item B<X-RateLimit-Limit>
+
+How many requests are authorized by hours
+
+=item B<X-RateLimit-Remaining>
+
+How many remaining requests 
+
+=item B<X-RateLimit-Reset>
+
+When will the counter be reseted (in epoch)
+
+=back
+
+=head2 VARIABLES
+
+=over 4
+
+=item B<backend>
+
+Which backend to use. Currently only Hash and Redis are supported. If no
+backend is specified, Hash is used by default.
+
+=item B<requests_per_hour>
+
+How many requests is allowed by hour.
+
+=back
 
 =head1 AUTHOR
 
-franck cuny E<lt>franck.cuny@rtgi.frE<gt>
+franck cuny E<lt>franck@linkfluence.netE<gt>
 
 =head1 SEE ALSO
 
diff --git a/lib/Plack/Middleware/APIRateLimit/Backend.pm b/lib/Plack/Middleware/APIRateLimit/Backend.pm
new file mode 100644
index 0000000..5628e86
--- /dev/null
+++ b/lib/Plack/Middleware/APIRateLimit/Backend.pm
@@ -0,0 +1,14 @@
+package Plack::Middleware::APIRateLimit::Backend;
+
+use Moose;
+use Carp;
+
+sub incr {
+    confess "Backend must implement an incr method";
+}
+
+sub get {
+    confess "Backend must implement a get method";
+}
+
+1;
diff --git a/lib/Plack/Middleware/APIRateLimit/Backend/Hash.pm b/lib/Plack/Middleware/APIRateLimit/Backend/Hash.pm
new file mode 100644
index 0000000..0e407d8
--- /dev/null
+++ b/lib/Plack/Middleware/APIRateLimit/Backend/Hash.pm
@@ -0,0 +1,21 @@
+package Plack::Middleware::APIRateLimit::Backend::Hash;
+
+use Moose;
+extends 'Plack::Middleware::APIRateLimit::Backend';
+
+has store => (
+    is      => 'rw',
+    isa     => 'HashRef',
+    traits  => ['Hash'],
+    lazy    => 1,
+    default => sub { {} },
+    handles => { get => 'get', set => 'set' }
+);
+
+sub incr {
+    my ( $self, $key ) = @_;
+    my $value = ($self->get($key) || 0) + 1;
+    $self->set( $key => $value );
+}
+
+1;
diff --git a/lib/Plack/Middleware/APIRateLimit/Backend/Redis.pm b/lib/Plack/Middleware/APIRateLimit/Backend/Redis.pm
new file mode 100644
index 0000000..e647f6a
--- /dev/null
+++ b/lib/Plack/Middleware/APIRateLimit/Backend/Redis.pm
@@ -0,0 +1,43 @@
+package Plack::Middleware::APIRateLimit::Backend::Redis;
+
+use Moose;
+extends 'Plack::Middleware::APIRateLimit::Backend';
+
+use AnyEvent::Redis;
+
+has store => (
+    is      => 'rw',
+    isa     => 'AnyEvent::Redis',
+    lazy    => 1,
+    default => sub {
+        return AnyEvent::Redis->new(
+            host => '127.0.0.1',
+            port => 6378,
+        );
+    }
+);
+
+sub BUILD {
+    my $self = shift;
+    if (scalar @_) {
+        $self->store(AnyEvent::Redis->new(@_));
+    }
+    return $self;
+}
+
+sub get {
+    my ( $self, $key ) = @_;
+    my $val = $self->store->get($key)->recv;
+    if ( !$val ) {
+        $self->store->set( $key => 1 )->recv;
+        $val = 1;
+    }
+    return $val;
+}
+
+sub incr {
+    my ($self, $key) = @_;
+    return $self->store->incr($key)->recv;
+}
+
+1;
diff --git a/t/01_basic.t b/t/01_basic.t
new file mode 100644
index 0000000..5fb3212
--- /dev/null
+++ b/t/01_basic.t
@@ -0,0 +1,43 @@
+use strict;
+use warnings;
+use Test::More;
+
+use Plack::Test;
+use Plack::Builder;
+use HTTP::Request::Common;
+
+#use AnyEvent::Redis;
+#my $redis = AnyEvent::Redis->new(port => 6379, server => '127.0.0.1');
+#$redis->flushall;
+
+my $handler = builder {
+    enable "APIRateLimit";
+    #enable "APIRateLimit", requests_per_hour => 2, backend => "Hash";
+    #enable "APIRateLimit",
+        #requests_per_hour => 2,
+        #backend => [ "Redis", { port => 6379, server => '127.0.0.1' } ];
+    sub { [ '200', [ 'Content-Type' => 'text/html' ], ['hello world'] ] };
+};
+
+test_psgi
+    app    => $handler,
+    client => sub {
+    my $cb = shift;
+    use YAML::Syck;
+    {
+        for ( 1 .. 2 ) {
+            my $req = GET "http://localhost/";
+            my $res = $cb->($req);
+            is $res->code, 200;
+	    ok $res->headers('X-RateLimit-Limit');
+#	    warn Dump $res;
+        }
+        my $req = GET "http://localhost/";
+        my $res = $cb->($req);
+        is $res->code, 503;
+	ok $res->headers('X-RateLimit-Reset');
+#	warn Dump $res;
+    }
+};
+
+done_testing;