From 21bcae9955d7aa623c9e917ff839e5d2379ced42 Mon Sep 17 00:00:00 2001 From: Ed Robinson Date: Fri, 31 May 2024 11:47:45 +0100 Subject: [PATCH] Initial implimentation --- .github/workflows/main.yml | 9 ++++ Gemfile | 3 ++ Gemfile.lock | 11 +++++ lib/millrace.rb | 4 +- lib/millrace/rate_limit.rb | 68 +++++++++++++++++++++++++++ lib/millrace/rate_limited.rb | 10 ++++ millrace.gemspec | 2 + spec/rate_limit_spec.rb | 89 ++++++++++++++++++++++++++++++++++++ 8 files changed, 194 insertions(+), 2 deletions(-) create mode 100644 lib/millrace/rate_limit.rb create mode 100644 lib/millrace/rate_limited.rb create mode 100644 spec/rate_limit_spec.rb diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index ffe7c79..edfcfe0 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -16,6 +16,14 @@ jobs: ruby: - '3.3.1' + services: + redis: + image: redis:7.2.5 + options: >- + --health-cmd "redis-cli ping" + --health-interval 10s + --health-timeout 5s + --health-retries 5 steps: - uses: actions/checkout@v3 - name: Set up Ruby @@ -25,3 +33,4 @@ jobs: bundler-cache: true - name: Run the default task run: bundle exec rake + env: MILLRACE_REDIS_URL: redis://redis:6379 diff --git a/Gemfile b/Gemfile index 8f80798..a59147e 100644 --- a/Gemfile +++ b/Gemfile @@ -11,3 +11,6 @@ gem "rubocop", "~> 1.21" gem "rubocop-performance" gem "rubocop-rails" gem "rubocop-rspec" + +gem "hiredis-client" +gem "redis" diff --git a/Gemfile.lock b/Gemfile.lock index eef14c1..c4e2f36 100644 --- a/Gemfile.lock +++ b/Gemfile.lock @@ -2,6 +2,7 @@ PATH remote: . specs: millrace (0.1.0) + prorate GEM remote: https://rubygems.org/ @@ -23,6 +24,8 @@ GEM connection_pool (2.4.1) diff-lcs (1.5.1) drb (2.2.1) + hiredis-client (0.22.2) + redis-client (= 0.22.2) i18n (1.14.5) concurrent-ruby (~> 1.0) json (2.7.2) @@ -33,10 +36,16 @@ GEM parser (3.3.1.0) ast (~> 2.4.1) racc + prorate (0.7.3) + redis (>= 2) racc (1.7.3) rack (3.0.11) rainbow (3.1.1) rake (13.2.1) + redis (5.2.0) + redis-client (>= 0.22.0) + redis-client (0.22.2) + connection_pool regexp_parser (2.9.0) rexml (3.2.8) strscan (>= 3.0.9) @@ -96,8 +105,10 @@ PLATFORMS ruby DEPENDENCIES + hiredis-client millrace! rake (~> 13.0) + redis rspec (~> 3.0) rubocop (~> 1.21) rubocop-performance diff --git a/lib/millrace.rb b/lib/millrace.rb index b4627dc..aa67065 100644 --- a/lib/millrace.rb +++ b/lib/millrace.rb @@ -1,8 +1,8 @@ # frozen_string_literal: true require_relative "millrace/version" +require_relative "millrace/rate_limited" +require_relative "millrace/rate_limit" module Millrace - class Error < StandardError; end - # Your code goes here... end diff --git a/lib/millrace/rate_limit.rb b/lib/millrace/rate_limit.rb new file mode 100644 index 0000000..8ef7f2a --- /dev/null +++ b/lib/millrace/rate_limit.rb @@ -0,0 +1,68 @@ +require "digest" +require "prorate" + +module Millrace + class RateLimit + def initialize(name:, rate:, window:, penalty: 0, redis_url: ENV.fetch("MILLRACE_REDIS_URL", {})) + @name = name + @rate = rate + @window = window + @penalty = penalty + @redis_url = redis_url + end + + attr_reader :name, :rate, :window, :redis_url + + def before(controller) + bucket = get_bucket(controller.request.remote_ip) + level = record_request(bucket) + + return unless level > threshold + + if level - 1 < threshold + level = bucket.fillup(penalty).level + end + + raise RateLimited.new(limit_name: name, retry_after: retry_after(level)) + end + + private + + def retry_after(level) + ((level - threshold) / rate).to_i + end + + def record_request(bucket) + bucket.fillup(1).level + end + + def get_bucket(ip) + Prorate::LeakyBucket.new( + redis: redis, + redis_key_prefix: key(ip), + leak_rate: rate, + bucket_capacity: capacity, + ) + end + + def key(ip) + "millrace.#{name}.#{Digest::SHA1.hexdigest(ip)}" + end + + def capacity + (threshold * 2) + penalty + end + + def threshold + window * rate + end + + def penalty + @penalty * rate + end + + def redis + Thread.current["millrace_#{name}_redis"] ||= Redis.new(redis_url) + end + end +end diff --git a/lib/millrace/rate_limited.rb b/lib/millrace/rate_limited.rb new file mode 100644 index 0000000..1c4f207 --- /dev/null +++ b/lib/millrace/rate_limited.rb @@ -0,0 +1,10 @@ +module Millrace + class RateLimited < StandardError + def initialize(limit_name:, retry_after:) + @limit_name = limit_name + @retry_after = retry_after + end + + attr_reader :limit_name, :retry_after + end +end diff --git a/millrace.gemspec b/millrace.gemspec index de3f05c..805a3a7 100644 --- a/millrace.gemspec +++ b/millrace.gemspec @@ -38,4 +38,6 @@ Gem::Specification.new do |spec| # For more information and examples about making a new gem, check out our # guide at: https://bundler.io/guides/creating_gem.html spec.metadata["rubygems_mfa_required"] = "true" + + spec.add_dependency "prorate" end diff --git a/spec/rate_limit_spec.rb b/spec/rate_limit_spec.rb new file mode 100644 index 0000000..d62764b --- /dev/null +++ b/spec/rate_limit_spec.rb @@ -0,0 +1,89 @@ +# frozen_string_literal: true + +RSpec.describe Millrace::RateLimit do + let(:subject) do + described_class.new( + name: "test", + rate: 10, + window: 2, + penalty: penalty, + ) + end + + let(:penalty) { 1 } + + let(:controller) do + double(:controller, request: double(:request, remote_ip: to_s)) + end + + describe "#before" do + it "rate limits" do + # Fill the bucket + 20.times { subject.before(controller) } + + # hit the threshold and get a penalty + expect { subject.before(controller) }.to raise_error Millrace::RateLimited + + sleep 1 + # Still blocked for the penalty duration + expect { subject.before(controller) }.to raise_error Millrace::RateLimited + + # Not blocked after the penalty duration is over + sleep 1 + subject.before(controller) + end + + it "returns an exeption with the correct name" do + # Fill the bucket + 20.times { subject.before(controller) } + + # hit the threshold and get an error + expect { subject.before(controller) }.to raise_error do |exception| + expect(exception.limit_name).to eq "test" + end + end + + it "returns an exeption with the correct retry time" do + # Fill the bucket + 20.times { subject.before(controller) } + + # hit the threshold and get an error + expect { subject.before(controller) }.to raise_error do |exception| + expect(exception.retry_after).to eq 1 + end + end + + context "a longer penalty" do + let(:penalty) { 10 } + + it "returns an exeption with the correct retry time" do + # Fill the bucket + 20.times { subject.before(controller) } + + # hit the threshold and get an error + expect { subject.before(controller) }.to raise_error do |exception| + expect(exception.retry_after).to eq 10 + end + end + end + + context "additional requests" do + let(:penalty) { 0 } + + it "returns an exeption with the correct retry time" do + # Fill the bucket + 40.times do + subject.before(controller) + # Keep making requests even though we are rate limited + rescue Millrace::RateLimited + nil + end + + # hit the threshold and get an error + expect { subject.before(controller) }.to raise_error do |exception| + expect(exception.retry_after).to eq 2 + end + end + end + end +end