improved-rack-throttle-w-expiry 0.8.0

data/lib/rack/throttle/limiters/sliding_window.rb

@@ -0,0 +1,86 @@
+module Rack; module Throttle
+  ##
+  # This rate limiter strategy throttles the application with
+  # a sliding window (implemented as a leaky bucket). It operates
+  # on second-level resolution. It takes :burst and :average 
+  # options, which correspond to the maximum size of a traffic 
+  # burst, and the maximum allowed average traffic level.
+  class SlidingWindow < Limiter
+    ##
+    # @param  [#call]                    app
+    # @param  [Hash{Symbol => Object}]   options
+    # @option options [Integer] :burst   5
+    # @option options [Float] :average   1
+    def initialize(app, options = {})
+      super
+      options[:burst] ||= 5
+      options[:average] ||= 1
+    end
+
+    ##
+    # Returns `true` if the request conforms to the 
+    # specified :average and :burst rules
+    #
+    # @param  [Rack::Request] request
+    # @return [Boolean]
+    def allowed?(request)
+      t1 = request_start_time(request)
+      key = cache_key(request)
+      bucket = cache_get(key) rescue nil
+      bucket ||= LeakyBucket.new(options[:burst], options[:average])
+      bucket.maximum, bucket.outflow = options[:burst], options[:average]
+      bucket.leak! 
+      bucket.increment!
+      allowed = !bucket.full?
+      begin
+        cache_set(key, bucket)
+        allowed
+      rescue StandardError => e
+        allowed = true
+        # If an error occurred while trying to update the timestamp stored
+        # in the cache, we will fall back to allowing the request through.
+        # This prevents the Rack application blowing up merely due to a
+        # backend cache server (Memcached, Redis, etc.) being offline.
+      end
+    end
+
+    ###
+    # LeakyBucket is an internal class used to implement the
+    # SlidingWindow limiter strategy. It is a (slightly tweaked)
+    # implementation of the {http://en.wikipedia.org/wiki/Leaky_bucket
+    # Leaky Bucket Algorithm}.
+    class LeakyBucket
+      attr_accessor :maximum, :outflow
+      attr_reader :count, :last_touched
+
+      ##
+      # @param [Integer] maximum
+      # @param [Float] outflow
+      def initialize(maximum, outflow)
+        @maximum, @outflow = maximum, outflow
+        @count, @last_touched = 0, Time.now
+      end
+
+      def leak!
+        t = Time.now
+        time = t - last_touched
+        loss = (outflow * time).to_f
+        if loss > 0
+          @count -= loss
+          @last_touched = t 
+        end
+      end
+
+      def increment!
+        @count = 0 if count < 0
+        @count += 1
+        @count = maximum if count > maximum
+      end
+
+      def full?
+        count == maximum
+      end
+    end
+
+  end
+end; end

data/lib/rack/throttle/limiters/time_window.rb

@@ -0,0 +1,21 @@
+module Rack; module Throttle
+  ##
+  class TimeWindow < Limiter
+    ##
+    # Returns `true` if fewer than the maximum number of requests permitted
+    # for the current window of time have been made.
+    #
+    # @param  [Rack::Request] request
+    # @return [Boolean]
+    def allowed?(request)
+      count = cache_get(key = cache_key(request)).to_i + 1 rescue 1
+      allowed = count <= max_per_window.to_i
+      begin
+        cache_set(key, count)
+        allowed
+      rescue => e
+        allowed = true
+      end
+    end
+  end
+end; end

data/lib/rack/throttle/matchers/matcher.rb

@@ -0,0 +1,32 @@
+module Rack; module Throttle
+  ###
+  # This is the base class for matcher implementations.
+  # Implementations are provided for User Agent, URL, and request
+  # method. Subclass Matcher if you want to provide a custom
+  # implementation.
+  class Matcher
+    attr_reader :rule
+
+    def initialize(rule)
+      @rule = rule
+    end
+
+    # Must be implemented in a subclass.
+    # MUST return true or false.
+    # @return [Boolean]
+    # @abstract
+    def match?
+      true
+    end
+
+    # Must be implemented in a subclass.
+    # Used to produce a unique key in our cache store.
+    # Typically of the form "xx-"
+    # @return [String]
+    # @abstract
+    def identifier
+      ""
+    end
+  end
+
+end; end

data/lib/rack/throttle/matchers/method_matcher.rb

@@ -0,0 +1,24 @@
+module Rack; module Throttle
+  ###
+  # MethodMatchers are used to restrict throttling based on the HTTP
+  # method used by the request. For instance, you may only care about
+  # throttling POST requests on a login form; GET requests are just
+  # fine.
+  # MethodMatchers take Symbol objects of :get, :put, :post, or :delete
+  class MethodMatcher < Matcher
+    ##
+    # @param [Rack::Request] request
+    # @return [Boolean]
+    def match?(request)
+      rack_method = :"#{@rule}?"
+      request.send(rack_method)
+    end
+
+    ##
+    # @return [String]
+    def identifier
+      "meth-#{@rule}"
+    end
+  end
+
+end; end

data/lib/rack/throttle/matchers/url_matcher.rb

@@ -0,0 +1,24 @@
+module Rack; module Throttle
+  ###
+  # UrlMatchers are used to restrict requests based on the URL
+  # requested. For instance, you may care about limiting requests
+  # to a machine-consumed API, but not be concerned about requests
+  # coming from browsers.
+  # UrlMatchers take Regexp object to matcha gainst the request path.
+  class UrlMatcher < Matcher
+    ###
+    # @param [Rack::Request] request
+    # @return [Boolean]
+    def match?(request)
+      !!(@rule =~ request.path)
+    end
+
+    ###
+    # @return [String]
+    def identifier
+      "url-" + @rule.inspect
+    end
+  end
+
+end; end
+

data/lib/rack/throttle/matchers/user_agent_matcher.rb

@@ -0,0 +1,23 @@
+module Rack; module Throttle
+  ###
+  # User Agent Matchers are used to restrict requests based on the User
+  # Agent supplied by the requester. For instance, you may care about
+  # limiting a specific API consumer who reliably uses a known User-Agent.
+  # UserAgentMatchers take Regexp objects to match against the
+  # User-Agent.
+  class UserAgentMatcher < Matcher
+    ###
+    # @param [Rack::Request] request
+    # @return [Boolean]
+    def match?(request)
+      !!(@rule =~ request.user_agent)
+    end
+
+    ###
+    # @return [String]
+    def identifier
+      "ua-" + @rule.inspect
+    end
+  end
+
+end; end

data/lib/rack/throttle/version.rb

@@ -0,0 +1,23 @@
+module Rack; module Throttle
+  module VERSION
+    MAJOR = 0
+    MINOR = 8
+    TINY  = 0
+    EXTRA = nil
+
+    STRING = [MAJOR, MINOR, TINY].join('.')
+    STRING << "-#{EXTRA}" if EXTRA
+
+    ##
+    # @return [String]
+    def self.to_s() STRING end
+
+    ##
+    # @return [String]
+    def self.to_str() STRING end
+
+    ##
+    # @return [Array(Integer, Integer, Integer)]
+    def self.to_a() [MAJOR, MINOR, TINY] end
+  end
+end; end

data/spec/limiters/daily_spec.rb

@@ -0,0 +1,31 @@
+require File.join(File.dirname(__FILE__), '..', 'spec_helper')
+
+describe Rack::Throttle::Daily do
+  include Rack::Test::Methods
+
+  def app
+    @target_app ||= example_target_app
+    @app ||= Rack::Throttle::Daily.new(@target_app, :max_per_day => 3)
+  end
+
+  it "should be allowed if not seen this day" do
+    get "/foo"
+    last_response.body.should show_allowed_response
+  end
+  
+  it "should be allowed if seen fewer than the max allowed per day" do
+    2.times { get "/foo" }
+    last_response.body.should show_allowed_response
+  end
+  
+  it "should not be allowed if seen more times than the max allowed per day" do
+    4.times { get "/foo" }
+    last_response.body.should show_throttled_response
+  end
+
+  it "should have an associated default_ttl" do
+    app.class.default_ttl.should be_an_instance_of Fixnum
+  end
+
+  # TODO mess with time travelling and requests to make sure no overlap
+end

data/spec/limiters/hourly_spec.rb

@@ -0,0 +1,32 @@
+require File.join(File.dirname(__FILE__), '..', 'spec_helper')
+
+describe Rack::Throttle::Hourly do
+  include Rack::Test::Methods
+
+  def app
+    @target_app ||= example_target_app
+    @app ||= Rack::Throttle::Hourly.new(@target_app, :max_per_hour => 3)
+  end
+
+  it "should be allowed if not seen this hour" do
+    get "/foo"
+    last_response.body.should show_allowed_response
+  end
+  
+  it "should be allowed if seen fewer than the max allowed per hour" do
+    2.times { get "/foo" }
+    last_response.body.should show_allowed_response
+  end
+  
+  it "should not be allowed if seen more times than the max allowed per hour" do
+    4.times { get "/foo" }
+    last_response.body.should show_throttled_response
+  end
+
+  it "should have an associated default_ttl" do
+    app.class.default_ttl.should be_an_instance_of Fixnum
+  end
+
+
+  # TODO mess with time travelling and requests to make sure no overlap
+end

data/spec/limiters/interval_spec.rb

@@ -0,0 +1,45 @@
+require File.join(File.dirname(__FILE__), '..', 'spec_helper')
+
+
+describe Rack::Throttle::Interval do
+  include Rack::Test::Methods
+  
+  def app
+    @target_app ||= example_target_app
+    @app ||= Rack::Throttle::Interval.new(@target_app, :min => 0.1)
+  end
+
+  it "should allow the request if the source has not been seen" do
+    get "/foo"
+    last_response.body.should show_allowed_response
+  end
+  
+  it "should allow the request if the source has not been seen in the current interval" do
+    Timecop.freeze do
+      get "/foo"
+      Timecop.freeze(1) do # Timecop.freeze won't do subsecond resolution
+        get "/foo"
+      end
+    end
+    last_response.body.should show_allowed_response
+  end
+  
+  it "should not allow the request if the source has been seen inside the current interval" do
+    Timecop.freeze do
+      2.times { get "/foo" }
+    end
+    last_response.body.should show_throttled_response
+  end
+  
+  it "should gracefully allow the request if the cache bombs on getting" do
+    app.should_receive(:cache_get).and_raise(StandardError)
+    get "/foo"
+    last_response.body.should show_allowed_response
+  end
+  
+  it "should gracefully allow the request if the cache bombs on setting" do
+    app.should_receive(:cache_set).and_raise(StandardError)
+    get "/foo"
+    last_response.body.should show_allowed_response
+  end
+end

data/spec/limiters/limiter_spec.rb

@@ -0,0 +1,51 @@
+require File.join(File.dirname(__FILE__), '..', 'spec_helper')
+
+describe Rack::Throttle::Limiter do
+  include Rack::Test::Methods
+
+  def app
+    @target_app ||= example_target_app
+    @app ||= Rack::Throttle::Limiter.new(@target_app)
+  end
+  
+  describe "basic calling" do
+    it "should return the example app" do
+      get "/foo"
+      last_response.body.should show_allowed_response
+    end
+  
+    it "should call the application if allowed" do
+      app.should_receive(:allowed?).and_return(true)
+      get "/foo"
+      last_response.body.should show_allowed_response
+    end
+  
+    it "should give a rate limit exceeded message if not allowed" do
+      app.should_receive(:allowed?).and_return(false)
+      get "/foo"
+      last_response.body.should show_throttled_response
+    end
+  end
+  
+  describe "allowed?" do
+    it "should return true if whitelisted" do
+      app.should_receive(:whitelisted?).and_return(true)
+      get "/foo"
+      last_response.body.should show_allowed_response
+    end
+    
+    it "should return false if blacklisted" do
+      app.should_receive(:blacklisted?).and_return(true)
+      get "/foo"
+      last_response.body.should show_throttled_response
+    end
+    
+    it "should return true if not whitelisted or blacklisted" do
+      app.should_receive(:whitelisted?).and_return(false)
+      app.should_receive(:blacklisted?).and_return(false)
+      get "/foo"
+      last_response.body.should show_allowed_response
+    end
+  end
+
+end

data/spec/limiters/sliding_window_spec.rb

@@ -0,0 +1,67 @@
+require File.join(File.dirname(__FILE__), '..', 'spec_helper')
+
+describe Rack::Throttle::SlidingWindow do
+  include Rack::Test::Methods
+
+  def app
+    @target_app ||= example_target_app
+    @app ||= Rack::Throttle::SlidingWindow.new(@target_app, :burst => 2, :average => 1)
+  end
+
+  before(:each) do
+    Timecop.freeze
+    @time = Time.now
+  end
+
+  after(:each) do
+    Timecop.return
+  end
+
+  it "should allow the request if the source has not been seen at all" do
+    get "/foo"
+    last_response.body.should show_allowed_response
+  end
+  
+  it "should allow the request if the rate is above-average but within the burst rule" do
+    Timecop.freeze(@time) { get "/foo" }
+    Timecop.freeze(@time + 0.5) { get "/foo" }
+    last_response.body.should show_allowed_response
+  end
+
+   it "should not allow the request if the rate is greater than the burst rule" do
+    Timecop.freeze(@time) { get "/foo" }
+    Timecop.freeze(@time + 0.3) { get "/foo" }
+    Timecop.freeze(@time + 0.6) { get "/foo" }
+    last_response.body.should show_throttled_response
+  end
+  
+  it "should allow the request if the rate is less than the average" do
+    Timecop.freeze(@time) { get "/foo" }
+    Timecop.freeze(@time + 0.5) { get "/foo" }
+    Timecop.freeze(@time + 2) { get "/foo" }
+
+    last_response.body.should show_allowed_response
+  end
+
+  it "should not allow the request if the rate is more than the average" do
+    Timecop.freeze(@time) { get "/foo" }
+    Timecop.freeze(@time + 0.5) { get "/foo" }
+    Timecop.freeze(@time + 1) { get "/foo" }
+    Timecop.freeze(@time + 1.5) { get "/foo" }
+
+    last_response.body.should show_throttled_response
+  end
+
+  it "should gracefully allow the request if the cache bombs on getting" do
+    app.should_receive(:cache_get).and_raise(StandardError)
+    get "/foo"
+    last_response.body.should show_allowed_response
+  end
+  
+  it "should gracefully allow the request if the cache bombs on setting" do
+    app.should_receive(:cache_set).and_raise(StandardError)
+    get "/foo"
+    last_response.body.should show_allowed_response
+  end
+end
+