From 58ba650e6b4a2bfa357dd07f04ec7a3bcbfbd62f Mon Sep 17 00:00:00 2001
From: Gonzalo Rodriguez <grzuy0@gmail.com>
Date: Wed, 10 Jul 2019 17:21:48 -0300
Subject: [PATCH] feat: provide discriminator in throttle_data

---
 README.md                         |  2 +-
 lib/rack/attack/throttle.rb       |  1 +
 spec/rack_attack_throttle_spec.rb | 37 +++++++++++++++++++++++++++----
 3 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index e7468fed..357a18ef 100644
--- a/README.md
+++ b/README.md
@@ -354,7 +354,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, :epoch_time => t }
+request.env['rack.attack.throttle_data'][name] # => { discriminator: d, count: n, period: p, limit: l, epoch_time: t }
 ```
 
 ## Logging & Instrumentation
diff --git a/lib/rack/attack/throttle.rb b/lib/rack/attack/throttle.rb
index 4d5316a6..672140a3 100644
--- a/lib/rack/attack/throttle.rb
+++ b/lib/rack/attack/throttle.rb
@@ -31,6 +31,7 @@ def matched_by?(request)
         epoch_time     = cache.last_epoch_time
 
         data = {
+          discriminator: discriminator,
           count: count,
           period: current_period,
           limit: current_limit,
diff --git a/spec/rack_attack_throttle_spec.rb b/spec/rack_attack_throttle_spec.rb
index ef47a7b0..d760d942 100644
--- a/spec/rack_attack_throttle_spec.rb
+++ b/spec/rack_attack_throttle_spec.rb
@@ -22,7 +22,14 @@
     end
 
     it 'should populate throttle data' do
-      data = { count: 1, limit: 1, period: @period, epoch_time: Rack::Attack.cache.last_epoch_time.to_i }
+      data = {
+        count: 1,
+        limit: 1,
+        period: @period,
+        epoch_time: Rack::Attack.cache.last_epoch_time.to_i,
+        discriminator: "1.2.3.4"
+      }
+
       last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data
     end
   end
@@ -39,7 +46,15 @@
     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, epoch_time: Rack::Attack.cache.last_epoch_time.to_i)
+
+      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,
+        discriminator: "1.2.3.4"
+      )
+
       last_request.env['rack.attack.match_discriminator'].must_equal('1.2.3.4')
     end
 
@@ -67,7 +82,14 @@
     end
 
     it 'should populate throttle data' do
-      data = { count: 1, limit: 1, period: @period, epoch_time: Rack::Attack.cache.last_epoch_time.to_i }
+      data = {
+        count: 1,
+        limit: 1,
+        period: @period,
+        epoch_time: Rack::Attack.cache.last_epoch_time.to_i,
+        discriminator: "1.2.3.4"
+      }
+
       last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data
     end
   end
@@ -91,7 +113,14 @@
     end
 
     it 'should populate throttle data' do
-      data = { count: 1, limit: 1, period: @period, epoch_time: Rack::Attack.cache.last_epoch_time.to_i }
+      data = {
+        count: 1,
+        limit: 1,
+        period: @period,
+        epoch_time: Rack::Attack.cache.last_epoch_time.to_i,
+        discriminator: "1.2.3.4"
+      }
+
       last_request.env['rack.attack.throttle_data']['ip/sec'].must_equal data
     end
   end