Use remote_ip in Rack::Attack (2nd attempt) (#18409)

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

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