webextension-rpc/src

import type { RpcServer } from "./RpcServer"
import { isRpcResponseMessage, RPC_CALL } from "./common"
import type { AsyncFunction, RpcCallMessage, RpcResponseMessage, ReplaceRpcInfo } from "./common"

export type { AsyncFunction }

/**
 * Error thrown when the remote function could not be found/executed.
 */
export class RpcError extends Error {
	constructor(message: string) {
		super(message)
		this.name = this.constructor.name
	}
}

/**
 * Error thrown when the remote function threw an error.
 */
export class RemoteError extends Error {
	constructor(message: string) {
		super(message)
		this.name = this.constructor.name
	}
}

/**
 * If the special symbol `injectRpcInfo` is passed as the first argument to a proxy function, this
 * argument will be replaced on the executing side by an `RpcInfo` object.
 */
export const injectRpcInfo = Symbol('RpcInfo')

export interface RpcOptions {
	/**
	 * The id of the tab whose content script is the remote side. Leave undefined
	 * to call the background script (from a content script).
	 */
	tabId?: number,
}

/**
 * RpcFunction<F> equals the function F, except for two tweaks:
 * - In the parameters, any RpcInfo is swapped for rpcInfoSymbol.
 * - In the return type, nested promises become one promise.
 */
export type RpcFunction<F extends AsyncFunction> = (
	...args: ReplaceRpcInfo<Parameters<F>, typeof injectRpcInfo>
) => Promise<Awaited<ReturnType<F>>>

type FunctionsOf<R extends RpcServer> = R['functions']

/**
 * Define a ‘connection’ for remote procedure calls.
 * @param options.tabId - The id of the tab whose content script is the remote side. Leave undefined
 * to call functions of the background script (from a content script).
 */
export class RpcClient<TheRpcServer extends RpcServer> {
	tabId: number | undefined

	constructor(options: RpcOptions = {}) {
		this.tabId = options.tabId;

	}

	/**
	 * Create a proxy function that invokes the specified remote function.
	 * @param funcName - Name of the function as registered on the remote side.
	 * @param options.tabId - Overrides the `tabId` passed to `RpcClient`.
	 * @returns The proxy function.
	 */
	func<FunctionName extends string & keyof FunctionsOf<TheRpcServer>>(
		funcName: FunctionName,
		{
			tabId = this.tabId,
		}: RpcOptions = {},
	): RpcFunction<FunctionsOf<TheRpcServer>[FunctionName]> {
		type RemoteFunction = FunctionsOf<TheRpcServer>[FunctionName]

		const otherSide = (tabId !== undefined)
			? "the tab's content script"
			: 'the background script'

		const f: RpcFunction<RemoteFunction> = async function (...args): ReturnType<RpcFunction<RemoteFunction>> {
			const message: RpcCallMessage = {
				__WEBEXTENSION_RPC_MESSAGE__: RPC_CALL,
				funcName,
				args,
				addRpcInfoAsArgument: false,
			}

			if (args.includes(injectRpcInfo)) {
				const argIndex = args.indexOf(injectRpcInfo)
				message.addRpcInfoAsArgument = argIndex
				message.args[argIndex] = null
			}

			// Try send the message and await the response.
			let response: RpcResponseMessage<RemoteFunction>
			try {
				response = (tabId !== undefined && tabId !== null)
					? await browser.tabs.sendMessage(tabId, message)
					: await browser.runtime.sendMessage(message)
			} catch (err) {}

			// Check if we got an error or no response.
			if (response === undefined) {
				throw new RpcError(
					`Got no response when trying to call '${funcName}'. `
					+ `Did you enable RPC in ${otherSide}?`
				)
			}

			// Check if it was *our* listener that responded.
			if (!isRpcResponseMessage(response)) {
				throw new RpcError(
					`RPC got a response from an interfering listener while calling '${funcName}' in `
					+ `${otherSide}`
				)
			}

			// If we could not invoke the function on the other side, throw an error.
			if ('rpcError' in response) {
				throw new RpcError(response.rpcError)
			}

			// Return the value or throw the error we received from the other side.
			if ('errorMessage' in response) {
				throw new RemoteError(response.errorMessage)
			} else {
				return response.returnValue
			}
		}

		// Give it a name, could be helpful in debugging
		Object.defineProperty(f, 'name', { value: `${funcName}_RPC` })
		return f
	}
}