[PM-7810] Handle Multithread Decryption Through Offscreen API (#8978)

* [PM-7810] Handle Multithread Decryption through Offscreen API

* [PM-7810] Handle Multithread Decryption through Offscreen API

* Use a service to track when to open and close offscreen document

There some strangeness around maintaining the offscreen document for more callbacks, that need not have the same reasons and justifications as the original.

We'd need to test, but perhaps the intent is something closer to maintaining a work queue ourselves and creating the offscreen page for only a single reason as it comes in, then waiting for that page to close before opening another.

* [PM-7810] Handle Multithread Decryption through Offscreen API

* [PM-7810] Handle Multithread Decryption through Offscreen API

* [PM-7810] Handle Multithread Decryption through Offscreen API

* [PM-7810] Handle Multithread Decryption through Offscreen API

* [PM-7810] Implementing jest tests for OffscreenDocument and BrowserMultithreadEncryptServiceImplementation

* [PM-7810] Separating out the process by which we get decrypted items from the web worker to ensure we do not do duplicate effort

* [PM-7810] Separating out the process by which we get decrypted items from the web worker to ensure we do not do duplicate effort

* Prefer builtin promise flattening

* [PM-7810] Introducing a fallback to the MultithreadEncryptServiceImplementation to ensure we can fallback to single thread decryption if necessary

* [PM-7810] Updating documentation

* [PM-7810] Fixing implementation to leverage the new OffscreenDocumentService

---------

Co-authored-by: Matt Gibson <mgibson@bitwarden.com>

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

@@ -106,7 +106,6 @@ import { DefaultConfigService } from "@bitwarden/common/platform/services/config
 import { ConsoleLogService } from "@bitwarden/common/platform/services/console-log.service";
 import { ContainerService } from "@bitwarden/common/platform/services/container.service";
 import { EncryptServiceImplementation } from "@bitwarden/common/platform/services/cryptography/encrypt.service.implementation";
-import { MultithreadEncryptServiceImplementation } from "@bitwarden/common/platform/services/cryptography/multithread-encrypt.service.implementation";
 import { FileUploadService } from "@bitwarden/common/platform/services/file-upload/file-upload.service";
 import { KeyGenerationService } from "@bitwarden/common/platform/services/key-generation.service";
 import { MigrationBuilderService } from "@bitwarden/common/platform/services/migration-builder.service";
@@ -219,6 +218,7 @@ import { BrowserCryptoService } from "../platform/services/browser-crypto.servic
 import { BrowserEnvironmentService } from "../platform/services/browser-environment.service";
 import BrowserLocalStorageService from "../platform/services/browser-local-storage.service";
 import BrowserMemoryStorageService from "../platform/services/browser-memory-storage.service";
+import { BrowserMultithreadEncryptServiceImplementation } from "../platform/services/browser-multithread-encrypt.service.implementation";
 import { BrowserScriptInjectorService } from "../platform/services/browser-script-injector.service";
 import { DefaultBrowserStateService } from "../platform/services/default-browser-state.service";
 import I18nService from "../platform/services/i18n.service";
@@ -475,14 +475,14 @@ export default class MainBackground {
       storageServiceProvider,
     );
 
-    this.encryptService =
-      flagEnabled("multithreadDecryption") && BrowserApi.isManifestVersion(2)
-        ? new MultithreadEncryptServiceImplementation(
-            this.cryptoFunctionService,
-            this.logService,
-            true,
-          )
-        : new EncryptServiceImplementation(this.cryptoFunctionService, this.logService, true);
+    this.encryptService = flagEnabled("multithreadDecryption")
+      ? new BrowserMultithreadEncryptServiceImplementation(
+          this.cryptoFunctionService,
+          this.logService,
+          true,
+          this.offscreenDocumentService,
+        )
+      : new EncryptServiceImplementation(this.cryptoFunctionService, this.logService, true);
 
     this.singleUserStateProvider = new DefaultSingleUserStateProvider(
       storageServiceProvider,

apps/browser/src/platform/offscreen-document/abstractions/offscreen-document.ts

@@ -2,6 +2,7 @@ export type OffscreenDocumentExtensionMessage = {
   [key: string]: any;
   command: string;
   text?: string;
+  decryptRequest?: string;
 };
 
 type OffscreenExtensionMessageEventParams = {
@@ -13,6 +14,7 @@ export type OffscreenDocumentExtensionMessageHandlers = {
   [key: string]: ({ message, sender }: OffscreenExtensionMessageEventParams) => any;
   offscreenCopyToClipboard: ({ message }: OffscreenExtensionMessageEventParams) => any;
   offscreenReadFromClipboard: () => any;
+  offscreenDecryptItems: ({ message }: OffscreenExtensionMessageEventParams) => Promise<string>;
 };
 
 export interface OffscreenDocument {

apps/browser/src/platform/offscreen-document/offscreen-document.spec.ts

@@ -1,7 +1,25 @@
+import { mock } from "jest-mock-extended";
+
+import { Decryptable } from "@bitwarden/common/platform/interfaces/decryptable.interface";
+import { InitializerMetadata } from "@bitwarden/common/platform/interfaces/initializer-metadata.interface";
+import { SymmetricCryptoKey } from "@bitwarden/common/platform/models/domain/symmetric-crypto-key";
+
 import { flushPromises, sendExtensionRuntimeMessage } from "../../autofill/spec/testing-utils";
 import { BrowserApi } from "../browser/browser-api";
 import BrowserClipboardService from "../services/browser-clipboard.service";
 
+jest.mock(
+  "@bitwarden/common/platform/services/cryptography/multithread-encrypt.service.implementation",
+  () => ({
+    MultithreadEncryptServiceImplementation: class MultithreadEncryptServiceImplementation {
+      getDecryptedItemsFromWorker = async <T extends InitializerMetadata>(
+        items: Decryptable<T>[],
+        _key: SymmetricCryptoKey,
+      ): Promise<string> => JSON.stringify(items);
+    },
+  }),
+);
+
 describe("OffscreenDocument", () => {
   const browserApiMessageListenerSpy = jest.spyOn(BrowserApi, "messageListener");
   const browserClipboardServiceCopySpy = jest.spyOn(BrowserClipboardService, "copy");
@@ -60,5 +78,37 @@ describe("OffscreenDocument", () => {
         expect(browserClipboardServiceReadSpy).toHaveBeenCalledWith(window);
       });
     });
+
+    describe("handleOffscreenDecryptItems", () => {
+      it("returns an empty array as a string if the decrypt request is not present in the message", async () => {
+        let response: string | undefined;
+        sendExtensionRuntimeMessage(
+          { command: "offscreenDecryptItems" },
+          mock<chrome.runtime.MessageSender>(),
+          (res: string) => (response = res),
+        );
+        await flushPromises();
+
+        expect(response).toBe("[]");
+      });
+
+      it("decrypts the items and sends back the response as a string", async () => {
+        const items = [{ id: "test" }];
+        const key = { id: "test" };
+        const decryptRequest = JSON.stringify({ items, key });
+        let response: string | undefined;
+
+        sendExtensionRuntimeMessage(
+          { command: "offscreenDecryptItems", decryptRequest },
+          mock<chrome.runtime.MessageSender>(),
+          (res: string) => {
+            response = res;
+          },
+        );
+        await flushPromises();
+
+        expect(response).toBe(JSON.stringify(items));
+      });
+    });
   });
 });

apps/browser/src/platform/offscreen-document/offscreen-document.ts

@@ -1,21 +1,35 @@
 import { ConsoleLogService } from "@bitwarden/common/platform/services/console-log.service";
+import { MultithreadEncryptServiceImplementation } from "@bitwarden/common/platform/services/cryptography/multithread-encrypt.service.implementation";
+import { WebCryptoFunctionService } from "@bitwarden/common/platform/services/web-crypto-function.service";
 
 import { BrowserApi } from "../browser/browser-api";
 import BrowserClipboardService from "../services/browser-clipboard.service";
 
 import {
+  OffscreenDocument as OffscreenDocumentInterface,
   OffscreenDocumentExtensionMessage,
   OffscreenDocumentExtensionMessageHandlers,
-  OffscreenDocument as OffscreenDocumentInterface,
 } from "./abstractions/offscreen-document";
 
 class OffscreenDocument implements OffscreenDocumentInterface {
-  private consoleLogService: ConsoleLogService = new ConsoleLogService(false);
+  private readonly consoleLogService: ConsoleLogService;
+  private encryptService: MultithreadEncryptServiceImplementation;
   private readonly extensionMessageHandlers: OffscreenDocumentExtensionMessageHandlers = {
     offscreenCopyToClipboard: ({ message }) => this.handleOffscreenCopyToClipboard(message),
     offscreenReadFromClipboard: () => this.handleOffscreenReadFromClipboard(),
+    offscreenDecryptItems: ({ message }) => this.handleOffscreenDecryptItems(message),
   };
 
+  constructor() {
+    const cryptoFunctionService = new WebCryptoFunctionService(self);
+    this.consoleLogService = new ConsoleLogService(false);
+    this.encryptService = new MultithreadEncryptServiceImplementation(
+      cryptoFunctionService,
+      this.consoleLogService,
+      true,
+    );
+  }
+
   /**
    * Initializes the offscreen document extension.
    */
@@ -39,6 +53,23 @@ class OffscreenDocument implements OffscreenDocumentInterface {
     return await BrowserClipboardService.read(self);
   }
 
+  /**
+   * Decrypts the items in the message using the encrypt service.
+   *
+   * @param message - The extension message containing the items to decrypt
+   */
+  private async handleOffscreenDecryptItems(
+    message: OffscreenDocumentExtensionMessage,
+  ): Promise<string> {
+    const { decryptRequest } = message;
+    if (!decryptRequest) {
+      return "[]";
+    }
+
+    const request = JSON.parse(decryptRequest);
+    return await this.encryptService.getDecryptedItemsFromWorker(request.items, request.key);
+  }
+
   /**
    * Sets up the listener for extension messages.
    */

apps/browser/src/platform/services/browser-multithread-encrypt.service.implementation.spec.ts

@@ -0,0 +1,97 @@
+import { mock, MockProxy } from "jest-mock-extended";
+
+import { CryptoFunctionService } from "@bitwarden/common/platform/abstractions/crypto-function.service";
+import { LogService } from "@bitwarden/common/platform/abstractions/log.service";
+import { EncryptionType } from "@bitwarden/common/platform/enums";
+import { Decryptable } from "@bitwarden/common/platform/interfaces/decryptable.interface";
+import { InitializerMetadata } from "@bitwarden/common/platform/interfaces/initializer-metadata.interface";
+import { SymmetricCryptoKey } from "@bitwarden/common/platform/models/domain/symmetric-crypto-key";
+import { InitializerKey } from "@bitwarden/common/platform/services/cryptography/initializer-key";
+import { makeStaticByteArray } from "@bitwarden/common/spec";
+
+import { BrowserApi } from "../browser/browser-api";
+import { OffscreenDocumentService } from "../offscreen-document/abstractions/offscreen-document";
+
+import { BrowserMultithreadEncryptServiceImplementation } from "./browser-multithread-encrypt.service.implementation";
+
+describe("BrowserMultithreadEncryptServiceImplementation", () => {
+  let cryptoFunctionServiceMock: MockProxy<CryptoFunctionService>;
+  let logServiceMock: MockProxy<LogService>;
+  let offscreenDocumentServiceMock: MockProxy<OffscreenDocumentService>;
+  let encryptService: BrowserMultithreadEncryptServiceImplementation;
+  const manifestVersionSpy = jest.spyOn(BrowserApi, "manifestVersion", "get");
+  const sendMessageWithResponseSpy = jest.spyOn(BrowserApi, "sendMessageWithResponse");
+  const encType = EncryptionType.AesCbc256_HmacSha256_B64;
+  const key = new SymmetricCryptoKey(makeStaticByteArray(64, 100), encType);
+  const items: Decryptable<InitializerMetadata>[] = [
+    {
+      decrypt: jest.fn(),
+      initializerKey: InitializerKey.Cipher,
+    },
+  ];
+
+  beforeEach(() => {
+    cryptoFunctionServiceMock = mock<CryptoFunctionService>();
+    logServiceMock = mock<LogService>();
+    offscreenDocumentServiceMock = mock<OffscreenDocumentService>({
+      withDocument: jest.fn((_, __, callback) => callback() as any),
+    });
+    encryptService = new BrowserMultithreadEncryptServiceImplementation(
+      cryptoFunctionServiceMock,
+      logServiceMock,
+      false,
+      offscreenDocumentServiceMock,
+    );
+    manifestVersionSpy.mockReturnValue(3);
+    sendMessageWithResponseSpy.mockResolvedValue(JSON.stringify([]));
+  });
+
+  afterEach(() => {
+    jest.clearAllMocks();
+  });
+
+  it("decrypts items using web workers if the chrome.offscreen API is not supported", async () => {
+    manifestVersionSpy.mockReturnValue(2);
+
+    await encryptService.decryptItems([], key);
+
+    expect(offscreenDocumentServiceMock.withDocument).not.toHaveBeenCalled();
+  });
+
+  it("decrypts items using the chrome.offscreen API if it is supported", async () => {
+    sendMessageWithResponseSpy.mockResolvedValue(JSON.stringify(items));
+
+    await encryptService.decryptItems(items, key);
+
+    expect(offscreenDocumentServiceMock.withDocument).toHaveBeenCalledWith(
+      [chrome.offscreen.Reason.WORKERS],
+      "Use web worker to decrypt items.",
+      expect.any(Function),
+    );
+    expect(BrowserApi.sendMessageWithResponse).toHaveBeenCalledWith("offscreenDecryptItems", {
+      decryptRequest: expect.any(String),
+    });
+  });
+
+  it("returns an empty array if the passed items are not defined", async () => {
+    const result = await encryptService.decryptItems(null, key);
+
+    expect(result).toEqual([]);
+  });
+
+  it("returns an empty array if the offscreen document message returns an empty value", async () => {
+    sendMessageWithResponseSpy.mockResolvedValue("");
+
+    const result = await encryptService.decryptItems(items, key);
+
+    expect(result).toEqual([]);
+  });
+
+  it("returns an empty array if the offscreen document message returns an empty array", async () => {
+    sendMessageWithResponseSpy.mockResolvedValue("[]");
+
+    const result = await encryptService.decryptItems(items, key);
+
+    expect(result).toEqual([]);
+  });
+});

apps/browser/src/platform/services/browser-multithread-encrypt.service.implementation.ts

@@ -0,0 +1,91 @@
+import { CryptoFunctionService } from "@bitwarden/common/platform/abstractions/crypto-function.service";
+import { LogService } from "@bitwarden/common/platform/abstractions/log.service";
+import { Decryptable } from "@bitwarden/common/platform/interfaces/decryptable.interface";
+import { InitializerMetadata } from "@bitwarden/common/platform/interfaces/initializer-metadata.interface";
+import { Utils } from "@bitwarden/common/platform/misc/utils";
+import { SymmetricCryptoKey } from "@bitwarden/common/platform/models/domain/symmetric-crypto-key";
+import { MultithreadEncryptServiceImplementation } from "@bitwarden/common/platform/services/cryptography/multithread-encrypt.service.implementation";
+
+import { BrowserApi } from "../browser/browser-api";
+import { OffscreenDocumentService } from "../offscreen-document/abstractions/offscreen-document";
+
+export class BrowserMultithreadEncryptServiceImplementation extends MultithreadEncryptServiceImplementation {
+  constructor(
+    cryptoFunctionService: CryptoFunctionService,
+    logService: LogService,
+    logMacFailures: boolean,
+    private offscreenDocumentService: OffscreenDocumentService,
+  ) {
+    super(cryptoFunctionService, logService, logMacFailures);
+  }
+
+  /**
+   * Handles decryption of items, will use the offscreen document if supported.
+   *
+   * @param items - The items to decrypt.
+   * @param key - The key to use for decryption.
+   */
+  async decryptItems<T extends InitializerMetadata>(
+    items: Decryptable<T>[],
+    key: SymmetricCryptoKey,
+  ): Promise<T[]> {
+    if (!this.isOffscreenDocumentSupported()) {
+      return await super.decryptItems(items, key);
+    }
+
+    return await this.decryptItemsInOffscreenDocument(items, key);
+  }
+
+  /**
+   * Decrypts items using the offscreen document api.
+   *
+   * @param items - The items to decrypt.
+   * @param key - The key to use for decryption.
+   */
+  private async decryptItemsInOffscreenDocument<T extends InitializerMetadata>(
+    items: Decryptable<T>[],
+    key: SymmetricCryptoKey,
+  ): Promise<T[]> {
+    if (items == null || items.length < 1) {
+      return [];
+    }
+
+    const request = {
+      id: Utils.newGuid(),
+      items: items,
+      key: key,
+    };
+
+    const response = await this.offscreenDocumentService.withDocument(
+      [chrome.offscreen.Reason.WORKERS],
+      "Use web worker to decrypt items.",
+      async () => {
+        return (await BrowserApi.sendMessageWithResponse("offscreenDecryptItems", {
+          decryptRequest: JSON.stringify(request),
+        })) as string;
+      },
+    );
+
+    if (!response) {
+      return [];
+    }
+
+    const responseItems = JSON.parse(response);
+    if (responseItems?.length < 1) {
+      return [];
+    }
+
+    return this.initializeItems(responseItems);
+  }
+
+  /**
+   * Checks if the offscreen document api is supported.
+   */
+  private isOffscreenDocumentSupported() {
+    return (
+      BrowserApi.isManifestVersion(3) &&
+      typeof chrome !== "undefined" &&
+      typeof chrome.offscreen !== "undefined"
+    );
+  }
+}

libs/common/src/platform/services/cryptography/multithread-encrypt.service.implementation.ts

@@ -19,17 +19,36 @@ export class MultithreadEncryptServiceImplementation extends EncryptServiceImple
   private clear$ = new Subject<void>();
 
   /**
-   * Sends items to a web worker to decrypt them.
-   * This utilises multithreading to decrypt items faster without interrupting other operations (e.g. updating UI).
+   * Decrypts items using a web worker if the environment supports it.
+   * Will fall back to the main thread if the window object is not available.
    */
   async decryptItems<T extends InitializerMetadata>(
     items: Decryptable<T>[],
     key: SymmetricCryptoKey,
   ): Promise<T[]> {
+    if (typeof window === "undefined") {
+      return super.decryptItems(items, key);
+    }
+
     if (items == null || items.length < 1) {
       return [];
     }
 
+    const decryptedItems = await this.getDecryptedItemsFromWorker(items, key);
+    const parsedItems = JSON.parse(decryptedItems);
+
+    return this.initializeItems(parsedItems);
+  }
+
+  /**
+   * Sends items to a web worker to decrypt them. This utilizes multithreading to decrypt items
+   * faster without interrupting other operations (e.g. updating UI). This method returns values
+   * prior to deserialization to support forwarding results to another party
+   */
+  async getDecryptedItemsFromWorker<T extends InitializerMetadata>(
+    items: Decryptable<T>[],
+    key: SymmetricCryptoKey,
+  ): Promise<string> {
     this.logService.info("Starting decryption using multithreading");
 
     this.worker ??= new Worker(
@@ -53,19 +72,20 @@ export class MultithreadEncryptServiceImplementation extends EncryptServiceImple
     return await firstValueFrom(
       fromEvent(this.worker, "message").pipe(
         filter((response: MessageEvent) => response.data?.id === request.id),
-        map((response) => JSON.parse(response.data.items)),
-        map((items) =>
-          items.map((jsonItem: Jsonify<T>) => {
-            const initializer = getClassInitializer<T>(jsonItem.initializerKey);
-            return initializer(jsonItem);
-          }),
-        ),
+        map((response) => response.data.items),
         takeUntil(this.clear$),
-        defaultIfEmpty([]),
+        defaultIfEmpty("[]"),
       ),
     );
   }
 
+  protected initializeItems<T extends InitializerMetadata>(items: Jsonify<T>[]): T[] {
+    return items.map((jsonItem: Jsonify<T>) => {
+      const initializer = getClassInitializer<T>(jsonItem.initializerKey);
+      return initializer(jsonItem);
+    });
+  }
+
   private clear() {
     this.clear$.next();
     this.worker?.terminate();