From 8ec75613dcb05a931db366b5600c3df8f9654b15 Mon Sep 17 00:00:00 2001
From: rr-bw <102181210+rr-bw@users.noreply.github.com>
Date: Thu, 12 Dec 2024 13:10:54 -0800
Subject: [PATCH] fix(passkeys): [PM-13932] Fix passkey flow incorrect routing
 (#12363)

This PR fixes a bug in the LockComponent refresh that affected the setup/save and use passkey flows. The user was wrongly directly to the /vault after unlock instead of to /fido2 (the passkey screen).

Feature Flag: ExtensionRefresh ON
---
 apps/browser/src/popup/app-routing.module.ts  |  9 ++-
 .../unauth-ui-refresh-redirect.spec.ts        | 13 ++--
 .../functions/unauth-ui-refresh-redirect.ts   |  2 +-
 .../utils/extension-refresh-redirect.spec.ts  | 62 +++++++++++++++++++
 .../src/utils/extension-refresh-redirect.ts   |  2 +-
 libs/auth/src/angular/lock/lock.component.ts  |  6 ++
 6 files changed, 85 insertions(+), 9 deletions(-)
 create mode 100644 libs/angular/src/utils/extension-refresh-redirect.spec.ts

diff --git a/apps/browser/src/popup/app-routing.module.ts b/apps/browser/src/popup/app-routing.module.ts
index 2893647f1a..38071a9e5c 100644
--- a/apps/browser/src/popup/app-routing.module.ts
+++ b/apps/browser/src/popup/app-routing.module.ts
@@ -659,7 +659,14 @@ const routes: Routes = [
           },
           showReadonlyHostname: true,
           showAcctSwitcher: true,
-        } satisfies ExtensionAnonLayoutWrapperData,
+          elevation: 1,
+          /**
+           * This ensures that in a passkey flow the `/fido2?<queryParams>` URL does not get
+           * overwritten in the `BrowserRouterService` by the `/lockV2` route. This way, after
+           * unlocking, the user can be redirected back to the `/fido2?<queryParams>` URL.
+           */
+          doNotSaveUrl: true,
+        } satisfies ExtensionAnonLayoutWrapperData & RouteDataProperties,
         children: [
           {
             path: "",
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 6a19f1ace7..887f528d54 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
@@ -40,15 +40,14 @@ describe("unauthUiRefreshRedirect", () => {
   it("returns UrlTree when UnauthenticatedExtensionUIRefresh flag is enabled and preserves query params", async () => {
     configService.getFeatureFlag.mockResolvedValue(true);
 
-    const queryParams = { test: "test" };
+    const urlTree = new UrlTree();
+    urlTree.queryParams = { test: "test" };
 
     const navigation: Navigation = {
-      extras: {
-        queryParams: queryParams,
-      },
+      extras: {},
       id: 0,
       initialUrl: new UrlTree(),
-      extractedUrl: new UrlTree(),
+      extractedUrl: urlTree,
       trigger: "imperative",
       previousNavigation: undefined,
     };
@@ -60,6 +59,8 @@ describe("unauthUiRefreshRedirect", () => {
     expect(configService.getFeatureFlag).toHaveBeenCalledWith(
       FeatureFlag.UnauthenticatedExtensionUIRefresh,
     );
-    expect(router.createUrlTree).toHaveBeenCalledWith(["/redirect"], { queryParams });
+    expect(router.createUrlTree).toHaveBeenCalledWith(["/redirect"], {
+      queryParams: urlTree.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 a54bad1147..2cb53d5324 100644
--- a/libs/angular/src/auth/functions/unauth-ui-refresh-redirect.ts
+++ b/libs/angular/src/auth/functions/unauth-ui-refresh-redirect.ts
@@ -17,7 +17,7 @@ export function unauthUiRefreshRedirect(redirectUrl: string): () => Promise<bool
     );
     if (shouldRedirect) {
       const currentNavigation = router.getCurrentNavigation();
-      const queryParams = currentNavigation?.extras?.queryParams || {};
+      const queryParams = currentNavigation?.extractedUrl?.queryParams || {};
 
       // Preserve query params when redirecting as it is likely that the refreshed component
       // will be consuming the same query params.
diff --git a/libs/angular/src/utils/extension-refresh-redirect.spec.ts b/libs/angular/src/utils/extension-refresh-redirect.spec.ts
new file mode 100644
index 0000000000..3291a4496f
--- /dev/null
+++ b/libs/angular/src/utils/extension-refresh-redirect.spec.ts
@@ -0,0 +1,62 @@
+import { TestBed } from "@angular/core/testing";
+import { Navigation, Router, UrlTree } from "@angular/router";
+import { mock, MockProxy } from "jest-mock-extended";
+
+import { FeatureFlag } from "@bitwarden/common/enums/feature-flag.enum";
+import { ConfigService } from "@bitwarden/common/platform/abstractions/config/config.service";
+
+import { extensionRefreshRedirect } from "./extension-refresh-redirect";
+
+describe("extensionRefreshRedirect", () => {
+  let configService: MockProxy<ConfigService>;
+  let router: MockProxy<Router>;
+
+  beforeEach(() => {
+    configService = mock<ConfigService>();
+    router = mock<Router>();
+
+    TestBed.configureTestingModule({
+      providers: [
+        { provide: ConfigService, useValue: configService },
+        { provide: Router, useValue: router },
+      ],
+    });
+  });
+
+  it("returns true when ExtensionRefresh flag is disabled", async () => {
+    configService.getFeatureFlag.mockResolvedValue(false);
+
+    const result = await TestBed.runInInjectionContext(() =>
+      extensionRefreshRedirect("/redirect")(),
+    );
+
+    expect(result).toBe(true);
+    expect(configService.getFeatureFlag).toHaveBeenCalledWith(FeatureFlag.ExtensionRefresh);
+    expect(router.parseUrl).not.toHaveBeenCalled();
+  });
+
+  it("returns UrlTree when ExtensionRefresh flag is enabled and preserves query params", async () => {
+    configService.getFeatureFlag.mockResolvedValue(true);
+
+    const urlTree = new UrlTree();
+    urlTree.queryParams = { test: "test" };
+
+    const navigation: Navigation = {
+      extras: {},
+      id: 0,
+      initialUrl: new UrlTree(),
+      extractedUrl: urlTree,
+      trigger: "imperative",
+      previousNavigation: undefined,
+    };
+
+    router.getCurrentNavigation.mockReturnValue(navigation);
+
+    await TestBed.runInInjectionContext(() => extensionRefreshRedirect("/redirect")());
+
+    expect(configService.getFeatureFlag).toHaveBeenCalledWith(FeatureFlag.ExtensionRefresh);
+    expect(router.createUrlTree).toHaveBeenCalledWith(["/redirect"], {
+      queryParams: urlTree.queryParams,
+    });
+  });
+});
diff --git a/libs/angular/src/utils/extension-refresh-redirect.ts b/libs/angular/src/utils/extension-refresh-redirect.ts
index 81c50ceca1..2baa3a3ec8 100644
--- a/libs/angular/src/utils/extension-refresh-redirect.ts
+++ b/libs/angular/src/utils/extension-refresh-redirect.ts
@@ -16,7 +16,7 @@ export function extensionRefreshRedirect(redirectUrl: string): () => Promise<boo
     const shouldRedirect = await configService.getFeatureFlag(FeatureFlag.ExtensionRefresh);
     if (shouldRedirect) {
       const currentNavigation = router.getCurrentNavigation();
-      const queryParams = currentNavigation?.extras?.queryParams || {};
+      const queryParams = currentNavigation?.extractedUrl?.queryParams || {};
 
       // Preserve query params when redirecting as it is likely that the refreshed component
       // will be consuming the same query params.
diff --git a/libs/auth/src/angular/lock/lock.component.ts b/libs/auth/src/angular/lock/lock.component.ts
index fb29bb731d..14a5553577 100644
--- a/libs/auth/src/angular/lock/lock.component.ts
+++ b/libs/auth/src/angular/lock/lock.component.ts
@@ -534,8 +534,14 @@ export class LockV2Component implements OnInit, OnDestroy {
 
     if (this.clientType === "browser") {
       const previousUrl = this.lockComponentService.getPreviousUrl();
+      /**
+       * In a passkey flow, the `previousUrl` will still be `/fido2?<queryParams>` at this point
+       * because the `/lockV2` route doesn't save the URL in the `BrowserRouterService`. This is
+       * handled by the `doNotSaveUrl` property on the `lockV2` route in `app-routing.module.ts`.
+       */
       if (previousUrl) {
         await this.router.navigateByUrl(previousUrl);
+        return;
       }
     }