comlink chrome extensions

endpoint.ts

import * as Comlink from 'comlink';

export const PORT_ID = '@@@PORT_ID';

const ports = new Map<string, chrome.runtime.Port>();
chrome.runtime.onConnect.addListener((port) => {
  if (port.name.startsWith(PORT_ID)) {
    ports.set(port.name, port);
  }
});

function deserialize(data: any): any {
  if (Array.isArray(data)) {
    data.forEach((value, i) => {
      data[i] = deserialize(value);
    });
  } else if (data && typeof data === 'object') {
    const id = data[PORT_ID];
    if (id) {
      const port = ports.get(id)!;
      const { port1, port2 } = new MessageChannel();
      forward(port2, port);
      return port1;
    }

    for (const key in data) {
      data[key] = deserialize(data[key]);
    }
  }
  return data;
}

export function extensionEndpoint(port: chrome.runtime.Port): Comlink.Endpoint {
  const listeners = new WeakMap();
  return {
    postMessage: (message, transfer: MessagePort[]) => {
      for (const port of transfer) {
        const id = PORT_ID + `${+new Date()}${Math.random()}`;
        // @ts-ignore
        port[PORT_ID] = id;
        forward(port, chrome.runtime.connect({ name: id }));
      }
      port.postMessage(message);
    },
    addEventListener: (_, handler) => {
      const listener = (data: any) => {
        const event = new MessageEvent('message', { data: deserialize(data) });
        if ('handleEvent' in handler) {
          handler.handleEvent(event);
        } else {
          handler(event);
        }
      };
      port.onMessage.addListener(listener);
      listeners.set(handler, listener);
    },
    removeEventListener: (_, handler) => {
      const listener = listeners.get(handler);
      if (!listener) {
        return;
      }
      port.onMessage.removeListener(listener);
      listeners.delete(handler);
    },
  };
}

export function forward(from: MessagePort, to: chrome.runtime.Port) {
  const port = extensionEndpoint(to);

  from.onmessage = ({ data, ports }) => port.postMessage(data, ports as any);
  port.addEventListener('message', ({ data, ports }: any) => {
    from.postMessage(data, ports as any);
  });
}