feat(browser-extension): support pick client for browser ext main tra…

…nport

examples/browser-extension/src/background/index.ts

@@ -27,11 +27,18 @@ class BackgroundTransport
   }
 }
 
-const backgroundTransport = new BackgroundTransport({
+const transport = new BackgroundTransport({
   browser: browser as any,
 });
 
-(global as any).backgroundTransport = backgroundTransport;
+(global as any).transport = transport;
+
+transport.onConnect((id) => {
+  console.log('connect:', id);
+});
+
+// @ts-ignore
+transport.listen('contentToBg', async (a, b) => a);
 
 // chrome.runtime.onMessage.addListener((message, sender, sendResponse) => {
 //   //

examples/browser-extension/src/content/index.tsx

@@ -1,16 +1,40 @@
-import React from "react";
-import {
-  createRoot
-} from "react-dom/client";
+import React from 'react';
+import { createRoot } from 'react-dom/client';
 
-import App from "./components/app";
+import browser from 'webextension-polyfill';
+import { createTransport } from 'data-transport';
 
-import "./index.scss";
+import App from './components/app';
 
-const container = document.createElement("popup");
+import './index.scss';
+
+const container = document.createElement('popup');
 document.body.appendChild(container);
 
 const root = createRoot(container);
 root.render(<App />);
 
-console.log("Content Script 👋");
+console.log('Content Script 👋');
+
+const transport = createTransport('BrowserExtensionsClient', {
+  browser: browser as any,
+});
+
+transport.onConnect(() => {
+  console.log('connect');
+});
+
+(global as any).transport = transport;
+
+console.log('transport:', transport);
+
+// @ts-ignore
+transport.listen('a', () => {
+  console.log('a event in content script');
+});
+
+setTimeout(async () => {
+  console.log('send contentToBg');
+  const a = await transport.emit('contentToBg', 1, 2);
+  console.log('a:', a);
+}, 1000 * 10);

examples/browser-extension/src/popup/index.tsx

@@ -10,18 +10,31 @@ class PopupTransport extends BrowserExtensionsClientTransport<{ emit: PopupToBac
         respond: false,
       },
       {
-        path: 'client.html',
+        path: 'option.html',
         features: 'width=300,height=600',
       }
     );
   }
 }
 
-const popupTransport = new PopupTransport({
+const transport = new PopupTransport({
   browser: browser as any,
+  verbose: true,
 });
 
 const button = document.getElementById('button');
 button!.addEventListener('click', () => {
-  popupTransport.openClient();
+  transport.openClient();
 });
+
+(global as any).transport = transport;
+
+transport.onConnect(() => {
+  console.log('connect');
+});
+
+// @ts-ignore
+transport.listen('a', () => {
+  console.log('a event in popup');
+});
+

src/transports/browserExtensionsTransport.ts

@@ -50,6 +50,9 @@ interface BrowserExtensionsMainPort {
   _port?: Port;
 }
 
+type ClientCallback = () => void | Promise<void>;
+type MainCallback = (clientId: string) => void | Promise<void>;
+
 export abstract class BrowserExtensionsGenericTransport<
   T extends BaseInteraction = any
 > extends Transport<T> {
@@ -97,10 +100,12 @@ export abstract class BrowserExtensionsGenericTransport<
   }
 }
 
+const connectEventName = 'sharedworker-connect';
+
 export abstract class BrowserExtensionsMainTransport<
   T extends BaseInteraction = any
 > extends Transport<T> {
-  protected ports = new Set<Port>();
+  protected ports = new Map<string, Port>();
 
   private [callbackKey]!: (
     options: ListenerOptions<BrowserExtensionsMainPort>
@@ -118,12 +123,19 @@ export abstract class BrowserExtensionsMainTransport<
         };
       },
       sender = function (this: BrowserExtensionsMainTransport, message) {
-        const port = message._port;
+        const port = message._extra?._port;
         if (port) {
           delete message._port;
           port.postMessage(message);
+        } else if (
+          message.type === 'response' &&
+          // @ts-ignore
+          this.ports.has(message.requestId)
+        ) {
+          // @ts-ignore
+          const port = this.ports.get(message.requestId)!;
+          port.postMessage(message);
         } else {
-          // TODO: select an assignable port
           this.ports.forEach((port) => {
             port.postMessage(message);
           });
@@ -136,21 +148,54 @@ export abstract class BrowserExtensionsMainTransport<
       listener,
       sender,
     });
-    browser.runtime.onConnect.addListener((port: Port) => {
+    browser.runtime.onConnect.addListener(async (port: Port) => {
       if (port.name === transportName) {
-        this.ports.add(port);
         const handler = (data: any) => {
-          data._port = port;
+          data._extra = data._extra ?? {};
+          data._extra._port = port;
           this[callbackKey](data as ListenerOptions<BrowserExtensionsMainPort>);
         };
         port.onMessage.addListener(handler);
         port.onDisconnect.addListener(() => {
           port.onMessage.removeListener(handler);
-          this.ports.delete(port);
+          this.ports.forEach((_port, id) => {
+            if (_port === port) {
+              this.ports.delete(id);
+            }
+          });
+        });
+
+        // @ts-ignore
+        const id: string = await this.emit({
+          // @ts-ignore
+          name: connectEventName,
+          _extra: { _port: port },
+        });
+        this.ports.set(id, port);
+        this._onConnectCallback.forEach((callback) => {
+          callback(id);
         });
       }
     });
   }
+
+  private _onConnectCallback = new Set<MainCallback>();
+
+  onConnect(callback: MainCallback) {
+    this._onConnectCallback.add(callback);
+    return () => {
+      this._onConnectCallback.delete(callback);
+    };
+  }
+
+  private _onDisconnectCallback = new Set<MainCallback>();
+
+  onDisconnect(callback: MainCallback) {
+    this._onDisconnectCallback.add(callback);
+    return () => {
+      this._onDisconnectCallback.delete(callback);
+    };
+  }
 }
 
 export abstract class BrowserExtensionsClientTransport<
@@ -179,6 +224,24 @@ export abstract class BrowserExtensionsClientTransport<
       listener,
       sender,
     });
+    // @ts-ignore
+    this.listen(connectEventName, async () => {
+      Promise.resolve().then(() => {
+        this._onConnectCallback.forEach((callback) => {
+          callback();
+        });
+      });
+      return this.id;
+    });
+  }
+
+  private _onConnectCallback = new Set<ClientCallback>();
+
+  onConnect(callback: ClientCallback) {
+    this._onConnectCallback.add(callback);
+    return () => {
+      this._onConnectCallback.delete(callback);
+    };
   }
 }
 

src/transports/sharedWorkerTransport.ts

@@ -21,8 +21,8 @@ export interface SharedWorkerClientTransportOptions
   worker: SharedWorker;
 }
 
-type WorkerCallback = () => void | Promise<void>;
-type ClientCallback = (clientId: string) => void | Promise<void>;
+type ClientCallback = () => void | Promise<void>;
+type WorkerCallback = (clientId: string) => void | Promise<void>;
 
 const connectEventName = 'sharedworker-connect';
 const disconnectEventName = 'sharedworker-disconnect';
@@ -75,9 +75,9 @@ export abstract class SharedWorkerClientTransport<
     });
   }
 
-  private _onConnectCallback = new Set<WorkerCallback>();
+  private _onConnectCallback = new Set<ClientCallback>();
 
-  onConnect(callback: WorkerCallback) {
+  onConnect(callback: ClientCallback) {
     this._onConnectCallback.add(callback);
     return () => {
       this._onConnectCallback.delete(callback);
@@ -178,18 +178,18 @@ export abstract class SharedWorkerInternalTransport<
     });
   }
 
-  private _onConnectCallback = new Set<ClientCallback>();
+  private _onConnectCallback = new Set<WorkerCallback>();
 
-  onConnect(callback: ClientCallback) {
+  onConnect(callback: WorkerCallback) {
     this._onConnectCallback.add(callback);
     return () => {
       this._onConnectCallback.delete(callback);
     };
   }
 
-  private _onDisconnectCallback = new Set<ClientCallback>();
+  private _onDisconnectCallback = new Set<WorkerCallback>();
 
-  onDisconnect(callback: ClientCallback) {
+  onDisconnect(callback: WorkerCallback) {
     this._onDisconnectCallback.add(callback);
     return () => {
       this._onDisconnectCallback.delete(callback);