Add ability to use multiple rate limits per controller (rails#52960)

jhawthorn · Sep 17, 2024 · c0807dc · c0807dc
1 parent 4aa7811
commit c0807dc
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 12 deletions.
diff --git a/actionpack/lib/action_controller/metal/rate_limiting.rb b/actionpack/lib/action_controller/metal/rate_limiting.rb
@@ -29,6 +29,9 @@ module ClassMethods
       # datastore as your general caches, you can pass a custom store in the `store`
       # parameter.
       #
+      # If you want to use multiple rate limits per controller, you need to give each of
+      # them and explicit name via the `name:` option.
+      #
       # Examples:
       #
       #     class SessionsController < ApplicationController
@@ -44,14 +47,19 @@ module ClassMethods
       #       RATE_LIMIT_STORE = ActiveSupport::Cache::RedisCacheStore.new(url: ENV["REDIS_URL"])
       #       rate_limit to: 10, within: 3.minutes, store: RATE_LIMIT_STORE
       #     end
-      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
+      #
+      #     class SessionsController < ApplicationController
+      #       rate_limit to: 3, within: 2.seconds, name: "short-term"
+      #       rate_limit to: 10, within: 5.minutes, name: "long-term"
+      #     end
+      def rate_limit(to:, within:, by: -> { request.remote_ip }, with: -> { head :too_many_requests }, store: cache_store, name: controller_path, **options)
+        before_action -> { rate_limiting(to: to, within: within, by: by, with: with, store: store, name: name) }, **options
       end
     end
 
     private
-      def rate_limiting(to:, within:, by:, with:, store:)
-        count = store.increment("rate-limit:#{controller_path}:#{instance_exec(&by)}", 1, expires_in: within)
+      def rate_limiting(to:, within:, by:, with:, store:, name:)
+        count = store.increment("rate-limit:#{name}:#{instance_exec(&by)}", 1, expires_in: within)
         if count && count > to
           ActiveSupport::Notifications.instrument("rate_limit.action_controller", request: request) do
             instance_exec(&with)

diff --git a/actionpack/test/controller/rate_limiting_test.rb b/actionpack/test/controller/rate_limiting_test.rb
@@ -4,9 +4,10 @@
 
 class RateLimitedController < ActionController::Base
   self.cache_store = ActiveSupport::Cache::MemoryStore.new
-  rate_limit to: 2, within: 2.seconds, only: :limited_to_two
+  rate_limit to: 2, within: 2.seconds, only: :limited
+  rate_limit to: 5, within: 1.minute, name: "long-term", only: :limited
 
-  def limited_to_two
+  def limited
     head :ok
   end
 
@@ -24,21 +25,39 @@ class RateLimitingTest < ActionController::TestCase
   end
 
   test "exceeding basic limit" do
-    get :limited_to_two
-    get :limited_to_two
+    get :limited
+    get :limited
     assert_response :ok
 
-    get :limited_to_two
+    get :limited
     assert_response :too_many_requests
   end
 
+  test "multiple rate limits" do
+    get :limited
+    get :limited
+    assert_response :ok
+
+    travel_to 3.seconds.from_now do
+      get :limited
+      get :limited
+      assert_response :ok
+    end
+
+    travel_to 3.seconds.from_now do
+      get :limited
+      get :limited
+      assert_response :too_many_requests
+    end
+  end
+
   test "limit resets after time" do
-    get :limited_to_two
-    get :limited_to_two
+    get :limited
+    get :limited
     assert_response :ok
 
     travel_to Time.now + 3.seconds do
-      get :limited_to_two
+      get :limited
       assert_response :ok
     end
   end