4 files changed, 119 insertions, 14 deletions
diff --git a/app/controllers/api_controller.rb b/app/controllers/api_controller.rb
index 1e72549bd..42b85865e 100644
--- a/app/controllers/api_controller.rb
+++ b/app/controllers/api_controller.rb
@@ -4,11 +4,11 @@ class ApiController < ApplicationController
   DEFAULT_STATUSES_LIMIT = 20
   DEFAULT_ACCOUNTS_LIMIT = 40
 
+  include RateLimitHeaders
+
   skip_before_action :verify_authenticity_token
   skip_before_action :store_current_location
 
-  before_action :set_rate_limit_headers
-
   rescue_from ActiveRecord::RecordInvalid, Mastodon::ValidationError do |e|
     render json: { error: e.to_s }, status: 422
   end
@@ -43,17 +43,6 @@ class ApiController < ApplicationController
 
   protected
 
-  def set_rate_limit_headers
-    return if request.env['rack.attack.throttle_data'].nil?
-
-    now        = Time.now.utc
-    match_data = request.env['rack.attack.throttle_data']['api']
-
-    response.headers['X-RateLimit-Limit']     = match_data[:limit].to_s
-    response.headers['X-RateLimit-Remaining'] = (match_data[:limit] - match_data[:count]).to_s
-    response.headers['X-RateLimit-Reset']     = (now + (match_data[:period] - now.to_i % match_data[:period])).iso8601(6)
-  end
-
   def set_pagination_headers(next_path = nil, prev_path = nil)
     links = []
     links << [next_path, [%w(rel next)]] if next_path
diff --git a/app/controllers/concerns/rate_limit_headers.rb b/app/controllers/concerns/rate_limit_headers.rb
new file mode 100644
index 000000000..36cb91075
--- /dev/null
+++ b/app/controllers/concerns/rate_limit_headers.rb
@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module RateLimitHeaders
+  extend ActiveSupport::Concern
+
+  included do
+    before_action :set_rate_limit_headers, if: :rate_limited_request?
+  end
+
+  private
+
+  def set_rate_limit_headers
+    apply_header_limit
+    apply_header_remaining
+    apply_header_reset
+  end
+
+  def rate_limited_request?
+    !request.env['rack.attack.throttle_data'].nil?
+  end
+
+  def apply_header_limit
+    response.headers['X-RateLimit-Limit'] = rate_limit_limit
+  end
+
+  def rate_limit_limit
+    api_throttle_data[:limit].to_s
+  end
+
+  def apply_header_remaining
+    response.headers['X-RateLimit-Remaining'] = rate_limit_remaining
+  end
+
+  def rate_limit_remaining
+    (api_throttle_data[:limit] - api_throttle_data[:count]).to_s
+  end
+
+  def apply_header_reset
+    response.headers['X-RateLimit-Reset'] = rate_limit_reset
+  end
+
+  def rate_limit_reset
+    (request_time + reset_period_offset).iso8601(6)
+  end
+
+  def api_throttle_data
+    request.env['rack.attack.throttle_data']['api']
+  end
+
+  def request_time
+    @_request_time ||= Time.now.utc
+  end
+
+  def reset_period_offset
+    api_throttle_data[:period] - request_time.to_i % api_throttle_data[:period]
+  end
+end
diff --git a/spec/controllers/api_controller_spec.rb b/spec/controllers/api_controller_spec.rb
index 1026afbbc..44be4276a 100644
--- a/spec/controllers/api_controller_spec.rb
+++ b/spec/controllers/api_controller_spec.rb
@@ -9,9 +9,12 @@ describe ApiController, type: :controller do
     end
   end
 
+  before do
+    routes.draw { post 'success' => 'api#success' }
+  end
+
   it 'does not protect from forgery' do
     ActionController::Base.allow_forgery_protection = true
-    routes.draw { post 'success' => 'api#success' }
     post 'success'
     expect(response).to have_http_status(:success)
   end
diff --git a/spec/controllers/concerns/rate_limit_headers_spec.rb b/spec/controllers/concerns/rate_limit_headers_spec.rb
new file mode 100644
index 000000000..719978dc2
--- /dev/null
+++ b/spec/controllers/concerns/rate_limit_headers_spec.rb
@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+require 'rails_helper'
+
+describe ApplicationController do
+  controller do
+    include RateLimitHeaders
+
+    def show
+      head 200
+    end
+  end
+
+  before do
+    routes.draw { get 'show' => 'anonymous#show' }
+  end
+
+  describe 'rate limiting' do
+    context 'throttling is off' do
+      before do
+        request.env['rack.attack.throttle_data'] = nil
+      end
+
+      it 'does not apply rate limiting' do
+        get 'show'
+
+        expect(response.headers['X-RateLimit-Limit']).to be_nil
+        expect(response.headers['X-RateLimit-Remaining']).to be_nil
+        expect(response.headers['X-RateLimit-Reset']).to be_nil
+      end
+    end
+
+    context 'throttling is on' do
+      let(:start_time) { DateTime.new(2017, 1, 1, 12, 0, 0).utc }
+
+      before do
+        request.env['rack.attack.throttle_data'] = { 'api' => { limit: 100, count: 20, period: 10 } }
+        travel_to start_time do
+          get 'show'
+        end
+      end
+
+      it 'applies rate limiting limit header' do
+        expect(response.headers['X-RateLimit-Limit']).to eq '100'
+      end
+
+      it 'applies rate limiting remaining header' do
+        expect(response.headers['X-RateLimit-Remaining']).to eq '80'
+      end
+
+      it 'applies rate limiting reset header' do
+        expect(response.headers['X-RateLimit-Reset']).to eq (start_time + 10.seconds).iso8601(6)
+      end
+    end
+  end
+end