Skip to content

Commit c0807dc

Browse files
authored
Add ability to use multiple rate limits per controller (rails#52960)
1 parent 4aa7811 commit c0807dc

File tree

2 files changed

+39
-12
lines changed

2 files changed

+39
-12
lines changed

actionpack/lib/action_controller/metal/rate_limiting.rb

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@ module ClassMethods
2929
# datastore as your general caches, you can pass a custom store in the `store`
3030
# parameter.
3131
#
32+
# If you want to use multiple rate limits per controller, you need to give each of
33+
# them and explicit name via the `name:` option.
34+
#
3235
# Examples:
3336
#
3437
# class SessionsController < ApplicationController
@@ -44,14 +47,19 @@ module ClassMethods
4447
# RATE_LIMIT_STORE = ActiveSupport::Cache::RedisCacheStore.new(url: ENV["REDIS_URL"])
4548
# rate_limit to: 10, within: 3.minutes, store: RATE_LIMIT_STORE
4649
# end
47-
def rate_limit(to:, within:, by: -> { request.remote_ip }, with: -> { head :too_many_requests }, store: cache_store, **options)
48-
before_action -> { rate_limiting(to: to, within: within, by: by, with: with, store: store) }, **options
50+
#
51+
# class SessionsController < ApplicationController
52+
# rate_limit to: 3, within: 2.seconds, name: "short-term"
53+
# rate_limit to: 10, within: 5.minutes, name: "long-term"
54+
# end
55+
def rate_limit(to:, within:, by: -> { request.remote_ip }, with: -> { head :too_many_requests }, store: cache_store, name: controller_path, **options)
56+
before_action -> { rate_limiting(to: to, within: within, by: by, with: with, store: store, name: name) }, **options
4957
end
5058
end
5159

5260
private
53-
def rate_limiting(to:, within:, by:, with:, store:)
54-
count = store.increment("rate-limit:#{controller_path}:#{instance_exec(&by)}", 1, expires_in: within)
61+
def rate_limiting(to:, within:, by:, with:, store:, name:)
62+
count = store.increment("rate-limit:#{name}:#{instance_exec(&by)}", 1, expires_in: within)
5563
if count && count > to
5664
ActiveSupport::Notifications.instrument("rate_limit.action_controller", request: request) do
5765
instance_exec(&with)

actionpack/test/controller/rate_limiting_test.rb

Lines changed: 27 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,10 @@
44

55
class RateLimitedController < ActionController::Base
66
self.cache_store = ActiveSupport::Cache::MemoryStore.new
7-
rate_limit to: 2, within: 2.seconds, only: :limited_to_two
7+
rate_limit to: 2, within: 2.seconds, only: :limited
8+
rate_limit to: 5, within: 1.minute, name: "long-term", only: :limited
89

9-
def limited_to_two
10+
def limited
1011
head :ok
1112
end
1213

@@ -24,21 +25,39 @@ class RateLimitingTest < ActionController::TestCase
2425
end
2526

2627
test "exceeding basic limit" do
27-
get :limited_to_two
28-
get :limited_to_two
28+
get :limited
29+
get :limited
2930
assert_response :ok
3031

31-
get :limited_to_two
32+
get :limited
3233
assert_response :too_many_requests
3334
end
3435

36+
test "multiple rate limits" do
37+
get :limited
38+
get :limited
39+
assert_response :ok
40+
41+
travel_to 3.seconds.from_now do
42+
get :limited
43+
get :limited
44+
assert_response :ok
45+
end
46+
47+
travel_to 3.seconds.from_now do
48+
get :limited
49+
get :limited
50+
assert_response :too_many_requests
51+
end
52+
end
53+
3554
test "limit resets after time" do
36-
get :limited_to_two
37-
get :limited_to_two
55+
get :limited
56+
get :limited
3857
assert_response :ok
3958

4059
travel_to Time.now + 3.seconds do
41-
get :limited_to_two
60+
get :limited
4261
assert_response :ok
4362
end
4463
end

0 commit comments

Comments
 (0)