Переглянути джерело

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
master
Matt Jankowski 7 роки тому
committed by Eugen Rochko
джерело
коміт
f0634ba876
4 змінених файлів з 119 додано та 14 видалено
  1. +2
    -13
      app/controllers/api_controller.rb
  2. +57
    -0
      app/controllers/concerns/rate_limit_headers.rb
  3. +4
    -1
      spec/controllers/api_controller_spec.rb
  4. +56
    -0
      spec/controllers/concerns/rate_limit_headers_spec.rb

+ 2
- 13
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


+ 57
- 0
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

+ 4
- 1
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


+ 56
- 0
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

Завантаження…
Відмінити
Зберегти