diff --git a/api/messaging/src/chunked-stream.ts b/api/messaging/src/chunked-stream.ts new file mode 100644 index 000000000..31237381e --- /dev/null +++ b/api/messaging/src/chunked-stream.ts @@ -0,0 +1,183 @@ +import type { + Chunk, + ChunkCollectionID, + InitChunk, + MessageEventCallback, + PortName +} from "./types" +import { getExtRuntime } from "./utils" + +const maxChunkSize = 4_200_000 + +/** + * Split large data into multiple chunks to + * bypass the browser's limit on runtime messages. + */ +export function createChunksFromData(data: unknown): Chunk[] { + // serialize data to buffer + const jsonObj = JSON.stringify(data) + const serialized = new TextEncoder().encode(jsonObj) + + // split serialized data + const bytes: number[][] = [] + + for (let i = 0; i < serialized.length; i++) { + const chunk = Math.floor(i / maxChunkSize) + + if (!bytes[chunk]) bytes[chunk] = [] + + bytes[chunk].push(serialized[i]) + } + + // create a chunk collection ID + const collectionID = Math.floor(Math.random() * 100) + + // create chunks + const chunks: Chunk[] = bytes.map((byteGroup, i) => ({ + name: "__PLASMO_MESSAGING_CHUNK__", + type: i === byteGroup.length - 1 ? "end" : i === 0 ? "init" : "data", + index: i, + chunkCollectionId: collectionID, + data: byteGroup + })) + + // add total chunk length + const initChunk = chunks.find((chunk) => chunk.type === "init") as InitChunk + + initChunk.totalChunks = chunks.length + initChunk.dataLength = serialized.length + + return chunks +} + +/** + * Reconstruct split data from "createChunksFromData()" + */ +export function buildDataFromChunks(chunks: Chunk[]): T { + // find the init chunk + const initChunk = chunks.find((chunk) => chunk.type === "init") as InitChunk + + // validate init chunk and check if + // the chunks are complete + if ( + !initChunk || + initChunk.totalChunks !== chunks.length || + typeof initChunk.dataLength === "undefined" + ) { + throw new Error( + "Failed to validate init chunk: incomplete chunk array / no data length / no init chunk" + ) + } + + // initialize the encoded data + const encoded = new Uint8Array(initChunk.dataLength) + + // sort chunks by their index + // this is to make sure we are + // setting the encoded bytes in + // the correct order + chunks.sort((a, b) => a.index - b.index) + + // set bytes + for (let i = 0; i < chunks.length; i++) { + encoded.set(chunks[i].data, chunks[i - 1]?.data?.length || 0) + } + + // decode the data + const serialized = new TextDecoder().decode(encoded) + const obj: T = JSON.parse(serialized) + + return obj +} + +/** + * Advanced chunked streaming port extending the default + * chrome.runtime.Port + */ +export const createChunkedStreamPort = ( + name: PortName +): chrome.runtime.Port => { + // connect to the port + const port = getExtRuntime().connect({ name }) + + // chunk map + const chunkMap = new Map() + + // intercepted event listeners map + // Map format: key - original handler, value - interceptor + const listenerMap = new Map() + + // setup interceptor + return { + ...port, + postMessage(message: unknown) { + // split chunks + const chunks = createChunksFromData(message) + + // get if chunks are needed + // if not, just send the message + if (chunks.length >= 1) { + return port.postMessage(message) + } + + // send chunks + for (let i = 0; i < chunks.length; i++) { + port.postMessage(chunks[i]) + } + }, + onMessage: { + ...port.onMessage, + addListener(callback: MessageEventCallback) { + // interceptor for the chunks + const interceptor: MessageEventCallback = (message: Chunk, port) => { + // only handle chunks + if (message?.name !== "__PLASMO_MESSAGING_CHUNK__") { + return callback(message, port) + } + + // check if a group exists for this + // chunk in the chunkMap + let group = chunkMap.get(message.chunkCollectionId) + + // if the group exists, add chunk to it + // otherwise create the group + if (!!group) group.push(message) + else chunkMap.set(message.chunkCollectionId, [message]) + + // update group (in case it was undefined before) + group = chunkMap.get(message.chunkCollectionId) + + // check if all chunks have been received + const initChunk = group.find( + (chunk) => chunk.type === "init" + ) as InitChunk + + if (group.length !== initChunk.totalChunks) return + + // check if the listener is present + if (!listenerMap.get(callback)) return + + // build message data + const data = buildDataFromChunks(group) + + // call original listener to handle + // the reconstructed message + return callback(data, port) + } + + // add listener + port.onMessage.addListener(interceptor) + + // map listener + listenerMap.set(callback, interceptor) + }, + removeListener(callback: MessageEventCallback) { + // remove listener from the original port + port.onMessage.removeListener(listenerMap.get(callback)) + + // remove listener from listener map + listenerMap.delete(callback) + } + } + } +} diff --git a/api/messaging/src/hook.ts b/api/messaging/src/hook.ts index 86d5b2ab6..df3b88489 100644 --- a/api/messaging/src/hook.ts +++ b/api/messaging/src/hook.ts @@ -27,18 +27,18 @@ export const useMessage = ( } } -export const usePort: PlasmoMessaging.PortHook = (name) => { +export const usePort: PlasmoMessaging.PortHook = (portKey) => { const portRef = useRef() const reconnectRef = useRef(0) const [data, setData] = useState() useEffect(() => { - if (!name) { + if (!portKey) { return null } const { port, disconnect } = portListen( - name, + portKey, (msg) => { setData(msg) }, @@ -50,7 +50,7 @@ export const usePort: PlasmoMessaging.PortHook = (name) => { portRef.current = port return disconnect }, [ - name, + portKey, reconnectRef.current // This is needed to force a new port ref ]) @@ -58,11 +58,11 @@ export const usePort: PlasmoMessaging.PortHook = (name) => { data, send: (body) => { portRef.current.postMessage({ - name, + name: portKey, body }) }, - listen: (handler) => portListen(name, handler) + listen: (handler) => portListen(portKey, handler) } } diff --git a/api/messaging/src/index.ts b/api/messaging/src/index.ts index 257075c4d..c515ba6cc 100644 --- a/api/messaging/src/index.ts +++ b/api/messaging/src/index.ts @@ -8,7 +8,8 @@ export type { PortName, PortsMetadata, MessagesMetadata, - OriginContext + OriginContext, + PortKey } from "./types" /** diff --git a/api/messaging/src/port.ts b/api/messaging/src/port.ts index 5172549af..0bdab96fe 100644 --- a/api/messaging/src/port.ts +++ b/api/messaging/src/port.ts @@ -1,31 +1,39 @@ -import type { PortName } from "./index" +import { createChunkedStreamPort } from "./chunked-stream" +import type { PortKey, PortName } from "./index" import { getExtRuntime } from "./utils" const portMap = new Map() -export const getPort = (name: PortName) => { - const port = portMap.get(name) +export const getPort = (portKey: PortKey) => { + const portName = typeof portKey === "string" ? portKey : portKey.name + const isChunked = typeof portKey !== "string" && portKey.isChunked + + const port = portMap.get(portName) + if (!!port) { return port } - const newPort = getExtRuntime().connect({ name }) - portMap.set(name, newPort) + const newPort = isChunked + ? createChunkedStreamPort(portName) + : getExtRuntime().connect({ name: portName }) + + portMap.set(portName, newPort) return newPort } -export const removePort = (name: PortName) => { - portMap.delete(name) +export const removePort = (portKey: PortKey) => { + portMap.delete(typeof portKey === "string" ? portKey : portKey.name) } export const listen = ( - name: PortName, + portKey: PortKey, handler: (msg: ResponseBody) => Promise | void, onReconnect?: () => void ) => { - const port = getPort(name) + const port = getPort(portKey) function reconnectHandler() { - removePort(name) + removePort(portKey) onReconnect?.() } diff --git a/api/messaging/src/types.ts b/api/messaging/src/types.ts index b8b7b98f8..d984793cb 100644 --- a/api/messaging/src/types.ts +++ b/api/messaging/src/types.ts @@ -92,7 +92,9 @@ export namespace PlasmoMessaging { } export interface PortHook { - , TResponseBody = any>(name: PortName): { + , TResponseBody = any>( + portKey: PortKey + ): { data?: TResponseBody send: (payload: TRequestBody) => void listen: ( @@ -111,3 +113,40 @@ export type OriginContext = | "sandbox-page" | "content-script" | "window" + +export type PortKey = + | PortName + | { + name: PortName + // Enable chunking of port data stream. This split the data into smaller chunk and stream them through the port, overcoming the port bandwidth limitation. + isChunked?: boolean + } + +export type ChunkCollectionID = number + +export type MessageEventCallback = ( + message: unknown, + port: chrome.runtime.Port +) => void + +export interface Chunk { + name: "__PLASMO_MESSAGING_CHUNK__" + type: "init" | "end" | "data" + index: number + chunkCollectionId: ChunkCollectionID + data: number[] +} + +export interface InitChunk extends Chunk { + type: "init" + dataLength: number + totalChunks: number +} + +export interface DataChunk extends Chunk { + type: "data" +} + +export interface EndChunk extends Chunk { + type: "end" +}