From 3593ee2e36284de71be8dc74c1772de7a7e1a7e3 Mon Sep 17 00:00:00 2001
From: Claire <claire.github-309c@sitedethib.com>
Date: Fri, 19 Jan 2024 13:19:49 +0100
Subject: [PATCH] Add rate-limit of TOTP authentication attempts at controller
 level (#28801)

---
 app/controllers/auth/sessions_controller.rb   | 22 +++++++++++++++++++
 .../auth/two_factor_authentication_concern.rb |  5 +++++
 config/locales/en.yml                         |  1 +
 .../auth/sessions_controller_spec.rb          | 20 +++++++++++++++++
 4 files changed, 48 insertions(+)

diff --git a/app/controllers/auth/sessions_controller.rb b/app/controllers/auth/sessions_controller.rb
index 148ad5375..6bc48a780 100644
--- a/app/controllers/auth/sessions_controller.rb
+++ b/app/controllers/auth/sessions_controller.rb
@@ -1,6 +1,10 @@
 # frozen_string_literal: true
 
 class Auth::SessionsController < Devise::SessionsController
+  include Redisable
+
+  MAX_2FA_ATTEMPTS_PER_HOUR = 10
+
   layout 'auth'
 
   skip_before_action :check_self_destruct!
@@ -130,9 +134,23 @@ class Auth::SessionsController < Devise::SessionsController
     session.delete(:attempt_user_updated_at)
   end
 
+  def clear_2fa_attempt_from_user(user)
+    redis.del(second_factor_attempts_key(user))
+  end
+
+  def check_second_factor_rate_limits(user)
+    attempts, = redis.multi do |multi|
+      multi.incr(second_factor_attempts_key(user))
+      multi.expire(second_factor_attempts_key(user), 1.hour)
+    end
+
+    attempts >= MAX_2FA_ATTEMPTS_PER_HOUR
+  end
+
   def on_authentication_success(user, security_measure)
     @on_authentication_success_called = true
 
+    clear_2fa_attempt_from_user(user)
     clear_attempt_from_session
 
     user.update_sign_in!(new_sign_in: true)
@@ -164,4 +182,8 @@ class Auth::SessionsController < Devise::SessionsController
       user_agent: request.user_agent
     )
   end
+
+  def second_factor_attempts_key(user)
+    "2fa_auth_attempts:#{user.id}:#{Time.now.utc.hour}"
+  end
 end
diff --git a/app/controllers/concerns/auth/two_factor_authentication_concern.rb b/app/controllers/concerns/auth/two_factor_authentication_concern.rb
index effdb8d21..404164751 100644
--- a/app/controllers/concerns/auth/two_factor_authentication_concern.rb
+++ b/app/controllers/concerns/auth/two_factor_authentication_concern.rb
@@ -66,6 +66,11 @@ module Auth::TwoFactorAuthenticationConcern
   end
 
   def authenticate_with_two_factor_via_otp(user)
+    if check_second_factor_rate_limits(user)
+      flash.now[:alert] = I18n.t('users.rate_limited')
+      return prompt_for_two_factor(user)
+    end
+
     if valid_otp_attempt?(user)
       on_authentication_success(user, :otp)
     else
diff --git a/config/locales/en.yml b/config/locales/en.yml
index 78820c3b5..89ca0ad72 100644
--- a/config/locales/en.yml
+++ b/config/locales/en.yml
@@ -1844,6 +1844,7 @@ en:
     go_to_sso_account_settings: Go to your identity provider's account settings
     invalid_otp_token: Invalid two-factor code
     otp_lost_help_html: If you lost access to both, you may get in touch with %{email}
+    rate_limited: Too many authentication attempts, try again later.
     seamless_external_login: You are logged in via an external service, so password and e-mail settings are not available.
     signed_in_as: 'Signed in as:'
   verification:
diff --git a/spec/controllers/auth/sessions_controller_spec.rb b/spec/controllers/auth/sessions_controller_spec.rb
index e3f2b278b..d238626c9 100644
--- a/spec/controllers/auth/sessions_controller_spec.rb
+++ b/spec/controllers/auth/sessions_controller_spec.rb
@@ -262,6 +262,26 @@ RSpec.describe Auth::SessionsController do
           end
         end
 
+        context 'when repeatedly using an invalid TOTP code before using a valid code' do
+          before do
+            stub_const('Auth::SessionsController::MAX_2FA_ATTEMPTS_PER_HOUR', 2)
+          end
+
+          it 'does not log the user in' do
+            # Travel to the beginning of an hour to avoid crossing rate-limit buckets
+            travel_to '2023-12-20T10:00:00Z'
+
+            Auth::SessionsController::MAX_2FA_ATTEMPTS_PER_HOUR.times do
+              post :create, params: { user: { otp_attempt: '1234' } }, session: { attempt_user_id: user.id, attempt_user_updated_at: user.updated_at.to_s }
+              expect(controller.current_user).to be_nil
+            end
+
+            post :create, params: { user: { otp_attempt: user.current_otp } }, session: { attempt_user_id: user.id, attempt_user_updated_at: user.updated_at.to_s }
+            expect(controller.current_user).to be_nil
+            expect(flash[:alert]).to match I18n.t('users.rate_limited')
+          end
+        end
+
         context 'when using a valid OTP' do
           before do
             post :create, params: { user: { otp_attempt: user.current_otp } }, session: { attempt_user_id: user.id, attempt_user_updated_at: user.updated_at.to_s }