diff --git a/CHANGELOG.md b/CHANGELOG.md index b5ccf9f..bb5e7c1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,8 @@ ## master (unreleased) * Implement proxy for Dalli with better Memcachier support. (thanks @hakanensari) + * Rack::Attack.new returns an instance to ease testing (thanks @stevehodgkiss) + * Use Rack::Attack::Request subclass of Rack::Request for easier extending (thanks @tristandunn) ## v3.0.0 - 15 March 2014 * Change default blacklisted response to 403 Forbidden (thanks @carpodaster). diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb index 91baafd..24a74b5 100644 --- a/lib/rack/attack.rb +++ b/lib/rack/attack.rb @@ -1,5 +1,7 @@ require 'rack' -module Rack::Attack +require 'forwardable' + +class Rack::Attack autoload :Cache, 'rack/attack/cache' autoload :Check, 'rack/attack/check' autoload :Throttle, 'rack/attack/throttle' @@ -8,7 +10,8 @@ module Rack::Attack autoload :Track, 'rack/attack/track' autoload :StoreProxy,'rack/attack/store_proxy' autoload :Fail2Ban, 'rack/attack/fail2ban' - autoload :Allow2Ban, 'rack/attack/allow2ban' + autoload :Allow2Ban, 'rack/attack/allow2ban' + autoload :Request, 'rack/attack/request' class << self @@ -35,35 +38,6 @@ module Rack::Attack def throttles; @throttles ||= {}; end def tracks; @tracks ||= {}; end - def new(app) - @app = app - - # Set defaults - @notifier ||= ActiveSupport::Notifications if defined?(ActiveSupport::Notifications) - @blacklisted_response ||= lambda {|env| [403, {}, ["Forbidden\n"]] } - @throttled_response ||= lambda {|env| - retry_after = env['rack.attack.match_data'][:period] rescue nil - [429, {'Retry-After' => retry_after.to_s}, ["Retry later\n"]] - } - - self - end - - def call(env) - req = Rack::Request.new(env) - - if whitelisted?(req) - @app.call(env) - elsif blacklisted?(req) - blacklisted_response[env] - elsif throttled?(req) - throttled_response[env] - else - tracked?(req) - @app.call(env) - end - end - def whitelisted?(req) whitelists.any? do |name, whitelist| whitelist[req] @@ -101,4 +75,37 @@ module Rack::Attack end end + + # Set defaults + @notifier = ActiveSupport::Notifications if defined?(ActiveSupport::Notifications) + @blacklisted_response = lambda {|env| [403, {}, ["Forbidden\n"]] } + @throttled_response = lambda {|env| + retry_after = env['rack.attack.match_data'][:period] rescue nil + [429, {'Retry-After' => retry_after.to_s}, ["Retry later\n"]] + } + + def initialize(app) + @app = app + end + + def call(env) + req = Rack::Attack::Request.new(env) + + if whitelisted?(req) + @app.call(env) + elsif blacklisted?(req) + self.class.blacklisted_response[env] + elsif throttled?(req) + self.class.throttled_response[env] + else + tracked?(req) + @app.call(env) + end + end + + extend Forwardable + def_delegators self, :whitelisted?, + :blacklisted?, + :throttled?, + :tracked? end diff --git a/lib/rack/attack/allow2ban.rb b/lib/rack/attack/allow2ban.rb index f5772c8..bf8c598 100644 --- a/lib/rack/attack/allow2ban.rb +++ b/lib/rack/attack/allow2ban.rb @@ -1,5 +1,5 @@ module Rack - module Attack + class Attack class Allow2Ban < Fail2Ban class << self protected diff --git a/lib/rack/attack/blacklist.rb b/lib/rack/attack/blacklist.rb index 4d2f666..c3a8341 100644 --- a/lib/rack/attack/blacklist.rb +++ b/lib/rack/attack/blacklist.rb @@ -1,5 +1,5 @@ module Rack - module Attack + class Attack class Blacklist < Check def initialize(name, block) super diff --git a/lib/rack/attack/cache.rb b/lib/rack/attack/cache.rb index 9c54505..e65fa84 100644 --- a/lib/rack/attack/cache.rb +++ b/lib/rack/attack/cache.rb @@ -1,5 +1,5 @@ module Rack - module Attack + class Attack class Cache attr_accessor :prefix diff --git a/lib/rack/attack/check.rb b/lib/rack/attack/check.rb index 79f3cd9..caea75a 100644 --- a/lib/rack/attack/check.rb +++ b/lib/rack/attack/check.rb @@ -1,5 +1,5 @@ module Rack - module Attack + class Attack class Check attr_reader :name, :block, :type def initialize(name, block) diff --git a/lib/rack/attack/fail2ban.rb b/lib/rack/attack/fail2ban.rb index be0cba9..443182d 100644 --- a/lib/rack/attack/fail2ban.rb +++ b/lib/rack/attack/fail2ban.rb @@ -1,5 +1,5 @@ module Rack - module Attack + class Attack class Fail2Ban class << self def filter(discriminator, options) diff --git a/lib/rack/attack/request.rb b/lib/rack/attack/request.rb new file mode 100644 index 0000000..603a126 --- /dev/null +++ b/lib/rack/attack/request.rb @@ -0,0 +1,6 @@ +module Rack + class Attack + class Request < ::Rack::Request + end + end +end diff --git a/lib/rack/attack/store_proxy.rb b/lib/rack/attack/store_proxy.rb index 177a3c3..136cbb5 100644 --- a/lib/rack/attack/store_proxy.rb +++ b/lib/rack/attack/store_proxy.rb @@ -1,7 +1,7 @@ require 'delegate' module Rack - module Attack + class Attack class StoreProxy def self.build(store) # RedisStore#increment needs different behavior, so detect that diff --git a/lib/rack/attack/throttle.rb b/lib/rack/attack/throttle.rb index fa82266..7ad6267 100644 --- a/lib/rack/attack/throttle.rb +++ b/lib/rack/attack/throttle.rb @@ -1,5 +1,5 @@ module Rack - module Attack + class Attack class Throttle MANDATORY_OPTIONS = [:limit, :period] attr_reader :name, :limit, :period, :block diff --git a/lib/rack/attack/track.rb b/lib/rack/attack/track.rb index c9d9152..3e516d7 100644 --- a/lib/rack/attack/track.rb +++ b/lib/rack/attack/track.rb @@ -1,5 +1,5 @@ module Rack - module Attack + class Attack class Track < Check def initialize(name, block) super diff --git a/lib/rack/attack/version.rb b/lib/rack/attack/version.rb index 8dc08be..1246847 100644 --- a/lib/rack/attack/version.rb +++ b/lib/rack/attack/version.rb @@ -1,5 +1,5 @@ module Rack - module Attack - VERSION = '3.0.0' + class Attack + VERSION = '3.1.0' end end diff --git a/lib/rack/attack/whitelist.rb b/lib/rack/attack/whitelist.rb index cd2699b..604268e 100644 --- a/lib/rack/attack/whitelist.rb +++ b/lib/rack/attack/whitelist.rb @@ -1,5 +1,5 @@ module Rack - module Attack + class Attack class Whitelist < Check def initialize(name, block) super diff --git a/spec/rack_attack_request_spec.rb b/spec/rack_attack_request_spec.rb new file mode 100644 index 0000000..ca3eadd --- /dev/null +++ b/spec/rack_attack_request_spec.rb @@ -0,0 +1,19 @@ +require_relative 'spec_helper' + +describe 'Rack::Attack' do + describe 'helpers' do + before do + class Rack::Attack::Request + def remote_ip + ip + end + end + + Rack::Attack.whitelist('valid IP') do |req| + req.remote_ip == "127.0.0.1" + end + end + + allow_ok_requests + end +end