From fa537638bfe7b9ece1a29522401d90ccc7afc89c Mon Sep 17 00:00:00 2001
From: Jared Snider <116684653+JaredSnider-Bitwarden@users.noreply.github.com>
Date: Fri, 25 Oct 2024 11:29:48 -0400
Subject: [PATCH] Auth/PM-13945 - Extension - Fix TDE User with MP can't
 navigate to Lock Screen (#11700)

* PM-13945 - Extension Refresh redirects should persist query params as we use query params to execute guard logic (e.g., lockGuard). The loss of the from: login-initiated query param prevented navigation to the lock screen.

* PM-13945 - Test updated
---
 .../unauth-ui-refresh-redirect.spec.ts        | 28 +++++++++++++------
 .../functions/unauth-ui-refresh-redirect.ts   |  7 ++++-
 .../src/utils/extension-refresh-redirect.ts   |  8 +++++-
 3 files changed, 32 insertions(+), 11 deletions(-)

diff --git a/libs/angular/src/auth/functions/unauth-ui-refresh-redirect.spec.ts b/libs/angular/src/auth/functions/unauth-ui-refresh-redirect.spec.ts
index 5fc904a5b9..6a19f1ace7 100644
--- a/libs/angular/src/auth/functions/unauth-ui-refresh-redirect.spec.ts
+++ b/libs/angular/src/auth/functions/unauth-ui-refresh-redirect.spec.ts
@@ -1,5 +1,5 @@
 import { TestBed } from "@angular/core/testing";
-import { Router, UrlTree } from "@angular/router";
+import { Navigation, Router, UrlTree } from "@angular/router";
 import { mock, MockProxy } from "jest-mock-extended";
 
 import { FeatureFlag } from "@bitwarden/common/enums/feature-flag.enum";
@@ -37,19 +37,29 @@ describe("unauthUiRefreshRedirect", () => {
     expect(router.parseUrl).not.toHaveBeenCalled();
   });
 
-  it("returns UrlTree when UnauthenticatedExtensionUIRefresh flag is enabled", async () => {
-    const mockUrlTree = mock<UrlTree>();
+  it("returns UrlTree when UnauthenticatedExtensionUIRefresh flag is enabled and preserves query params", async () => {
     configService.getFeatureFlag.mockResolvedValue(true);
-    router.parseUrl.mockReturnValue(mockUrlTree);
 
-    const result = await TestBed.runInInjectionContext(() =>
-      unauthUiRefreshRedirect("/redirect")(),
-    );
+    const queryParams = { test: "test" };
+
+    const navigation: Navigation = {
+      extras: {
+        queryParams: queryParams,
+      },
+      id: 0,
+      initialUrl: new UrlTree(),
+      extractedUrl: new UrlTree(),
+      trigger: "imperative",
+      previousNavigation: undefined,
+    };
+
+    router.getCurrentNavigation.mockReturnValue(navigation);
+
+    await TestBed.runInInjectionContext(() => unauthUiRefreshRedirect("/redirect")());
 
-    expect(result).toBe(mockUrlTree);
     expect(configService.getFeatureFlag).toHaveBeenCalledWith(
       FeatureFlag.UnauthenticatedExtensionUIRefresh,
     );
-    expect(router.parseUrl).toHaveBeenCalledWith("/redirect");
+    expect(router.createUrlTree).toHaveBeenCalledWith(["/redirect"], { queryParams });
   });
 });
diff --git a/libs/angular/src/auth/functions/unauth-ui-refresh-redirect.ts b/libs/angular/src/auth/functions/unauth-ui-refresh-redirect.ts
index 22ed23273b..a54bad1147 100644
--- a/libs/angular/src/auth/functions/unauth-ui-refresh-redirect.ts
+++ b/libs/angular/src/auth/functions/unauth-ui-refresh-redirect.ts
@@ -16,7 +16,12 @@ export function unauthUiRefreshRedirect(redirectUrl: string): () => Promise<bool
       FeatureFlag.UnauthenticatedExtensionUIRefresh,
     );
     if (shouldRedirect) {
-      return router.parseUrl(redirectUrl);
+      const currentNavigation = router.getCurrentNavigation();
+      const queryParams = currentNavigation?.extras?.queryParams || {};
+
+      // Preserve query params when redirecting as it is likely that the refreshed component
+      // will be consuming the same query params.
+      return router.createUrlTree([redirectUrl], { queryParams });
     } else {
       return true;
     }
diff --git a/libs/angular/src/utils/extension-refresh-redirect.ts b/libs/angular/src/utils/extension-refresh-redirect.ts
index f7399d9c27..81c50ceca1 100644
--- a/libs/angular/src/utils/extension-refresh-redirect.ts
+++ b/libs/angular/src/utils/extension-refresh-redirect.ts
@@ -12,9 +12,15 @@ export function extensionRefreshRedirect(redirectUrl: string): () => Promise<boo
   return async () => {
     const configService = inject(ConfigService);
     const router = inject(Router);
+
     const shouldRedirect = await configService.getFeatureFlag(FeatureFlag.ExtensionRefresh);
     if (shouldRedirect) {
-      return router.parseUrl(redirectUrl);
+      const currentNavigation = router.getCurrentNavigation();
+      const queryParams = currentNavigation?.extras?.queryParams || {};
+
+      // Preserve query params when redirecting as it is likely that the refreshed component
+      // will be consuming the same query params.
+      return router.createUrlTree([redirectUrl], { queryParams });
     } else {
       return true;
     }