From 9dbece52729ba0c5170caad24b60874ee953eb8c Mon Sep 17 00:00:00 2001
From: Kyle d'Oliveira <kyle.doliveira@clio.com>
Date: Thu, 15 Feb 2018 12:50:42 -0800
Subject: [PATCH] Add an reader for the epoch_time variable in the cache so
 that it can also be returned in the data from the throttle. This is allows
 access to the same time that the cache uses for the count. This can be
 important for clients that want to provide rate limit information for
 well-behaved clients

---
 README.md                         | 6 +++---
 lib/rack/attack/cache.rb          | 7 ++++---
 lib/rack/attack/throttle.rb       | 5 ++++-
 spec/rack_attack_throttle_spec.rb | 8 ++++----
 4 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index b5381e7..8fe6d36 100644
--- a/README.md
+++ b/README.md
@@ -254,13 +254,13 @@ Here's an example response that includes conventional `X-RateLimit-*` headers:
 
 ```ruby
 Rack::Attack.throttled_response = lambda do |env|
-  now = Time.now
   match_data = env['rack.attack.match_data']
+  now = match_data[:epoch_time]
 
   headers = {
     'X-RateLimit-Limit' => match_data[:limit].to_s,
     'X-RateLimit-Remaining' => '0',
-    'X-RateLimit-Reset' => (now + (match_data[:period] - now.to_i % match_data[:period])).to_s
+    'X-RateLimit-Reset' => (now + (match_data[:period] - now % match_data[:period])).to_s
   }
 
   [ 429, headers, ["Throttled\n"]]
@@ -271,7 +271,7 @@ end
 For responses that did not exceed a throttle limit, Rack::Attack annotates the env with match data:
 
 ```ruby
-request.env['rack.attack.throttle_data'][name] # => { :count => n, :period => p, :limit => l }
+request.env['rack.attack.throttle_data'][name] # => { :count => n, :period => p, :limit => l, :epoch_time => t }
 ```
 
 ## Logging & Instrumentation
diff --git a/lib/rack/attack/cache.rb b/lib/rack/attack/cache.rb
index c2dd06c..73f0ca2 100644
--- a/lib/rack/attack/cache.rb
+++ b/lib/rack/attack/cache.rb
@@ -3,6 +3,7 @@ module Rack
     class Cache
 
       attr_accessor :prefix
+      attr_reader :last_epoch_time
 
       def initialize
         self.store = ::Rails.cache if defined?(::Rails.cache)
@@ -39,10 +40,10 @@ module Rack
       private
 
       def key_and_expiry(unprefixed_key, period)
-        epoch_time = Time.now.to_i
+        @last_epoch_time = Time.now.to_i
         # Add 1 to expires_in to avoid timing error: http://git.io/i1PHXA
-        expires_in = (period - (epoch_time % period) + 1).to_i
-        ["#{prefix}:#{(epoch_time / period).to_i}:#{unprefixed_key}", expires_in]
+        expires_in = (period - (@last_epoch_time % period) + 1).to_i
+        ["#{prefix}:#{(@last_epoch_time / period).to_i}:#{unprefixed_key}", expires_in]
       end
 
       def do_count(key, expires_in)
diff --git a/lib/rack/attack/throttle.rb b/lib/rack/attack/throttle.rb
index 19a6519..4dc999d 100644
--- a/lib/rack/attack/throttle.rb
+++ b/lib/rack/attack/throttle.rb
@@ -26,12 +26,15 @@ module Rack
         current_limit  = limit.respond_to?(:call) ? limit.call(req) : limit
         key            = "#{name}:#{discriminator}"
         count          = cache.count(key, current_period)
+        epoch_time     = cache.last_epoch_time
 
         data = {
           :count => count,
           :period => current_period,
-          :limit => current_limit
+          :limit => current_limit,
+          :epoch_time => epoch_time
         }
+
         (req.env['rack.attack.throttle_data'] ||= {})[name] = data
 
         (count > current_limit).tap do |throttled|
diff --git a/spec/rack_attack_throttle_spec.rb b/spec/rack_attack_throttle_spec.rb
index 28ab26c..5befbe9 100644
--- a/spec/rack_attack_throttle_spec.rb
+++ b/spec/rack_attack_throttle_spec.rb
@@ -20,7 +20,7 @@ describe 'Rack::Attack.throttle' do
     end
 
     it 'should populate throttle data' do
-      data = { :count => 1, :limit => 1, :period => @period }
+      data = { :count => 1, :limit => 1, :period => @period, epoch_time: Rack::Attack.cache.last_epoch_time.to_i }
       last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data
     end
   end
@@ -37,7 +37,7 @@ describe 'Rack::Attack.throttle' do
     it 'should tag the env' do
       last_request.env['rack.attack.matched'].must_equal 'ip/sec'
       last_request.env['rack.attack.match_type'].must_equal :throttle
-      last_request.env['rack.attack.match_data'].must_equal({:count => 2, :limit => 1, :period => @period})
+      last_request.env['rack.attack.match_data'].must_equal({:count => 2, :limit => 1, :period => @period, epoch_time: Rack::Attack.cache.last_epoch_time.to_i})
       last_request.env['rack.attack.match_discriminator'].must_equal('1.2.3.4')
     end
 
@@ -65,7 +65,7 @@ describe 'Rack::Attack.throttle with limit as proc' do
     end
 
     it 'should populate throttle data' do
-      data = { :count => 1, :limit => 1, :period => @period }
+      data = { :count => 1, :limit => 1, :period => @period, epoch_time: Rack::Attack.cache.last_epoch_time.to_i }
       last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data
     end
   end
@@ -89,7 +89,7 @@ describe 'Rack::Attack.throttle with period as proc' do
     end
 
     it 'should populate throttle data' do
-      data = { :count => 1, :limit => 1, :period => @period }
+      data = { :count => 1, :limit => 1, :period => @period, epoch_time: Rack::Attack.cache.last_epoch_time.to_i }
       last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data
     end
   end