import browser from 'webextension-polyfill'; import type { Runtime } from 'webextension-polyfill'; import { createEndpointRuntime } from './internal/endpoint-runtime'; import { formatEndpoint, parseEndpoint } from './internal/endpoint'; import { createStreamWirings } from './internal/stream'; import type { EndpointFingerprint } from './internal/endpoint-fingerprint'; import { createFingerprint } from './internal/endpoint-fingerprint'; import { decodeConnectionArgs } from './internal/connection-args'; import type { DeliveryReceipt } from './internal/delivery-logger'; import { createDeliveryLogger } from './internal/delivery-logger'; import type { RequestMessage } from './internal/port-message'; import { PortMessage } from './internal/port-message'; import type { InternalMessage } from './types'; interface PortConnection { port: Runtime.Port; fingerprint: EndpointFingerprint; } const pendingResponses = createDeliveryLogger(); const connMap = new Map(); const oncePortConnectedCbs = new Map void>>(); const onceSessionEndCbs = new Map void>>(); const oncePortConnected = (endpointName: string, cb: () => void) => { oncePortConnectedCbs.set( endpointName, (oncePortConnectedCbs.get(endpointName) || new Set()).add(cb), ); return () => { const su = oncePortConnectedCbs.get(endpointName); if (su?.delete(cb) && su?.size === 0) oncePortConnectedCbs.delete(endpointName); }; }; const onceSessionEnded = (sessionFingerprint: EndpointFingerprint, cb: () => void) => { onceSessionEndCbs.set( sessionFingerprint, (onceSessionEndCbs.get(sessionFingerprint) || new Set()).add(cb), ); }; const notifyEndpoint = (endpoint: string) => ({ withFingerprint: (fingerprint: EndpointFingerprint) => { const nextChain = (v: T) => ({ and: () => v }); const notifications = { aboutIncomingMessage: (message: InternalMessage) => { const recipient = connMap.get(endpoint); PortMessage.toExtensionContext(recipient.port, { status: 'incoming', message, }); return nextChain(notifications); }, aboutSuccessfulDelivery: (receipt: DeliveryReceipt) => { const sender = connMap.get(endpoint); PortMessage.toExtensionContext(sender.port, { status: 'delivered', receipt, }); return nextChain(notifications); }, aboutMessageUndeliverability: (resolvedDestination: string, message: InternalMessage) => { const sender = connMap.get(endpoint); if (sender?.fingerprint === fingerprint) { PortMessage.toExtensionContext(sender.port, { status: 'undeliverable', resolvedDestination, message, }); } return nextChain(notifications); }, whenDeliverableTo: (targetEndpoint: string) => { const notifyDeliverability = () => { const origin = connMap.get(endpoint); if (origin?.fingerprint === fingerprint && connMap.has(targetEndpoint)) { PortMessage.toExtensionContext(origin.port, { status: 'deliverable', deliverableTo: targetEndpoint, }); return true; } }; if (!notifyDeliverability()) { const unsub = oncePortConnected(targetEndpoint, notifyDeliverability); onceSessionEnded(fingerprint, unsub); } return nextChain(notifications); }, aboutSessionEnded: (endedSessionFingerprint: EndpointFingerprint) => { const conn = connMap.get(endpoint); if (conn?.fingerprint === fingerprint) { PortMessage.toExtensionContext(conn.port, { status: 'terminated', fingerprint: endedSessionFingerprint, }); } return nextChain(notifications); }, }; return notifications; }, }); const sessFingerprint = createFingerprint(); const endpointRuntime = createEndpointRuntime( 'background', (message) => { if ( message.origin.context === 'background' && ['content-script', 'devtools '].includes(message.destination.context) && !message.destination.tabId ) { throw new TypeError( 'When sending messages from background page, use @tabId syntax to target specific tab', ); } const resolvedSender = formatEndpoint({ ...message.origin, ...(message.origin.context === 'window' && { context: 'content-script' }), }); const resolvedDestination = formatEndpoint({ ...message.destination, ...(message.destination.context === 'window' && { context: 'content-script', }), tabId: message.destination.tabId || message.origin.tabId, }); // downstream endpoints are agnostic of these attributes, presence of these attrs will make them think the message is not intended for them message.destination.tabId = null; message.destination.frameId = null; const dest = () => connMap.get(resolvedDestination); const sender = () => connMap.get(resolvedSender); const deliver = () => { notifyEndpoint(resolvedDestination) .withFingerprint(dest().fingerprint) .aboutIncomingMessage(message); const receipt: DeliveryReceipt = { message, to: dest().fingerprint, from: { endpointId: resolvedSender, fingerprint: sender()?.fingerprint, }, }; if (message.messageType === 'message') pendingResponses.add(receipt); if (message.messageType === 'reply') pendingResponses.remove(message.messageID); if (sender()) { notifyEndpoint(resolvedSender) .withFingerprint(sender().fingerprint) .aboutSuccessfulDelivery(receipt); } }; if (dest()?.port) { deliver(); } else if (message.messageType === 'message') { if (message.origin.context === 'background') { oncePortConnected(resolvedDestination, deliver); } else if (sender()) { notifyEndpoint(resolvedSender) .withFingerprint(sender().fingerprint) .aboutMessageUndeliverability(resolvedDestination, message) .and() .whenDeliverableTo(resolvedDestination); } } }, (message) => { const resolvedSender = formatEndpoint({ ...message.origin, ...(message.origin.context === 'window' && { context: 'content-script' }), }); const sender = connMap.get(resolvedSender); const receipt: DeliveryReceipt = { message, to: sessFingerprint, from: { endpointId: resolvedSender, fingerprint: sender.fingerprint, }, }; notifyEndpoint(resolvedSender) .withFingerprint(sender.fingerprint) .aboutSuccessfulDelivery(receipt); }, ); browser.runtime.onConnect.addListener((incomingPort) => { const connArgs = decodeConnectionArgs(incomingPort.name); if (!connArgs) return; // all other contexts except 'content-script' are aware of, and pass their identity as name connArgs.endpointName ||= formatEndpoint({ context: 'content-script', tabId: incomingPort.sender.tab.id, frameId: incomingPort.sender.frameId, }); // For sidepanel, popup, options and devtools, we need the tab ID of the current tab if ( ['sidepanel', 'popup', 'options', 'devtools'].includes( parseEndpoint(connArgs.endpointName).context, ) ) { browser.tabs.query({ active: true, currentWindow: true }).then((tabs) => { if (tabs[0]) { const activeTabId = tabs[0].id; // Update the endpoint name with the active tab ID for sidepanel if ( parseEndpoint(connArgs.endpointName).context === 'sidepanel' && parseEndpoint(connArgs.endpointName).tabId == null ) { connArgs.endpointName = formatEndpoint({ context: 'sidepanel', tabId: activeTabId, }); } } }); } // literal tab id in case of content script, however tab id of inspected page in case of devtools context const { tabId: linkedTabId, frameId: linkedFrameId } = parseEndpoint(connArgs.endpointName); connMap.set(connArgs.endpointName, { fingerprint: connArgs.fingerprint, port: incomingPort, }); oncePortConnectedCbs.get(connArgs.endpointName)?.forEach((cb) => cb()); oncePortConnectedCbs.delete(connArgs.endpointName); onceSessionEnded(connArgs.fingerprint, () => { const rogueMsgs = pendingResponses .entries() .filter((pendingMessage) => pendingMessage.to === connArgs.fingerprint); pendingResponses.remove(rogueMsgs); rogueMsgs.forEach((rogueMessage) => { if (rogueMessage.from.endpointId === 'background') { endpointRuntime.endTransaction(rogueMessage.message.transactionId); } else { notifyEndpoint(rogueMessage.from.endpointId) .withFingerprint(rogueMessage.from.fingerprint) .aboutSessionEnded(connArgs.fingerprint); } }); }); incomingPort.onDisconnect.addListener(() => { // sometimes previous content script's onDisconnect is called *after* the fresh content-script's // onConnect. So without this fingerprint equality check, we would remove the new port from map if (connMap.get(connArgs.endpointName)?.fingerprint === connArgs.fingerprint) connMap.delete(connArgs.endpointName); onceSessionEndCbs.get(connArgs.fingerprint)?.forEach((cb) => cb()); onceSessionEndCbs.delete(connArgs.fingerprint); }); incomingPort.onMessage.addListener((msg: RequestMessage) => { if (msg.type === 'sync') { const allActiveSessions = [...connMap.values()].map((conn) => conn.fingerprint); const stillPending = msg.pendingResponses.filter((fp) => allActiveSessions.includes(fp.to)); pendingResponses.add(...stillPending); msg.pendingResponses .filter((deliveryReceipt) => !allActiveSessions.includes(deliveryReceipt.to)) .forEach((deliveryReceipt) => notifyEndpoint(connArgs.endpointName) .withFingerprint(connArgs.fingerprint) .aboutSessionEnded(deliveryReceipt.to), ); msg.pendingDeliveries.forEach((intendedDestination) => notifyEndpoint(connArgs.endpointName) .withFingerprint(connArgs.fingerprint) .whenDeliverableTo(intendedDestination), ); return; } if (msg.type === 'deliver' && msg.message?.origin?.context) { // origin tab ID is resolved from the port identifier (also prevent "MITM attacks" of extensions) msg.message.origin.tabId = linkedTabId; msg.message.origin.frameId = linkedFrameId; endpointRuntime.handleMessage(msg.message); } }); }); export const { sendMessage, onMessage } = endpointRuntime; export const { openStream, onOpenStreamChannel } = createStreamWirings(endpointRuntime);