Merge pull request #50781 from Shopify/as-cache-limitter

Refactor `ActionController::RateLimiting` to use `AS::Cache`
This commit is contained in:
Jean Boussier 2024-01-17 17:49:27 +01:00 committed by GitHub
commit 1205a751c5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 49 additions and 53 deletions

View File

@ -87,7 +87,6 @@ else
gem "rack", git: "https://github.com/rack/rack.git", branch: "main"
end
gem "kredis", ">= 1.7.0", require: false
gem "useragent", require: false
# Active Job

View File

@ -299,10 +299,6 @@ GEM
rexml
kramdown-parser-gfm (1.1.0)
kramdown (~> 2.0)
kredis (1.7.0)
activemodel (>= 6.0.0)
activesupport (>= 6.0.0)
redis (>= 4.2, < 6)
language_server-protocol (3.17.0.3)
libxml-ruby (5.0.0)
listen (3.8.0)
@ -611,7 +607,6 @@ DEPENDENCIES
jbuilder
jsbundling-rails
json (>= 2.0.0, != 2.7.0)
kredis (>= 1.7.0)
libxml-ruby
listen (~> 3.3)
mdl (!= 0.13.0)

View File

@ -21,7 +21,7 @@
*DHH*
* Add rate limiting API using Redis and the [Kredis limiter type](https://github.com/rails/kredis/blob/main/lib/kredis/types/limiter.rb).
* Add rate limiting API.
```ruby
class SessionsController < ApplicationController
@ -34,7 +34,7 @@
end
```
*DHH*
*DHH*, *Jean Boussier*
* Add `image/svg+xml` to the compressible content types of ActionDispatch::Static

View File

@ -15,6 +15,10 @@ module ActionController # :nodoc:
# Requests that exceed the rate limit are refused with a <tt>429 Too Many Requests</tt> response. You can specialize this by passing a callable
# in the <tt>with:</tt> parameter. It's evaluated within the context of the controller processing the request.
#
# Rate limiting relies on a backing <tt>ActiveSupport::Cache</tt> store and defaults to <tt>config.action_controller.cache_store</tt>, which
# itself default to the global `config.cache_store`. If you don't want to store rate limits in the same datastore than your general caches
# you can pass a custom store in the <tt>store</tt> parameter.
#
# Examples:
#
# class SessionsController < ApplicationController
@ -26,38 +30,24 @@ module ActionController # :nodoc:
# by: -> { request.domain }, with: -> { redirect_to busy_controller_url, alert: "Too many signups on domain!" }, only: :new
# end
#
# Note: Rate limiting relies on the application having an accessible Redis server and on Kredis 1.7.0+ being available in the bundle.
# This uses the Kredis limiter type underneath, which is failsafe, so in case Redis is inaccessible, the rate limit will not refuse action execution.
def rate_limit(to:, within:, by: -> { request.remote_ip }, with: -> { head :too_many_requests }, **options)
require_compatible_kredis
before_action -> { rate_limiting(to: to, within: within, by: by, with: with) }, **options
# class APIController < ApplicationController
# RATE_LIMIT_STORE = ActiveSupport::Cache::RedisCacheStore.new(url: ENV["REDIS_URL"])
# rate_limit to: 10, within: 3.minutes, store: RATE_LIMIT_STORE
# end
#
# TODO Note
def rate_limit(to:, within:, by: -> { request.remote_ip }, with: -> { head :too_many_requests }, store: cache_store, **options)
before_action -> { rate_limiting(to: to, within: within, by: by, with: with, store: store) }, **options
end
private
def require_compatible_kredis
require "kredis"
if Kredis::VERSION < "1.7.0"
raise StandardError, \
"Rate limiting requires Kredis 1.7.0+. Please update by calling `bundle update kredis`."
end
rescue LoadError
raise LoadError, \
"Rate limiting requires Redis and Kredis. " +
"Please ensure you have Redis installed on your system and the Kredis gem in your Gemfile."
end
end
private
def rate_limiting(to:, within:, by:, with:)
limiter = Kredis.limiter "rate-limit:#{controller_path}:#{instance_exec(&by)}", limit: to, expires_in: within
if limiter.exceeded?
def rate_limiting(to:, within:, by:, with:, store:)
count = store.increment("rate-limit:#{controller_path}:#{instance_exec(&by)}", 1, expires_in: within)
if count && count > to
ActiveSupport::Notifications.instrument("rate_limit.action_controller", request: request) do
instance_exec(&with)
end
else
limiter.poke
end
end
end

View File

@ -1,17 +1,9 @@
# frozen_string_literal: true
require "abstract_unit"
require "kredis"
Kredis.configurator = Class.new do
def config_for(name) { db: "2" } end
def root() Pathname.new(Dir.pwd) end
end.new
# Enable Kredis logging
# ActiveSupport::LogSubscriber.logger = ActiveSupport::Logger.new(STDOUT)
class RateLimitedController < ActionController::Base
self.cache_store = ActiveSupport::Cache::MemoryStore.new
rate_limit to: 2, within: 2.seconds, by: -> { Thread.current[:redis_test_seggregation] }, only: :limited_to_two
def limited_to_two
@ -29,7 +21,7 @@ class RateLimitingTest < ActionController::TestCase
setup do
Thread.current[:redis_test_seggregation] = Random.hex(10)
Kredis.counter("rate-limit:rate_limited:#{Thread.current[:redis_test_seggregation]}").del
RateLimitedController.cache_store.clear
end
test "exceeding basic limit" do
@ -46,9 +38,10 @@ class RateLimitingTest < ActionController::TestCase
get :limited_to_two
assert_response :ok
sleep 3
get :limited_to_two
assert_response :ok
travel_to Time.now + 3.seconds do
get :limited_to_two
assert_response :ok
end
end
test "limited with" do

View File

@ -211,17 +211,22 @@ module ActiveSupport
# If the key is not found it is created and set to +amount+.
def modify_value(name, amount, options)
file_name = normalize_key(name, options)
options = merged_options(options)
key = normalize_key(name, options)
version = normalize_version(name, options)
amount = Integer(amount)
lock_file(file_name) do
options = merged_options(options)
entry = read_entry(key, **options)
if num = read(name, options)
num = num.to_i + amount
write(name, num, options)
num
else
write(name, Integer(amount), options)
if !entry || entry.expired? || entry.mismatched?(version)
write(name, amount, options)
amount
else
num = entry.value.to_i + amount
entry = Entry.new(num, expires_at: entry.expires_at, version: entry.version)
write_entry(key, entry)
num
end
end
end

View File

@ -30,4 +30,18 @@ module CacheIncrementDecrementBehavior
missing = @cache.decrement(SecureRandom.alphanumeric, 100)
assert_equal @cache.is_a?(ActiveSupport::Cache::MemCacheStore) ? 0 : -100, missing
end
def test_ttl_isnt_updated
key = SecureRandom.uuid
assert_equal 1, @cache.increment(key, 1, expires_in: 1)
assert_equal 2, @cache.increment(key, 1, expires_in: 5000)
# having to sleep two seconds in a test is bad, but we're testing
# a wide range of backends with different TTL mecanisms, most without
# subsecond granularity, so this is the only reliable way.
sleep 2
assert_nil @cache.read(key, raw: true)
end
end