From bbdf704763da7141a65f3ec3800c7efd62c0a4d5 Mon Sep 17 00:00:00 2001 From: Matt Gibson Date: Thu, 21 Dec 2023 14:01:48 -0500 Subject: [PATCH] Target messages to only the interested ports (#7301) --- .../background-memory-storage.service.ts | 28 ++++++++++++------ ...emory-storage-service-interactions.spec.ts | 17 +++++++++-- .../platform/storage/mock-port.spec-util.ts | 25 ---------------- .../platform/storage/mock-ports.spec-util.ts | 29 +++++++++++++++++++ 4 files changed, 62 insertions(+), 37 deletions(-) delete mode 100644 apps/browser/src/platform/storage/mock-port.spec-util.ts create mode 100644 apps/browser/src/platform/storage/mock-ports.spec-util.ts diff --git a/apps/browser/src/platform/storage/background-memory-storage.service.ts b/apps/browser/src/platform/storage/background-memory-storage.service.ts index 9fb8cb7162..458863ff55 100644 --- a/apps/browser/src/platform/storage/background-memory-storage.service.ts +++ b/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) { + private broadcastMessage(data: Omit) { this._ports.forEach((port) => { - port.postMessage({ - ...data, - originator: "background", - }); + this.sendMessageTo(port, data); + }); + } + + private sendMessageTo( + port: chrome.runtime.Port, + data: Omit, + ) { + port.postMessage({ + ...data, + originator: "background", }); } } diff --git a/apps/browser/src/platform/storage/memory-storage-service-interactions.spec.ts b/apps/browser/src/platform/storage/memory-storage-service-interactions.spec.ts index 0f60403524..77aa096073 100644 --- a/apps/browser/src/platform/storage/memory-storage-service-interactions.spec.ts +++ b/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 { portName } from "./port-name"; +import { mockPorts } from "./mock-ports.spec-util"; 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(); + }); }); diff --git a/apps/browser/src/platform/storage/mock-port.spec-util.ts b/apps/browser/src/platform/storage/mock-port.spec-util.ts deleted file mode 100644 index 0ab5a0ae70..0000000000 --- a/apps/browser/src/platform/storage/mock-port.spec-util.ts +++ /dev/null @@ -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(); - // 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; -} diff --git a/apps/browser/src/platform/storage/mock-ports.spec-util.ts b/apps/browser/src/platform/storage/mock-ports.spec-util.ts new file mode 100644 index 0000000000..b5f7825d8e --- /dev/null +++ b/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(); + 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; + }); +}