fix race conditions, added tests

apps/browser/src/dirt/phishing-detection/services/phishing-detection.service.spec.ts

@@ -1,5 +1,5 @@
 import { mock, MockProxy } from "jest-mock-extended";
-import { Observable, of } from "rxjs";
+import { firstValueFrom, Observable, of, skip, Subject } from "rxjs";
 
 import { AccountService } from "@bitwarden/common/auth/abstractions/account.service";
 import { BillingAccountProfileStateService } from "@bitwarden/common/billing/abstractions";
@@ -8,8 +8,17 @@ import { LogService } from "@bitwarden/common/platform/abstractions/log.service"
 import { FakeGlobalStateProvider } from "@bitwarden/common/spec";
 import { MessageListener } from "@bitwarden/messaging";
 
+import { BrowserApi } from "../../../platform/browser/browser-api";
+
 import { PhishingDataService } from "./phishing-data.service";
-import { PhishingDetectionService } from "./phishing-detection.service";
+import {
+  IGNORED_PHISHING_HOSTNAMES_KEY,
+  PHISHING_DETECTION_CANCEL_COMMAND,
+  PHISHING_DETECTION_CONTINUE_COMMAND,
+  PhishingDetectionService,
+} from "./phishing-detection.service";
+
+jest.mock("../../../platform/browser/browser-api");
 
 describe("PhishingDetectionService", () => {
   let accountService: AccountService;
@@ -19,24 +28,59 @@ describe("PhishingDetectionService", () => {
   let phishingDataService: MockProxy<PhishingDataService>;
   let messageListener: MockProxy<MessageListener>;
   let globalStateProvider: FakeGlobalStateProvider;
+  let continueCommandSubject: Subject<{ tabId: number; url: string }>;
+  let cancelCommandSubject: Subject<{ tabId: number }>;
+
+  let cleanupFn: (() => void) | undefined;
 
   beforeEach(() => {
-    accountService = { getAccount$: jest.fn(() => of(null)) } as any;
-    billingAccountProfileStateService = {} as any;
-    configService = { getFeatureFlag$: jest.fn(() => of(false)) } as any;
+    // Mock a premium account with access to phishing detection
+    const mockAccount = { id: "test-user-id" };
+    accountService = {
+      getAccount$: jest.fn(() => of(mockAccount)),
+      activeAccount$: of(mockAccount),
+    } as any;
+    billingAccountProfileStateService = {
+      hasPremiumFromAnySource$: jest.fn(() => of(true)),
+    } as any;
+    configService = { getFeatureFlag$: jest.fn(() => of(true)) } as any;
     logService = { info: jest.fn(), debug: jest.fn(), warning: jest.fn(), error: jest.fn() } as any;
     phishingDataService = mock();
-    messageListener = mock<MessageListener>({
-      messages$(_commandDefinition) {
-        return new Observable();
-      },
+    phishingDataService.update$ = of(undefined);
+
+    continueCommandSubject = new Subject();
+    cancelCommandSubject = new Subject();
+
+    messageListener = mock<MessageListener>();
+    messageListener.messages$.mockImplementation((commandDefinition: any) => {
+      if (commandDefinition === PHISHING_DETECTION_CONTINUE_COMMAND) {
+        return continueCommandSubject.asObservable();
+      } else if (commandDefinition === PHISHING_DETECTION_CANCEL_COMMAND) {
+        return cancelCommandSubject.asObservable();
+      }
+      return new Observable();
     });
+
     globalStateProvider = new FakeGlobalStateProvider();
+
+    jest.spyOn(BrowserApi, "addListener").mockImplementation(() => {});
+    jest.spyOn(BrowserApi, "removeListener").mockImplementation(() => {});
+    jest.spyOn(BrowserApi, "navigateTabToUrl").mockResolvedValue(undefined);
+    jest.spyOn(BrowserApi, "closeTab").mockResolvedValue(undefined);
+    jest.spyOn(BrowserApi, "getRuntimeURL").mockReturnValue("chrome-extension://test/");
+  });
+
+  afterEach(() => {
+    if (cleanupFn) {
+      cleanupFn();
+      cleanupFn = undefined;
+    }
+    (PhishingDetectionService as any)._didInit = false;
   });
 
   it("should initialize without errors", () => {
     expect(() => {
-      PhishingDetectionService.initialize(
+      cleanupFn = PhishingDetectionService.initialize(
         accountService,
         billingAccountProfileStateService,
         configService,
@@ -48,6 +92,175 @@ describe("PhishingDetectionService", () => {
     }).not.toThrow();
   });
 
+  describe("Persistent Storage", () => {
+    it("loads ignored hostnames from storage on initialization", async () => {
+      const state = globalStateProvider.get(IGNORED_PHISHING_HOSTNAMES_KEY);
+      await state.update(() => ["malicious.com", "phishing.net"]);
+
+      cleanupFn = PhishingDetectionService.initialize(
+        accountService,
+        billingAccountProfileStateService,
+        configService,
+        logService,
+        phishingDataService,
+        messageListener,
+        globalStateProvider,
+      );
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+
+      expect(logService.debug).toHaveBeenCalledWith(
+        expect.stringContaining("Loaded/updated 2 ignored hostnames"),
+      );
+    });
+
+    it("persists ignored hostname when user continues to phishing site", async () => {
+      const state = globalStateProvider.get(IGNORED_PHISHING_HOSTNAMES_KEY);
+
+      cleanupFn = PhishingDetectionService.initialize(
+        accountService,
+        billingAccountProfileStateService,
+        configService,
+        logService,
+        phishingDataService,
+        messageListener,
+        globalStateProvider,
+      );
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+
+      const stateUpdatePromise = firstValueFrom(state.state$.pipe(skip(1)));
+
+      continueCommandSubject.next({
+        tabId: 123,
+        url: "https://malicious.com/login",
+      });
+
+      const storedHostnames = await stateUpdatePromise;
+      expect(storedHostnames).toContain("malicious.com");
+    });
+
+    it("handles multiple ignored hostnames", async () => {
+      const state = globalStateProvider.get(IGNORED_PHISHING_HOSTNAMES_KEY);
+
+      cleanupFn = PhishingDetectionService.initialize(
+        accountService,
+        billingAccountProfileStateService,
+        configService,
+        logService,
+        phishingDataService,
+        messageListener,
+        globalStateProvider,
+      );
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+
+      const firstUpdatePromise = firstValueFrom(state.state$.pipe(skip(1)));
+      continueCommandSubject.next({
+        tabId: 123,
+        url: "https://malicious1.com/page",
+      });
+      await firstUpdatePromise;
+
+      const secondUpdatePromise = firstValueFrom(state.state$.pipe(skip(1)));
+      continueCommandSubject.next({
+        tabId: 124,
+        url: "https://malicious2.com/page",
+      });
+      const storedHostnames = await secondUpdatePromise;
+
+      expect(storedHostnames).toContain("malicious1.com");
+      expect(storedHostnames).toContain("malicious2.com");
+      expect((storedHostnames as string[]).length).toBe(2);
+    });
+
+    it("does not duplicate hostnames when same site is ignored multiple times", async () => {
+      const state = globalStateProvider.get(IGNORED_PHISHING_HOSTNAMES_KEY);
+
+      cleanupFn = PhishingDetectionService.initialize(
+        accountService,
+        billingAccountProfileStateService,
+        configService,
+        logService,
+        phishingDataService,
+        messageListener,
+        globalStateProvider,
+      );
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+
+      const firstUpdatePromise = firstValueFrom(state.state$.pipe(skip(1)));
+      continueCommandSubject.next({
+        tabId: 123,
+        url: "https://malicious.com/page1",
+      });
+      await firstUpdatePromise;
+
+      const secondUpdatePromise = firstValueFrom(state.state$.pipe(skip(1)));
+      continueCommandSubject.next({
+        tabId: 124,
+        url: "https://malicious.com/page2",
+      });
+      const storedHostnames = await secondUpdatePromise;
+
+      const hostnameArray = storedHostnames as string[];
+      const count = hostnameArray.filter((h) => h === "malicious.com").length;
+      expect(count).toBe(1);
+      expect(hostnameArray.length).toBe(1);
+    });
+
+    it("handles storage initialization with empty state", async () => {
+      const state = globalStateProvider.get(IGNORED_PHISHING_HOSTNAMES_KEY);
+
+      expect(() => {
+        cleanupFn = PhishingDetectionService.initialize(
+          accountService,
+          billingAccountProfileStateService,
+          configService,
+          logService,
+          phishingDataService,
+          messageListener,
+          globalStateProvider,
+        );
+      }).not.toThrow();
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+
+      const stateUpdatePromise = firstValueFrom(state.state$.pipe(skip(1)));
+      continueCommandSubject.next({
+        tabId: 123,
+        url: "https://test.com/page",
+      });
+
+      const storedHostnames = await stateUpdatePromise;
+      expect(storedHostnames).toContain("test.com");
+    });
+
+    it("syncs state updates from other contexts", async () => {
+      const state = globalStateProvider.get(IGNORED_PHISHING_HOSTNAMES_KEY);
+
+      cleanupFn = PhishingDetectionService.initialize(
+        accountService,
+        billingAccountProfileStateService,
+        configService,
+        logService,
+        phishingDataService,
+        messageListener,
+        globalStateProvider,
+      );
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+
+      await state.update(() => ["external.com", "updated.com"]);
+
+      await new Promise((resolve) => setTimeout(resolve, 0));
+
+      expect(logService.debug).toHaveBeenCalledWith(
+        expect.stringContaining("Loaded/updated 2 ignored hostnames"),
+      );
+    });
+  });
+
   // TODO
   // it("should enable phishing detection for premium account", (done) => {
   //   const premiumAccount = { id: "user1" };

apps/browser/src/dirt/phishing-detection/services/phishing-detection.service.ts

@@ -4,7 +4,6 @@ import {
   distinctUntilChanged,
   EMPTY,
   filter,
-  firstValueFrom,
   map,
   merge,
   of,
@@ -78,25 +77,21 @@ export class PhishingDetectionService {
 
     logService.debug("[PhishingDetectionService] Initialize called. Checking prerequisites...");
 
-    // Load previously ignored hostnames from storage (synchronously to avoid race conditions)
     const ignoredHostnamesState = globalStateProvider.get(IGNORED_PHISHING_HOSTNAMES_KEY);
-    firstValueFrom(ignoredHostnamesState.state$.pipe(map((hostnames) => hostnames ?? [])))
-      .then((initialHostnames) => {
-        this._ignoredHostnames = new Set(initialHostnames);
-        logService.debug(
-          `[PhishingDetectionService] Loaded ${initialHostnames.length} ignored hostnames from storage`,
-        );
-      })
-      .catch((error) => {
-        logService.error("[PhishingDetectionService] Failed to load ignored hostnames", error);
-      });
 
-    // Subscribe to future state changes
+    // Subscribe to state changes (including initial state) to keep in-memory Set in sync
+    // This handles both the initial load and any future updates from other contexts
     const ignoredHostnamesSub = ignoredHostnamesState.state$
-      .pipe(map((hostnames) => hostnames ?? []))
-      .subscribe((hostnames) => {
-        this._ignoredHostnames = new Set(hostnames);
-      });
+      .pipe(
+        map((hostnames) => hostnames ?? []),
+        tap((hostnames) => {
+          this._ignoredHostnames = new Set(hostnames);
+          logService.debug(
+            `[PhishingDetectionService] Loaded/updated ${hostnames.length} ignored hostnames`,
+          );
+        }),
+      )
+      .subscribe();
 
     BrowserApi.addListener(chrome.tabs.onUpdated, this._handleTabUpdated.bind(this));
 
@@ -164,7 +159,7 @@ export class PhishingDetectionService {
 
     const onCancelCommand$ = messageListener
       .messages$(PHISHING_DETECTION_CANCEL_COMMAND)
-      .pipe(switchMap((message) => BrowserApi.closeTab(message.tabId)));
+      .pipe(concatMap(async (message) => await BrowserApi.closeTab(message.tabId)));
 
     const activeAccountHasAccess$ = combineLatest([
       accountService.activeAccount$,

apps/browser/src/platform/browser/browser-api.ts

@@ -292,10 +292,8 @@ export class BrowserApi {
           throw new Error("Failed to navigate tab to URL: " + error.message);
         });
       } else if (BrowserApi.isChromeApi) {
-        chrome.tabs.update(tabId, { url: url.href }, () => {
-          if (chrome.runtime.lastError) {
-            throw new Error("Failed to navigate tab to URL: " + chrome.runtime.lastError.message);
-          }
+        await chrome.tabs.update(tabId, { url: url.href }).catch((error) => {
+          throw new Error("Failed to navigate tab to URL: " + error.message);
         });
       }
     }