From c0d866e0e425bcb4694a562e7680cffe8a686d0f Mon Sep 17 00:00:00 2001
From: Eebs Kobeissi <ebrahim.kobeissi@gmail.com>
Date: Tue, 29 May 2018 09:51:07 -0400
Subject: [PATCH] Set the 'X-RateLimit-Reset' header

The X-RateLimit-Reset' header will be set to the timestamp when the
current period will reset.
---
 CHANGELOG.md                        |  6 +++++-
 lib/rack/attack/rate-limit.rb       | 12 ++++++++++++
 spec/rack/attack/rate-limit_spec.rb | 21 +++++++++++++++++++--
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 75bdcbc..a349c3c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,9 @@
 # rack-attack-rate-limit changelog
 
+## Unreleased
+
+* Set the 'X-RateLimit-Reset' header to the timestamp when the current throttle period expires.
+
 ## 1.1.0
 
 * Add support for multiple throttles.
@@ -13,4 +17,4 @@
 
 ## 0.1.0
 
-* Initial release.
\ No newline at end of file
+* Initial release.
diff --git a/lib/rack/attack/rate-limit.rb b/lib/rack/attack/rate-limit.rb
index 5664dbd..88a9720 100644
--- a/lib/rack/attack/rate-limit.rb
+++ b/lib/rack/attack/rate-limit.rb
@@ -53,6 +53,7 @@ def add_rate_limit_headers!(headers, env)
         throttle_data = throttle_data_closest_to_limit(env)
         headers['X-RateLimit-Limit']      = rate_limit_limit(throttle_data).to_s
         headers['X-RateLimit-Remaining']  = rate_limit_remaining(throttle_data).to_s
+        headers['X-RateLimit-Reset']      = rate_limit_reset(throttle_data).to_s
         headers
       end
 
@@ -76,6 +77,17 @@ def rate_limit_remaining(throttle_data)
         rate_limit_limit(throttle_data) - throttle_data[:count]
       end
 
+      # RateLimit timestamp when the current period expires resets from Rack::Attack
+      #
+      # env - Hash
+      #
+      # Returns Fixnum
+      def rate_limit_reset(throttle_data)
+        now = Time.now.to_i
+        throttle_period = throttle_data[:period]
+        now + (throttle_period - now % throttle_period)
+      end
+
       # Rate Limit available method for Rack::Attack provider
       # Checks that at least one of the keys provided by the user are in the rack.attack.throttle_data env hash key
       #
diff --git a/spec/rack/attack/rate-limit_spec.rb b/spec/rack/attack/rate-limit_spec.rb
index d68190a..f28ae53 100644
--- a/spec/rack/attack/rate-limit_spec.rb
+++ b/spec/rack/attack/rate-limit_spec.rb
@@ -25,6 +25,7 @@
     it 'should not create RateLimit headers' do
       last_response.header.key?('X-RateLimit-Limit').should be false
       last_response.header.key?('X-RateLimit-Remaining').should be false
+      last_response.header.key?('X-RateLimit-Reset').should be false
     end
 
   end
@@ -36,16 +37,18 @@
 
     let(:request_limit) { (1..10_000).to_a.sample }
     let(:request_count) { (1..(request_limit - 10)).to_a.sample }
+    let(:request_period) { rand(60..3600) }
 
     context 'one throttle only' do
 
       let(:rack_attack_throttle_data) do
-        { "#{throttle_one}" => { count: request_count, limit: request_limit } }
+        { "#{throttle_one}" => { count: request_count, limit: request_limit, period: request_period} }
       end
 
       it 'should include RateLimit headers' do
         last_response.header.key?('X-RateLimit-Limit').should be true
         last_response.header.key?('X-RateLimit-Remaining').should be true
+        last_response.header.key?('X-RateLimit-Reset').should be true
       end
 
       it 'should return correct rate limit in header' do
@@ -55,6 +58,12 @@
       it 'should return correct remaining calls in header' do
         last_response.header['X-RateLimit-Remaining'].to_i.should eq(request_limit - request_count)
       end
+
+      it 'should return the timestamp when the current throttle period resets' do
+        current_epoch_time = Time.now.to_i
+        timestamp_for_next_period = current_epoch_time + (request_period - current_epoch_time % request_period)
+        last_response.header['X-RateLimit-Reset'].to_i.should eq(timestamp_for_next_period)
+      end
     end
 
     context 'multiple throttles' do
@@ -69,17 +78,19 @@
 
       let(:request_limits) { 3.times.map { (1..10_000).to_a.sample } }
       let(:request_counts) { 3.times.map { |index| (1..(request_limits[index] - 10)).to_a.sample } }
+      let(:request_periods) { 3.times.map { rand(60..3600) } }
 
       let(:rack_attack_throttle_data) do
         data = {}
         [throttle_one, throttle_two, throttle_three].each_with_index do |thr, thr_index|
-          data["#{thr}"] = { count: request_counts[thr_index], limit: request_limits[thr_index] }
+          data["#{thr}"] = { count: request_counts[thr_index], limit: request_limits[thr_index], period: request_periods[thr_index] }
         end
         data
       end
       it 'should include RateLimit headers' do
         last_response.header.key?('X-RateLimit-Limit').should be true
         last_response.header.key?('X-RateLimit-Remaining').should be true
+        last_response.header.key?('X-RateLimit-Reset').should be true
       end
 
       describe 'header values' do
@@ -95,6 +106,12 @@
         it 'should return correct remaining calls' do
           last_response.header['X-RateLimit-Remaining'].to_i.should eq(request_differences[min_index])
         end
+
+        it 'should return the timestamp when the current throttle period resets' do
+          current_epoch_time = Time.now.to_i
+          timestamp_for_next_period = current_epoch_time + (request_periods[min_index] - current_epoch_time % request_periods[min_index])
+          last_response.header['X-RateLimit-Reset'].to_i.should eq(timestamp_for_next_period)
+        end
       end
     end
   end