rack-ratelimit 1.0.0 → 1.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
-SHA1:
-  metadata.gz: 73b500a4dca2655841490da8c02106a49f69f62b
-  data.tar.gz: 265204cbba7c818c0190eee355d55a7147083d1a
+SHA256:
+  metadata.gz: f495171388825bee8e7248a597406cad160d6d0fe2adfb03f9d33617ab136832
+  data.tar.gz: 4b6256bf56d9384a80265ac5b78a9f63809e29e08bc33918b87c22ae3717cc5d
 SHA512:
-  metadata.gz: da84278754fd180b780687f5819f4847d2e8fcf49378029e023b95706324b3b06eedd5cedf1c3b453e26723ed8b409ea78799e190497076db3411be9f9f55081
-  data.tar.gz: 2af1a03a36288b8cbec652185ff1207b3e9052b58d3c0737c7a07159f349eeadaacf42c79acfde8ad2ae8a46f011b61fd04619d21b629117ae1f059ec4e9d7b3
+  metadata.gz: c8eb47b3e77c43c16df7c984463366f78d0be4c5f46b425e15d027932d5d43be4771e16a8338325c631408229bc3011576daf888aaff0d021ba505a080e06a9f
+  data.tar.gz: 467d0a44fab4f3b6490a912ad1ed8b699b315d4158a2ce709ce2743bfa16ba9c614e151d941674c2b830e4ac406ff1305fc09ee6fb2612bf7176d21477852d55

data/lib/rack/ratelimit.rb

@@ -1,4 +1,3 @@
-require 'dalli'
 require 'logger'
 require 'time'
 
@@ -10,9 +9,9 @@ module Rack
   # * Apply each rate limit by request characteristics: IP, subdomain, OAuth2 token, etc.
   # * Flexible time window to limit burst traffic vs hourly or daily traffic:
   #     100 requests per 10 sec, 500 req/minute, 10000 req/hour, etc.
-  # * Fast, low-overhead implementation using memcache counters per time window:
+  # * Fast, low-overhead implementation using counters per time window:
   #     timeslice = window * ceiling(current time / window)
-  #     memcache.incr(counter for timeslice)
+  #     store.incr(timeslice)
   class Ratelimit
     # Takes a block that classifies requests for rate limiting. Given a
     # Rack env, return a string such as IP address, API token, etc. If the
@@ -21,7 +20,12 @@ module Rack
     #
     # Required configuration:
     #   rate: an array of [max requests, period in seconds]: [500, 5.minutes]
-    #   cache: a Dalli::Client instance, or an object that quacks like it.
+    # and one of
+    #   cache: a Dalli::Client instance
+    #   redis: a Redis instance
+    #   counter: Your own custom counter. Must respond to
+    #     `#increment(classification_string, end_of_time_window_epoch_timestamp)`
+    #     and return the counter value after increment.
     #
     # Optional configuration:
     #   name: name of the rate limiter. Defaults to 'HTTP'. Used in messages.
@@ -35,7 +39,10 @@ module Rack
     #     subsequently blocked requests.
     #   error_message: the message returned in the response body when the rate
     #     limit is exceeded. Defaults to "<name> rate limit exceeded. Please
-    #     wait <period> seconds then retry your request."
+    #     wait %d seconds then retry your request." The number of seconds
+    #     until the end of the rate-limiting window is interpolated into the
+    #     message string, but the %d placeholder is optional if you wish to
+    #     omit it.
     #
     # Example:
     #
@@ -51,7 +58,7 @@ module Rack
     #   use(Rack::Ratelimit, name: 'API',
     #     conditions: ->(env) { env['REMOTE_USER'] },
     #     rate:   [1000, 1.hour],
-    #     cache:  Dalli::Client.new,
+    #     redis:  Redis.new(ratelimit_redis_config),
     #     logger: Rails.logger) { |env| env['REMOTE_USER'] }
     def initialize(app, options, &classifier)
       @app, @classifier = app, classifier
@@ -61,10 +68,20 @@ module Rack
       @max, @period = options.fetch(:rate)
       @status = options.fetch(:status, 429)
 
-      @counter = Counter.new(options.fetch(:cache), @name, @period)
+      @counter =
+        if counter = options[:counter]
+          raise ArgumentError, 'Counter must respond to #increment' unless counter.respond_to?(:increment)
+          counter
+        elsif cache = options[:cache]
+          MemcachedCounter.new(cache, @name, @period)
+        elsif redis = options[:redis]
+          RedisCounter.new(redis, @name, @period)
+        else
+          raise ArgumentError, ':cache, :redis, or :counter is required'
+        end
 
       @logger = options[:logger]
-      @error_message = options.fetch(:error_message, "#{@name} rate limit exceeded. Please wait #{@period} seconds then retry your request.")
+      @error_message = options.fetch(:error_message, "#{@name} rate limit exceeded. Please wait %d seconds then retry your request.")
 
       @conditions = Array(options[:conditions])
       @exceptions = Array(options[:exceptions])
@@ -92,6 +109,11 @@ module Rack
       @exceptions.none? { |e| e.call(env) } && @conditions.all? { |c| c.call(env) }
     end
 
+    # Give subclasses an opportunity to specialize classification.
+    def classify(env)
+      @classifier.call env
+    end
+
     # Handle a Rack request:
     #   * Check whether the rate limit applies to the request.
     #   * Classify the request by IP, API token, etc.
@@ -101,33 +123,34 @@ module Rack
     #   * If it's the first request that exceeds the limit, log it.
     #   * If the count doesn't exceed the limit, pass through the request.
     def call(env)
-      if apply_rate_limit?(env) && classification = @classifier.call(env)
-
-        # Marks the end of the current rate-limiting window.
-        timestamp = @period * (Time.now.to_f / @period).ceil
-        time = Time.at(timestamp).utc.xmlschema
+      # Accept an optional start-of-request timestamp from the Rack env for
+      # upstream timing and for testing.
+      now = env.fetch('ratelimit.timestamp', Time.now).to_f
 
+      if apply_rate_limit?(env) && classification = classify(env)
         # Increment the request counter.
-        count = @counter.increment(classification, timestamp)
-        remaining = @max - count + 1
-
-        json = %({"name":"#{@name}","period":#{@period},"limit":#{@max},"remaining":#{remaining},"until":"#{time}"})
+        epoch = ratelimit_epoch(now)
+        count = @counter.increment(classification, epoch)
+        remaining = @max - count
 
         # If exceeded, return a 429 Rate Limit Exceeded response.
-        if remaining <= 0
+        if remaining < 0
           # Only log the first hit that exceeds the limit.
-          if @logger && remaining == 0
-            @logger.info '%s: %s exceeded %d request limit for %s' % [@name, classification, @max, time]
+          if @logger && remaining == -1
+            @logger.info '%s: %s exceeded %d request limit for %s' % [@name, classification, @max, format_epoch(epoch)]
           end
 
+          retry_after = seconds_until_epoch(epoch)
+
           [ @status,
-            { 'X-Ratelimit' => json, 'Retry-After' => @period.to_s },
-            [@error_message] ]
+            { 'X-Ratelimit' => ratelimit_json(remaining, epoch),
+              'Retry-After' => retry_after.to_s },
+            [ @error_message % retry_after ] ]
 
         # Otherwise, pass through then add some informational headers.
         else
           @app.call(env).tap do |status, headers, body|
-            headers['X-Ratelimit'] = [headers['X-Ratelimit'], json].compact.join("\n")
+            amend_headers headers, 'X-Ratelimit', ratelimit_json(remaining, epoch)
           end
         end
       else
@@ -135,14 +158,39 @@ module Rack
       end
     end
 
-    class Counter
+    private
+      # Calculate the end of the current rate-limiting window.
+      def ratelimit_epoch(timestamp)
+        @period * (timestamp / @period).ceil
+      end
+
+      def ratelimit_json(remaining, epoch)
+        %({"name":"#{@name}","period":#{@period},"limit":#{@max},"remaining":#{remaining < 0 ? 0 : remaining},"until":"#{format_epoch(epoch)}"})
+      end
+
+      def format_epoch(epoch)
+        Time.at(epoch).utc.xmlschema
+      end
+
+      # Clamp negative durations in case we're in a new rate-limiting window.
+      def seconds_until_epoch(epoch)
+        sec = (epoch - Time.now.to_f).ceil
+        sec = 0 if sec < 0
+        sec
+      end
+
+      def amend_headers(headers, name, value)
+        headers[name] = [headers[name], value].compact.join("\n")
+      end
+
+    class MemcachedCounter
       def initialize(cache, name, period)
         @cache, @name, @period = cache, name, period
       end
 
       # Increment the request counter and return the current count.
-      def increment(classification, timestamp)
-        key = 'rack-ratelimit/%s/%s/%i' % [@name, classification, timestamp]
+      def increment(classification, epoch)
+        key = 'rack-ratelimit/%s/%s/%i' % [@name, classification, epoch]
 
         # Try to increment the counter if it's present.
         if count = @cache.incr(key, 1)
@@ -158,5 +206,23 @@ module Rack
         end
       end
     end
+
+    class RedisCounter
+      def initialize(redis, name, period)
+        @redis, @name, @period = redis, name, period
+      end
+
+      # Increment the request counter and return the current count.
+      def increment(classification, epoch)
+        key = 'rack-ratelimit/%s/%s/%i' % [@name, classification, epoch]
+
+        # Returns [count, expire_ok] response for each multi command.
+        # Return the first, the count.
+        @redis.multi do |redis|
+          redis.incr key
+          redis.expire key, @period
+        end.first
+      end
+    end
   end
 end

metadata

@@ -1,84 +1,28 @@
 --- !ruby/object:Gem::Specification
 name: rack-ratelimit
 version: !ruby/object:Gem::Version
-  version: 1.0.0
+  version: 1.2.0
 platform: ruby
 authors:
-- Jeremy Kemper
-autorequire: 
+- Jeremy Daer
+autorequire:
 bindir: bin
 cert_chain: []
-date: 2014-03-21 00:00:00.000000000 Z
-dependencies:
-- !ruby/object:Gem::Dependency
-  name: rack
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-  type: :runtime
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-- !ruby/object:Gem::Dependency
-  name: dalli
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-  type: :runtime
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-- !ruby/object:Gem::Dependency
-  name: rake
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-  type: :development
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - ">="
-      - !ruby/object:Gem::Version
-        version: '0'
-- !ruby/object:Gem::Dependency
-  name: minitest
-  requirement: !ruby/object:Gem::Requirement
-    requirements:
-    - - "~>"
-      - !ruby/object:Gem::Version
-        version: 5.3.0
-  type: :development
-  prerelease: false
-  version_requirements: !ruby/object:Gem::Requirement
-    requirements:
-    - - "~>"
-      - !ruby/object:Gem::Version
-        version: 5.3.0
-description: 
-email: jeremy@bitsweat.net
+date: 2021-01-22 00:00:00.000000000 Z
+dependencies: []
+description:
+email: jeremydaer@gmail.com
 executables: []
 extensions: []
 extra_rdoc_files: []
 files:
 - "./lib/rack-ratelimit.rb"
 - "./lib/rack/ratelimit.rb"
-homepage: 
+homepage: https://github.com/jeremy/rack-ratelimit
 licenses:
 - MIT
 metadata: {}
-post_install_message: 
+post_install_message:
 rdoc_options: []
 require_paths:
 - lib
@@ -86,16 +30,15 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '1.8'
+      version: '2.0'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubyforge_project: 
-rubygems_version: 2.2.0
-signing_key: 
+rubygems_version: 3.1.4
+signing_key:
 specification_version: 4
 summary: Flexible rate limits for your Rack apps
 test_files: []