Permalink
Browse files

Merge branch 'master' of git://github.com/bemurphy/rack-throttle

  • Loading branch information...
2 parents e4a6fc7 + a9440ac commit dfac45dc24de401016c78499c56b880d5f11b3e2 Arto Bendiken committed Mar 22, 2010
@@ -2,6 +2,7 @@
module Rack
module Throttle
+ autoload :TimeWindow, 'rack/throttle/time_window'
autoload :Daily, 'rack/throttle/daily'
autoload :Hourly, 'rack/throttle/hourly'
autoload :Interval, 'rack/throttle/interval'
@@ -5,8 +5,22 @@ module Rack; module Throttle
# requests per 24 hours, which works out to an average of 1 request per
# second).
#
- # _Not yet implemented in the current release._
- class Daily < Limiter
- # TODO
+ # This is rough as it doesn't have a sliding window, but rather tracks per
+ # calendar day. I can't think of a way to not have a gazillion timestamps
+ # in the cache value, otherwise
+ class Daily < TimeWindow
+ def max_per_day
+ @max_per_hour ||= @options[:max_per_day] || 86400
+ end
+ alias_method :max_per_window, :max_per_day
+
+ protected
+
+ ##
+ # @param [Rack::Request] request
+ # @return [String]
+ def cache_key(request)
+ super + "_" + Time.now.strftime("%Y-%m-%d")
+ end
end
end; end
@@ -5,8 +5,19 @@ module Rack; module Throttle
# requests per 60 minutes, which works out to an average of 1 request per
# second).
#
- # _Not yet implemented in the current release._
- class Hourly < Limiter
- # TODO
+ class Hourly < TimeWindow
+ def max_per_hour
+ @max_per_hour ||= @options[:max_per_hour] || 3600
+ end
+ alias_method :max_per_window, :max_per_hour
+
+ protected
+
+ ##
+ # @param [Rack::Request] request
+ # @return [String]
+ def cache_key(request)
+ super + "_" + Time.now.strftime("%Y-%m-%d-%H")
+ end
end
end; end
@@ -0,0 +1,20 @@
+module Rack; module Throttle
+ class TimeWindow < Limiter
+ ##
+ # Returns `true` if fewer than the max 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
+ begin
+ cache_set(key, count)
+ allowed
+ rescue => e
+ allowed = true
+ end
+ end
+ end
+end; end
View
@@ -0,0 +1,27 @@
+require File.dirname(__FILE__) + '/spec_helper'
+
+def app
+ @target_app ||= example_target_app
+ @app ||= Rack::Throttle::Daily.new(@target_app, :max_per_day => 3)
+end
+
+describe Rack::Throttle::Daily do
+ include Rack::Test::Methods
+
+ 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
+
+ # TODO mess with time travelling and requests to make sure no overlap
+end
View
@@ -0,0 +1,27 @@
+require File.dirname(__FILE__) + '/spec_helper'
+
+def app
+ @target_app ||= example_target_app
+ @app ||= Rack::Throttle::Hourly.new(@target_app, :max_per_hour => 3)
+end
+
+describe Rack::Throttle::Hourly do
+ include Rack::Test::Methods
+
+ 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
+
+ # TODO mess with time travelling and requests to make sure no overlap
+end
View
@@ -0,0 +1,39 @@
+require File.dirname(__FILE__) + '/spec_helper'
+
+def app
+ @target_app ||= example_target_app
+ @app ||= Rack::Throttle::Interval.new(@target_app, :min => 0.1)
+end
+
+describe Rack::Throttle::Interval do
+ include Rack::Test::Methods
+
+ 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
+ get "/foo"
+ sleep 0.2 # Should time travel this instead?
+ get "/foo"
+ last_response.body.should show_allowed_response
+ end
+
+ it "should not all the request if the source has been seen inside the current interval" do
+ 2.times { 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
View
@@ -0,0 +1,50 @@
+require File.dirname(__FILE__) + '/spec_helper'
+
+def app
+ @target_app ||= example_target_app
+ @app ||= Rack::Throttle::Limiter.new(@target_app)
+end
+
+describe Rack::Throttle::Limiter do
+ include Rack::Test::Methods
+
+ 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
View
@@ -1 +1,44 @@
-require 'rack/throttle'
+require "spec"
+require "rack/test"
+require "rack/throttle"
+
+def example_target_app
+ @target_app ||= mock("Example Rack App")
+ @target_app.stub!(:call).and_return([200, {}, "Example App Body"])
+end
+
+Spec::Matchers.define :show_allowed_response do
+ match do |body|
+ body.include?("Example App Body")
+ end
+
+ failure_message_for_should do
+ "expected response to show the allowed response"
+ end
+
+ failure_message_for_should_not do
+ "expected response not to show the allowed response"
+ end
+
+ description do
+ "expected the allowed response"
+ end
+end
+
+Spec::Matchers.define :show_throttled_response do
+ match do |body|
+ body.include?("Rate Limit Exceeded")
+ end
+
+ failure_message_for_should do
+ "expected response to show the throttled response"
+ end
+
+ failure_message_for_should_not do
+ "expected response not to show the throttled response"
+ end
+
+ description do
+ "expected the throttled response"
+ end
+end

0 comments on commit dfac45d

Please sign in to comment.