Fix: ensure the service worker is awake before every port message (#1433)

* fix: ensure there is an open port connection before sending a port message

* chore: remove unnecessary console.log

* fix: prevent Error: Invalid url messages when setting a non http active tab

* fix test

Author: F-OBrien

packages/extension-base/src/utils/portUtils.ts

@@ -0,0 +1,65 @@
+// Copyright 2019-2024 @polkadot/extension-base authors & contributors
+// SPDX-License-Identifier: Apache-2.0
+
+import type { Message } from '@polkadot/extension-base/types';
+
+import { chrome } from '@polkadot/extension-inject/chrome';
+
+export function setupPort (portName: string, onMessageHandler: (data: Message['data']) => void, onDisconnectHandler: () => void): chrome.runtime.Port {
+  const port = chrome.runtime.connect({ name: portName });
+
+  port.onMessage.addListener(onMessageHandler);
+
+  port.onDisconnect.addListener(() => {
+    console.log(`Disconnected from ${portName}`);
+    onDisconnectHandler();
+  });
+
+  return port;
+}
+
+export async function wakeUpServiceWorker (): Promise<{ status: string }> {
+  return new Promise((resolve, reject) => {
+    chrome.runtime.sendMessage({ type: 'wakeup' }, (response: { status: string }) => {
+      if (chrome.runtime.lastError) {
+        reject(new Error(chrome.runtime.lastError.message));
+      } else {
+        resolve(response);
+      }
+    });
+  });
+}
+
+// This object is required to allow jest.spyOn to be used to create a mock Implementation for testing
+export const wakeUpServiceWorkerWrapper = { wakeUpServiceWorker };
+
+export async function ensurePortConnection (
+  portRef: chrome.runtime.Port | undefined,
+  portConfig: {
+    portName: string,
+    onPortMessageHandler: (data: Message['data']) => void,
+    onPortDisconnectHandler: () => void
+  }
+): Promise<chrome.runtime.Port> {
+  const maxAttempts = 5;
+  const delayMs = 1000;
+
+  for (let attempt = 0; attempt < maxAttempts; attempt++) {
+    try {
+      const response = await wakeUpServiceWorkerWrapper.wakeUpServiceWorker();
+
+      if (response?.status === 'awake') {
+        if (!portRef) {
+          return setupPort(portConfig.portName, portConfig.onPortMessageHandler, portConfig.onPortDisconnectHandler);
+        }
+
+        return portRef;
+      }
+    } catch (error) {
+      console.error(`Attempt ${attempt + 1} failed: ${(error as Error).message}`);
+      await new Promise((resolve) => setTimeout(resolve, delayMs));
+    }
+  }
+
+  throw new Error('Failed to wake up the service worker and setup the port after multiple attempts');
+}

packages/extension-ui/src/Popup/index.tsx

@@ -79,8 +79,6 @@ export default function Popup (): React.ReactElement {
   const [settingsCtx, setSettingsCtx] = useState<SettingsStruct>(startSettings);
   const history = useHistory();
 
-  console.log('WINDOW; ', window);
-
   const _onAction = useCallback(
     (to?: string): void => {
       setWelcomeDone(window.localStorage.getItem('welcome_read') === 'ok');

packages/extension-ui/src/messaging.spec.ts

@@ -10,17 +10,31 @@ import type * as _ from '@polkadot/dev-test/globals.d.ts';
 import Adapter from '@wojtekmaj/enzyme-adapter-react-17';
 import enzyme from 'enzyme';
 
+import { wakeUpServiceWorkerWrapper } from '../../extension-base/src/utils/portUtils.js';
 import { exportAccount } from './messaging.js';
 
 // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment,@typescript-eslint/no-unsafe-call
 enzyme.configure({ adapter: new Adapter() });
 
 describe('messaging sends message to background via extension port for', () => {
-  it('exportAccount', () => {
+  beforeEach(() => {
+    jest.spyOn(wakeUpServiceWorkerWrapper, 'wakeUpServiceWorker').mockImplementation(() => Promise.resolve({ status: 'awake' }));
+  });
+
+  afterEach(() => {
+    jest.restoreAllMocks();
+  });
+
+  it('exportAccount', async () => {
     const callback = jest.fn();
 
     chrome.runtime.connect().onMessage.addListener(callback);
-    exportAccount('HjoBp62cvsWDA3vtNMWxz6c9q13ReEHi9UGHK7JbZweH5g5', 'passw0rd').catch(console.error);
+
+    try {
+      await exportAccount('HjoBp62cvsWDA3vtNMWxz6c9q13ReEHi9UGHK7JbZweH5g5', 'passw0rd');
+    } catch (error) {
+      console.error(error);
+    }
 
     expect(callback).toHaveBeenCalledWith(
       expect.objectContaining({

packages/extension-ui/src/messaging.ts

@@ -15,6 +15,7 @@ import type { KeypairType } from '@polkadot/util-crypto/types';
 
 import { PORT_EXTENSION } from '@polkadot/extension-base/defaults';
 import { getId } from '@polkadot/extension-base/utils/getId';
+import { ensurePortConnection } from '@polkadot/extension-base/utils/portUtils';
 import { metadataExpand } from '@polkadot/extension-chains';
 
 import allChains from './util/chains.js';
@@ -30,41 +31,11 @@ interface Handler {
 
 type Handlers = Record<string, Handler>;
 
-async function wakeupBackground (): Promise<Error | null> {
-  try {
-    await chrome.runtime.sendMessage({ type: 'wakeup' });
-
-    return null;
-  } catch (cause) {
-    return cause instanceof Error ? cause : new Error(String(cause));
-  }
-}
-
-async function createPort (name: string, maxAttempts: number, delayMs: number): Promise<chrome.runtime.Port> {
-  let lastError: Error | null = null;
-
-  for (let attempt = 0; attempt < maxAttempts; attempt++) {
-    const error = await wakeupBackground();
-
-    if (error) {
-      lastError = error;
-      await new Promise((resolve) => setTimeout(resolve, delayMs));
-      continue;
-    }
-
-    const port = chrome.runtime.connect({ name });
-
-    return port;
-  }
-
-  throw new Error('Failed to create port after multiple attempts', { cause: lastError });
-}
-
-const port = await createPort(PORT_EXTENSION, 5, 1000);
 const handlers: Handlers = {};
 
-// setup a listener for messages, any incoming resolves the promise
-port.onMessage.addListener((data: Message['data']): void => {
+let port: chrome.runtime.Port | undefined;
+
+function onPortMessageHandler (data: Message['data']): void {
   const handler = handlers[data.id];
 
   if (!handler) {
@@ -85,7 +56,17 @@ port.onMessage.addListener((data: Message['data']): void => {
   } else {
     handler.resolve(data.response);
   }
-});
+}
+
+function onPortDisconnectHandler (): void {
+  port = undefined;
+}
+
+const portConfig = {
+  onPortDisconnectHandler,
+  onPortMessageHandler,
+  portName: PORT_EXTENSION
+};
 
 function sendMessage<TMessageType extends MessageTypesWithNullRequest>(message: TMessageType): Promise<ResponseTypes[TMessageType]>;
 function sendMessage<TMessageType extends MessageTypesWithNoSubscriptions>(message: TMessageType, request: RequestTypes[TMessageType]): Promise<ResponseTypes[TMessageType]>;
@@ -96,7 +77,13 @@ function sendMessage<TMessageType extends MessageTypes> (message: TMessageType,
 
     handlers[id] = { reject, resolve, subscriber };
 
-    port.postMessage({ id, message, request: request || {} });
+    ensurePortConnection(port, portConfig).then((connectedPort) => {
+      connectedPort.postMessage({ id, message, request: request || {} });
+      port = connectedPort;
+    }).catch((error) => {
+      console.error(`Failed to send message: ${(error as Error).message}`);
+      reject(error);
+    });
   });
 }
 

packages/extension/src/background.ts

@@ -29,15 +29,26 @@ chrome.runtime.onConnect.addListener((port): void => {
   port.onDisconnect.addListener(() => console.log(`Disconnected from ${port.name}`));
 });
 
+function isValidUrl (url: string) {
+  try {
+    const urlObj = new URL(url);
+
+    return urlObj.protocol === 'http:' || urlObj.protocol === 'https:';
+  } catch (_e) {
+    return false;
+  }
+}
+
 function getActiveTabs () {
   // queriing the current active tab in the current window should only ever return 1 tab
   // although an array is specified here
   chrome.tabs.query({ active: true, currentWindow: true }, (tabs) => {
-    // get the urls of the active tabs. In the case of new tab the url may be empty or undefined
+    // get the urls of the active tabs. Only http or https urls are supported. Other urls will be filtered out.
+    // e.g. browser tabs like chrome://newtab/, chrome://extensions/, about:addons etc will be filtered out
     // we filter these out
     const urls: string[] = tabs
       .map(({ url }) => url)
-      .filter((url) => !!url) as string[];
+      .filter((url) => !!url && isValidUrl(url)) as string[];
 
     const request: TransportRequestMessage<'pri(activeTabsUrl.update)'> = {
       id: 'background',

packages/extension/src/content.ts

@@ -4,15 +4,24 @@
 import type { Message } from '@polkadot/extension-base/types';
 
 import { MESSAGE_ORIGIN_CONTENT, MESSAGE_ORIGIN_PAGE, PORT_CONTENT } from '@polkadot/extension-base/defaults';
+import { ensurePortConnection } from '@polkadot/extension-base/utils/portUtils';
 import { chrome } from '@polkadot/extension-inject/chrome';
 
-// connect to the extension
-const port = chrome.runtime.connect({ name: PORT_CONTENT });
+let port: chrome.runtime.Port | undefined;
 
-// send any messages from the extension back to the page
-port.onMessage.addListener((data): void => {
+function onPortMessageHandler (data: Message['data']): void {
   window.postMessage({ ...data, origin: MESSAGE_ORIGIN_CONTENT }, '*');
-});
+}
+
+function onPortDisconnectHandler (): void {
+  port = undefined;
+}
+
+const portConfig = {
+  onPortDisconnectHandler,
+  onPortMessageHandler,
+  portName: PORT_CONTENT
+};
 
 // all messages from the page, pass them to the extension
 window.addEventListener('message', ({ data, source }: Message): void => {
@@ -21,7 +30,10 @@ window.addEventListener('message', ({ data, source }: Message): void => {
     return;
   }
 
-  port.postMessage(data);
+  ensurePortConnection(port, portConfig).then((connectedPort) => {
+    connectedPort.postMessage(data);
+    port = connectedPort;
+  }).catch((error) => console.error(`Failed to send message: ${(error as Error).message}`));
 });
 
 // inject our data injector