limiter.rb in rack-throttle-0.4.0

- old
+ new

@@ -29,11 +29,11 @@
     # @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)
-      allowed?(request) ? app.call(env) : rate_limit_exceeded
+      allowed?(request) ? app.call(env) : rate_limit_exceeded(request)
     end
 
     ##
     # Returns `false` if the rate limit has been exceeded for the given
     # `request`, or `true` otherwise.
@@ -122,11 +122,11 @@
     def cache_set(key, value)
       case
         when cache.respond_to?(:[]=)
           begin
             cache[key] = value
-          rescue TypeError => e
+          rescue TypeError
             # 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.
@@ -160,23 +160,25 @@
 
     ##
     # @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
+      # Check whether HTTP_X_REQUEST_START or HTTP_X_QUEUE_START exist and parse its value (for
+      # example, when having nginx in your stack, it's going to be in the "t=\d+" format).
+      if val = (request.env['HTTP_X_REQUEST_START'] || request.env['HTTP_X_QUEUE_START'])
+        val[/(?:^t=)?(\d+)/, 1].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
+    def rate_limit_exceeded(request)
+      options[:rate_limit_exceeded_callback].call(request) if options[:rate_limit_exceeded_callback]
       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
 
     ##
@@ -186,10 +188,10 @@
     # @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),
-        http_status(code) + (message.nil? ? "\n" : " (#{message})\n")]
+        [http_status(code), (message.nil? ? "\n" : " (#{message})\n")]]
     end
 
     ##
     # Returns the standard HTTP status message for the given status `code`.
     #