From f0634ba876639fcd7e506466683bf71ae81362d4 Mon Sep 17 00:00:00 2001 From: Matt Jankowski Date: Wed, 7 Jun 2017 11:23:26 -0400 Subject: [PATCH] Coverage improvement and concern extraction for rate limit headers in API controller (#3625) * Coverage for rate limit headers * Move rate limit headers methods to concern * Move throttle check to condition on before_action * Move match_data variable into method * Move utc timestamp to separate method * Move header setting into smaller methods * specs cleanup --- app/controllers/api_controller.rb | 15 +---- .../concerns/rate_limit_headers.rb | 57 +++++++++++++++++++ spec/controllers/api_controller_spec.rb | 5 +- .../concerns/rate_limit_headers_spec.rb | 56 ++++++++++++++++++ 4 files changed, 119 insertions(+), 14 deletions(-) create mode 100644 app/controllers/concerns/rate_limit_headers.rb create mode 100644 spec/controllers/concerns/rate_limit_headers_spec.rb 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