prorate 0.4.0 → 0.5.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: 415042776ee7b18bf44e586fe148c822a44746bd
-  data.tar.gz: 7e249159cd2b18b699fc1f7b4fcb0f3923f2082d
+  metadata.gz: 89ad7bfadc58561263e35de3de2a0eef9c4b5aca
+  data.tar.gz: 51591b0aabe76204c4a9029a095150faf434c84d
 SHA512:
-  metadata.gz: f84dc40f3b0d789a2a8a3b0d9739b958958d895c2a43d2a60265a325465c46923ec8b5ce39fbf94d7942ce806315ffcb19bf39d4de9039853285aadfa8ecb933
-  data.tar.gz: 9335ecd1f3c25f082515dc5a7396aee2670748e6f287b481ef723699baaf355e5b67d0e23123721aefca894f0a2dd2e03065700ef581cad1b6bf9a64f963a1f9
+  metadata.gz: 5393f7eb13d4bd9236f3809c90215ae1a9700dd20a17c195db4180a526502278055365fc8e5dfce4644bba737a808cd738e7e6d3120f07ff733d2a8d214a7749
+  data.tar.gz: 604f4fd70fd8d9506cdccff6e819d91d2e26eae94e70ba00f3cf212118ca5d2fb2eeab8052c0702bea4a76d423be93c860dded6de64cd140ad1ddacad9ea1e2e

data/README.md

@@ -4,6 +4,9 @@ Provides a low-level time-based throttle. Is mainly meant for situations where u
 useful since you need access to more variables. Under the hood, this uses a Lua script that implements the
 [Leaky Bucket](https://en.wikipedia.org/wiki/Leaky_bucket) algorithm in a single threaded and race condition safe way.
 
+[![Build Status](https://travis-ci.org/WeTransfer/prorate.svg?branch=master)](https://travis-ci.org/WeTransfer/prorate)
+[![Gem Version](https://badge.fury.io/rb/prorate.svg)](https://badge.fury.io/rb/prorate)
+
 ## Installation
 
 Add this line to your application's Gemfile:
@@ -27,14 +30,14 @@ Within your Rails controller:
     t = Prorate::Throttle.new(redis: Redis.new, logger: Rails.logger,
         name: "throttle-login-email", limit: 20, period: 5.seconds)
     # Add all the parameters that function as a discriminator
-    t << request.ip
-    t << params.require(:email)
+    t << request.ip << params.require(:email)
     # ...and call the throttle! method
     t.throttle! # Will raise a Prorate::Throttled exception if the limit has been reached
 
 To capture that exception, in the controller
 
     rescue_from Prorate::Throttled do |e|
+      response.set_header('Retry-After', e.retry_in_seconds.to_s)
       render nothing: true, status: 429
     end
 

data/Rakefile

@@ -1,6 +1,13 @@
 require "bundler/gem_tasks"
 require "rspec/core/rake_task"
 require 'rubocop/rake_task'
+require 'yard'
+
+YARD::Rake::YardocTask.new(:doc) do |t|
+  # The dash has to be between the two to "divide" the source files and
+  # miscellaneous documentation files that contain no code
+  t.files = ['lib/**/*.rb', '-', 'LICENSE.txt']
+end
 
 RSpec::Core::RakeTask.new(:spec)
 RuboCop::RakeTask.new(:rubocop)

data/lib/prorate/rate_limit.lua

@@ -14,6 +14,7 @@ local block_key = ARGV[1] .. ".block"
 local max_bucket_capacity = tonumber(ARGV[2])
 local leak_rate = tonumber(ARGV[3])
 local block_duration = tonumber(ARGV[4])
+local n_tokens = tonumber(ARGV[5]) -- How many tokens this call adds to the bucket. Defaults to 1
 local now = tonumber(redis.call("TIME")[1]) --unix timestamp, will be required in all paths
 
 local key_lifetime = math.ceil(max_bucket_capacity / leak_rate)
@@ -23,28 +24,29 @@ if blocked_until then
   return {(tonumber(blocked_until) - now), 0}
 end
 
--- get current bucket level
-local bucket_level = tonumber(redis.call("GET", bucket_level_key))
-if not bucket_level then
-  -- this throttle/identifier combo does not exist yet, so much calculation can be skipped
-  redis.call("SETEX", bucket_level_key, key_lifetime, 1) -- set bucket with initial value
-  retval =  {0, 1}
+-- get current bucket level. The throttle key might not exist yet in which
+-- case we default to 0
+local bucket_level = tonumber(redis.call("GET", bucket_level_key)) or 0
+
+-- ...and then perform the leaky bucket fillup/leak. We need to do this also when the bucket has
+-- just been created because the initial n_tokens to add might be so high that it will
+-- immediately overflow the bucket and trigger the throttle, on the first call.
+local last_updated = tonumber(redis.call("GET", last_updated_key)) or now -- use sensible default of 'now' if the key does not exist
+local new_bucket_level = math.max(0, bucket_level - (leak_rate * (now - last_updated)))
+
+if (new_bucket_level + n_tokens) <= max_bucket_capacity then
+  new_bucket_level = math.max(0, new_bucket_level + n_tokens)
+  retval = {0, math.ceil(new_bucket_level)}
 else
-  -- if it already exists, do the leaky bucket thing
-  local last_updated = tonumber(redis.call("GET", last_updated_key)) or now -- use sensible default of 'now' if the key does not exist
-  local new_bucket_level = math.max(0, bucket_level - (leak_rate * (now - last_updated)))
-
-  if (new_bucket_level + 1) <= max_bucket_capacity then
-    new_bucket_level = new_bucket_level + 1
-    retval = {0, math.ceil(new_bucket_level)}
-  else
-    redis.call("SETEX", block_key, block_duration, now + block_duration)
-    retval = {block_duration, 0}
-  end
-  redis.call("SETEX", bucket_level_key, key_lifetime, new_bucket_level) --still needs to be saved
+  redis.call("SETEX", block_key, block_duration, now + block_duration)
+  retval = {block_duration, 0}
 end
 
--- update last_updated for this bucket, required in all branches
+-- Save the new bucket level
+redis.call("SETEX", bucket_level_key, key_lifetime, new_bucket_level)
+
+-- Record when we updated the bucket so that the amount of tokens leaked
+-- can be correctly determined on the next invocation
 redis.call("SETEX", last_updated_key, key_lifetime, now)
 
 return retval

data/lib/prorate/throttle.rb

@@ -1,20 +1,12 @@
 require 'digest'
 
 module Prorate
-  class ScriptHashMismatch < StandardError
-  end
-
   class MisconfiguredThrottle < StandardError
   end
 
   class Throttle < Ks.strict(:name, :limit, :period, :block_for, :redis, :logger)
-    def self.lua_script_hash
-      script_filepath = File.join(__dir__, "rate_limit.lua")
-      script = File.read(script_filepath)
-      Digest::SHA1.hexdigest(script)
-    end
-
-    CURRENT_SCRIPT_HASH = lua_script_hash
+    LUA_SCRIPT_CODE = File.read(File.join(__dir__, "rate_limit.lua"))
+    LUA_SCRIPT_HASH = Digest::SHA1.hexdigest(LUA_SCRIPT_CODE)
 
     def initialize(*)
       super
@@ -24,17 +16,77 @@ module Prorate
       @leak_rate = limit.to_f / period # tokens per second;
     end
 
+    # Add a value that will be used to distinguish this throttle from others.
+    # It has to be something user- or connection-specific, and multiple
+    # discriminators can be combined:
+    #
+    #    throttle << ip_address << user_agent_fingerprint
+    #
+    # @param discriminator[Object] a Ruby object that can be marshaled
+    #    in an equivalent way between requests, using `Marshal.dump
     def <<(discriminator)
       @discriminators << discriminator
     end
 
-    def throttle!
+    # Applies the throttle and raises a {Throttled} exception if it has been triggered
+    #
+    # Accepts an optional number of tokens to put in the bucket (default is 1).
+    # The effect of `n_tokens:` set to 0 is a "ping".
+    # It makes sure the throttle keys in Redis get created and adjusts the
+    # last invoked time of the leaky bucket. Can be used when a throttle
+    # is applied in a "shadow" fashion. For example, imagine you
+    # have a cascade of throttles with the following block times:
+    #
+    #   Throttle A: [-------]
+    #   Throttle B: [----------]
+    #
+    # You apply Throttle A: and it fires, but when that happens you also
+    # want to enable a throttle that is applied to "repeat offenders" only -
+    # - for instance ones that probe for tokens and/or passwords.
+    #
+    #   Throttle C: [-------------------------------]
+    #
+    # If your "Throttle A" fires, you can trigger Throttle C
+    #
+    #    Throttle A: [-----|-]
+    #    Throttle C: [-----|-------------------------]
+    #
+    # because you know that Throttle A has fired and thus Throttle C comes
+    # into effect.  What you want to do, however, is to fire Throttle C
+    # even though Throttle A: would have unlatched, which would create this
+    # call sequence:
+    #
+    #    Throttle A: [-------]    *(A not triggered)
+    #    Throttle C: [------------|------------------]
+    #
+    # To achieve that you can keep Throttle C alive using `throttle!(n_tokens: 0)`,
+    # on every check that touches Throttle A and/or Throttle C. It keeps the leaky bucket
+    # updated but does not add any tokens to it:
+    #
+    #    Throttle A: [------]    *(A not triggered since block period has ended)
+    #    Throttle C: [-----------|(ping)------------------]  C is still blocking
+    #
+    # So you can effectively "keep a throttle alive" without ever triggering it,
+    # or keep it alive in combination with other throttles.
+    #
+    # @param n_tokens[Integer] the number of tokens to put in the bucket. If you are
+    #   using Prorate for rate limiting, and a single request is adding N objects to your
+    #   database for example, you can "top up" the bucket with a set number of tokens
+    #   with a arbitrary ratio - like 1 token per inserted row. Once the bucket fills up
+    #   the Throttled exception is going to be raised. Defaults to 1.
+    def throttle!(n_tokens: 1)
       discriminator = Digest::SHA1.hexdigest(Marshal.dump(@discriminators))
       identifier = [name, discriminator].join(':')
 
       redis.with do |r|
-        logger.info { "Applying throttle counter %s" % name }
-        remaining_block_time, bucket_level = run_lua_throttler(redis: r, identifier: identifier, bucket_capacity: limit, leak_rate: @leak_rate, block_for: block_for)
+        logger.debug { "Applying throttle counter %s" % name }
+        remaining_block_time, bucket_level = run_lua_throttler(
+          redis: r,
+          identifier: identifier,
+          bucket_capacity: limit,
+          leak_rate: @leak_rate,
+          block_for: block_for,
+          n_tokens: n_tokens)
 
         if remaining_block_time > 0
           logger.warn { "Throttle %s exceeded limit of %d in %d seconds and is blocked for the next %d seconds" % [name, limit, period, remaining_block_time] }
@@ -44,16 +96,16 @@ module Prorate
       end
     end
 
-    def run_lua_throttler(redis:, identifier:, bucket_capacity:, leak_rate:, block_for:)
-      redis.evalsha(CURRENT_SCRIPT_HASH, [], [identifier, bucket_capacity, leak_rate, block_for])
+    private
+
+    def run_lua_throttler(redis:, identifier:, bucket_capacity:, leak_rate:, block_for:, n_tokens:)
+      redis.evalsha(LUA_SCRIPT_HASH, [], [identifier, bucket_capacity, leak_rate, block_for, n_tokens])
     rescue Redis::CommandError => e
       if e.message.include? "NOSCRIPT"
-        # The Redis server has never seen this script before. Needs to run only once in the entire lifetime of the Redis server (unless the script changes)
-        script_filepath = File.join(__dir__, "rate_limit.lua")
-        script = File.read(script_filepath)
-        raise ScriptHashMismatch if Digest::SHA1.hexdigest(script) != CURRENT_SCRIPT_HASH
-        redis.script(:load, script)
-        redis.evalsha(CURRENT_SCRIPT_HASH, [], [identifier, bucket_capacity, leak_rate, block_for])
+        # The Redis server has never seen this script before. Needs to run only once in the entire lifetime
+        # of the Redis server, until the script changes - in which case it will be loaded under a different SHA
+        redis.script(:load, LUA_SCRIPT_CODE)
+        retry
       else
         raise e
       end

data/lib/prorate/throttled.rb

@@ -1,5 +1,17 @@
+# The Throttled exception gets raised when a throttle is triggered.
+#
+# The exception carries additional attributes which can be used for
+# error tracking and for creating a correct Retry-After HTTP header for
+# a 429 response
 class Prorate::Throttled < StandardError
-  attr_reader :throttle_name, :retry_in_seconds
+  # @attr [String] the name of the throttle (like "shpongs-per-ip").
+  #   Can be used to detect which throttle has fired when multiple
+  #   throttles are used within the same block.
+  attr_reader :throttle_name
+
+  # @attr [Integer] for how long the caller will be blocked, in seconds.
+  attr_reader :retry_in_seconds
+
   def initialize(throttle_name, try_again_in)
     @throttle_name = throttle_name
     @retry_in_seconds = try_again_in

data/lib/prorate/version.rb

@@ -1,3 +1,3 @@
 module Prorate
-  VERSION = "0.4.0"
+  VERSION = "0.5.0"
 end

data/prorate.gemspec

@@ -34,4 +34,5 @@ Gem::Specification.new do |spec|
   spec.add_development_dependency "rake", "~> 12.3"
   spec.add_development_dependency "rspec", "~> 3.0"
   spec.add_development_dependency 'wetransfer_style', '0.6.0'
+  spec.add_development_dependency 'yard', '~> 0.9'
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: prorate
 version: !ruby/object:Gem::Version
-  version: 0.4.0
+  version: 0.5.0
 platform: ruby
 authors:
 - Julik Tarkhanov
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2019-08-06 00:00:00.000000000 Z
+date: 2019-08-13 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: ks
@@ -108,6 +108,20 @@ dependencies:
     - - '='
       - !ruby/object:Gem::Version
         version: 0.6.0
+- !ruby/object:Gem::Dependency
+  name: yard
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.9'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - "~>"
+      - !ruby/object:Gem::Version
+        version: '0.9'
 description: Can be used to implement all kinds of throttles
 email:
 - me@julik.nl