webextension-rpc/src

// Our secret tokens to recognise our messages
const RPC_CALL = '__RPC_CALL__'
const RPC_RESPONSE = '__RPC_RESPONSE__'

export class RpcError extends Error {
	constructor(message) {
		super(message)
		this.name = this.constructor.name
	}
}

export class RemoteError extends Error {
	constructor(message) {
		super(message)
		this.name = this.constructor.name
	}
}

export const injectRpcInfo = Symbol('RpcInfo')

export interface RpcInfo {
	tab: browser.tabs.Tab,
}

type AsyncFunction = Function & ((...args: any[]) => Promise<any>)

interface RpcMessage {
	__WEBEXTENSION_RPC_MESSAGE__: typeof RPC_CALL | typeof RPC_RESPONSE,
	funcName: string,
}

interface RpcCallMessage<F extends AsyncFunction = AsyncFunction> extends RpcMessage {
	__WEBEXTENSION_RPC_MESSAGE__: typeof RPC_CALL,
	funcName: F['name'],
	args: ParametersWithRpcInfo<F>,
	addRpcInfoAsArgument: number | false,
}

type RpcResponseMessage<F extends AsyncFunction = AsyncFunction> =
	| RpcResponseResolve<F>
	| RpcResponseReject
	| RpcResponseRpcError

interface RpcResponseMessage_base {
__WEBEXTENSION_RPC_MESSAGE__: typeof RPC_RESPONSE,
}

interface RpcResponseResolve<F extends AsyncFunction> extends RpcResponseMessage_base {
	returnValue: ReturnType<F>,
}

interface RpcResponseReject extends RpcResponseMessage_base {
	errorMessage: string,
}

interface RpcResponseRpcError extends RpcResponseMessage_base {
	rpcError: string,
}

function isRpcCallMessage(message: any): message is RpcCallMessage {
	return !!(message && message['__WEBEXTENSION_RPC_MESSAGE__'] === RPC_CALL)
}

function isRpcResponseMessage(message: any): message is RpcResponseMessage {
	return !!(message && message['__WEBEXTENSION_RPC_MESSAGE__'] === RPC_RESPONSE)
}

type Tail<T extends any[]> = T extends [infer _Head, ...infer Tail] ? Tail : [];

type ParametersWithRpcInfo<F extends AsyncFunction> = Parameters<F> extends [RpcInfo, ...any]
	? [typeof injectRpcInfo, ...Tail<Parameters<F>>]
	: Parameters<F>

// I thought the type below should allow putting the RpcInfo argument at any position, but it does not work.
// type ParametersWithRpcInfo<F extends AsyncFunction> = Array<any> & {
// 	[N in keyof Parameters<F>]: Parameters<F>[N] extends RpcInfo
// 		? typeof injectRpcInfo
// 		: Parameters<F>[N]
// }

/**
 * 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: ParametersWithRpcInfo<F>) => Promise<Awaited<ReturnType<F>>>


// === Initiating side ===

/**
 * 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 - The id of the tab whose content script is the remote side. Leave undefined
   to call the background script (from a content script).
 * @returns The proxy function.
 */
export function remoteFunction<F extends AsyncFunction>(
	funcName: string,
	{ tabId }: { tabId?: number } = {},
): RpcFunction<F> {
	const otherSide = (tabId !== undefined)
		? "the tab's content script"
		: 'the background script'

	const f: RpcFunction<F> = async function (...args): Promise<Awaited<ReturnType<F>>> {
		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<F>
		try {
			response = (tabId !== undefined)
				? 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
}


// === Executing side ===


/**
 * Register one or more functions to enable remote scripts to call them.
 *
 * @param functions - A `{ functionName: function }` mapping. Each function will be remotely
 * callable using the given name.
 * @returns The passed `functions` object.
 */
export function makeRemotelyCallable<Fs extends Record<string, AsyncFunction>>(
	functions: Fs,
): typeof functions {
	browser.runtime.onMessage.addListener(incomingRPCListener)
	return functions

	async function incomingRPCListener(message: any, sender: browser.runtime.MessageSender): Promise<RpcResponseMessage> {
		if (!isRpcCallMessage(message)) return

		const funcName = message.funcName
		const func = functions[funcName]
		if (func === undefined) {
			console.error(`Received RPC for unknown function: ${funcName}`)
			return {
				rpcError: `No such function registered for RPC: ${funcName}`,
				__WEBEXTENSION_RPC_MESSAGE__: RPC_RESPONSE,
			}
		}

		const args = message.args

		if (message.addRpcInfoAsArgument !== false) {
			const rpcInfo: RpcInfo = {
				tab: sender.tab,
			}
			args[message.addRpcInfoAsArgument] = rpcInfo
		}

		// Run the function, return the result.
		try {
			const returnValue = await func(...args)
			return {
				returnValue,
				__WEBEXTENSION_RPC_MESSAGE__: RPC_RESPONSE,
			}
		} catch (error) {
			return {
				errorMessage: error.message,
				__WEBEXTENSION_RPC_MESSAGE__: RPC_RESPONSE,
			}
		}
	}
}