limiter.rb in rack-throttle-0.2.0

- old
+ new

@@ -120,11 +120,20 @@
     # @param  [Object] value
     # @return [void]
     def cache_set(key, value)
       case
         when cache.respond_to?(:[]=)
-          cache[key] = value
+          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
 
@@ -162,25 +171,24 @@
     end
 
     ##
     # Outputs a `Rate Limit Exceeded` error.
     #
-    # @param  [Integer] code
-    # @param  [String]  message
     # @return [Array(Integer, Hash, #each)]
-    def rate_limit_exceeded(code = nil, message = nil)
-      http_error(code || options[:code] || 403,
-        message || options[:message] || 'Rate Limit Exceeded')
+    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  [Integer]                code
+    # @param  [String, #to_s]          message
+    # @param  [Hash{String => String}] headers
     # @return [Array(Integer, Hash, #each)]
-    def http_error(code, message = nil)
-      [code, {'Content-Type' => 'text/plain; charset=utf-8'},
+    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")]
     end
 
     ##
     # Returns the standard HTTP status message for the given status `code`.