1 files changed, 15 insertions, 1 deletions
diff --git a/app/controllers/concerns/rate_limit_headers.rb b/app/controllers/concerns/rate_limit_headers.rb
index b79c558d8..86fe58a71 100644
--- a/app/controllers/concerns/rate_limit_headers.rb
+++ b/app/controllers/concerns/rate_limit_headers.rb
@@ -3,6 +3,20 @@
 module RateLimitHeaders
   extend ActiveSupport::Concern
 
+  class_methods do
+    def override_rate_limit_headers(method_name, options = {})
+      around_action(only: method_name, if: :current_account) do |_controller, block|
+        begin
+          block.call
+        ensure
+          rate_limiter = RateLimiter.new(current_account, options)
+          rate_limit_headers = rate_limiter.to_headers
+          response.headers.merge!(rate_limit_headers) unless response.headers['X-RateLimit-Remaining'].present? && rate_limit_headers['X-RateLimit-Remaining'].to_i > response.headers['X-RateLimit-Remaining'].to_i
+        end
+      end
+    end
+  end
+
   included do
     before_action :set_rate_limit_headers, if: :rate_limited_request?
   end
@@ -44,7 +58,7 @@ module RateLimitHeaders
   end
 
   def api_throttle_data
-    most_limited_type, = request.env['rack.attack.throttle_data'].min_by { |_, v| v[:limit] }
+    most_limited_type, = request.env['rack.attack.throttle_data'].min_by { |_, v| v[:limit] - v[:count] }
     request.env['rack.attack.throttle_data'][most_limited_type]
   end