[PM-24101] Switching to use the orgKeys$ from the key service instead of getOrgKey (#15781)

* Switching to use the orgKeys$ from the key service instead of getOrgKey

* Using account service instead of state provider

* First try for fixing test cases

* fixing test cases

* PM-24101 fix identified by failing test

* Error checking on the orgId

* Private method did not need error check

* Setting OrganizationId type

* Fixing test cases for setting org id

* Moving the get of critical apps to the init

* The critical apps component was being set again

---------

Co-authored-by: voommen-livefront <voommen@livefront.com>

bitwarden_license/bit-common/src/dirt/reports/risk-insights/services/critical-apps.service.spec.ts

@@ -1,14 +1,13 @@
 import { randomUUID } from "crypto";
 
-import { fakeAsync, flush } from "@angular/core/testing";
 import { mock } from "jest-mock-extended";
-import { of } from "rxjs";
+import { of, BehaviorSubject } from "rxjs";
 
 import { EncryptService } from "@bitwarden/common/key-management/crypto/abstractions/encrypt.service";
 import { EncString } from "@bitwarden/common/key-management/crypto/models/enc-string";
 import { SymmetricCryptoKey } from "@bitwarden/common/platform/models/domain/symmetric-crypto-key";
 import { CsprngArray } from "@bitwarden/common/types/csprng";
-import { OrganizationId } from "@bitwarden/common/types/guid";
+import { UserId, OrganizationId } from "@bitwarden/common/types/guid";
 import { OrgKey } from "@bitwarden/common/types/key";
 import { KeyService } from "@bitwarden/key-management";
 
@@ -21,6 +20,17 @@ import {
 import { CriticalAppsApiService } from "./critical-apps-api.service";
 import { CriticalAppsService } from "./critical-apps.service";
 
+const SomeCsprngArray = new Uint8Array(64) as CsprngArray;
+const SomeUser = "some user" as UserId;
+const SomeOrganization = "some organization" as OrganizationId;
+const AnotherOrganization = "another organization" as OrganizationId;
+const SomeOrgKey = new SymmetricCryptoKey(SomeCsprngArray) as OrgKey;
+const AnotherOrgKey = new SymmetricCryptoKey(SomeCsprngArray) as OrgKey;
+const OrgRecords: Record<OrganizationId, OrgKey> = {
+  [SomeOrganization]: SomeOrgKey,
+  [AnotherOrganization]: AnotherOrgKey,
+};
+
 describe("CriticalAppsService", () => {
   let service: CriticalAppsService;
   const keyService = mock<KeyService>();
@@ -35,10 +45,6 @@ describe("CriticalAppsService", () => {
 
     // reset mocks
     jest.resetAllMocks();
-
-    const mockRandomBytes = new Uint8Array(64) as CsprngArray;
-    const mockOrgKey = new SymmetricCryptoKey(mockRandomBytes) as OrgKey;
-    keyService.getOrgKey.mockResolvedValue(mockOrgKey);
   });
 
   it("should be created", () => {
@@ -50,23 +56,27 @@ describe("CriticalAppsService", () => {
     const criticalApps = ["https://example.com", "https://example.org"];
 
     const request = [
-      { organizationId: "org1", url: "encryptedUrlName" },
+      { organizationId: SomeOrganization, url: "encryptedUrlName" },
-      { organizationId: "org1", url: "encryptedUrlName" },
+      { organizationId: SomeOrganization, url: "encryptedUrlName" },
     ] as PasswordHealthReportApplicationsRequest[];
 
     const response = [
-      { id: "id1", organizationId: "org1", uri: "https://example.com" },
+      { id: "id1", organizationId: SomeOrganization, uri: "https://example.com" },
-      { id: "id2", organizationId: "org1", uri: "https://example.org" },
+      { id: "id2", organizationId: SomeOrganization, uri: "https://example.org" },
     ] as PasswordHealthReportApplicationsResponse[];
 
     encryptService.encryptString.mockResolvedValue(new EncString("encryptedUrlName"));
     criticalAppsApiService.saveCriticalApps.mockReturnValue(of(response));
+    const orgKey$ = new BehaviorSubject(OrgRecords);
+    keyService.orgKeys$.mockReturnValue(orgKey$);
+
+    service.setOrganizationId(SomeOrganization, SomeUser);
 
     // act
-    await service.setCriticalApps("org1", criticalApps);
+    await service.setCriticalApps(SomeOrganization, criticalApps);
 
     // expectations
-    expect(keyService.getOrgKey).toHaveBeenCalledWith("org1");
+    expect(keyService.orgKeys$).toHaveBeenCalledWith(SomeUser);
     expect(encryptService.encryptString).toHaveBeenCalledTimes(2);
     expect(criticalAppsApiService.saveCriticalApps).toHaveBeenCalledWith(request);
   });
@@ -77,7 +87,7 @@ describe("CriticalAppsService", () => {
     service.setAppsInListForOrg([
       {
         id: randomUUID() as PasswordHealthReportApplicationId,
-        organizationId: "org1" as OrganizationId,
+        organizationId: SomeOrganization,
         uri: "https://example.com",
       },
     ]);
@@ -87,59 +97,65 @@ describe("CriticalAppsService", () => {
 
     // expect only one record to be sent to the server
     const request = [
-      { organizationId: "org1", url: "encryptedUrlName" },
+      { organizationId: SomeOrganization, url: "encryptedUrlName" },
     ] as PasswordHealthReportApplicationsRequest[];
 
     // mocked response
     const response = [
-      { id: "id1", organizationId: "org1", uri: "test" },
+      { id: "id1", organizationId: SomeOrganization, uri: "test" },
     ] as PasswordHealthReportApplicationsResponse[];
 
     encryptService.encryptString.mockResolvedValue(new EncString("encryptedUrlName"));
     criticalAppsApiService.saveCriticalApps.mockReturnValue(of(response));
 
+    // mock org keys
+    const orgKey$ = new BehaviorSubject(OrgRecords);
+    keyService.orgKeys$.mockReturnValue(orgKey$);
+
+    service.setOrganizationId(SomeOrganization, SomeUser);
+
     // act
-    await service.setCriticalApps("org1", selectedUrls);
+    await service.setCriticalApps(SomeOrganization, selectedUrls);
 
     // expectations
-    expect(keyService.getOrgKey).toHaveBeenCalledWith("org1");
+    expect(keyService.orgKeys$).toHaveBeenCalledWith(SomeUser);
     expect(encryptService.encryptString).toHaveBeenCalledTimes(1);
     expect(criticalAppsApiService.saveCriticalApps).toHaveBeenCalledWith(request);
   });
 
-  it("should get critical apps", fakeAsync(() => {
+  it("should get critical apps", () => {
-    const orgId = "org1" as OrganizationId;
     const response = [
-      { id: "id1", organizationId: "org1", uri: "https://example.com" },
+      { id: "id1", organizationId: SomeOrganization, uri: "https://example.com" },
-      { id: "id2", organizationId: "org1", uri: "https://example.org" },
+      { id: "id2", organizationId: SomeOrganization, uri: "https://example.org" },
     ] as PasswordHealthReportApplicationsResponse[];
 
     encryptService.decryptString.mockResolvedValue("https://example.com");
     criticalAppsApiService.getCriticalApps.mockReturnValue(of(response));
 
-    const mockRandomBytes = new Uint8Array(64) as CsprngArray;
+    // mock org keys
-    const mockOrgKey = new SymmetricCryptoKey(mockRandomBytes) as OrgKey;
+    const orgKey$ = new BehaviorSubject(OrgRecords);
-    keyService.getOrgKey.mockResolvedValue(mockOrgKey);
+    keyService.orgKeys$.mockReturnValue(orgKey$);
 
-    service.setOrganizationId(orgId as OrganizationId);
+    service.setOrganizationId(SomeOrganization, SomeUser);
-    flush();
 
-    expect(keyService.getOrgKey).toHaveBeenCalledWith(orgId.toString());
+    expect(keyService.orgKeys$).toHaveBeenCalledWith(SomeUser);
     expect(encryptService.decryptString).toHaveBeenCalledTimes(2);
-    expect(criticalAppsApiService.getCriticalApps).toHaveBeenCalledWith(orgId);
+    expect(criticalAppsApiService.getCriticalApps).toHaveBeenCalledWith(SomeOrganization);
-  }));
+  });
 
   it("should get by org id", () => {
-    const orgId = "org1" as OrganizationId;
+    const orgId = "some organization" as OrganizationId;
     const response = [
-      { id: "id1", organizationId: "org1", uri: "https://example.com" },
+      { id: "id1", organizationId: "some organization", uri: "https://example.com" },
-      { id: "id2", organizationId: "org1", uri: "https://example.org" },
+      { id: "id2", organizationId: "some organization", uri: "https://example.org" },
-      { id: "id3", organizationId: "org2", uri: "https://example.org" },
+      { id: "id3", organizationId: "another organization", uri: "https://example.org" },
-      { id: "id4", organizationId: "org2", uri: "https://example.org" },
+      { id: "id4", organizationId: "another organization", uri: "https://example.org" },
     ] as PasswordHealthReportApplicationsResponse[];
 
+    const orgKey$ = new BehaviorSubject(OrgRecords);
+    keyService.orgKeys$.mockReturnValue(orgKey$);
+    service.setOrganizationId(SomeOrganization, SomeUser);
     service.setAppsInListForOrg(response);
-
     service.getAppsListForOrg(orgId as OrganizationId).subscribe((res) => {
       expect(res).toHaveLength(2);
     });
@@ -147,26 +163,30 @@ describe("CriticalAppsService", () => {
 
   it("should drop a critical app", async () => {
     // arrange
-    const orgId = "org1" as OrganizationId;
     const selectedUrl = "https://example.com";
 
     const initialList = [
-      { id: "id1", organizationId: "org1", uri: "https://example.com" },
+      { id: "id1", organizationId: SomeOrganization, uri: "https://example.com" },
-      { id: "id2", organizationId: "org1", uri: "https://example.org" },
+      { id: "id2", organizationId: SomeOrganization, uri: "https://example.org" },
     ] as PasswordHealthReportApplicationsResponse[];
 
+    const orgKey$ = new BehaviorSubject(OrgRecords);
+    keyService.orgKeys$.mockReturnValue(orgKey$);
+
+    service.setOrganizationId(SomeOrganization, SomeUser);
+
     service.setAppsInListForOrg(initialList);
 
     // act
-    await service.dropCriticalApp(orgId, selectedUrl);
+    await service.dropCriticalApp(SomeOrganization, selectedUrl);
 
     // expectations
     expect(criticalAppsApiService.dropCriticalApp).toHaveBeenCalledWith({
-      organizationId: orgId,
+      organizationId: SomeOrganization,
       passwordHealthReportApplicationIds: ["id1"],
     });
-    expect(service.getAppsListForOrg(orgId)).toBeTruthy();
+    expect(service.getAppsListForOrg(SomeOrganization)).toBeTruthy();
-    service.getAppsListForOrg(orgId).subscribe((res) => {
+    service.getAppsListForOrg(SomeOrganization).subscribe((res) => {
       expect(res).toHaveLength(1);
       expect(res[0].uri).toBe("https://example.org");
     });
@@ -174,23 +194,27 @@ describe("CriticalAppsService", () => {
 
   it("should not drop a critical app if it does not exist", async () => {
     // arrange
-    const orgId = "org1" as OrganizationId;
     const selectedUrl = "https://nonexistent.com";
 
     const initialList = [
-      { id: "id1", organizationId: "org1", uri: "https://example.com" },
+      { id: "id1", organizationId: SomeOrganization, uri: "https://example.com" },
-      { id: "id2", organizationId: "org1", uri: "https://example.org" },
+      { id: "id2", organizationId: SomeOrganization, uri: "https://example.org" },
     ] as PasswordHealthReportApplicationsResponse[];
 
+    const orgKey$ = new BehaviorSubject(OrgRecords);
+    keyService.orgKeys$.mockReturnValue(orgKey$);
+
+    service.setOrganizationId(SomeOrganization, SomeUser);
+
     service.setAppsInListForOrg(initialList);
 
     // act
-    await service.dropCriticalApp(orgId, selectedUrl);
+    await service.dropCriticalApp(SomeOrganization, selectedUrl);
 
     // expectations
     expect(criticalAppsApiService.dropCriticalApp).not.toHaveBeenCalled();
-    expect(service.getAppsListForOrg(orgId)).toBeTruthy();
+    expect(service.getAppsListForOrg(SomeOrganization)).toBeTruthy();
-    service.getAppsListForOrg(orgId).subscribe((res) => {
+    service.getAppsListForOrg(SomeOrganization).subscribe((res) => {
       expect(res).toHaveLength(2);
     });
   });

bitwarden_license/bit-common/src/dirt/reports/risk-insights/services/critical-apps.service.ts

@@ -1,9 +1,9 @@
 import {
   BehaviorSubject,
+  filter,
   first,
   firstValueFrom,
   forkJoin,
-  from,
   map,
   Observable,
   of,
@@ -15,7 +15,7 @@ import {
 
 import { EncryptService } from "@bitwarden/common/key-management/crypto/abstractions/encrypt.service";
 import { EncString } from "@bitwarden/common/key-management/crypto/models/enc-string";
-import { OrganizationId } from "@bitwarden/common/types/guid";
+import { OrganizationId, UserId } from "@bitwarden/common/types/guid";
 import { OrgKey } from "@bitwarden/common/types/key";
 import { KeyService } from "@bitwarden/key-management";
 
@@ -31,6 +31,7 @@ import { CriticalAppsApiService } from "./critical-apps-api.service";
  */
 export class CriticalAppsService {
   private orgId = new BehaviorSubject<OrganizationId | null>(null);
+  private orgKey$ = new Observable<OrgKey>();
   private criticalAppsList = new BehaviorSubject<PasswordHealthReportApplicationsResponse[]>([]);
   private teardown = new Subject<void>();
 
@@ -48,7 +49,11 @@ export class CriticalAppsService {
   ) {}
 
   // Get a list of critical apps for a given organization
-  getAppsListForOrg(orgId: string): Observable<PasswordHealthReportApplicationsResponse[]> {
+  getAppsListForOrg(orgId: OrganizationId): Observable<PasswordHealthReportApplicationsResponse[]> {
+    if (orgId != this.orgId.value) {
+      throw new Error("Organization ID mismatch");
+    }
+
     return this.criticalAppsList
       .asObservable()
       .pipe(map((apps) => apps.filter((app) => app.organizationId === orgId)));
@@ -60,17 +65,22 @@ export class CriticalAppsService {
   }
 
   // Save the selected critical apps for a given organization
-  async setCriticalApps(orgId: string, selectedUrls: string[]) {
+  async setCriticalApps(orgId: OrganizationId, selectedUrls: string[]) {
-    const key = await this.keyService.getOrgKey(orgId);
+    if (orgId != this.orgId.value) {
-    if (key == null) {
+      throw new Error("Organization ID mismatch");
+    }
+
+    const orgKey = await firstValueFrom(this.orgKey$);
+
+    if (orgKey == null) {
       throw new Error("Organization key not found");
     }
 
     // only save records that are not already in the database
     const newEntries = await this.filterNewEntries(orgId as OrganizationId, selectedUrls);
     const criticalAppsRequests = await this.encryptNewEntries(
-      orgId as OrganizationId,
+      this.orgId.value as OrganizationId,
-      key,
+      orgKey,
       newEntries,
     );
 
@@ -83,7 +93,7 @@ export class CriticalAppsService {
     for (const responseItem of dbResponse) {
       const decryptedUrl = await this.encryptService.decryptString(
         new EncString(responseItem.uri),
-        key,
+        orgKey,
       );
       if (!updatedList.some((f) => f.uri === decryptedUrl)) {
         updatedList.push({
@@ -97,13 +107,21 @@ export class CriticalAppsService {
   }
 
   // Get the critical apps for a given organization
-  setOrganizationId(orgId: OrganizationId) {
+  setOrganizationId(orgId: OrganizationId, userId: UserId) {
+    this.orgKey$ = this.keyService.orgKeys$(userId).pipe(
+      filter((OrgKeys) => !!OrgKeys),
+      map((organizationKeysById) => organizationKeysById[orgId as OrganizationId]),
+    );
     this.orgId.next(orgId);
   }
 
   // Drop a critical app for a given organization
   // Only one app may be dropped at a time
   async dropCriticalApp(orgId: OrganizationId, selectedUrl: string) {
+    if (orgId != this.orgId.value) {
+      throw new Error("Organization ID mismatch");
+    }
+
     const app = this.criticalAppsList.value.find(
       (f) => f.organizationId === orgId && f.uri === selectedUrl,
     );
@@ -127,10 +145,7 @@ export class CriticalAppsService {
       return of([]);
     }
 
-    const result$ = zip(
+    const result$ = zip(this.criticalAppsApiService.getCriticalApps(orgId), this.orgKey$).pipe(
-      this.criticalAppsApiService.getCriticalApps(orgId),
-      from(this.keyService.getOrgKey(orgId)),
-    ).pipe(
       switchMap(([response, key]) => {
         if (key == null) {
           throw new Error("Organization key not found");

bitwarden_license/bit-web/src/app/dirt/access-intelligence/all-applications.component.ts

@@ -25,6 +25,7 @@ import { AccountService } from "@bitwarden/common/auth/abstractions/account.serv
 import { getUserId } from "@bitwarden/common/auth/services/account.service";
 import { ConfigService } from "@bitwarden/common/platform/abstractions/config/config.service";
 import { I18nService } from "@bitwarden/common/platform/abstractions/i18n.service";
+import { OrganizationId } from "@bitwarden/common/types/guid";
 import { CipherService } from "@bitwarden/common/vault/abstractions/cipher.service";
 import {
   IconButtonModule,
@@ -86,7 +87,7 @@ export class AllApplicationsComponent implements OnInit {
 
       combineLatest([
         this.dataService.applications$,
-        this.criticalAppsService.getAppsListForOrg(organizationId),
+        this.criticalAppsService.getAppsListForOrg(organizationId as OrganizationId),
         organization$,
       ])
         .pipe(
@@ -168,7 +169,7 @@ export class AllApplicationsComponent implements OnInit {
 
     try {
       await this.criticalAppsService.setCriticalApps(
-        this.organization.id,
+        this.organization.id as OrganizationId,
         Array.from(this.selectedUrls),
       );
 

bitwarden_license/bit-web/src/app/dirt/access-intelligence/critical-applications.component.ts

@@ -4,7 +4,7 @@ import { Component, DestroyRef, inject, OnInit } from "@angular/core";
 import { takeUntilDestroyed } from "@angular/core/rxjs-interop";
 import { FormControl } from "@angular/forms";
 import { ActivatedRoute, Router } from "@angular/router";
-import { combineLatest, debounceTime, map, switchMap } from "rxjs";
+import { combineLatest, debounceTime, firstValueFrom, map, switchMap } from "rxjs";
 
 import { Security } from "@bitwarden/assets/svg";
 import {
@@ -17,6 +17,8 @@ import {
   ApplicationHealthReportDetailWithCriticalFlagAndCipher,
   ApplicationHealthReportSummary,
 } from "@bitwarden/bit-common/dirt/reports/risk-insights/models/password-health";
+import { AccountService } from "@bitwarden/common/auth/abstractions/account.service";
+import { getUserId } from "@bitwarden/common/auth/services/account.service";
 import { ConfigService } from "@bitwarden/common/platform/abstractions/config/config.service";
 import { I18nService } from "@bitwarden/common/platform/abstractions/i18n.service";
 import { CipherId, OrganizationId } from "@bitwarden/common/types/guid";
@@ -61,10 +63,11 @@ export class CriticalApplicationsComponent implements OnInit {
 
   async ngOnInit() {
     this.organizationId = this.activatedRoute.snapshot.paramMap.get("organizationId") ?? "";
-
+    const userId = await firstValueFrom(getUserId(this.accountService.activeAccount$));
+    this.criticalAppsService.setOrganizationId(this.organizationId as OrganizationId, userId);
     combineLatest([
       this.dataService.applications$,
-      this.criticalAppsService.getAppsListForOrg(this.organizationId),
+      this.criticalAppsService.getAppsListForOrg(this.organizationId as OrganizationId),
     ])
       .pipe(
         takeUntilDestroyed(this.destroyRef),
@@ -168,6 +171,7 @@ export class CriticalApplicationsComponent implements OnInit {
     protected i18nService: I18nService,
     private configService: ConfigService,
     private adminTaskService: DefaultAdminTaskService,
+    private accountService: AccountService,
   ) {
     this.searchControl.valueChanges
       .pipe(debounceTime(200), takeUntilDestroyed())

bitwarden_license/bit-web/src/app/dirt/access-intelligence/risk-insights.component.ts

@@ -2,7 +2,7 @@ import { CommonModule } from "@angular/common";
 import { Component, DestroyRef, OnInit, inject } from "@angular/core";
 import { takeUntilDestroyed } from "@angular/core/rxjs-interop";
 import { ActivatedRoute, Router } from "@angular/router";
-import { EMPTY, Observable } from "rxjs";
+import { EMPTY, firstValueFrom, Observable } from "rxjs";
 import { map, switchMap } from "rxjs/operators";
 
 import { JslibModule } from "@bitwarden/angular/jslib.module";
@@ -15,6 +15,8 @@ import {
   DrawerType,
   PasswordHealthReportApplicationsResponse,
 } from "@bitwarden/bit-common/dirt/reports/risk-insights/models/password-health";
+import { AccountService } from "@bitwarden/common/auth/abstractions/account.service";
+import { getUserId } from "@bitwarden/common/auth/services/account.service";
 import { ConfigService } from "@bitwarden/common/platform/abstractions/config/config.service";
 import { OrganizationId } from "@bitwarden/common/types/guid";
 import {
@@ -78,15 +80,16 @@ export class RiskInsightsComponent implements OnInit {
     private configService: ConfigService,
     protected dataService: RiskInsightsDataService,
     private criticalAppsService: CriticalAppsService,
+    private accountService: AccountService,
   ) {
     this.route.queryParams.pipe(takeUntilDestroyed()).subscribe(({ tabIndex }) => {
       this.tabIndex = !isNaN(Number(tabIndex)) ? Number(tabIndex) : RiskInsightsTabType.AllApps;
     });
-    const orgId = this.route.snapshot.paramMap.get("organizationId") ?? "";
-    this.criticalApps$ = this.criticalAppsService.getAppsListForOrg(orgId);
   }
 
   async ngOnInit() {
+    const userId = await firstValueFrom(getUserId(this.accountService.activeAccount$));
+
     this.route.paramMap
       .pipe(
         takeUntilDestroyed(this.destroyRef),
@@ -109,7 +112,11 @@ export class RiskInsightsComponent implements OnInit {
           if (applications) {
             this.appsCount = applications.length;
           }
-          this.criticalAppsService.setOrganizationId(this.organizationId as OrganizationId);
+
+          this.criticalAppsService.setOrganizationId(this.organizationId as OrganizationId, userId);
+          this.criticalApps$ = this.criticalAppsService.getAppsListForOrg(
+            this.organizationId as OrganizationId,
+          );
         },
       });
   }