FIX: Add Retry-Header to rate limited responses (#11736)

It returned a 429 error code with a 'Retry-After' header if a
RateLimiter::LimitExceeded was raised and unhandled, but the header was
missing if the request was limited in the 'RequestTracker' middleware.

lib/middleware/request_tracker.rb

@@ -20,8 +20,7 @@ class Middleware::RequestTracker
   end
 
   def self.unregister_detailed_request_logger(callback)
-    @@detailed_request_loggers.delete callback
-
+    @@detailed_request_loggers.delete(callback)
     if @@detailed_request_loggers.length == 0
       @detailed_request_loggers = nil
     end
@@ -78,10 +77,9 @@ class Middleware::RequestTracker
       ApplicationRequest.increment!(:http_4xx)
     elsif status >= 300
       ApplicationRequest.increment!(:http_3xx)
-    elsif status >= 200 && status < 300
+    elsif status >= 200
       ApplicationRequest.increment!(:http_2xx)
     end
-
   end
 
   def self.get_data(env, result, timing)
@@ -120,11 +118,11 @@ class Middleware::RequestTracker
     if cache = headers["X-Discourse-Cached"]
       h[:cache] = cache
     end
+
     h
   end
 
   def log_request_info(env, result, info)
-
     # we got to skip this on error ... its just logging
     data = self.class.get_data(env, result, info) rescue nil
 
@@ -139,7 +137,6 @@ class Middleware::RequestTracker
 
       log_later(data)
     end
-
   end
 
   def self.populate_request_queue_seconds!(env)
@@ -166,15 +163,20 @@ class Middleware::RequestTracker
 
     request = Rack::Request.new(env)
 
-    if rate_limit(request)
-      result = [429, {}, ["Slow down, too many requests from this IP address"]]
-      return result
+    if available_in = rate_limit(request)
+      return [
+        429,
+        { "Retry-After" => available_in },
+        ["Slow down, too many requests from this IP address"]
+      ]
     end
 
     env["discourse.request_tracker"] = self
+
     MethodProfiler.start
     result = @app.call(env)
     info = MethodProfiler.stop
+
     # possibly transferred?
     if info && (headers = result[1])
       headers["X-Runtime"] = "%0.6f" % info[:total_duration]
@@ -221,7 +223,6 @@ class Middleware::RequestTracker
   end
 
   def rate_limit(request)
-
     if (
       GlobalSetting.max_reqs_per_ip_mode == "block" ||
       GlobalSetting.max_reqs_per_ip_mode == "warn" ||
@@ -265,29 +266,38 @@ class Middleware::RequestTracker
       request.env['DISCOURSE_RATE_LIMITERS'] = [limiter10, limiter60]
       request.env['DISCOURSE_ASSET_RATE_LIMITERS'] = [limiter_assets10]
 
-      warn = GlobalSetting.max_reqs_per_ip_mode == "warn" ||
-        GlobalSetting.max_reqs_per_ip_mode == "warn+block"
+      warn = GlobalSetting.max_reqs_per_ip_mode == "warn" || GlobalSetting.max_reqs_per_ip_mode == "warn+block"
 
       if !limiter_assets10.can_perform?
         if warn
           Discourse.warn("Global asset IP rate limit exceeded for #{ip}: 10 second rate limit", uri: request.env["REQUEST_URI"])
         end
 
-        return !(GlobalSetting.max_reqs_per_ip_mode == "warn")
+        if GlobalSetting.max_reqs_per_ip_mode != "warn"
+          return limiter_assets10.seconds_to_wait(Time.now.to_i)
+        else
+          return false
+        end
       end
 
-      type = 10
       begin
+        type = 10
         limiter10.performed!
+
         type = 60
         limiter60.performed!
+
         false
-      rescue RateLimiter::LimitExceeded
+      rescue RateLimiter::LimitExceeded => e
         if warn
           Discourse.warn("Global IP rate limit exceeded for #{ip}: #{type} second rate limit", uri: request.env["REQUEST_URI"])
-          !(GlobalSetting.max_reqs_per_ip_mode == "warn")
+          if GlobalSetting.max_reqs_per_ip_mode != "warn"
+            e.available_in
+          else
+            false
+          end
         else
-          true
+          e.available_in
         end
       end
     end

lib/rate_limiter.rb

@@ -55,6 +55,10 @@ class RateLimiter
     rate_unlimited? || is_under_limit?
   end
 
+  def seconds_to_wait(now)
+    @secs - age_of_oldest(now)
+  end
+
   # reloader friendly
   unless defined? PERFORM_LUA
     PERFORM_LUA = <<~LUA
@@ -173,10 +177,6 @@ class RateLimiter
     Discourse.redis.without_namespace
   end
 
-  def seconds_to_wait(now)
-    @secs - age_of_oldest(now)
-  end
-
   def age_of_oldest(now)
     # age of oldest event in buffer, in seconds
     now - redis.lrange(prefixed_key, -1, -1).first.to_i

spec/components/middleware/request_tracker_spec.rb

@@ -237,10 +237,11 @@ describe Middleware::RequestTracker do
       global_setting :max_reqs_per_ip_mode, 'warn+block'
 
       status, _ = middleware.call(env)
-      status, _ = middleware.call(env)
+      status, headers = middleware.call(env)
 
       expect(Rails.logger.warnings).to eq(1)
       expect(status).to eq(429)
+      expect(headers["Retry-After"]).to eq(10)
     end
 
     it "does warn if rate limiter is enabled" do
@@ -267,13 +268,15 @@ describe Middleware::RequestTracker do
       expect(status).to eq(200)
       status, _ = middleware.call(env1)
       expect(status).to eq(200)
-      status, _ = middleware.call(env1)
+      status, headers = middleware.call(env1)
       expect(status).to eq(429)
+      expect(headers["Retry-After"]).to eq(10)
 
       env2 = env("REMOTE_ADDR" => "1.1.1.1")
 
-      status, _ = middleware.call(env2)
+      status, headers = middleware.call(env2)
       expect(status).to eq(429)
+      expect(headers["Retry-After"]).to eq(10)
     end
 
     it "does block if rate limiter is enabled" do
@@ -286,8 +289,9 @@ describe Middleware::RequestTracker do
       status, _ = middleware.call(env1)
       expect(status).to eq(200)
 
-      status, _ = middleware.call(env1)
+      status, headers = middleware.call(env1)
       expect(status).to eq(429)
+      expect(headers["Retry-After"]).to eq(10)
 
       status, _ = middleware.call(env2)
       expect(status).to eq(200)