johnandersen777.bsky.social/atlibs

Switch repository

Loading…

// XRPC subscriptions over WebSocket for hosts:
// the browser WebSocket or Node 22+'s global one, no deps. atproto event streams send binary
// frames, a DAG-CBOR header { op, t } then a DAG-CBOR body; jetstream-style streams send JSON text.

import { dagCbor, toBase32 } from "../encoding/index.js";

// fromCharCode in chunks: spreading a large byte string overflows the stack
const toBase64 = (bytes) => {
	let text = "";
	for (let i = 0; i < bytes.length; i += 0x8000) text += String.fromCharCode(...bytes.subarray(i, i + 0x8000));
	return btoa(text).replace(/=+$/, "");
};
const fromBase64 = (text) => Uint8Array.from(atob(text.replace(/-/g, "+").replace(/_/g, "/")), (c) => c.charCodeAt(0));

/** Arrays, maps and tags nest at most this deep: frames come from the network, and each level recurses. */
const MAX_DEPTH = 128;

/**
 * Decode consecutive DAG-CBOR items into the atproto JSON data model: bytes become { $bytes: base64 },
 * CIDs (tag 42) { $link: base32 CIDv1 }. Returns the items.
 * @param {Uint8Array} bytes
 */
export function decodeCbor(bytes) {
	const view = new DataView(bytes.buffer, bytes.byteOffset, bytes.byteLength);
	let at = 0;
	const length = (info) => {
		if (info < 24) return info;
		if (info === 24) return view.getUint8(at++);
		if (info === 25) return (at += 2), view.getUint16(at - 2);
		if (info === 26) return (at += 4), view.getUint32(at - 4);
		if (info === 27) return (at += 8), Number(view.getBigUint64(at - 8));
		throw new Error("CBOR: indefinite lengths are not DAG-CBOR");
	};
	const span = (n) => {
		if (at + n > bytes.length) throw new Error("CBOR: truncated");
		return bytes.subarray(at, (at += n));
	};
	const item = (depth = 0) => {
		if (depth > MAX_DEPTH) throw new Error(`CBOR: nested deeper than ${MAX_DEPTH}`);
		const next = () => item(depth + 1);
		const head = view.getUint8(at++);
		const major = head >> 5;
		const info = head & 31;
		if (major === 7) {
			if (info === 20) return false;
			if (info === 21) return true;
			if (info === 22 || info === 23) return null;
			if (info === 25) return (at += 2), view.getFloat16?.(at - 2) ?? NaN;
			if (info === 26) return (at += 4), view.getFloat32(at - 4);
			if (info === 27) return (at += 8), view.getFloat64(at - 8);
			throw new Error(`CBOR: simple value ${info}`);
		}
		const n = length(info);
		if (major === 0) return n;
		if (major === 1) return -1 - n;
		if (major === 2) return { $bytes: toBase64(span(n)) };
		if (major === 3) return new TextDecoder().decode(span(n));
		if (major === 4) return Array.from({ length: n }, next);
		if (major === 5) return Object.fromEntries(Array.from({ length: n }, () => [next(), next()]));
		// major 6: tag 42 is a CID, bytes with a 0x00 multibase-identity prefix.
		const tagged = next();
		if (n !== 42 || !tagged.$bytes) throw new Error(`CBOR: tag ${n}`);
		return { $link: `b${toBase32(fromBase64(tagged.$bytes).subarray(1))}` };
	};
	const items = [];
	while (at < bytes.length) items.push(item());
	return items;
}

/** Encode values (the atproto JSON data model, as decodeCbor returns it) as consecutive DAG-CBOR items. */
export function encodeCbor(...values) {
	const items = values.map((value) => dagCbor(value, { links: true }));
	const out = new Uint8Array(items.reduce((n, item) => n + item.length, 0));
	items.reduce((at, item) => (out.set(item, at), at + item.length), 0);
	return out;
}

/**
 * One event-stream frame as a message: { $type: nsid + t, ...body } for op 1; an error frame (op -1)
 * throws its error. Text is JSON.
 */
export function decodeFrame(data, nsid = "") {
	if (typeof data === "string") return JSON.parse(data);
	const [header, body] = decodeCbor(new Uint8Array(data));
	if (header?.op === -1) throw new Error(`${body?.error ?? "error frame"}${body?.message ? `: ${body.message}` : ""}`);
	return { $type: `${nsid}${header?.t ?? ""}`, ...body };
}

/**
 * What send(data) writes: a string as a text frame; { $bytes } as a binary frame; with framing "cbor"
 * an object { $type: "...#t", ...body } as an event-stream frame (header { op: 1, t }); otherwise JSON text.
 */
export function encodeFrame(data, framing) {
	if (typeof data === "string") return data;
	if (data && typeof data.$bytes === "string" && Object.keys(data).length === 1) return fromBase64(data.$bytes);
	if (framing === "cbor") {
		const { $type = "", ...body } = data ?? {};
		// t is the "#fragment" of $type; with no "#" the header has no t ($type is the nsid itself)
		const hash = String($type).indexOf("#");
		return encodeCbor(hash < 0 ? { op: 1 } : { op: 1, t: String($type).slice(hash) }, body);
	}
	return JSON.stringify(data);
}

/**
 * Resolve a subscribe request's serviceAuth: mint(service, nsid)
 * resolves to a service-auth token (aud the service's did:web, lxm nsid). { param } puts it in that query
 * parameter; {} sends it as Authorization, which only hosts whose WebSocket takes headers can do
 * (canHeader false in browsers). Returns the request without serviceAuth, with params or headers set.
 */
export async function withServiceAuth(nsid, request, mint, { canHeader = true } = {}) {
	const { serviceAuth, ...rest } = request;
	if (rest.headers && !canHeader) throw new Error(`${nsid}: browsers cannot set WebSocket headers`);
	if (!serviceAuth) return rest;
	if (!rest.service?.startsWith("https://")) throw new Error(`${nsid} serviceAuth needs an https service`);
	if (serviceAuth.param === undefined && !canHeader) throw new Error(`${nsid}: browsers cannot set WebSocket headers; serviceAuth needs { param }, for a service that takes the token in a query parameter`);
	const token = await mint(rest.service, nsid);
	if (typeof token !== "string" || !token) throw new Error(`${nsid}: no service-auth token`);
	return serviceAuth.param === undefined ? { ...rest, headers: { ...rest.headers, authorization: `Bearer ${token}` } } : { ...rest, params: { ...rest.params, [serviceAuth.param]: token } };
}

export const wsUrl = (base, nsid, params) => {
	const url = new URL(`/xrpc/${nsid}`, base);
	url.protocol = url.protocol === "http:" ? "ws:" : "wss:";
	if (params) url.search = new URLSearchParams(params).toString();
	return url.href;
};

const HIGH_WATER = 1 << 20; // bytes buffered before send() waits for the socket to drain

/**
 * Open a WebSocket subscription. The result is an async iterable of decoded messages with
 * send(data) (see encodeFrame; resolves once the frame is buffered under HIGH_WATER) and close().
 * The iterable ends when the socket closes cleanly and throws when it fails or an error frame arrives.
 * @param {string} url  wss://host/xrpc/nsid?params (wsUrl)
 * @param {{ nsid?: string, signal?: AbortSignal, framing?: "cbor" | "json", headers?: Record<string, string>,
 *   keepalive?: { ms: number, data: unknown }, WebSocket?: typeof WebSocket }} options
 *   framing only changes what send() writes; incoming binary is always CBOR frames and text JSON.
 *   headers: Node's WebSocket only (browsers cannot set them). keepalive: send data every ms.
 */
export function subscribeXrpc(url, { nsid, signal, framing, headers, keepalive, WebSocket: WS = globalThis.WebSocket } = {}) {
	const ws = headers ? new WS(url, { headers }) : new WS(url);
	ws.binaryType = "arraybuffer";
	const queue = [];
	let finished = false;
	let failure;
	let wake;
	const notify = () => wake?.();
	const opened = new Promise((resolve, reject) => {
		ws.onopen = resolve;
		ws.onclose = (event) => {
			finished = true;
			if (!failure && event.code !== 1000 && event.code !== 1005) failure = new Error(`${nsid ?? url} closed: ${event.code}${event.reason ? ` ${event.reason}` : ""}`);
			clearInterval(ping);
			reject(failure ?? new Error(`${nsid ?? url} closed`));
			notify();
		};
	});
	opened.catch(() => {});
	ws.onerror = () => (failure ??= new Error(`${nsid ?? url}: WebSocket error`));
	ws.onmessage = ({ data }) => {
		try {
			queue.push(decodeFrame(data, nsid));
		} catch (error) {
			failure = error;
			ws.close();
		}
		notify();
	};
	const close = () => ws.readyState < 2 && ws.close(1000);
	signal?.addEventListener("abort", close, { once: true });
	if (signal?.aborted) close();
	const send = async (data) => {
		await opened;
		const frame = encodeFrame(data, framing);
		while (ws.bufferedAmount > HIGH_WATER && ws.readyState === 1) await new Promise((r) => setTimeout(r, 10));
		if (ws.readyState !== 1) throw new Error(`${nsid ?? url} is closed`);
		ws.send(frame);
	};
	// clamped to setTimeout's limit: past it Node fires every 1 ms
	const ping = keepalive?.ms > 0 ? setInterval(() => send(keepalive.data).catch(() => {}), Math.min(keepalive.ms, 2 ** 31 - 1)) : undefined;
	return {
		send,
		close,
		[Symbol.asyncIterator]() {
			return {
				next: async () => {
					while (!queue.length && !finished) await new Promise((resolve) => (wake = resolve));
					if (queue.length) return { value: queue.shift(), done: false };
					if (failure) throw failure;
					return { value: undefined, done: true };
				},
				return: async () => (close(), { value: undefined, done: true }),
			};
		},
	};
}