Rate limiting whitelist based on headers provided by client. Stashing initial...

Rate limiting whitelist based on headers provided by client. Stashing initial changes to switch to master branch.

lib/Plack/Middleware/EnsThrottle.pm

@@ -45,6 +45,8 @@ enable 'EnsThrottle::Second',
   },
   blacklist => '192.168.2.1', #blacklist a single IP
   whitelist => ['192.168.0.0-192.169.255.255'], #whitelist a range
+  whitelist_hdr => 'Token',
+  whitelist_hdr_values => ['loveme'],
   client_id_prefix => 'second', #make the generated key more unique
   message => 'custom rate exceeded message',
   # If specified once limit has been hit add this value to the Retry-After header.
@@ -100,12 +102,13 @@ In all cases they limit up-to their time unit. For example a user is first seen
 use strict;
 use warnings;
 use parent 'Plack::Middleware';
-use Plack::Util::Accessor qw(max_requests backend path blacklist whitelist client_id_prefix message retry_after_addition);
+use Plack::Util::Accessor qw(max_requests backend path blacklist whitelist whitelist_hdr whitelist_hdr_values client_id_prefix message retry_after_addition);
 use Plack::Util;
 use Carp;
 use Net::CIDR::Lite;
 use Readonly;
 use Plack::Middleware::EnsThrottle::SimpleBackend;
+use Plack::Request;
 
 Readonly::Scalar my $CLIENT_ID_PREFIX => 'throttle';
 Readonly::Scalar my $MESSAGE => 'Too many requests';
@@ -125,6 +128,7 @@ sub prepare_app {
   croak "Cannot continue. No max_requests given" unless $self->max_requests();
   $self->blacklist($self->_populate_cidr($self->blacklist));
   $self->whitelist($self->_populate_cidr($self->whitelist));
+  $self->whitelist_hdr_values($self->_populate_array($self->whitelist_hdr_values));
   my $path = $self->path();
   croak "Cannot continue. No path given" unless $path;
   croak "Cannot continue. path must be an CODE ref" if ref($path) ne 'CODE';
@@ -139,6 +143,8 @@ sub call {
   my $res;
   my $remaining_requests = 0;
 
+  my $request = Plack::Request->new($env);
+
   if($self->_throttle_path($env)) {
 
     #If IP was in the blacklist then reject & allow no headers
@@ -151,6 +157,10 @@ sub call {
       $res = $self->app->($env);
       $add_headers = 0;
     }
+    elsif($self->_whitelisted_hdr($request->headers)) {
+      $res = $self->app->($env);
+      $add_headers = 0;
+    }
     else {
       $remaining_requests = $self->remaining_requests($env);
       #If we are throttled then we have to prepare the throttle response
@@ -267,6 +277,20 @@ sub _whitelisted {
   return $list->find($env->{REMOTE_ADDR});
 }
 
+sub _whitelisted_hdr {
+    my ($self, $headers) = @_;
+
+    my $header = $self->whitelist_hdr();
+    my @values = $self->whitelist_hdr_values();
+    return 0 unless $header && @values;
+    # Since this is a special header just for us we're
+    # going to assume there's only one value, to speed
+    # the checking code below
+    my $hdr_value = $headers->header($header);
+    return 0 unless $hdr_value;
+    return grep $hdr_value, @values
+}
+
 sub _client_id {
   my ($self, $env) = @_;
   my $id;
@@ -290,4 +314,14 @@ sub _populate_cidr {
   return $cidr;
 }
 
-1;
\ No newline at end of file
+# Ensure value is an array if a scalar is passed
+sub _populate_array {
+    my ($self, $input) = @_;
+
+    if($input) {
+	$input = (ref($input) eq 'ARRAY') ? $input : [$input];
+    }
+    return $input;
+}
+
+1;

t/ratelimit.t

@@ -85,6 +85,8 @@ assert_basic_rate_limit('EnsThrottle::Second', 1, 'second', 'You have done too m
       max_requests => 1,
       whitelist => '192.168.2.1',
       blacklist => ['192.167.0.0-192.167.255.255'],
+#      whitelist_hdr => 'Token',
+#      whitelist_hdr_values => ['loveme', 'imacool'],
       client_id_prefix => 'second', backend => Plack::Middleware::EnsThrottle::SimpleBackend->new();
 
     sub {
@@ -297,4 +299,4 @@ sub assert_basic_rate_limit {
   return;
 }
 
-done_testing();
\ No newline at end of file
+done_testing();