From 1247463e299641368ffd2cde21c3826e066b6964 Mon Sep 17 00:00:00 2001
From: Jared Snider <116684653+JaredSnider-Bitwarden@users.noreply.github.com>
Date: Mon, 29 May 2023 13:55:37 -0400
Subject: [PATCH] PM-2403 - Persist SSO flag between 2FA component and 2FA
 Options component and back so that the correct onSuccessful login logic can
 run which closes the tab extension and sidesteps Safari master password
 invalid issues due to null KDF config / iterations again. Tested on Chrome,
 Firefox + sidebar, Edge, Opera + sidebar, and Safari (#5535)

---
 .../popup/two-factor-options.component.html   |  2 +-
 .../popup/two-factor-options.component.ts     | 25 ++++++++++++++++---
 .../src/auth/popup/two-factor.component.ts    | 10 +++++++-
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/apps/browser/src/auth/popup/two-factor-options.component.html b/apps/browser/src/auth/popup/two-factor-options.component.html
index 3fc510bb27..f25944aba6 100644
--- a/apps/browser/src/auth/popup/two-factor-options.component.html
+++ b/apps/browser/src/auth/popup/two-factor-options.component.html
@@ -1,6 +1,6 @@
 <header>
   <div class="left">
-    <button type="button" routerLink="/2fa">{{ "close" | i18n }}</button>
+    <button type="button" (click)="close()">{{ "close" | i18n }}</button>
   </div>
   <h1 class="center">
     <span class="title">{{ "twoStepOptions" | i18n }}</span>
diff --git a/apps/browser/src/auth/popup/two-factor-options.component.ts b/apps/browser/src/auth/popup/two-factor-options.component.ts
index 6aa8109f05..b0aced9dd9 100644
--- a/apps/browser/src/auth/popup/two-factor-options.component.ts
+++ b/apps/browser/src/auth/popup/two-factor-options.component.ts
@@ -1,5 +1,5 @@
 import { Component } from "@angular/core";
-import { Router } from "@angular/router";
+import { ActivatedRoute, Router } from "@angular/router";
 
 import { TwoFactorOptionsComponent as BaseTwoFactorOptionsComponent } from "@bitwarden/angular/auth/components/two-factor-options.component";
 import { I18nService } from "@bitwarden/common/abstractions/i18n.service";
@@ -15,14 +15,33 @@ export class TwoFactorOptionsComponent extends BaseTwoFactorOptionsComponent {
     twoFactorService: TwoFactorService,
     router: Router,
     i18nService: I18nService,
-    platformUtilsService: PlatformUtilsService
+    platformUtilsService: PlatformUtilsService,
+    private activatedRoute: ActivatedRoute
   ) {
     super(twoFactorService, router, i18nService, platformUtilsService, window);
   }
 
+  close() {
+    this.navigateTo2FA();
+  }
+
   choose(p: any) {
     super.choose(p);
     this.twoFactorService.setSelectedProvider(p.type);
-    this.router.navigate(["2fa"]);
+
+    this.navigateTo2FA();
+  }
+
+  navigateTo2FA() {
+    const sso = this.activatedRoute.snapshot.queryParamMap.get("sso") === "true";
+
+    if (sso) {
+      // Persist SSO flag back to the 2FA comp if it exists
+      // in order for successful login logic to work properly for
+      // SSO + 2FA in browser extension
+      this.router.navigate(["2fa"], { queryParams: { sso: true } });
+    } else {
+      this.router.navigate(["2fa"]);
+    }
   }
 }
diff --git a/apps/browser/src/auth/popup/two-factor.component.ts b/apps/browser/src/auth/popup/two-factor.component.ts
index 99a850d3c0..e6652e7ba3 100644
--- a/apps/browser/src/auth/popup/two-factor.component.ts
+++ b/apps/browser/src/auth/popup/two-factor.component.ts
@@ -147,7 +147,15 @@ export class TwoFactorComponent extends BaseTwoFactorComponent {
   }
 
   anotherMethod() {
-    this.router.navigate(["2fa-options"]);
+    const sso = this.route.snapshot.queryParamMap.get("sso") === "true";
+
+    if (sso) {
+      // We must persist this so when the user returns to the 2FA comp, the
+      // proper onSuccessfulLogin logic is executed.
+      this.router.navigate(["2fa-options"], { queryParams: { sso: true } });
+    } else {
+      this.router.navigate(["2fa-options"]);
+    }
   }
 
   async isLinux() {