├── .gitignore ├── .rubocop.yml ├── .travis.yml ├── CHANGELOG.md ├── Gemfile ├── LICENSE.txt ├── README.md ├── Rakefile ├── lib └── rack │ └── attack │ ├── rate-limit.rb │ └── rate-limit │ └── version.rb ├── rack-attack-rate-limit.gemspec └── spec ├── rack └── attack │ └── rate-limit_spec.rb └── spec_helper.rb /.gitignore: -------------------------------------------------------------------------------- 1 | *.gem 2 | *.rbc 3 | .bundle 4 | .config 5 | .yardoc 6 | 7 | Gemfile.lock 8 | InstalledFiles 9 | _yardoc 10 | coverage 11 | doc/ 12 | lib/bundler/man 13 | pkg 14 | rdoc 15 | spec/reports 16 | test/tmp 17 | test/version_tmp 18 | tmp 19 | -------------------------------------------------------------------------------- /.rubocop.yml: -------------------------------------------------------------------------------- 1 | Encoding: 2 | Enabled: false 3 | Documentation: 4 | Enabled: false 5 | LineLength: 6 | Max: 120 -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | language: ruby 2 | rvm: 3 | - "1.9.3" 4 | - "2.0.0" 5 | script: bundle exec rspec -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # rack-attack-rate-limit changelog 2 | 3 | ## 1.1.0 4 | 5 | * Add support for multiple throttles. 6 | 7 | ## 1.0.0 8 | 9 | * Support for Rails 4 10 | * Rack::Attack should be *class* not a *module*. 11 | * Define Rack::Attack class if not already defined. 12 | * Style changes. 13 | 14 | ## 0.1.0 15 | 16 | * Initial release. -------------------------------------------------------------------------------- /Gemfile: -------------------------------------------------------------------------------- 1 | source 'https://rubygems.org' 2 | 3 | # Specify your gem's dependencies in rack-attack-ratelimit.gemspec 4 | gemspec 5 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | Copyright (c) 2013 Jason Byck 2 | 3 | MIT License 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining 6 | a copy of this software and associated documentation files (the 7 | "Software"), to deal in the Software without restriction, including 8 | without limitation the rights to use, copy, modify, merge, publish, 9 | distribute, sublicense, and/or sell copies of the Software, and to 10 | permit persons to whom the Software is furnished to do so, subject to 11 | the following conditions: 12 | 13 | The above copyright notice and this permission notice shall be 14 | included in all copies or substantial portions of the Software. 15 | 16 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, 17 | EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF 18 | MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND 19 | NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE 20 | LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION 21 | OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION 22 | WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 23 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Rack::Attack::RateLimit 2 | 3 | [![Build Status](https://travis-ci.org/jbyck/rack-attack-rate-limit.png?branch=master)](https://travis-ci.org/jbyck/rack-attack-rate-limit) 4 | 5 | Add rate limit headers for [Rack::Attack](https://github.com/kickstarter/rack-attack) throttles. 6 | 7 | ## Installation 8 | 9 | Install the gem: 10 | 11 | ```shell 12 | gem install 'rack-attack-rate-limit' 13 | ``` 14 | 15 | In your gemfile: 16 | ```ruby 17 | gem 'rack-attack-rate-limit', require: 'rack/attack/rate-limit' 18 | ``` 19 | 20 | 21 | 22 | And then execute: 23 | ```shell 24 | bundle 25 | ``` 26 | 27 | ## Usage 28 | 29 | Rack::Attack::RateLimit expects at least one Rack::Attack throttle to be defined: 30 | 31 | ```ruby 32 | Rack::Attack.throttle('my_throttle') do |req| 33 | req.ip 34 | end 35 | ``` 36 | 37 | To include rate limit headers for throttles, include the Rack::Attack::RateLimit middleware, and provide it with the names of the throttles you want to add rate limit headers for. A single throttle name can be provided as a string, while multiple throttle names must be provided as an array of strings. 38 | 39 | For Rails 3+: 40 | 41 | ```ruby 42 | config.middleware.use Rack::Attack::RateLimit, throttle: ['my_throttle', 'my_other_throttle'] 43 | ``` 44 | 45 | Rate limit headers are: 46 | 47 | * 'X-RateLimit-Limit' - The total number of requests allowed. 48 | * 'X-RateLimit-Remaining' - The number of remaining requests. 49 | 50 | If a request triggers multiple throttles, the gem will add headers for the throttle with the lowest number of remaining requests. 51 | 52 | ## Contributing 53 | 54 | 1. Fork it 55 | 2. Create your feature branch (`git checkout -b my-new-feature`) 56 | 3. Commit your changes (`git commit -am 'Add some feature'`) 57 | 4. Push to the branch (`git push origin my-new-feature`) 58 | 5. Create new Pull Request 59 | -------------------------------------------------------------------------------- /Rakefile: -------------------------------------------------------------------------------- 1 | require "bundler/gem_tasks" 2 | -------------------------------------------------------------------------------- /lib/rack/attack/rate-limit.rb: -------------------------------------------------------------------------------- 1 | unless defined?(Rack::Attack) 2 | module Rack 3 | class Attack 4 | end 5 | end 6 | end 7 | 8 | module Rack 9 | class Attack 10 | class RateLimit 11 | RACK_ATTACK_KEY = 'rack.attack.throttle_data' 12 | 13 | attr_reader :app, :options 14 | 15 | def initialize(app, options = {}) 16 | @app = app 17 | @options = default_options.merge(options) 18 | end 19 | 20 | def call(env) 21 | # If env does not have necessary data to extract rate limit data for the provider, then app.call 22 | return app.call(env) unless rate_limit_available?(env) 23 | # Otherwise, add rate limit headers 24 | status, headers, body = app.call(env) 25 | add_rate_limit_headers!(headers, env) 26 | [status, headers, body] 27 | end 28 | 29 | # Returns env key used by Rack::Attack to namespace data 30 | # 31 | # Returns string 32 | def rack_attack_key 33 | RACK_ATTACK_KEY 34 | end 35 | 36 | # Default options to configure Rack::RateLimit 37 | # 38 | # Returns hash 39 | def default_options 40 | { throttle: 'throttle' } 41 | end 42 | 43 | def throttle 44 | Array(options[:throttle]) || [] 45 | end 46 | 47 | # Return hash of headers with Rate Limiting data 48 | # 49 | # headers - Hash of headers 50 | # 51 | # Returns hash 52 | def add_rate_limit_headers!(headers, env) 53 | throttle_data = throttle_data_closest_to_limit(env) 54 | headers['X-RateLimit-Limit'] = rate_limit_limit(throttle_data).to_s 55 | headers['X-RateLimit-Remaining'] = rate_limit_remaining(throttle_data).to_s 56 | headers 57 | end 58 | 59 | protected 60 | 61 | # RateLimit upper limit from Rack::Attack 62 | # 63 | # env - Hash 64 | # 65 | # Returns Fixnum 66 | def rate_limit_limit(throttle_data) 67 | throttle_data[:limit] 68 | end 69 | 70 | # RateLimit remaining request from Rack::Attack 71 | # 72 | # env - Hash 73 | # 74 | # Returns Fixnum 75 | def rate_limit_remaining(throttle_data) 76 | rate_limit_limit(throttle_data) - throttle_data[:count] 77 | end 78 | 79 | # Rate Limit available method for Rack::Attack provider 80 | # Checks that at least one of the keys provided by the user are in the rack.attack.throttle_data env hash key 81 | # 82 | # env - Hash 83 | # 84 | # Returns boolean 85 | def rate_limit_available?(env) 86 | env.key?(rack_attack_key) && (env[rack_attack_key].keys & throttle).any? 87 | end 88 | 89 | # Throttle Data of Interest 90 | # Filters the rack.attack.throttle_data env hash key for the throttle names provided by the user 91 | # 92 | # env - Hash 93 | # 94 | # Returns Hash 95 | def throttle_data_of_interest(env) 96 | env[rack_attack_key].select { |k, _v| throttle.include?(k) } 97 | end 98 | 99 | # Throttle Data Closest to Limit 100 | # Selects the hash in throttle_data_of_interest where the user is closest to the limit 101 | # 102 | # env - Hash 103 | # 104 | # Returns Hash 105 | def throttle_data_closest_to_limit(env) 106 | min_array = throttle_data_of_interest(env).min_by { |_k, v| v[:limit] - v[:count] } 107 | # The min_by method returns an array of the form [key, value] 108 | # We only need the values 109 | min_array.last 110 | end 111 | end 112 | end 113 | end 114 | -------------------------------------------------------------------------------- /lib/rack/attack/rate-limit/version.rb: -------------------------------------------------------------------------------- 1 | module Rack 2 | class Attack 3 | class RateLimit 4 | VERSION = '1.1.0' 5 | end 6 | end 7 | end 8 | -------------------------------------------------------------------------------- /rack-attack-rate-limit.gemspec: -------------------------------------------------------------------------------- 1 | # coding: utf-8 2 | lib = File.expand_path('../lib', __FILE__) 3 | $LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib) 4 | 5 | require 'rack/attack/rate-limit/version' 6 | 7 | Gem::Specification.new do |spec| 8 | spec.name = "rack-attack-rate-limit" 9 | spec.version = Rack::Attack::RateLimit::VERSION 10 | spec.authors = ["Jason Byck"] 11 | spec.email = ["jasonbyck@gmail.com"] 12 | spec.description = %q{ Add RateLimit headers for Rack::Attack throttling } 13 | spec.summary = %q{ Add RateLimit headers for Rack::Attack throttling } 14 | spec.homepage = "https://github.com/jbyck/rack-attack-rate-limit" 15 | spec.license = "MIT" 16 | 17 | spec.files = `git ls-files`.split($/) 18 | spec.executables = spec.files.grep(%r{^bin/}) { |f| File.basename(f) } 19 | spec.test_files = spec.files.grep(%r{^(test|spec|features)/}) 20 | spec.require_paths = ["lib"] 21 | 22 | spec.add_dependency 'rack' 23 | 24 | spec.add_development_dependency "bundler", "~> 1.3" 25 | spec.add_development_dependency 'rake' 26 | spec.add_development_dependency 'rspec' 27 | spec.add_development_dependency 'rack-test' 28 | spec.add_development_dependency 'rubocop' 29 | 30 | end 31 | -------------------------------------------------------------------------------- /spec/rack/attack/rate-limit_spec.rb: -------------------------------------------------------------------------------- 1 | require 'spec_helper' 2 | 3 | describe Rack::Attack::RateLimit do 4 | 5 | include Rack::Test::Methods 6 | 7 | let(:throttle_one) { 'foo_throttle' } 8 | let(:throttle_two) { 'bar_throttle' } 9 | let(:throttle_three) { 'baz_throttle' } 10 | 11 | let(:app) do 12 | use_throttle = throttle_one 13 | Rack::Builder.new do 14 | use Rack::Attack::RateLimit, throttle: use_throttle 15 | run ->(_env) { [200, {}, 'Hello, World!'] } 16 | end.to_app 17 | end 18 | 19 | context 'Throttle data not present from Rack::Attack' do 20 | 21 | before(:each) do 22 | get '/' 23 | end 24 | 25 | it 'should not create RateLimit headers' do 26 | last_response.header.key?('X-RateLimit-Limit').should be false 27 | last_response.header.key?('X-RateLimit-Remaining').should be false 28 | end 29 | 30 | end 31 | 32 | context 'Throttle data present from Rack::Attack' do 33 | before(:each) do 34 | get '/', {}, "#{Rack::Attack::RateLimit::RACK_ATTACK_KEY}" => rack_attack_throttle_data 35 | end 36 | 37 | let(:request_limit) { (1..10_000).to_a.sample } 38 | let(:request_count) { (1..(request_limit - 10)).to_a.sample } 39 | 40 | context 'one throttle only' do 41 | 42 | let(:rack_attack_throttle_data) do 43 | { "#{throttle_one}" => { count: request_count, limit: request_limit } } 44 | end 45 | 46 | it 'should include RateLimit headers' do 47 | last_response.header.key?('X-RateLimit-Limit').should be true 48 | last_response.header.key?('X-RateLimit-Remaining').should be true 49 | end 50 | 51 | it 'should return correct rate limit in header' do 52 | last_response.header['X-RateLimit-Limit'].to_i.should eq request_limit 53 | end 54 | 55 | it 'should return correct remaining calls in header' do 56 | last_response.header['X-RateLimit-Remaining'].to_i.should eq(request_limit - request_count) 57 | end 58 | end 59 | 60 | context 'multiple throttles' do 61 | 62 | let(:app) do 63 | use_throttle = [throttle_one, throttle_two, throttle_three] 64 | Rack::Builder.new do 65 | use Rack::Attack::RateLimit, throttle: use_throttle 66 | run ->(_env) { [200, {}, 'Hello, World!'] } 67 | end.to_app 68 | end 69 | 70 | let(:request_limits) { 3.times.map { (1..10_000).to_a.sample } } 71 | let(:request_counts) { 3.times.map { |index| (1..(request_limits[index] - 10)).to_a.sample } } 72 | 73 | let(:rack_attack_throttle_data) do 74 | data = {} 75 | [throttle_one, throttle_two, throttle_three].each_with_index do |thr, thr_index| 76 | data["#{thr}"] = { count: request_counts[thr_index], limit: request_limits[thr_index] } 77 | end 78 | data 79 | end 80 | it 'should include RateLimit headers' do 81 | last_response.header.key?('X-RateLimit-Limit').should be true 82 | last_response.header.key?('X-RateLimit-Remaining').should be true 83 | end 84 | 85 | describe 'header values' do 86 | let(:request_differences) do 87 | request_limits.map.each_with_index { |limit, index| limit - request_counts[index] } 88 | end 89 | let(:min_index) { request_differences.each_with_index.min.last } 90 | 91 | it 'should return correct rate limit' do 92 | last_response.header['X-RateLimit-Limit'].to_i.should eq request_limits[min_index] 93 | end 94 | 95 | it 'should return correct remaining calls' do 96 | last_response.header['X-RateLimit-Remaining'].to_i.should eq(request_differences[min_index]) 97 | end 98 | end 99 | end 100 | end 101 | end 102 | -------------------------------------------------------------------------------- /spec/spec_helper.rb: -------------------------------------------------------------------------------- 1 | $LOAD_PATH.unshift File.expand_path(File.dirname(__FILE__) + "../lib") 2 | 3 | require 'rspec' 4 | require 'rack/test' 5 | require 'rack/attack/rate-limit' 6 | 7 | RSpec::configure do |config| 8 | 9 | config.formatter = :documentation 10 | config.tty = true 11 | config.color = true 12 | end 13 | --------------------------------------------------------------------------------