From a34c187dda0243673198e20314e7fe5280ea09e8 Mon Sep 17 00:00:00 2001
From: fatkodima <fatkodima123@gmail.com>
Date: Fri, 11 Oct 2019 23:06:44 +0300
Subject: [PATCH] Allow to configure Retry-After header for default
 throttled_response handler

---
 README.md                          |  5 +++++
 lib/rack/attack.rb                 |  2 ++
 lib/rack/attack/configuration.rb   | 14 +++++++++++---
 spec/acceptance/throttling_spec.rb | 20 +++++++++++++++++++-
 spec/rack_attack_throttle_spec.rb  |  4 ----
 5 files changed, 37 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index f18f9f3..866cf35 100644
--- a/README.md
+++ b/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
diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb
index be9fd74..f533ca6 100644
--- a/lib/rack/attack.rb
+++ b/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,
diff --git a/lib/rack/attack/configuration.rb b/lib/rack/attack/configuration.rb
index 45c579d..5a15c74 100644
--- a/lib/rack/attack/configuration.rb
+++ b/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
diff --git a/spec/acceptance/throttling_spec.rb b/spec/acceptance/throttling_spec.rb
index bf3a79b..0db89dd 100644
--- a/spec/acceptance/throttling_spec.rb
+++ b/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
diff --git a/spec/rack_attack_throttle_spec.rb b/spec/rack_attack_throttle_spec.rb
index 8346611..dcb1e41 100644
--- a/spec/rack_attack_throttle_spec.rb
+++ b/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