From 739d1d761a94b82e738600852e5a0a0350000140 Mon Sep 17 00:00:00 2001
From: baarkerlounger <5101747+baarkerlounger@users.noreply.github.com>
Date: Thu, 10 Mar 2022 08:52:17 +0000
Subject: [PATCH] CLDC-1016: Fix 2FA bypass (#367)

* Failing test

* Simplest thing to make spec pass

* Extract to method

* Set condition based on class having the 2fa module rather than hardcoding class name
---
 app/controllers/auth/passwords_controller.rb    | 8 ++++++++
 spec/requests/auth/passwords_controller_spec.rb | 4 +++-
 2 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/app/controllers/auth/passwords_controller.rb b/app/controllers/auth/passwords_controller.rb
index ec709c653..acbe63e07 100644
--- a/app/controllers/auth/passwords_controller.rb
+++ b/app/controllers/auth/passwords_controller.rb
@@ -37,6 +37,7 @@ class Auth::PasswordsController < Devise::PasswordsController
         set_flash_message!(:notice, password_update_flash_message)
         resource.after_database_authentication
         sign_in(resource_name, resource)
+        set_2fa_required
       else
         set_flash_message!(:notice, :updated_not_active)
       end
@@ -49,6 +50,13 @@ class Auth::PasswordsController < Devise::PasswordsController
 
 protected
 
+  def set_2fa_required
+    return unless resource.respond_to?(:need_two_factor_authentication?) &&
+      resource.need_two_factor_authentication?(request)
+
+    warden.session(resource_class.name.underscore)[TwoFactorAuthentication::NEED_AUTHENTICATION] = true
+  end
+
   def password_update_flash_message
     resource_class == AdminUser ? :updated_2FA : :updated
   end
diff --git a/spec/requests/auth/passwords_controller_spec.rb b/spec/requests/auth/passwords_controller_spec.rb
index c84cc850e..4dd4dccdd 100644
--- a/spec/requests/auth/passwords_controller_spec.rb
+++ b/spec/requests/auth/passwords_controller_spec.rb
@@ -130,9 +130,11 @@ RSpec.describe Auth::PasswordsController, type: :request do
           }.to change(admin_user, :encrypted_password)
         end
 
-        it "sends you to the 2FA page" do
+        it "sends you to the 2FA page and does not allow bypassing 2FA code" do
           put "/admin/password", headers: headers, params: params
           expect(response).to redirect_to("/admin/two-factor-authentication")
+          get "/admin/case_logs", headers: headers
+          expect(response).to redirect_to("/admin/two-factor-authentication")
         end
 
         it "triggers an SMS" do