From 1ebe1c3517363cf88f50b22764be3c011fc0c197 Mon Sep 17 00:00:00 2001 From: Paul Coates Date: Mon, 19 May 2014 11:11:01 -0700 Subject: [PATCH 1/4] Added limit and period options to track. Delegates [] to Throttle if they are present otherwise Check. --- lib/rack/attack.rb | 4 ++-- lib/rack/attack/check.rb | 4 ++-- lib/rack/attack/throttle.rb | 5 +++-- lib/rack/attack/track.rb | 19 +++++++++++++++---- spec/rack_attack_track_spec.rb | 14 ++++++++++++++ 5 files changed, 36 insertions(+), 10 deletions(-) diff --git a/lib/rack/attack.rb b/lib/rack/attack.rb index 0ec982c..eb9043e 100644 --- a/lib/rack/attack.rb +++ b/lib/rack/attack.rb @@ -31,8 +31,8 @@ class Rack::Attack self.throttles[name] = Throttle.new(name, options, block) end - def track(name, &block) - self.tracks[name] = Track.new(name, block) + def track(name, options = {}, &block) + self.tracks[name] = Track.new(name, options, block) end def whitelists; @whitelists ||= {}; end diff --git a/lib/rack/attack/check.rb b/lib/rack/attack/check.rb index caea75a..21451de 100644 --- a/lib/rack/attack/check.rb +++ b/lib/rack/attack/check.rb @@ -2,9 +2,9 @@ module Rack class Attack class Check attr_reader :name, :block, :type - def initialize(name, block) + def initialize(name, options = {}, block) @name, @block = name, block - @type = nil + @type = options.fetch(:type, nil) end def [](req) diff --git a/lib/rack/attack/throttle.rb b/lib/rack/attack/throttle.rb index 2824014..6bf91c9 100644 --- a/lib/rack/attack/throttle.rb +++ b/lib/rack/attack/throttle.rb @@ -2,7 +2,7 @@ module Rack class Attack class Throttle MANDATORY_OPTIONS = [:limit, :period] - attr_reader :name, :limit, :period, :block + attr_reader :name, :limit, :period, :block, :type def initialize(name, options, block) @name, @block = name, block MANDATORY_OPTIONS.each do |opt| @@ -10,6 +10,7 @@ module Rack end @limit = options[:limit] @period = options[:period].to_i + @type = options.fetch(:type, :throttle) end def cache @@ -34,7 +35,7 @@ module Rack if throttled req.env['rack.attack.matched'] = name req.env['rack.attack.match_discriminator'] = discriminator - req.env['rack.attack.match_type'] = :throttle + req.env['rack.attack.match_type'] = type req.env['rack.attack.match_data'] = data Rack::Attack.instrument(req) end diff --git a/lib/rack/attack/track.rb b/lib/rack/attack/track.rb index 3e516d7..4bfc145 100644 --- a/lib/rack/attack/track.rb +++ b/lib/rack/attack/track.rb @@ -1,10 +1,21 @@ module Rack class Attack - class Track < Check - def initialize(name, block) - super - @type = :track + class Track + extend Forwardable + + attr_reader :checker + + def initialize(name, options = {}, block) + options[:type] = :track + + if options[:limit] && options[:period] + @checker = Throttle.new(name, options, block) + else + @checker = Check.new(name, options, block) + end end + + def_delegator :@checker, :[], :[] end end end diff --git a/spec/rack_attack_track_spec.rb b/spec/rack_attack_track_spec.rb index 3f0e56e..8603266 100644 --- a/spec/rack_attack_track_spec.rb +++ b/spec/rack_attack_track_spec.rb @@ -41,4 +41,18 @@ describe 'Rack::Attack.track' do Counter.check.must_equal 2 end end + + describe "without limit and period options" do + it "should delegate [] to check" do + tracker = Rack::Attack.track("homepage") { |req| req.path == "/"} + tracker.checker.class.must_equal Rack::Attack::Check + end + end + + describe "with limit and period options" do + it "should delegate [] method to throttle" do + tracker = Rack::Attack.track("homepage", :limit => 10, :period => 10) { |req| req.path == "/"} + tracker.checker.class.must_equal Rack::Attack::Throttle + end + end end From c51676d29501a6df828a1b5c6d5f065734a37f89 Mon Sep 17 00:00:00 2001 From: Paul Coates Date: Mon, 19 May 2014 11:14:50 -0700 Subject: [PATCH 2/4] Updated README.md to reflect optional limit and period parmas for tracks. --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 25bc84f..278c224 100644 --- a/README.md +++ b/README.md @@ -188,7 +188,7 @@ end ### Tracks ```ruby -# Track requests from a special user agent +# Track requests from a special user agent, supports limit and period options like throttle. Rack::Attack.track("special_agent") do |req| req.user_agent == "SpecialAgent" end From 1857f8dd579367cf9c09d74bde0649a2f0239884 Mon Sep 17 00:00:00 2001 From: Paul Coates Date: Mon, 19 May 2014 12:30:18 -0700 Subject: [PATCH 3/4] Added track with limit/period options example to the README.md for clarity. --- README.md | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 278c224..2431616 100644 --- a/README.md +++ b/README.md @@ -188,11 +188,16 @@ end ### Tracks ```ruby -# Track requests from a special user agent, supports limit and period options like throttle. +# Track requests from a special user agent. Rack::Attack.track("special_agent") do |req| req.user_agent == "SpecialAgent" end +# Supports optional limit and period, triggers the notification only when the limit is reached. +Rack::Attack.track("special_agent", :limit 6, :period => 60.seconds) do |req| + req.user_agent == "SpecialAgent" +end + # Track it using ActiveSupport::Notification ActiveSupport::Notifications.subscribe("rack.attack") do |name, start, finish, request_id, req| if req.env['rack.attack.matched'] == "special_agent" && req.env['rack.attack.match_type'] == :track From e8d98a7ad3b5981660600c7ffffbb6e66b0bb15d Mon Sep 17 00:00:00 2001 From: Paul Coates Date: Thu, 22 May 2014 10:11:23 -0700 Subject: [PATCH 4/4] Changed track checker to track filter. Made track filter tests more clear. --- lib/rack/attack/track.rb | 8 ++++---- spec/rack_attack_track_spec.rb | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/lib/rack/attack/track.rb b/lib/rack/attack/track.rb index 4bfc145..e0039e3 100644 --- a/lib/rack/attack/track.rb +++ b/lib/rack/attack/track.rb @@ -3,19 +3,19 @@ module Rack class Track extend Forwardable - attr_reader :checker + attr_reader :filter def initialize(name, options = {}, block) options[:type] = :track if options[:limit] && options[:period] - @checker = Throttle.new(name, options, block) + @filter = Throttle.new(name, options, block) else - @checker = Check.new(name, options, block) + @filter = Check.new(name, options, block) end end - def_delegator :@checker, :[], :[] + def_delegator :@filter, :[] end end end diff --git a/spec/rack_attack_track_spec.rb b/spec/rack_attack_track_spec.rb index 8603266..9f9ca9c 100644 --- a/spec/rack_attack_track_spec.rb +++ b/spec/rack_attack_track_spec.rb @@ -43,16 +43,16 @@ describe 'Rack::Attack.track' do end describe "without limit and period options" do - it "should delegate [] to check" do + it "should assign the track filter to a Check instance" do tracker = Rack::Attack.track("homepage") { |req| req.path == "/"} - tracker.checker.class.must_equal Rack::Attack::Check + tracker.filter.class.must_equal Rack::Attack::Check end end describe "with limit and period options" do - it "should delegate [] method to throttle" do + it "should assign the track filter to a Throttle instance" do tracker = Rack::Attack.track("homepage", :limit => 10, :period => 10) { |req| req.path == "/"} - tracker.checker.class.must_equal Rack::Attack::Throttle + tracker.filter.class.must_equal Rack::Attack::Throttle end end end