Target messages to only the interested ports (#7301)

apps/browser/src/platform/storage/background-memory-storage.service.ts

@@ -25,20 +25,23 @@ export class BackgroundMemoryStorageService extends MemoryStorageService {
       });
       port.onMessage.addListener(listenerCallback);
       // Initialize the new memory storage service with existing data
-      this.sendMessage({
+      this.sendMessageTo(port, {
         action: "initialization",
         data: Array.from(this.store.keys()),
       });
     });
     this.updates$.subscribe((update) => {
-      this.sendMessage({
+      this.broadcastMessage({
         action: "subject_update",
         data: update,
       });
     });
   }
 
-  private async onMessageFromForeground(message: MemoryStoragePortMessage) {
+  private async onMessageFromForeground(
+    message: MemoryStoragePortMessage,
+    port: chrome.runtime.Port,
+  ) {
     if (message.originator === "background") {
       return;
     }
@@ -60,19 +63,26 @@ export class BackgroundMemoryStorageService extends MemoryStorageService {
         break;
     }
 
-    this.sendMessage({
+    this.sendMessageTo(port, {
       id: message.id,
       key: message.key,
       data: JSON.stringify(result),
     });
   }
 
-  private async sendMessage(data: Omit<MemoryStoragePortMessage, "originator">) {
+  private broadcastMessage(data: Omit<MemoryStoragePortMessage, "originator">) {
     this._ports.forEach((port) => {
+      this.sendMessageTo(port, data);
+    });
+  }
+
+  private sendMessageTo(
+    port: chrome.runtime.Port,
+    data: Omit<MemoryStoragePortMessage, "originator">,
+  ) {
     port.postMessage({
       ...data,
       originator: "background",
     });
-    });
   }
 }

apps/browser/src/platform/storage/memory-storage-service-interactions.spec.ts

@@ -7,15 +7,14 @@ import { trackEmissions } from "@bitwarden/common/../spec/utils";
 
 import { BackgroundMemoryStorageService } from "./background-memory-storage.service";
 import { ForegroundMemoryStorageService } from "./foreground-memory-storage.service";
-import { mockPort } from "./mock-port.spec-util";
+import { mockPorts } from "./mock-ports.spec-util";
-import { portName } from "./port-name";
 
 describe("foreground background memory storage interaction", () => {
   let foreground: ForegroundMemoryStorageService;
   let background: BackgroundMemoryStorageService;
 
   beforeEach(() => {
-    mockPort(portName(chrome.storage.session));
+    mockPorts();
 
     background = new BackgroundMemoryStorageService();
     foreground = new ForegroundMemoryStorageService();
@@ -63,4 +62,16 @@ describe("foreground background memory storage interaction", () => {
 
     expect(emissions).toEqual([{ key, updateType }]);
   });
+
+  test("background should message only the requesting foreground", async () => {
+    const secondForeground = new ForegroundMemoryStorageService();
+    const secondPort = secondForeground["_port"];
+    const secondPost = secondPort.postMessage as jest.Mock;
+    secondPost.mockClear();
+
+    const key = "key";
+    await foreground.get(key);
+
+    expect(secondPost).not.toHaveBeenCalled();
+  });
 });

apps/browser/src/platform/storage/mock-port.spec-util.ts

@@ -1,25 +0,0 @@
-import { mockDeep } from "jest-mock-extended";
-
-/**
- * Mocks a chrome.runtime.Port set up to send messages through `postMessage` to `onMessage.addListener` callbacks.
- * @returns a mock chrome.runtime.Port
- */
-export function mockPort(name: string) {
-  const port = mockDeep<chrome.runtime.Port>();
-  // notify listeners of a new port
-  (chrome.runtime.connect as jest.Mock).mockImplementation((portInfo) => {
-    port.name = portInfo.name;
-    (chrome.runtime.onConnect.addListener as jest.Mock).mock.calls.forEach(([callbackFn]) => {
-      callbackFn(port);
-    });
-    return port;
-  });
-
-  // set message broadcast
-  (port.postMessage as jest.Mock).mockImplementation((message) => {
-    (port.onMessage.addListener as jest.Mock).mock.calls.forEach(([callbackFn]) => {
-      callbackFn(message);
-    });
-  });
-  return port;
-}

apps/browser/src/platform/storage/mock-ports.spec-util.ts

@@ -0,0 +1,29 @@
+import { mockDeep } from "jest-mock-extended";
+
+/**
+ * Mocks a chrome.runtime.Port set up to send messages through `postMessage` to `onMessage.addListener` callbacks.
+ * @param name - The name of the port.
+ * @param immediateOnConnectExecution - Whether to immediately execute the onConnect callbacks against the new port.
+ * Defaults to false. If true, the creator of the port will not have had a chance to set up listeners yet.
+ * @returns a mock chrome.runtime.Port
+ */
+export function mockPorts() {
+  // notify listeners of a new port
+  (chrome.runtime.connect as jest.Mock).mockImplementation((portInfo) => {
+    const port = mockDeep<chrome.runtime.Port>();
+    port.name = portInfo.name;
+
+    // set message broadcast
+    (port.postMessage as jest.Mock).mockImplementation((message) => {
+      (port.onMessage.addListener as jest.Mock).mock.calls.forEach(([callbackFn]) => {
+        callbackFn(message, port);
+      });
+    });
+
+    (chrome.runtime.onConnect.addListener as jest.Mock).mock.calls.forEach(([callbackFn]) => {
+      callbackFn(port);
+    });
+
+    return port;
+  });
+}