From f7bd05e534acb7d05d29e93536dae6140bcf4324 Mon Sep 17 00:00:00 2001
From: Kyle Zhao <kzhao.sw@gmail.com>
Date: Tue, 13 Mar 2018 11:12:41 -0400
Subject: [PATCH] FEATURE: set 'Retry-After' header for 429 responses (#5659)

---
 app/controllers/application_controller.rb | 19 +++++++++++++++----
 config/initializers/004-message_bus.rb    |  2 +-
 spec/integration/rate_limiting_spec.rb    |  3 ++-
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/app/controllers/application_controller.rb b/app/controllers/application_controller.rb
index 566c79be634..8a64fb843a4 100644
--- a/app/controllers/application_controller.rb
+++ b/app/controllers/application_controller.rb
@@ -107,7 +107,15 @@ class ApplicationController < ActionController::Base
   end
 
   def render_rate_limit_error(e)
-    render_json_error e.description, type: :rate_limit, status: 429, extras: { wait_seconds: e&.available_in }
+    retry_time_in_seconds = e&.available_in
+
+    render_json_error(
+      e.description,
+      type: :rate_limit,
+      status: 429,
+      extras: { wait_seconds: retry_time_in_seconds },
+      headers: { 'Retry-After': retry_time_in_seconds },
+    )
   end
 
   # If they hit the rate limiter
@@ -523,12 +531,15 @@ class ApplicationController < ActionController::Base
 
     # Render action for a JSON error.
     #
-    # obj      - a translated string, an ActiveRecord model, or an array of translated strings
+    # obj       - a translated string, an ActiveRecord model, or an array of translated strings
     # opts:
-    #   type   - a machine-readable description of the error
-    #   status - HTTP status code to return
+    #   type    - a machine-readable description of the error
+    #   status  - HTTP status code to return
+    #   headers - extra headers for the response
     def render_json_error(obj, opts = {})
       opts = { status: opts } if opts.is_a?(Integer)
+      opts.fetch(:headers, {}).each { |name, value| headers[name.to_s] = value }
+
       render json: MultiJson.dump(create_errors_json(obj, opts)), status: opts[:status] || 422
     end
 
diff --git a/config/initializers/004-message_bus.rb b/config/initializers/004-message_bus.rb
index d9adc7a8c13..58ba09b6b88 100644
--- a/config/initializers/004-message_bus.rb
+++ b/config/initializers/004-message_bus.rb
@@ -70,7 +70,7 @@ MessageBus.on_middleware_error do |env, e|
   if Discourse::InvalidAccess === e
     [403, {}, ["Invalid Access"]]
   elsif RateLimiter::LimitExceeded === e
-    [429, {}, [e.description]]
+    [429, { 'Retry-After' => e.available_in }, [e.description]]
   end
 end
 
diff --git a/spec/integration/rate_limiting_spec.rb b/spec/integration/rate_limiting_spec.rb
index 93ad7bc53a9..25a389cf9b0 100644
--- a/spec/integration/rate_limiting_spec.rb
+++ b/spec/integration/rate_limiting_spec.rb
@@ -25,7 +25,7 @@ describe 'rate limiter integration' do
     }
   end
 
-  it 'can cleanly limit requests' do
+  it 'can cleanly limit requests and sets a Retry-After header' do
     freeze_time
     #request.set_header("action_dispatch.show_exceptions", true)
 
@@ -50,6 +50,7 @@ describe 'rate limiter integration' do
 
     data = JSON.parse(response.body)
 
+    expect(response.headers['Retry-After']).to eq(60)
     expect(data["extras"]["wait_seconds"]).to eq(60)
   end
 end