import old Plack::Middleware::APIRateLimit (renamed as miyagawa's suggestion) and stole API from rack::throttle

7 files changed, 250 insertions, 6 deletions

diff --git a/lib/Plack/Middleware/Throttle.pm b/lib/Plack/Middleware/Throttle.pm
index 6fd034a..96a66f7 100644
--- a/lib/Plack/Middleware/Throttle.pm
+++ b/lib/Plack/Middleware/Throttle.pm
@@ -1,23 +1,145 @@
 package Plack::Middleware::Throttle;
 
-use strict;
-use warnings;
+use Moose;
+use Carp;
+use Scalar::Util;
+use DateTime;
+use Plack::Util;
+
 our $VERSION = '0.01';
 
+extends 'Plack::Middleware';
+
+has code => ( is => 'rw', isa => 'Int', lazy => 1, default => '503' );
+has message =>
+    ( is => 'rw', isa => 'Str', lazy => 1, default => 'Over rate limit' );
+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 );
+
+sub prepare_app {
+    my $self = shift;
+    $self->backend( $self->_create_backend( $self->backend ) );
+}
+
+sub _create_backend {
+    my ( $self, $backend ) = @_;
+
+    if ( defined !$backend ) {
+        Plack::Util::load_class("Plack::Middleware::Throttle::Backend::Hash");
+    }
+
+    return $backend if defined $backend && Scalar::Util::blessed $backend;
+    die "backend must be a cache objectn";
+}
+
+sub call {
+    my ( $self, $env ) = @_;
+
+    my $res          = $self->app->($env);
+    my $request_done = $self->request_done($env);
+
+    if ( $request_done > $self->max ) {
+        $self->over_rate_limit();
+    }
+    else {
+        $self->response_cb(
+            $res,
+            sub {
+                my $res = shift;
+                $self->add_headers( $res, $request_done );
+            }
+        );
+    }
+}
+
+sub request_done {
+    return 1;
+}
+
+sub over_rate_limit {
+    my $self = shift;
+    return [
+        $self->code,
+        [
+            'Content-Type'      => 'text/plain',
+            'X-RateLimit-Reset' => $self->reset_time
+        ],
+        [ $self->message ]
+    ];
+}
+
+sub add_headers {
+    my ( $self, $res, $request_done ) = @_;
+    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-Reset',
+        $self->reset_time );
+    return $res;
+}
+
+sub client_identifier {
+    my ( $self, $env ) = @_;
+    if ( $env->{REMOTE_USER} ) {
+        return $self->key_prefix."_".$env->{REMOTE_USER};
+    }
+    else {
+        return $self->key_prefix."_".$env->{REMOTE_ADDR};
+    }
+}
+
 1;
 __END__
 
 =head1 NAME
 
-Plack::Middleware::Throttle -
+Plack::Middleware::Throttle - A Plack Middleware for rate-limiting incoming HTTP requests.
 
 =head1 SYNOPSIS
 
-  use Plack::Middleware::Throttle;
-
 =head1 DESCRIPTION
 
-Plack::Middleware::Throttle is
+Set a limit on how many requests per hour is allowed on your API. In 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. Backend must implement B<set>,
+B<get> and B<incr>.
+
+=item B<code>
+
+HTTP code that will be returned when too many connections have been reached.
+
+=item B<message>
+
+HTTP message that will be returned when too many connections have been reached.
+
+=back
 
 =head1 AUTHOR
 
diff --git a/lib/Plack/Middleware/Throttle/Backend/Hash.pm b/lib/Plack/Middleware/Throttle/Backend/Hash.pm
new file mode 100644
index 0000000..9144e36
--- /dev/null
+++ b/lib/Plack/Middleware/Throttle/Backend/Hash.pm
@@ -0,0 +1,20 @@
+package Plack::Middleware::Throttle::Backend::Hash;
+
+use Moose;
+
+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/Throttle/Daily.pm b/lib/Plack/Middleware/Throttle/Daily.pm
new file mode 100644
index 0000000..d28d1d7
--- /dev/null
+++ b/lib/Plack/Middleware/Throttle/Daily.pm
@@ -0,0 +1,17 @@
+package Plack::Middleware::Throttle::Daily;
+
+use Moose;
+extends 'Plack::Middleware::Throttle::Limiter';
+
+sub cache_key {
+    my ( $self, $env ) = @_;
+    $self->client_identifier($env) . "_"
+        . DateTime->now->strftime("%Y-%m-%d");
+}
+
+sub reset_time {
+    my $dt = DateTime->now;
+    (24 * 3600) - (( 60 * $dt->minute ) + $dt->second);
+}
+
+1;
diff --git a/lib/Plack/Middleware/Throttle/Hourly.pm b/lib/Plack/Middleware/Throttle/Hourly.pm
new file mode 100644
index 0000000..818d70b
--- /dev/null
+++ b/lib/Plack/Middleware/Throttle/Hourly.pm
@@ -0,0 +1,17 @@
+package Plack::Middleware::Throttle::Hourly;
+
+use Moose;
+extends 'Plack::Middleware::Throttle::Limiter';
+
+sub cache_key {
+    my ( $self, $env ) = @_;
+    $self->client_identifier($env) . "_"
+        . DateTime->now->strftime("%Y-%m-%d-%H");
+}
+
+sub reset_time {
+    my $dt = DateTime->now;
+    3600 - (( 60 * $dt->minute ) + $dt->second);
+}
+
+1;
diff --git a/lib/Plack/Middleware/Throttle/Interval.pm b/lib/Plack/Middleware/Throttle/Interval.pm
new file mode 100644
index 0000000..cbe7d59
--- /dev/null
+++ b/lib/Plack/Middleware/Throttle/Interval.pm
@@ -0,0 +1,12 @@
+package Plack::Middleware::Throttle::Interval;
+
+use Moose;
+extends 'Plack::Middleware::Throttle';
+
+sub allowed {
+}
+
+sub cache_key {
+}
+
+1;
diff --git a/lib/Plack/Middleware/Throttle/Limiter.pm b/lib/Plack/Middleware/Throttle/Limiter.pm
new file mode 100644
index 0000000..626732d
--- /dev/null
+++ b/lib/Plack/Middleware/Throttle/Limiter.pm
@@ -0,0 +1,21 @@
+package Plack::Middleware::Throttle::Limiter;
+
+use Moose;
+extends 'Plack::Middleware::Throttle';
+
+sub request_done {
+    my ( $self, $env ) = @_;
+    my $key = $self->cache_key($env);
+
+    $self->backend->incr($key);
+
+    my $request_done = $self->backend->get($key);
+
+    if ( !$request_done ) {
+        $self->backend->set( $key, 1 );
+    }
+
+    $request_done;
+}
+
+1;
diff --git a/t/01_basic.t b/t/01_basic.t
new file mode 100644
index 0000000..31a4e98
--- /dev/null
+++ b/t/01_basic.t
@@ -0,0 +1,35 @@
+use strict;
+use warnings;
+use Test::More;
+
+use Plack::Test;
+use Plack::Builder;
+use HTTP::Request::Common;
+use Plack::Middleware::Throttle::Backend::Hash;
+
+my $handler = builder {
+    enable "Throttle::Hourly",
+        max     => 2,
+        backend => Plack::Middleware::Throttle::Backend::Hash->new();
+    sub { [ '200', [ 'Content-Type' => 'text/html' ], ['hello world'] ] };
+};
+
+test_psgi
+    app    => $handler,
+    client => sub {
+    my $cb = shift;
+    {
+        for ( 1 .. 2 ) {
+            my $req = GET "http://localhost/";
+            my $res = $cb->($req);
+            is $res->code, 200, 'http response is 200';
+            ok $res->headers('X-RateLimit-Limit'), 'header ratelimit';
+        }
+        my $req = GET "http://localhost/";
+        my $res = $cb->($req);
+        is $res->code, 503, 'http response is 503';
+        ok $res->headers('X-RateLimit-Reset'), 'header reset';
+    }
+    };
+
+done_testing;