Merge pull request #440 from fatkodima/retry-after-header

Allow to configure Retry-After header for default throttled_response handler

README.md

@@ -342,6 +342,11 @@ end
 While Rack::Attack's primary focus is minimizing harm from abusive clients, it
 can also be used to return rate limit data that's helpful for well-behaved clients.
 
+If you want to return to user how many seconds to wait until he can start sending requests again, this can be done through enabling `Retry-After` header:
+```ruby
+Rack::Attack.throttled_response_retry_after_header = true
+```
+
 Here's an example response that includes conventional `RateLimit-*` headers:
 
 ```ruby

lib/rack/attack.rb

@@ -67,6 +67,8 @@ module Rack
         :blocklisted_response=,
         :throttled_response,
         :throttled_response=,
+        :throttled_response_retry_after_header,
+        :throttled_response_retry_after_header=,
         :clear_configuration,
         :safelists,
         :blocklists,

lib/rack/attack/configuration.rb

@@ -4,7 +4,7 @@ module Rack
   class Attack
     class Configuration
       attr_reader :safelists, :blocklists, :throttles, :anonymous_blocklists, :anonymous_safelists
-      attr_accessor :blocklisted_response, :throttled_response
+      attr_accessor :blocklisted_response, :throttled_response, :throttled_response_retry_after_header
 
       def initialize
         @safelists = {}
@@ -13,11 +13,18 @@ module Rack
         @tracks = {}
         @anonymous_blocklists = []
         @anonymous_safelists = []
+        @throttled_response_retry_after_header = false
 
         @blocklisted_response = lambda { |_env| [403, { 'Content-Type' => 'text/plain' }, ["Forbidden\n"]] }
         @throttled_response   = lambda do |env|
-          retry_after = (env['rack.attack.match_data'] || {})[:period]
-          [429, { 'Content-Type' => 'text/plain', 'Retry-After' => retry_after.to_s }, ["Retry later\n"]]
+          if throttled_response_retry_after_header
+            match_data = env['rack.attack.match_data']
+            now = match_data[:epoch_time]
+            retry_after = match_data[:period] - (now % match_data[:period])
+            [429, { 'Content-Type' => 'text/plain', 'Retry-After' => retry_after.to_s }, ["Retry later\n"]]
+          else
+            [429, { 'Content-Type' => 'text/plain' }, ["Retry later\n"]]
+          end
         end
       end
 
@@ -86,6 +93,7 @@ module Rack
         @tracks = {}
         @anonymous_blocklists = []
         @anonymous_safelists = []
+        @throttled_response_retry_after_header = false
       end
     end
   end

spec/acceptance/throttling_spec.rb

@@ -20,7 +20,7 @@ describe "#throttle" do
     get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
 
     assert_equal 429, last_response.status
-    assert_equal "60", last_response.headers["Retry-After"]
+    assert_nil last_response.headers["Retry-After"]
     assert_equal "Retry later\n", last_response.body
 
     get "/", {}, "REMOTE_ADDR" => "5.6.7.8"
@@ -34,6 +34,24 @@ describe "#throttle" do
     end
   end
 
+  it "returns correct Retry-After header if enabled" do
+    Rack::Attack.throttled_response_retry_after_header = true
+
+    Rack::Attack.throttle("by ip", limit: 1, period: 60) do |request|
+      request.ip
+    end
+
+    Timecop.freeze(Time.at(0)) do
+      get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
+      assert_equal 200, last_response.status
+    end
+
+    Timecop.freeze(Time.at(25)) do
+      get "/", {}, "REMOTE_ADDR" => "1.2.3.4"
+      assert_equal "35", last_response.headers["Retry-After"]
+    end
+  end
+
   it "supports limit to be dynamic" do
     # Could be used to have different rate limits for authorized
     # vs general requests

spec/rack_attack_throttle_spec.rb

@@ -57,10 +57,6 @@ describe 'Rack::Attack.throttle' do
 
       _(last_request.env['rack.attack.match_discriminator']).must_equal('1.2.3.4')
     end
-
-    it 'should set a Retry-After header' do
-      _(last_response.headers['Retry-After']).must_equal @period.to_s
-    end
   end
 end