rivetkit 2.0.23 → 2.0.24

package/src/actor/router-websocket-endpoints.ts

@@ -0,0 +1,408 @@
+import type { WSContext } from "hono/ws";
+import invariant from "invariant";
+import type { AnyConn } from "@/actor/conn/mod";
+import type { AnyActorInstance } from "@/actor/instance/mod";
+import type { InputData } from "@/actor/protocol/serde";
+import { type Encoding, EncodingSchema } from "@/actor/protocol/serde";
+import {
+	PATH_CONNECT,
+	PATH_INSPECTOR_CONNECT,
+	PATH_WEBSOCKET_BASE,
+	PATH_WEBSOCKET_PREFIX,
+	WS_PROTOCOL_CONN_PARAMS,
+	WS_PROTOCOL_ENCODING,
+} from "@/common/actor-router-consts";
+import { deconstructError } from "@/common/utils";
+import type {
+	RivetMessageEvent,
+	UniversalWebSocket,
+} from "@/common/websocket-interface";
+import type { RunnerConfig } from "@/registry/run-config";
+import { promiseWithResolvers } from "@/utils";
+import type { ConnDriver } from "./conn/driver";
+import { createRawWebSocketDriver } from "./conn/drivers/raw-websocket";
+import { createWebSocketDriver } from "./conn/drivers/websocket";
+import type { ActorDriver } from "./driver";
+import { loggerWithoutContext } from "./log";
+import { parseMessage } from "./protocol/old";
+import { getRequestExposeInternalError } from "./router-endpoints";
+
+// TODO: Merge with ConnectWebSocketOutput interface
+export interface UpgradeWebSocketArgs {
+	conn?: AnyConn;
+	actor?: AnyActorInstance;
+	onRestore?: (ws: WSContext) => void;
+	onOpen: (event: any, ws: WSContext) => void;
+	onMessage: (event: any, ws: WSContext) => void;
+	onClose: (event: any, ws: WSContext) => void;
+	onError: (error: any, ws: WSContext) => void;
+}
+
+interface WebSocketHandlerOpts {
+	runConfig: RunnerConfig;
+	request: Request | undefined;
+	encoding: Encoding;
+	actor: AnyActorInstance;
+	closePromiseResolvers: ReturnType<typeof promiseWithResolvers<void>>;
+	conn: AnyConn;
+	exposeInternalError: boolean;
+}
+
+/** Handler for a specific WebSocket route. Used in routeWebSocket. */
+type WebSocketHandler = (
+	opts: WebSocketHandlerOpts,
+) => Promise<UpgradeWebSocketArgs>;
+
+export async function routeWebSocket(
+	request: Request | undefined,
+	requestPath: string,
+	requestHeaders: Record<string, string>,
+	runConfig: RunnerConfig,
+	actorDriver: ActorDriver,
+	actorId: string,
+	encoding: Encoding,
+	parameters: unknown,
+	gatewayId: ArrayBuffer | undefined,
+	requestId: ArrayBuffer | undefined,
+	isHibernatable: boolean,
+	isRestoringHibernatable: boolean,
+): Promise<UpgradeWebSocketArgs> {
+	const exposeInternalError = request
+		? getRequestExposeInternalError(request)
+		: false;
+
+	let createdConn: AnyConn | undefined;
+	try {
+		const actor = await actorDriver.loadActor(actorId);
+
+		actor.rLog.debug({
+			msg: "new websocket connection",
+			actorId,
+			requestPath,
+			isHibernatable,
+		});
+
+		// Promise used to wait for the websocket close in `disconnect`
+		const closePromiseResolvers = promiseWithResolvers<void>();
+
+		// Route WebSocket & create driver
+		let handler: WebSocketHandler;
+		let connDriver: ConnDriver;
+		if (requestPath === PATH_CONNECT) {
+			const { driver, setWebSocket } = createWebSocketDriver(
+				isHibernatable
+					? { gatewayId: gatewayId!, requestId: requestId! }
+					: undefined,
+				encoding,
+				closePromiseResolvers.promise,
+			);
+			handler = handleWebSocketConnect.bind(undefined, setWebSocket);
+			connDriver = driver;
+		} else if (
+			requestPath === PATH_WEBSOCKET_BASE ||
+			requestPath.startsWith(PATH_WEBSOCKET_PREFIX)
+		) {
+			const { driver, setWebSocket } = createRawWebSocketDriver(
+				isHibernatable
+					? { gatewayId: gatewayId!, requestId: requestId! }
+					: undefined,
+				closePromiseResolvers.promise,
+			);
+			handler = handleRawWebSocket.bind(undefined, setWebSocket);
+			connDriver = driver;
+		} else if (requestPath === PATH_INSPECTOR_CONNECT) {
+			// This returns raw UpgradeWebSocketArgs instead of accepting a
+			// Conn since this does not need a Conn
+			return await handleWebSocketInspectorConnect();
+		} else {
+			throw `WebSocket Path Not Found: ${requestPath}`;
+		}
+
+		// Prepare connection
+		const conn = await actor.connectionManager.prepareConn(
+			connDriver,
+			parameters,
+			request,
+			requestPath,
+			requestHeaders,
+			isHibernatable,
+			isRestoringHibernatable,
+		);
+		createdConn = conn;
+
+		// Create handler
+		//
+		// This must call actor.connectionManager.connectConn in onOpen.
+		return await handler({
+			runConfig,
+			request,
+			encoding,
+			actor,
+			closePromiseResolvers,
+			conn,
+			exposeInternalError,
+		});
+	} catch (error) {
+		const { group, code } = deconstructError(
+			error,
+			loggerWithoutContext(),
+			{},
+			exposeInternalError,
+		);
+
+		// Clean up connection
+		if (createdConn) {
+			createdConn.disconnect(`${group}.${code}`);
+		}
+
+		// Return handler that immediately closes with error
+		// Note: createdConn should always exist here, but we use a type assertion for safety
+		return {
+			conn: createdConn!,
+			onOpen: (_evt: any, ws: WSContext) => {
+				ws.close(1011, code);
+			},
+			onMessage: (_evt: { data: any }, ws: WSContext) => {
+				ws.close(1011, "actor.not_loaded");
+			},
+			onClose: (_event: any, _ws: WSContext) => {},
+			onError: (_error: unknown) => {},
+		};
+	}
+}
+
+/**
+ * Creates a WebSocket connection handler
+ */
+export async function handleWebSocketConnect(
+	setWebSocket: (ws: WSContext) => void,
+	{
+		runConfig,
+		encoding,
+		actor,
+		closePromiseResolvers,
+		conn,
+		exposeInternalError,
+	}: WebSocketHandlerOpts,
+): Promise<UpgradeWebSocketArgs> {
+	return {
+		conn,
+		actor,
+		onRestore: (ws: WSContext) => {
+			setWebSocket(ws);
+		},
+		// NOTE: onOpen cannot be async since this messes up the open event listener order
+		onOpen: (_evt: any, ws: WSContext) => {
+			actor.rLog.debug("actor websocket open");
+
+			setWebSocket(ws);
+
+			// This will not be called by restoring hibernatable
+			// connections. All restoration is done in prepareConn.
+			actor.connectionManager.connectConn(conn);
+		},
+		onMessage: (evt: RivetMessageEvent, ws: WSContext) => {
+			// Handle message asynchronously
+			actor.rLog.debug({ msg: "received message" });
+
+			const value = evt.data.valueOf() as InputData;
+			parseMessage(value, {
+				encoding: encoding,
+				maxIncomingMessageSize: runConfig.maxIncomingMessageSize,
+			})
+				.then((message) => {
+					actor.processMessage(message, conn).catch((error) => {
+						const { code } = deconstructError(
+							error,
+							actor.rLog,
+							{
+								wsEvent: "message",
+							},
+							exposeInternalError,
+						);
+						ws.close(1011, code);
+					});
+				})
+				.catch((error) => {
+					const { code } = deconstructError(
+						error,
+						actor.rLog,
+						{
+							wsEvent: "message",
+						},
+						exposeInternalError,
+					);
+					ws.close(1011, code);
+				});
+		},
+		onClose: (
+			event: {
+				wasClean: boolean;
+				code: number;
+				reason: string;
+			},
+			ws: WSContext,
+		) => {
+			closePromiseResolvers.resolve();
+
+			if (event.wasClean) {
+				actor.rLog.info({
+					msg: "websocket closed",
+					code: event.code,
+					reason: event.reason,
+					wasClean: event.wasClean,
+				});
+			} else {
+				actor.rLog.warn({
+					msg: "websocket closed",
+					code: event.code,
+					reason: event.reason,
+					wasClean: event.wasClean,
+				});
+			}
+
+			// HACK: Close socket in order to fix bug with Cloudflare leaving WS in closing state
+			// https://github.com/cloudflare/workerd/issues/2569
+			ws.close(1000, "hack_force_close");
+
+			// Wait for actor.createConn to finish before removing the connection
+			conn.disconnect(event?.reason);
+		},
+		onError: (_error: unknown) => {
+			try {
+				// Actors don't need to know about this, since it's abstracted away
+				actor.rLog.warn({ msg: "websocket error" });
+			} catch (error) {
+				deconstructError(
+					error,
+					actor.rLog,
+					{ wsEvent: "error" },
+					exposeInternalError,
+				);
+			}
+		},
+	};
+}
+
+export async function handleRawWebSocket(
+	setWebSocket: (ws: UniversalWebSocket) => void,
+	{ request, actor, closePromiseResolvers, conn }: WebSocketHandlerOpts,
+): Promise<UpgradeWebSocketArgs> {
+	return {
+		conn,
+		actor,
+		onRestore: (wsContext: WSContext) => {
+			const ws = wsContext.raw as UniversalWebSocket;
+			invariant(ws, "missing wsContext.raw");
+
+			setWebSocket(ws);
+		},
+		// NOTE: onOpen cannot be async since this will cause the client's open
+		// event to be called before this completes. Do all async work in
+		// handleRawWebSocket root.
+		onOpen: (_evt: any, wsContext: WSContext) => {
+			const ws = wsContext.raw as UniversalWebSocket;
+			invariant(ws, "missing wsContext.raw");
+
+			setWebSocket(ws);
+
+			// This will not be called by restoring hibernatable
+			// connections. All restoration is done in prepareConn.
+			actor.connectionManager.connectConn(conn);
+
+			// Call the actor's onWebSocket handler with the adapted WebSocket
+			//
+			// NOTE: onWebSocket is called inside this function. Make sure
+			// this is called synchronously within onOpen.
+			actor.handleRawWebSocket(conn, ws, request);
+		},
+		onMessage: (event: any, ws: any) => {
+			// Find the adapter for this WebSocket
+			const adapter = (ws as any).__adapter;
+			if (adapter) {
+				adapter._handleMessage(event);
+			}
+		},
+		onClose: (evt: any, ws: any) => {
+			// Resolve the close promise
+			closePromiseResolvers.resolve();
+
+			// Clean up the connection
+			conn.disconnect(evt?.reason);
+		},
+		onError: (error: any, ws: any) => {},
+	};
+}
+
+export async function handleWebSocketInspectorConnect(): Promise<UpgradeWebSocketArgs> {
+	return {
+		// NOTE: onOpen cannot be async since this messes up the open event listener order
+		onOpen: (_evt: any, ws: WSContext) => {
+			ws.send("Hello world");
+		},
+		onMessage: (evt: RivetMessageEvent, ws: WSContext) => {
+			ws.send("Pong");
+		},
+		onClose: (
+			event: {
+				wasClean: boolean;
+				code: number;
+				reason: string;
+			},
+			ws: WSContext,
+		) => {
+			// TODO:
+		},
+		onError: (_error: unknown) => {
+			// TODO:
+		},
+	};
+}
+
+/**
+ * Parse encoding and connection parameters from WebSocket Sec-WebSocket-Protocol header
+ */
+export function parseWebSocketProtocols(protocols: string | null | undefined): {
+	encoding: Encoding;
+	connParams: unknown;
+} {
+	let encodingRaw: string | undefined;
+	let connParamsRaw: string | undefined;
+
+	if (protocols) {
+		const protocolList = protocols.split(",").map((p) => p.trim());
+		for (const protocol of protocolList) {
+			if (protocol.startsWith(WS_PROTOCOL_ENCODING)) {
+				encodingRaw = protocol.substring(WS_PROTOCOL_ENCODING.length);
+			} else if (protocol.startsWith(WS_PROTOCOL_CONN_PARAMS)) {
+				connParamsRaw = decodeURIComponent(
+					protocol.substring(WS_PROTOCOL_CONN_PARAMS.length),
+				);
+			}
+		}
+	}
+
+	const encoding = EncodingSchema.parse(encodingRaw);
+	const connParams = connParamsRaw ? JSON.parse(connParamsRaw) : undefined;
+
+	return { encoding, connParams };
+}
+
+/**
+ * Truncase the PATH_WEBSOCKET_PREFIX path prefix in order to pass a clean
+ * path to the onWebSocket handler.
+ *
+ * Example:
+ * - `/websocket/foo` -> `/foo`
+ * - `/websocket` -> `/`
+ */
+export function truncateRawWebSocketPathPrefix(path: string): string {
+	// Extract the path after prefix and preserve query parameters
+	// Use URL API for cleaner parsing
+	const url = new URL(path, "http://actor");
+	const pathname = url.pathname.replace(/^\/websocket\/?/, "") || "/";
+	const normalizedPath =
+		(pathname.startsWith("/") ? pathname : `/${pathname}`) + url.search;
+
+	return normalizedPath;
+}