From 2413ece281d01130fc6fed955c45c016f7ce17f3 Mon Sep 17 00:00:00 2001
From: Mac Siri <mac@forem.com>
Date: Tue, 6 Sep 2022 10:06:29 -0400
Subject: [PATCH] Use remote_ip in Rack::Attack (2nd attempt) (#18409)

---
 .../{rack/attack.rb => rack_attack.rb}        | 36 +++++++++----------
 .../attack_spec.rb => rack_attack_spec.rb}    |  8 ++++-
 2 files changed, 24 insertions(+), 20 deletions(-)
 rename config/initializers/{rack/attack.rb => rack_attack.rb} (56%)
 rename spec/initializers/{rack/attack_spec.rb => rack_attack_spec.rb} (97%)

diff --git a/config/initializers/rack/attack.rb b/config/initializers/rack_attack.rb
similarity index 56%
rename from config/initializers/rack/attack.rb
rename to config/initializers/rack_attack.rb
index 3174d8dab..0df287598 100644
--- a/config/initializers/rack/attack.rb
+++ b/config/initializers/rack_attack.rb
@@ -2,22 +2,33 @@ Rack::Attack.throttled_response_retry_after_header = true
 
 module Rack
   class Attack
+    class Request < ::Rack::Request
+      def track_and_return_ip
+        if ApplicationConfig["FASTLY_API_KEY"].present?
+          Honeycomb.add_field("fastly_client_ip", env["HTTP_FASTLY_CLIENT_IP"])
+          env["HTTP_FASTLY_CLIENT_IP"]
+        else
+          ActionDispatch::Request.new(env).remote_ip
+        end
+      end
+    end
+
     throttle("search_throttle", limit: 5, period: 1) do |request|
       if request.path.starts_with?("/search/")
-        track_and_return_ip(request.env["HTTP_FASTLY_CLIENT_IP"])
+        request.track_and_return_ip
       end
     end
 
     throttle("api_throttle", limit: 3, period: 1) do |request|
       if request.path.starts_with?("/api/") && request.get?
-        track_and_return_ip(request.env["HTTP_FASTLY_CLIENT_IP"])
+        request.track_and_return_ip
       end
     end
 
     throttle("api_write_throttle", limit: 1, period: 1) do |request|
       if request.path.starts_with?("/api/") && (request.put? || request.post? || request.delete?)
         Honeycomb.add_field("user_api_key", request.env["HTTP_API_KEY"])
-        ip_address = track_and_return_ip(request.env["HTTP_FASTLY_CLIENT_IP"])
+        ip_address = request.track_and_return_ip
         if request.env["HTTP_API_KEY"].present?
           "#{ip_address}-#{request.env['HTTP_API_KEY']}"
         elsif ip_address.present?
@@ -26,25 +37,12 @@ module Rack
       end
     end
 
-    throttle("site_hits", limit: 40, period: 2) do |request|
-      track_and_return_ip(request.env["HTTP_FASTLY_CLIENT_IP"])
-    end
+    throttle("site_hits", limit: 40, period: 2, &:track_and_return_ip)
 
     throttle("tag_throttle", limit: 2, period: 1) do |request|
-      if tag_request?(request)
-        track_and_return_ip(request.env["HTTP_FASTLY_CLIENT_IP"])
+      if request.path.include?("/t/")
+        request.track_and_return_ip
       end
     end
-
-    def self.track_and_return_ip(ip_address)
-      return if ip_address.blank?
-
-      Honeycomb.add_field("fastly_client_ip", ip_address)
-      ip_address.to_s
-    end
-
-    def self.tag_request?(request)
-      request.path.include?("/t/")
-    end
   end
 end
diff --git a/spec/initializers/rack/attack_spec.rb b/spec/initializers/rack_attack_spec.rb
similarity index 97%
rename from spec/initializers/rack/attack_spec.rb
rename to spec/initializers/rack_attack_spec.rb
index e34d7fc59..af1c8c6a8 100644
--- a/spec/initializers/rack/attack_spec.rb
+++ b/spec/initializers/rack_attack_spec.rb
@@ -1,11 +1,17 @@
 require "rails_helper"
 
-describe Rack::Attack, type: :request, throttle: true do
+describe Rack, ".attack", type: :request, throttle: true do
   before do
     cache_db = ActiveSupport::Cache.lookup_store(:redis_cache_store)
     allow(Rails).to receive(:cache) { cache_db }
     cache_db.redis.flushdb
     allow(Honeycomb).to receive(:add_field)
+
+    ENV["FASTLY_API_KEY"] = "12345"
+  end
+
+  after do
+    ENV["FASTLY_API_KEY"] = nil
   end
 
   describe "search_throttle" do