Initial implimentation

.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

Gemfile

@@ -11,3 +11,6 @@ gem "rubocop", "~> 1.21"
 gem "rubocop-performance"
 gem "rubocop-rails"
 gem "rubocop-rspec"
+
+gem "hiredis-client"
+gem "redis"

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

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

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

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

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

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