From 2a68ed3d0138827f4fed8a3520f7859493da4e1b Mon Sep 17 00:00:00 2001 From: franck cuny Date: Wed, 31 Mar 2010 22:30:23 +0200 Subject: import old Plack::Middleware::APIRateLimit (renamed as miyagawa's suggestion) and stole API from rack::throttle --- lib/Plack/Middleware/Throttle.pm | 134 ++++++++++++++++++++++++-- lib/Plack/Middleware/Throttle/Backend/Hash.pm | 20 ++++ lib/Plack/Middleware/Throttle/Daily.pm | 17 ++++ lib/Plack/Middleware/Throttle/Hourly.pm | 17 ++++ lib/Plack/Middleware/Throttle/Interval.pm | 12 +++ lib/Plack/Middleware/Throttle/Limiter.pm | 21 ++++ t/01_basic.t | 35 +++++++ 7 files changed, 250 insertions(+), 6 deletions(-) create mode 100644 lib/Plack/Middleware/Throttle/Backend/Hash.pm create mode 100644 lib/Plack/Middleware/Throttle/Daily.pm create mode 100644 lib/Plack/Middleware/Throttle/Hourly.pm create mode 100644 lib/Plack/Middleware/Throttle/Interval.pm create mode 100644 lib/Plack/Middleware/Throttle/Limiter.pm create mode 100644 t/01_basic.t 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 + +How many requests are authorized by hours + +=item B + +How many remaining requests + +=item B + +When will the counter be reseted (in epoch) + +=back + +=head2 VARIABLES + +=over 4 + +=item B + +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, +B and B. + +=item B + +HTTP code that will be returned when too many connections have been reached. + +=item B + +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; -- cgit 1.4.1