From 3f4892fcc83db133a2e661d1cefb5090362efb67 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Bispo?= <abispo@bitwarden.com>
Date: Fri, 11 Aug 2023 17:32:51 +0100
Subject: [PATCH] [PM-3381] Fix TDE login 2FA flow (#2678)

* [PM-3381] Check for vault lock on 2FA screen

* [PM-3381] Move logic to ViewModel

* [PM-3381] Fix null vm error
---
 src/App/Pages/Accounts/TwoFactorPage.xaml.cs  | 29 ++++++++++---------
 .../Pages/Accounts/TwoFactorPageViewModel.cs  | 21 ++++++++++++--
 2 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/src/App/Pages/Accounts/TwoFactorPage.xaml.cs b/src/App/Pages/Accounts/TwoFactorPage.xaml.cs
index e3c1ec755..c999fe259 100644
--- a/src/App/Pages/Accounts/TwoFactorPage.xaml.cs
+++ b/src/App/Pages/Accounts/TwoFactorPage.xaml.cs
@@ -4,6 +4,7 @@ using Bit.App.Controls;
 using Bit.App.Models;
 using Bit.App.Utilities;
 using Bit.Core.Abstractions;
+using Bit.Core.Services;
 using Bit.Core.Utilities;
 using Xamarin.Forms;
 
@@ -24,17 +25,19 @@ namespace Bit.App.Pages
         {
             InitializeComponent();
             SetActivityIndicator();
-            _authingWithSso = authingWithSso ?? false;
             _appOptions = appOptions;
             _orgIdentifier = orgIdentifier;
             _broadcasterService = ServiceContainer.Resolve<IBroadcasterService>("broadcasterService");
             _messagingService = ServiceContainer.Resolve<IMessagingService>("messagingService");
             _vm = BindingContext as TwoFactorPageViewModel;
             _vm.Page = this;
+            _vm.AuthingWithSso = authingWithSso ?? false;
             _vm.StartSetPasswordAction = () =>
                 Device.BeginInvokeOnMainThread(async () => await StartSetPasswordAsync());
             _vm.TwoFactorAuthSuccessAction = () =>
-                Device.BeginInvokeOnMainThread(async () => await TwoFactorAuthSuccessAsync());
+                Device.BeginInvokeOnMainThread(async () => await TwoFactorAuthSuccessToMainAsync());
+            _vm.LockAction = () =>
+                Device.BeginInvokeOnMainThread(TwoFactorAuthSuccessWithSSOLocked);
             _vm.UpdateTempPasswordAction =
                 () => Device.BeginInvokeOnMainThread(async () => await UpdateTempPasswordAsync());
             _vm.StartDeviceApprovalOptionsAction =
@@ -188,21 +191,19 @@ namespace Bit.App.Pages
             await Navigation.PushModalAsync(new NavigationPage(page));
         }
 
-        private async Task TwoFactorAuthSuccessAsync()
+        private void TwoFactorAuthSuccessWithSSOLocked()
         {
-            if (_authingWithSso)
+            Application.Current.MainPage = new NavigationPage(new LockPage(_appOptions));
+        }
+
+        private async Task TwoFactorAuthSuccessToMainAsync()
+        {
+            if (AppHelpers.SetAlternateMainPage(_appOptions))
             {
-                Application.Current.MainPage = new NavigationPage(new LockPage(_appOptions));
-            }
-            else
-            {
-                if (AppHelpers.SetAlternateMainPage(_appOptions))
-                {
-                    return;
-                }
-                var previousPage = await AppHelpers.ClearPreviousPage();
-                Application.Current.MainPage = new TabsPage(_appOptions, previousPage);
+                return;
             }
+            var previousPage = await AppHelpers.ClearPreviousPage();
+            Application.Current.MainPage = new TabsPage(_appOptions, previousPage);
         }
 
         private void Token_TextChanged(object sender, TextChangedEventArgs e)
diff --git a/src/App/Pages/Accounts/TwoFactorPageViewModel.cs b/src/App/Pages/Accounts/TwoFactorPageViewModel.cs
index 54bd90983..f54123012 100644
--- a/src/App/Pages/Accounts/TwoFactorPageViewModel.cs
+++ b/src/App/Pages/Accounts/TwoFactorPageViewModel.cs
@@ -33,6 +33,7 @@ namespace Bit.App.Pages
         private readonly IStateService _stateService;
         private readonly II18nService _i18nService;
         private readonly IAppIdService _appIdService;
+        private readonly IVaultTimeoutService _vaultTimeoutService;
         private readonly ILogger _logger;
         private readonly IDeviceTrustCryptoService _deviceTrustCryptoService;
         private TwoFactorProviderType? _selectedProviderType;
@@ -55,6 +56,7 @@ namespace Bit.App.Pages
             _stateService = ServiceContainer.Resolve<IStateService>("stateService");
             _i18nService = ServiceContainer.Resolve<II18nService>("i18nService");
             _appIdService = ServiceContainer.Resolve<IAppIdService>("appIdService");
+            _vaultTimeoutService = ServiceContainer.Resolve<IVaultTimeoutService>();
             _logger = ServiceContainer.Resolve<ILogger>();
             _deviceTrustCryptoService = ServiceContainer.Resolve<IDeviceTrustCryptoService>();
 
@@ -71,6 +73,8 @@ namespace Bit.App.Pages
 
         public bool Remember { get; set; }
 
+        public bool AuthingWithSso { get; set; }
+
         public string Token { get; set; }
 
         public bool DuoMethod => SelectedProviderType == TwoFactorProviderType.Duo ||
@@ -120,6 +124,7 @@ namespace Bit.App.Pages
         public Command SubmitCommand { get; }
         public ICommand MoreCommand { get; }
         public Action TwoFactorAuthSuccessAction { get; set; }
+        public Action LockAction { get; set; }
         public Action StartDeviceApprovalOptionsAction { get; set; }
         public Action StartSetPasswordAction { get; set; }
         public Action CloseAction { get; set; }
@@ -344,7 +349,7 @@ namespace Bit.App.Pages
                     }
                     else if (await _deviceTrustCryptoService.IsDeviceTrustedAsync())
                     {
-                        TwoFactorAuthSuccessAction?.Invoke();
+                        await TwoFactorAuthSuccessAsync();
                     }
                     else
                     {
@@ -353,7 +358,7 @@ namespace Bit.App.Pages
                 }
                 else
                 {
-                    TwoFactorAuthSuccessAction?.Invoke();
+                    await TwoFactorAuthSuccessAsync();
                 }
             }
             catch (ApiException e)
@@ -447,5 +452,17 @@ namespace Bit.App.Pages
                 return false;
             }
         }
+
+        public async Task TwoFactorAuthSuccessAsync()
+        {
+            if (AuthingWithSso && await _vaultTimeoutService.IsLockedAsync())
+            {
+                LockAction?.Invoke();
+            }
+            else
+            {
+                TwoFactorAuthSuccessAction?.Invoke();
+            }
+        }
     }
 }