From 37414ade28da0185bf650c3f6793bda209739b65 Mon Sep 17 00:00:00 2001
From: Irfan Omur <irfan.omur@automattic.com>
Date: Thu, 1 Feb 2024 02:33:05 +0300
Subject: [PATCH] Fix Smart Lock Login with 2FA issue

---
 .../login/LoginUsernamePasswordFragment.java  | 27 ++++++++++++++-----
 1 file changed, 20 insertions(+), 7 deletions(-)

diff --git a/WordPressLoginFlow/src/main/java/org/wordpress/android/login/LoginUsernamePasswordFragment.java b/WordPressLoginFlow/src/main/java/org/wordpress/android/login/LoginUsernamePasswordFragment.java
index f8012864bf..32ba26c091 100644
--- a/WordPressLoginFlow/src/main/java/org/wordpress/android/login/LoginUsernamePasswordFragment.java
+++ b/WordPressLoginFlow/src/main/java/org/wordpress/android/login/LoginUsernamePasswordFragment.java
@@ -30,6 +30,7 @@
 import org.wordpress.android.fluxc.store.AccountStore.AuthenticatePayload;
 import org.wordpress.android.fluxc.store.AccountStore.AuthenticationErrorType;
 import org.wordpress.android.fluxc.store.AccountStore.OnAuthenticationChanged;
+import org.wordpress.android.fluxc.store.AccountStore.OnTwoFactorAuthStarted;
 import org.wordpress.android.fluxc.store.SiteStore.OnProfileFetched;
 import org.wordpress.android.fluxc.store.SiteStore.OnSiteChanged;
 import org.wordpress.android.fluxc.store.SiteStore.RefreshSitesXMLRPCPayload;
@@ -504,13 +505,7 @@ private void handleAuthError(AuthenticationErrorType error, XmlRpcErrorType xmlR
             case INVALID_TOKEN:
             case AUTHORIZATION_REQUIRED:
             case NEEDS_2FA:
-                if (mIsWpcom) {
-                    if (mLoginListener != null) {
-                        mLoginListener.needs2fa(mRequestedUsername, mRequestedPassword);
-                    }
-                } else {
-                    showError("2FA not supported for self-hosted sites. Please use an app-password.");
-                }
+                handle2fa();
                 break;
             default:
                 AppLog.e(T.NUX, "Server response: " + errorMessage);
@@ -521,8 +516,26 @@ private void handleAuthError(AuthenticationErrorType error, XmlRpcErrorType xmlR
         }
     }
 
+    private void handle2fa() {
+        if (mIsWpcom) {
+            if (mLoginListener != null) {
+                mLoginListener.needs2fa(mRequestedUsername, mRequestedPassword);
+            }
+        } else {
+            showError("2FA not supported for self-hosted sites. Please use an app-password.");
+        }
+    }
+
     // OnChanged events
 
+    @SuppressWarnings("unused")
+    @Subscribe(threadMode = ThreadMode.MAIN)
+    public void onTwoFactorAuthStarted(OnTwoFactorAuthStarted event) {
+        mLoginStarted = false;
+        handle2fa();
+        endProgress();
+    }
+
     @SuppressWarnings("unused")
     @Subscribe(threadMode = ThreadMode.MAIN)
     public void onAuthenticationChanged(OnAuthenticationChanged event) {