rivetkit 2.0.6 → 2.0.7-rc.1

package/src/manager/gateway.ts

@@ -0,0 +1,407 @@
+import type { Context as HonoContext, Next } from "hono";
+import type { WSContext } from "hono/ws";
+import { MissingActorHeader, WebSocketsNotEnabled } from "@/actor/errors";
+import type { Encoding, Transport } from "@/client/mod";
+import {
+	HEADER_RIVET_ACTOR,
+	HEADER_RIVET_TARGET,
+	WS_PROTOCOL_ACTOR,
+	WS_PROTOCOL_CONN_ID,
+	WS_PROTOCOL_CONN_PARAMS,
+	WS_PROTOCOL_CONN_TOKEN,
+	WS_PROTOCOL_ENCODING,
+	WS_PROTOCOL_TARGET,
+} from "@/common/actor-router-consts";
+import { deconstructError, noopNext } from "@/common/utils";
+import type { UniversalWebSocket, UpgradeWebSocketArgs } from "@/mod";
+import type { RunConfig } from "@/registry/run-config";
+import { promiseWithResolvers, stringifyError } from "@/utils";
+import type { ManagerDriver } from "./driver";
+import { logger } from "./log";
+
+/**
+ * Provides an endpoint to connect to individual actors.
+ *
+ * Routes requests based on the Upgrade header:
+ * - WebSocket requests: Uses sec-websocket-protocol for routing (target.actor, actor.{id})
+ * - HTTP requests: Uses x-rivet-target and x-rivet-actor headers for routing
+ */
+export async function actorGateway(
+	runConfig: RunConfig,
+	managerDriver: ManagerDriver,
+	c: HonoContext,
+	next: Next,
+) {
+	// Skip test routes - let them be handled by their specific handlers
+	if (c.req.path.startsWith("/.test/")) {
+		return next();
+	}
+
+	// Check if this is a WebSocket upgrade request
+	if (c.req.header("upgrade") === "websocket") {
+		return await handleWebSocketGateway(runConfig, managerDriver, c);
+	}
+
+	// Handle regular HTTP requests
+	return await handleHttpGateway(managerDriver, c, next);
+}
+
+/**
+ * Handle WebSocket requests using sec-websocket-protocol for routing
+ */
+async function handleWebSocketGateway(
+	runConfig: RunConfig,
+	managerDriver: ManagerDriver,
+	c: HonoContext,
+) {
+	const upgradeWebSocket = runConfig.getUpgradeWebSocket?.();
+	if (!upgradeWebSocket) {
+		throw new WebSocketsNotEnabled();
+	}
+
+	// Parse configuration from Sec-WebSocket-Protocol header
+	const protocols = c.req.header("sec-websocket-protocol");
+	let target: string | undefined;
+	let actorId: string | undefined;
+	let encodingRaw: string | undefined;
+	let connParamsRaw: string | undefined;
+	let connIdRaw: string | undefined;
+	let connTokenRaw: string | undefined;
+
+	if (protocols) {
+		const protocolList = protocols.split(",").map((p) => p.trim());
+		for (const protocol of protocolList) {
+			if (protocol.startsWith(WS_PROTOCOL_TARGET)) {
+				target = protocol.substring(WS_PROTOCOL_TARGET.length);
+			} else if (protocol.startsWith(WS_PROTOCOL_ACTOR)) {
+				actorId = protocol.substring(WS_PROTOCOL_ACTOR.length);
+			} else 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),
+				);
+			} else if (protocol.startsWith(WS_PROTOCOL_CONN_ID)) {
+				connIdRaw = protocol.substring(WS_PROTOCOL_CONN_ID.length);
+			} else if (protocol.startsWith(WS_PROTOCOL_CONN_TOKEN)) {
+				connTokenRaw = protocol.substring(WS_PROTOCOL_CONN_TOKEN.length);
+			}
+		}
+	}
+
+	if (target !== "actor") {
+		return c.text("WebSocket upgrade requires target.actor protocol", 400);
+	}
+
+	if (!actorId) {
+		throw new MissingActorHeader();
+	}
+
+	logger().debug({
+		msg: "proxying websocket to actor",
+		actorId,
+		path: c.req.path,
+		encoding: encodingRaw,
+	});
+
+	const encoding = encodingRaw || "json";
+	const connParams = connParamsRaw ? JSON.parse(connParamsRaw) : undefined;
+
+	// Include query string if present
+	const pathWithQuery = c.req.url.includes("?")
+		? c.req.path + c.req.url.substring(c.req.url.indexOf("?"))
+		: c.req.path;
+
+	return await managerDriver.proxyWebSocket(
+		c,
+		pathWithQuery,
+		actorId,
+		encoding as any, // Will be validated by driver
+		connParams,
+		connIdRaw,
+		connTokenRaw,
+	);
+}
+
+/**
+ * Handle HTTP requests using x-rivet headers for routing
+ */
+async function handleHttpGateway(
+	managerDriver: ManagerDriver,
+	c: HonoContext,
+	next: Next,
+) {
+	const target = c.req.header(HEADER_RIVET_TARGET);
+	const actorId = c.req.header(HEADER_RIVET_ACTOR);
+
+	if (target !== "actor") {
+		return next();
+	}
+
+	if (!actorId) {
+		throw new MissingActorHeader();
+	}
+
+	logger().debug({
+		msg: "proxying request to actor",
+		actorId,
+		path: c.req.path,
+		method: c.req.method,
+	});
+
+	// Preserve all headers except the routing headers
+	const proxyHeaders = new Headers(c.req.raw.headers);
+	proxyHeaders.delete(HEADER_RIVET_TARGET);
+	proxyHeaders.delete(HEADER_RIVET_ACTOR);
+
+	// Build the proxy request with the actor URL format
+	const url = new URL(c.req.url);
+	const proxyUrl = new URL(`http://actor${url.pathname}${url.search}`);
+
+	const proxyRequest = new Request(proxyUrl, {
+		method: c.req.raw.method,
+		headers: proxyHeaders,
+		body: c.req.raw.body,
+		signal: c.req.raw.signal,
+	});
+
+	return await managerDriver.proxyRequest(c, proxyRequest, actorId);
+}
+
+/**
+ * Creates a WebSocket proxy for test endpoints that forwards messages between server and client WebSockets
+ */
+export async function createTestWebSocketProxy(
+	clientWsPromise: Promise<UniversalWebSocket>,
+): Promise<UpgradeWebSocketArgs> {
+	// Store a reference to the resolved WebSocket
+	let clientWs: UniversalWebSocket | null = null;
+	const {
+		promise: serverWsPromise,
+		resolve: serverWsResolve,
+		reject: serverWsReject,
+	} = promiseWithResolvers<WSContext>();
+	try {
+		// Resolve the client WebSocket promise
+		logger().debug({ msg: "awaiting client websocket promise" });
+		const ws = await clientWsPromise;
+		clientWs = ws;
+		logger().debug({
+			msg: "client websocket promise resolved",
+			constructor: ws?.constructor.name,
+		});
+
+		// Wait for ws to open
+		await new Promise<void>((resolve, reject) => {
+			const onOpen = () => {
+				logger().debug({ msg: "test websocket connection to actor opened" });
+				resolve();
+			};
+			const onError = (error: any) => {
+				logger().error({ msg: "test websocket connection failed", error });
+				reject(
+					new Error(`Failed to open WebSocket: ${error.message || error}`),
+				);
+				serverWsReject();
+			};
+
+			ws.addEventListener("open", onOpen);
+
+			ws.addEventListener("error", onError);
+
+			ws.addEventListener("message", async (clientEvt: MessageEvent) => {
+				const serverWs = await serverWsPromise;
+
+				logger().debug({
+					msg: `test websocket connection message from client`,
+					dataType: typeof clientEvt.data,
+					isBlob: clientEvt.data instanceof Blob,
+					isArrayBuffer: clientEvt.data instanceof ArrayBuffer,
+					dataConstructor: clientEvt.data?.constructor?.name,
+					dataStr:
+						typeof clientEvt.data === "string"
+							? clientEvt.data.substring(0, 100)
+							: undefined,
+				});
+
+				if (serverWs.readyState === 1) {
+					// OPEN
+					// Handle Blob data
+					if (clientEvt.data instanceof Blob) {
+						clientEvt.data
+							.arrayBuffer()
+							.then((buffer) => {
+								logger().debug({
+									msg: "converted client blob to arraybuffer, sending to server",
+									bufferSize: buffer.byteLength,
+								});
+								serverWs.send(buffer as any);
+							})
+							.catch((error) => {
+								logger().error({
+									msg: "failed to convert blob to arraybuffer",
+									error,
+								});
+							});
+					} else {
+						logger().debug({
+							msg: "sending client data directly to server",
+							dataType: typeof clientEvt.data,
+							dataLength:
+								typeof clientEvt.data === "string"
+									? clientEvt.data.length
+									: undefined,
+						});
+						serverWs.send(clientEvt.data as any);
+					}
+				}
+			});
+
+			ws.addEventListener("close", async (clientEvt: any) => {
+				const serverWs = await serverWsPromise;
+
+				logger().debug({
+					msg: `test websocket connection closed`,
+				});
+
+				if (serverWs.readyState !== 3) {
+					// Not CLOSED
+					serverWs.close(clientEvt.code, clientEvt.reason);
+				}
+			});
+
+			ws.addEventListener("error", async () => {
+				const serverWs = await serverWsPromise;
+
+				logger().debug({
+					msg: `test websocket connection error`,
+				});
+
+				if (serverWs.readyState !== 3) {
+					// Not CLOSED
+					serverWs.close(1011, "Error in client websocket");
+				}
+			});
+		});
+	} catch (error) {
+		logger().error({
+			msg: `failed to establish client websocket connection`,
+			error,
+		});
+		return {
+			onOpen: (_evt, serverWs) => {
+				serverWs.close(1011, "Failed to establish connection");
+			},
+			onMessage: () => {},
+			onError: () => {},
+			onClose: () => {},
+		};
+	}
+
+	// Create WebSocket proxy handlers to relay messages between client and server
+	return {
+		onOpen: (_evt: any, serverWs: WSContext) => {
+			logger().debug({
+				msg: `test websocket connection from client opened`,
+			});
+
+			// Check WebSocket type
+			logger().debug({
+				msg: "clientWs info",
+				constructor: clientWs.constructor.name,
+				hasAddEventListener: typeof clientWs.addEventListener === "function",
+				readyState: clientWs.readyState,
+			});
+
+			serverWsResolve(serverWs);
+		},
+		onMessage: (evt: { data: any }) => {
+			logger().debug({
+				msg: "received message from server",
+				dataType: typeof evt.data,
+				isBlob: evt.data instanceof Blob,
+				isArrayBuffer: evt.data instanceof ArrayBuffer,
+				dataConstructor: evt.data?.constructor?.name,
+				dataStr:
+					typeof evt.data === "string" ? evt.data.substring(0, 100) : undefined,
+			});
+
+			// Forward messages from server websocket to client websocket
+			if (clientWs.readyState === 1) {
+				// OPEN
+				// Handle Blob data
+				if (evt.data instanceof Blob) {
+					evt.data
+						.arrayBuffer()
+						.then((buffer) => {
+							logger().debug({
+								msg: "converted blob to arraybuffer, sending",
+								bufferSize: buffer.byteLength,
+							});
+							clientWs.send(buffer);
+						})
+						.catch((error) => {
+							logger().error({
+								msg: "failed to convert blob to arraybuffer",
+								error,
+							});
+						});
+				} else {
+					logger().debug({
+						msg: "sending data directly",
+						dataType: typeof evt.data,
+						dataLength:
+							typeof evt.data === "string" ? evt.data.length : undefined,
+					});
+					clientWs.send(evt.data);
+				}
+			}
+		},
+		onClose: (
+			event: {
+				wasClean: boolean;
+				code: number;
+				reason: string;
+			},
+			serverWs: WSContext,
+		) => {
+			logger().debug({
+				msg: `server websocket closed`,
+				wasClean: event.wasClean,
+				code: event.code,
+				reason: event.reason,
+			});
+
+			// HACK: Close socket in order to fix bug with Cloudflare leaving WS in closing state
+			// https://github.com/cloudflare/workerd/issues/2569
+			serverWs.close(1000, "hack_force_close");
+
+			// Close the client websocket when the server websocket closes
+			if (
+				clientWs &&
+				clientWs.readyState !== clientWs.CLOSED &&
+				clientWs.readyState !== clientWs.CLOSING
+			) {
+				// Don't pass code/message since this may affect how close events are triggered
+				clientWs.close(1000, event.reason);
+			}
+		},
+		onError: (error: unknown) => {
+			logger().error({
+				msg: `error in server websocket`,
+				error,
+			});
+
+			// Close the client websocket on error
+			if (
+				clientWs &&
+				clientWs.readyState !== clientWs.CLOSED &&
+				clientWs.readyState !== clientWs.CLOSING
+			) {
+				clientWs.close(1011, "Error in server websocket");
+			}
+
+			serverWsReject();
+		},
+	};
+}