Add premium guard to phishing detection service (#16602)

apps/browser/src/background/main.background.ts

@@ -1415,12 +1415,14 @@ export default class MainBackground {
     this.inlineMenuFieldQualificationService = new InlineMenuFieldQualificationService();
 
     PhishingDetectionService.initialize(
-      this.configService,
+      this.accountService,
       this.auditService,
+      this.billingAccountProfileStateService,
+      this.configService,
+      this.eventCollectionService,
       this.logService,
       this.storageService,
       this.taskSchedulerService,
-      this.eventCollectionService,
     );
 
     this.ipcContentScriptManagerService = new IpcContentScriptManagerService(this.configService);

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

@@ -2,6 +2,8 @@ import { of } from "rxjs";
 
 import { AuditService } from "@bitwarden/common/abstractions/audit.service";
 import { EventCollectionService } from "@bitwarden/common/abstractions/event/event-collection.service";
+import { AccountService } from "@bitwarden/common/auth/abstractions/account.service";
+import { BillingAccountProfileStateService } from "@bitwarden/common/billing/abstractions";
 import { ConfigService } from "@bitwarden/common/platform/abstractions/config/config.service";
 import { LogService } from "@bitwarden/common/platform/abstractions/log.service";
 import { AbstractStorageService } from "@bitwarden/common/platform/abstractions/storage.service";
@@ -10,35 +12,109 @@ import { TaskSchedulerService } from "@bitwarden/common/platform/scheduling/task
 import { PhishingDetectionService } from "./phishing-detection.service";
 
 describe("PhishingDetectionService", () => {
+  let accountService: AccountService;
   let auditService: AuditService;
+  let billingAccountProfileStateService: BillingAccountProfileStateService;
+  let configService: ConfigService;
+  let eventCollectionService: EventCollectionService;
   let logService: LogService;
   let storageService: AbstractStorageService;
   let taskSchedulerService: TaskSchedulerService;
-  let configService: ConfigService;
-  let eventCollectionService: EventCollectionService;
 
   beforeEach(() => {
+    accountService = { getAccount$: jest.fn(() => of(null)) } as any;
     auditService = { getKnownPhishingDomains: jest.fn() } as any;
+    billingAccountProfileStateService = {} as any;
+    configService = { getFeatureFlag$: jest.fn(() => of(false)) } as any;
+    eventCollectionService = {} as any;
     logService = { info: jest.fn(), debug: jest.fn(), warning: jest.fn(), error: jest.fn() } as any;
     storageService = { get: jest.fn(), save: jest.fn() } as any;
     taskSchedulerService = { registerTaskHandler: jest.fn(), setInterval: jest.fn() } as any;
-    configService = { getFeatureFlag$: jest.fn(() => of(false)) } as any;
-    eventCollectionService = {} as any;
   });
 
   it("should initialize without errors", () => {
     expect(() => {
       PhishingDetectionService.initialize(
-        configService,
+        accountService,
         auditService,
+        billingAccountProfileStateService,
+        configService,
+        eventCollectionService,
         logService,
         storageService,
         taskSchedulerService,
-        eventCollectionService,
       );
     }).not.toThrow();
   });
 
+  it("should enable phishing detection for premium account", (done) => {
+    const premiumAccount = { id: "user1" };
+    accountService = { activeAccount$: of(premiumAccount) } as any;
+    configService = { getFeatureFlag$: jest.fn(() => of(true)) } as any;
+    billingAccountProfileStateService = {
+      hasPremiumFromAnySource$: jest.fn(() => of(true)),
+    } as any;
+
+    // Patch _setup to call done
+    const setupSpy = jest
+      .spyOn(PhishingDetectionService as any, "_setup")
+      .mockImplementation(async () => {
+        expect(setupSpy).toHaveBeenCalled();
+        done();
+      });
+
+    // Run the initialization
+    PhishingDetectionService.initialize(
+      accountService,
+      auditService,
+      billingAccountProfileStateService,
+      configService,
+      eventCollectionService,
+      logService,
+      storageService,
+      taskSchedulerService,
+    );
+  });
+
+  it("should not enable phishing detection for non-premium account", (done) => {
+    const nonPremiumAccount = { id: "user2" };
+    accountService = { activeAccount$: of(nonPremiumAccount) } as any;
+    configService = { getFeatureFlag$: jest.fn(() => of(true)) } as any;
+    billingAccountProfileStateService = {
+      hasPremiumFromAnySource$: jest.fn(() => of(false)),
+    } as any;
+
+    // Patch _setup to fail if called
+    // [FIXME] This test needs to check if the setupSpy fails or is called
+    // Refactor initialize in PhishingDetectionService to return a Promise or Observable that resolves/completes when initialization is done
+    // So that spy setups can be properly verified after initialization
+    // const setupSpy = jest
+    //   .spyOn(PhishingDetectionService as any, "_setup")
+    //   .mockImplementation(async () => {
+    //     throw new Error("Should not call _setup");
+    //   });
+
+    // Patch _cleanup to call done
+    const cleanupSpy = jest
+      .spyOn(PhishingDetectionService as any, "_cleanup")
+      .mockImplementation(() => {
+        expect(cleanupSpy).toHaveBeenCalled();
+        done();
+      });
+
+    // Run the initialization
+    PhishingDetectionService.initialize(
+      accountService,
+      auditService,
+      billingAccountProfileStateService,
+      configService,
+      eventCollectionService,
+      logService,
+      storageService,
+      taskSchedulerService,
+    );
+  });
+
   it("should detect phishing domains", () => {
     PhishingDetectionService["_knownPhishingDomains"].add("phishing.com");
     const url = new URL("https://phishing.com");

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

@@ -1,7 +1,18 @@
-import { concatMap, delay, Subject, Subscription } from "rxjs";
+import {
+  combineLatest,
+  concatMap,
+  delay,
+  EMPTY,
+  map,
+  Subject,
+  Subscription,
+  switchMap,
+} from "rxjs";
 
 import { AuditService } from "@bitwarden/common/abstractions/audit.service";
 import { EventCollectionService } from "@bitwarden/common/abstractions/event/event-collection.service";
+import { AccountService } from "@bitwarden/common/auth/abstractions/account.service";
+import { BillingAccountProfileStateService } from "@bitwarden/common/billing/abstractions";
 import { FeatureFlag } from "@bitwarden/common/enums/feature-flag.enum";
 import { ConfigService } from "@bitwarden/common/platform/abstractions/config/config.service";
 import { LogService } from "@bitwarden/common/platform/abstractions/log.service";
@@ -41,31 +52,44 @@ export class PhishingDetectionService {
   private static _lastUpdateTime: number = 0;
 
   static initialize(
-    configService: ConfigService,
+    accountService: AccountService,
     auditService: AuditService,
+    billingAccountProfileStateService: BillingAccountProfileStateService,
+    configService: ConfigService,
+    eventCollectionService: EventCollectionService,
     logService: LogService,
     storageService: AbstractStorageService,
     taskSchedulerService: TaskSchedulerService,
-    eventCollectionService: EventCollectionService,
   ): void {
     this._auditService = auditService;
     this._logService = logService;
     this._storageService = storageService;
     this._taskSchedulerService = taskSchedulerService;
 
-    logService.info("[PhishingDetectionService] Initialize called");
+    logService.info("[PhishingDetectionService] Initialize called. Checking prerequisites...");
 
-    configService
-      .getFeatureFlag$(FeatureFlag.PhishingDetection)
+    combineLatest([
+      accountService.activeAccount$,
+      configService.getFeatureFlag$(FeatureFlag.PhishingDetection),
+    ])
       .pipe(
-        concatMap(async (enabled) => {
-          if (!enabled) {
+        switchMap(([account, featureEnabled]) => {
+          if (!account) {
+            logService.info("[PhishingDetectionService] No active account.");
+            this._cleanup();
+            return EMPTY;
+          }
+          return billingAccountProfileStateService
+            .hasPremiumFromAnySource$(account.id)
+            .pipe(map((hasPremium) => ({ hasPremium, featureEnabled })));
+        }),
+        concatMap(async ({ hasPremium, featureEnabled }) => {
+          if (!hasPremium || !featureEnabled) {
             logService.info(
-              "[PhishingDetectionService] Phishing detection feature flag is disabled.",
+              "[PhishingDetectionService] User does not have access to phishing detection service.",
             );
             this._cleanup();
           } else {
-            // Enable phishing detection service
             logService.info("[PhishingDetectionService] Enabling phishing detection service");
             await this._setup();
           }