improved-rack-throttle 0.5.0

data/lib/rack/throttle/daily.rb

@@ -0,0 +1,44 @@
+module Rack; module Throttle
+  ##
+  # This rate limiter strategy throttles the application by defining a
+  # maximum number of allowed HTTP requests per day (by default, 86,400
+  # requests per 24 hours, which works out to an average of 1 request per
+  # second).
+  #
+  # Note that this strategy doesn't use a sliding time window, but rather
+  # tracks requests per calendar day. This means that the throttling counter
+  # is reset at midnight (according to the server's local timezone) every
+  # night.
+  #
+  # @example Allowing up to 86,400 requests per day
+  #   use Rack::Throttle::Daily
+  #
+  # @example Allowing up to 1,000 requests per day
+  #   use Rack::Throttle::Daily, :max => 1000
+  #
+  class Daily < TimeWindow
+    ##
+    # @param  [#call]                  app
+    # @param  [Hash{Symbol => Object}] options
+    # @option options [Integer] :max   (86400)
+    def initialize(app, options = {})
+      super
+    end
+
+    ##
+    def max_per_day
+      @max_per_hour ||= options[:max_per_day] || options[:max] || 86_400
+    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')].join(':')
+    end
+  end
+end; end

data/lib/rack/throttle/hourly.rb

@@ -0,0 +1,44 @@
+module Rack; module Throttle
+  ##
+  # This rate limiter strategy throttles the application by defining a
+  # maximum number of allowed HTTP requests per hour (by default, 3,600
+  # requests per 60 minutes, which works out to an average of 1 request per
+  # second).
+  #
+  # Note that this strategy doesn't use a sliding time window, but rather
+  # tracks requests per distinct hour. This means that the throttling
+  # counter is reset every hour on the hour (according to the server's local
+  # timezone).
+  #
+  # @example Allowing up to 3,600 requests per hour
+  #   use Rack::Throttle::Hourly
+  #
+  # @example Allowing up to 100 requests per hour
+  #   use Rack::Throttle::Hourly, :max => 100
+  #
+  class Hourly < TimeWindow
+    ##
+    # @param  [#call]                  app
+    # @param  [Hash{Symbol => Object}] options
+    # @option options [Integer] :max   (3600)
+    def initialize(app, options = {})
+      super
+    end
+
+    ##
+    def max_per_hour
+      @max_per_hour ||= options[:max_per_hour] || options[:max] || 3_600
+    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-%dT%H')].join(':')
+    end
+  end
+end; end

data/lib/rack/throttle/interval.rb

@@ -0,0 +1,63 @@
+module Rack; module Throttle
+  ##
+  # This rate limiter strategy throttles the application by enforcing a
+  # minimum interval (by default, 1 second) between subsequent allowed HTTP
+  # requests.
+  #
+  # @example Allowing up to two requests per second
+  #   use Rack::Throttle::Interval, :min => 0.5   #  500 ms interval
+  #
+  # @example Allowing a request every two seconds
+  #   use Rack::Throttle::Interval, :min => 2.0   # 2000 ms interval
+  #
+  class Interval < Limiter
+    ##
+    # @param  [#call]                  app
+    # @param  [Hash{Symbol => Object}] options
+    # @option options [Float] :min     (1.0)
+    def initialize(app, options = {})
+      super
+    end
+
+    ##
+    # Returns `true` if sufficient time (equal to or more than
+    # {#minimum_interval}) has passed since the last request and the given
+    # present `request`.
+    #
+    # @param  [Rack::Request] request
+    # @return [Boolean]
+    def allowed?(request)
+      t1 = request_start_time(request)
+      t0 = cache_get(key = cache_key(request)) rescue nil
+      allowed = !t0 || (dt = t1 - t0.to_f) >= minimum_interval
+      begin
+        cache_set(key, t1)
+        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
+
+    ##
+    # Returns the number of seconds before the client is allowed to retry an
+    # HTTP request.
+    #
+    # @return [Float]
+    def retry_after
+      minimum_interval
+    end
+
+    ##
+    # Returns the required minimal interval (in terms of seconds) that must
+    # elapse between two subsequent HTTP requests.
+    #
+    # @return [Float]
+    def minimum_interval
+      @min ||= (@options[:min] || 1.0).to_f
+    end
+  end
+end; end

data/lib/rack/throttle/limiter.rb

@@ -0,0 +1,228 @@
+module Rack; module Throttle
+  ##
+  # This is the base class for rate limiter implementations.
+  #
+  # @example Defining a rate limiter subclass
+  #   class MyLimiter < Limiter
+  #     def allowed?(request)
+  #       # TODO: custom logic goes here
+  #     end
+  #   end
+  #
+  class Limiter
+    attr_reader :app, :options, :matchers
+
+    ##
+    # @param  [#call]                       app
+    # @param  [Hash{Symbol => Object}]      options
+    # @option options [String]  :cache      (Hash.new)
+    # @option options [String]  :key        (nil)
+    # @option options [String]  :key_prefix (nil)
+    # @option options [Integer] :code       (403)
+    # @option options [String]  :message    ("Rate Limit Exceeded")
+    # @option options [Regexp]  :url_rule   (nil)
+    def initialize(app, options = {})
+      rules = options.delete(:rules) || {}
+      @app, @options, @matchers = app, options, []
+      @matchers += Array(rules[:ip]).map { |rule| IpMatcher.new(rule) } if rules[:ip]
+      @matchers += Array(rules[:url]).map { |rule| UrlMatcher.new(rule) } if rules[:url]
+      @matchers += Array(rules[:user_agent]).map { |rule| UserAgentMatcher.new(rule) } if rules[:user_agent]
+      @matchers += Array(rules[:method]).map { |rule| MethodMatcher.new(rule) } if rules[:method]
+    end
+
+    ##
+    # @param  [Hash{String => String}] env
+    # @return [Array(Integer, Hash, #each)]
+    # @see    http://rack.rubyforge.org/doc/SPEC.html
+    def call(env)
+      request = Rack::Request.new(env)
+      match_results = @matchers.map { |m| m.match?(request) }.uniq
+      applicable = @matchers.empty? || match_results == [true]
+      if applicable and !allowed?(request)
+        rate_limit_exceeded
+      else
+        app.call(env)
+      end
+    end
+
+    ##
+    # Returns `true` if no :url_rule regex or if the request path
+    # matches the :url regex, `false` otherwise.
+    #
+    # You can override this class, though that might be weird.
+    #
+    # @param  [String] path
+    # @return [Boolean]
+    def restricted_url?(path)
+      options[:url_rule].nil? || options[:url_rule].match(path)
+    end
+
+    ##
+    # Returns `false` if the rate limit has been exceeded for the given
+    # `request`, or `true` otherwise.
+    #
+    # Override this method in subclasses that implement custom rate limiter
+    # strategies.
+    #
+    # @param  [Rack::Request] request
+    # @return [Boolean]
+    def allowed?(request)
+      case
+        when whitelisted?(request) then true
+        when blacklisted?(request) then false
+        else true # override in subclasses
+      end
+    end
+
+    ##
+    # Returns `true` if the originator of the given `request` is whitelisted
+    # (not subject to further rate limits).
+    #
+    # The default implementation always returns `false`. Override this
+    # method in a subclass to implement custom whitelisting logic.
+    #
+    # @param  [Rack::Request] request
+    # @return [Boolean]
+    # @abstract
+    def whitelisted?(request)
+      false
+    end
+
+    ##
+    # Returns `true` if the originator of the given `request` is blacklisted
+    # (not honoring rate limits, and thus permanently forbidden access
+    # without the need to maintain further rate limit counters).
+    #
+    # The default implementation always returns `false`. Override this
+    # method in a subclass to implement custom blacklisting logic.
+    #
+    # @param  [Rack::Request] request
+    # @return [Boolean]
+    # @abstract
+    def blacklisted?(request)
+      false
+    end
+
+    protected
+
+    ##
+    # @return [Hash]
+    def cache
+      case cache = (options[:cache] ||= {})
+        when Proc then cache.call
+        else cache
+      end
+    end
+
+    ##
+    # @param  [String] key
+    def cache_has?(key)
+      case
+        when cache.respond_to?(:has_key?)
+          cache.has_key?(key)
+        when cache.respond_to?(:get)
+          cache.get(key) rescue false
+        else false
+      end
+    end
+
+    ##
+    # @param  [String] key
+    # @return [Object]
+    def cache_get(key, default = nil)
+      case
+      when cache.respond_to?(:[])
+        cache[key] || default
+      when cache.respond_to?(:get)
+        cache.get(key) || default
+      end
+    end
+
+    ##
+    # @param  [String] key
+    # @param  [Object] value
+    # @return [void]
+    def cache_set(key, value)
+      case
+      when cache.respond_to?(:[]=)
+        begin
+          cache[key] = value
+        rescue TypeError => e
+          # GDBM throws a "TypeError: can't convert Float into String"
+          # exception when trying to store a Float. On the other hand, we
+          # don't want to unnecessarily coerce the value to a String for
+          # any stores that do support other data types (e.g. in-memory
+          # hash objects). So, this is a compromise.
+          cache[key] = value.to_s
+        end
+      when cache.respond_to?(:set)
+        cache.set(key, value)
+      end
+    end
+
+    ##
+    # @param  [Rack::Request] request
+    # @return [String]
+    def cache_key(request)
+      id = client_identifier(request)
+      id = options[:key].call(request) if options.has_key?(:key)
+      id = [options[:key_prefix], id].join(':') if options.has_key?(:key_prefix)
+      @matchers.each do |matcher|
+        id += ":#{matcher.identifier}"
+      end
+
+      id
+    end
+
+    ##
+    # @param  [Rack::Request] request
+    # @return [String]
+    def client_identifier(request)
+      request.ip.to_s
+    end
+
+    ##
+    # @param  [Rack::Request] request
+    # @return [Float]
+    def request_start_time(request)
+      case
+      when request.env.has_key?('HTTP_X_REQUEST_START')
+        request.env['HTTP_X_REQUEST_START'].to_f / 1000
+      else
+        Time.now.to_f
+      end
+    end
+
+    ##
+    # Outputs a `Rate Limit Exceeded` error.
+    #
+    # @return [Array(Integer, Hash, #each)]
+    def rate_limit_exceeded
+      headers = respond_to?(:retry_after) ? {'Retry-After' => retry_after.to_f.ceil.to_s} : {}
+      http_error(options[:code] || 403, options[:message] || 'Rate Limit Exceeded', headers)
+    end
+
+    ##
+    # Outputs an HTTP `4xx` or `5xx` response.
+    #
+    # @param  [Integer]                code
+    # @param  [String, #to_s]          message
+    # @param  [Hash{String => String}] headers
+    # @return [Array(Integer, Hash, #each)]
+    def http_error(code, message = nil, headers = {})
+      [ code,
+        { 'Content-Type' => 'text/plain; charset=utf-8' }.merge(headers),
+        Array( http_status(code) + (message.nil? ? "\n" : " (#{message})\n") )
+      ]
+    end
+
+    ##
+    # Returns the standard HTTP status message for the given status `code`.
+    #
+    # @param  [Integer] code
+    # @return [String]
+    def http_status(code)
+      [code, Rack::Utils::HTTP_STATUS_CODES[code]].join(' ')
+    end
+  end
+end; end

data/lib/rack/throttle/matcher.rb

@@ -0,0 +1,24 @@
+module Rack; module Throttle
+  ###
+  # This is the base class for matcher implementations
+  # Implements IP, user agent, url, and request method based matching
+  # e.g. implement throttling rules that discriminate by ip, user agent, url, request method,
+  # or any combination thereof
+  class Matcher
+    attr_reader :rule
+
+    def initialize(rule)
+      @rule = rule
+    end
+
+    # MUST return true or false
+    def match?
+      true
+    end
+
+    def identifier
+      ""
+    end
+  end
+
+end; end

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

@@ -0,0 +1,14 @@
+module Rack; module Throttle
+  ###
+  # IpMatchers take RegExp objects and compare the request ip against them
+  class IpMatcher < Matcher
+    def match?(request)
+      !!(@rule =~ request.ip)
+    end
+
+    def identifier
+      "ip" + @rule.inspect
+    end
+  end
+
+end; end

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

@@ -0,0 +1,15 @@
+module Rack; module Throttle
+  ###
+  # MethodMatchers take Symbol objects of :get, :put, :post, or :delete
+  class MethodMatcher < Matcher
+    def match?(request)
+      rack_method = :"#{@rule}?"
+      request.send(rack_method)
+    end
+
+    def identifier
+      "meth-#{@rule}"
+    end
+  end
+
+end; end

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

@@ -0,0 +1,15 @@
+module Rack; module Throttle
+  ###
+  # UrlMatchers take Regexp objects and compare the request path against them
+  class UrlMatcher < Matcher
+    def match?(request)
+      !!(@rule =~ request.path)
+    end
+
+    def identifier
+      "url" + @rule.inspect
+    end
+  end
+
+end; end
+