rivetkit 2.0.2 → 2.0.4-rc.1

package/src/serverless/router.ts

@@ -0,0 +1,215 @@
+import invariant from "invariant";
+import {
+	EndpointMismatch,
+	InvalidRequest,
+	NamespaceMismatch,
+} from "@/actor/errors";
+import { convertRegistryConfigToClientConfig } from "@/client/config";
+import { handleHealthRequest, handleMetadataRequest } from "@/common/router";
+import { ServerlessStartHeadersSchema } from "@/manager/router-schema";
+import { createClientWithDriver } from "@/mod";
+import type { DriverConfig, RegistryConfig } from "@/registry/config";
+import { RemoteManagerDriver } from "@/remote-manager-driver/mod";
+import { createRouter } from "@/utils/router";
+import { logger } from "./log";
+
+export function buildServerlessRouter(
+	driverConfig: DriverConfig,
+	config: RegistryConfig,
+) {
+	return createRouter(config.serverless.basePath, (router) => {
+		// GET /
+		router.get("/", (c) => {
+			return c.text(
+				"This is a RivetKit server.\n\nLearn more at https://rivetkit.org",
+			);
+		});
+
+		// Serverless start endpoint
+		router.get("/start", async (c) => {
+			// Parse headers
+			const parseResult = ServerlessStartHeadersSchema.safeParse({
+				endpoint: c.req.header("x-rivet-endpoint"),
+				token: c.req.header("x-rivet-token") ?? undefined,
+				totalSlots: c.req.header("x-rivet-total-slots"),
+				runnerName: c.req.header("x-rivet-runner-name"),
+				namespace: c.req.header("x-rivet-namespace-name"),
+			});
+			if (!parseResult.success) {
+				throw new InvalidRequest(
+					parseResult.error.issues[0]?.message ??
+						"invalid serverless start headers",
+				);
+			}
+			const { endpoint, token, totalSlots, runnerName, namespace } =
+				parseResult.data;
+
+			logger().debug({
+				msg: "received serverless runner start request",
+				endpoint,
+				totalSlots,
+				runnerName,
+				namespace,
+			});
+
+			// Validate endpoint and namespace match config to catch
+			// misconfiguration or malicious requests.
+			//
+			// Only verify if namespace matches if endpoint configured since
+			// configuring an endpoint indicates you want to assert the
+			// incoming serverless requests.
+			if (config.endpoint) {
+				if (!endpointsMatch(endpoint, config.endpoint)) {
+					throw new EndpointMismatch(config.endpoint, endpoint);
+				}
+
+				if (namespace !== config.namespace) {
+					throw new NamespaceMismatch(config.namespace, namespace);
+				}
+			}
+
+			// Convert config to runner config
+			const newConfig: RegistryConfig = {
+				...config,
+				endpoint: endpoint,
+				namespace: namespace,
+				token: token,
+				runner: {
+					...config.runner,
+					totalSlots: totalSlots,
+					runnerName: runnerName,
+					// Not supported on serverless
+					runnerKey: undefined,
+				},
+			};
+
+			// Create manager driver on demand based on the properties provided
+			// by headers
+			//
+			// NOTE: This relies on the `newConfig.runner.runnerName` to
+			// configure which runner to create actors on.
+			const managerDriver = new RemoteManagerDriver(
+				convertRegistryConfigToClientConfig(newConfig),
+			);
+			const client = createClientWithDriver(managerDriver);
+
+			// Create new actor driver with updated config
+			const actorDriver = driverConfig.actor(
+				newConfig,
+				managerDriver,
+				client,
+			);
+			invariant(
+				actorDriver.serverlessHandleStart,
+				"missing serverlessHandleStart on ActorDriver",
+			);
+
+			return await actorDriver.serverlessHandleStart(c);
+		});
+
+		router.get("/health", (c) => handleHealthRequest(c));
+
+		router.get("/metadata", (c) =>
+			handleMetadataRequest(
+				c,
+				config,
+				{ serverless: {} },
+				config.publicEndpoint,
+				config.publicNamespace,
+				config.publicToken,
+			),
+		);
+	});
+}
+
+/**
+ * Normalizes a URL for comparison by extracting protocol, host, port, and pathname.
+ * Normalizes loopback addresses (127.0.0.1, 0.0.0.0, ::1) to localhost for consistent comparison.
+ * Normalizes regional endpoints (api-*.domain) to base endpoints (api.domain).
+ * Returns null if the URL is invalid.
+ */
+export function normalizeEndpointUrl(url: string): string | null {
+	try {
+		const parsed = new URL(url);
+		// Normalize pathname by removing trailing slash (except for root)
+		const pathname =
+			parsed.pathname === "/" ? "/" : parsed.pathname.replace(/\/+$/, "");
+
+		// Normalize loopback addresses to localhost
+		let hostname = isLoopbackAddress(parsed.hostname)
+			? "localhost"
+			: parsed.hostname;
+
+		// Normalize regional endpoints (api-region.domain) to base endpoints (api.domain)
+		// HACK: This is specific to Rivet Cloud and will not work for self-hosted
+		// engines with different regional endpoint naming conventions.
+		hostname = normalizeRegionalHostname(hostname);
+
+		// Reconstruct host with normalized hostname and port
+		const host = parsed.port ? `${hostname}:${parsed.port}` : hostname;
+
+		// Reconstruct normalized URL with protocol, host, and pathname
+		return `${parsed.protocol}//${host}${pathname}`;
+	} catch {
+		return null;
+	}
+}
+
+/**
+ * Normalizes regional hostnames (api-region.domain) to base hostnames (api.domain).
+ * Only applies to rivet.dev domains.
+ *
+ * Examples:
+ * - api-us-west-1.rivet.dev -> api.rivet.dev
+ * - api-lax.staging.rivet.dev -> api.staging.rivet.dev
+ * - api.rivet.dev -> api.rivet.dev (unchanged)
+ * - api-us-west-1.example.com -> api-us-west-1.example.com (unchanged, not rivet.dev)
+ * - foo-bar.rivet.dev -> foo-bar.rivet.dev (unchanged, not api- prefix)
+ */
+function normalizeRegionalHostname(hostname: string): string {
+	// Only apply to rivet.dev domains
+	if (!hostname.endsWith(".rivet.dev")) {
+		return hostname;
+	}
+
+	if (!hostname.startsWith("api-")) {
+		return hostname;
+	}
+
+	// Find the first dot after "api-"
+	const withoutPrefix = hostname.slice(4); // Remove "api-"
+	const firstDotIndex = withoutPrefix.indexOf(".");
+	if (firstDotIndex === -1) {
+		return hostname;
+	}
+
+	// Extract the domain part and prepend "api."
+	const domain = withoutPrefix.slice(firstDotIndex + 1);
+	return `api.${domain}`;
+}
+
+/**
+ * Compares two endpoint URLs after normalization.
+ * Returns true if they match (same protocol, host, port, and path).
+ */
+export function endpointsMatch(a: string, b: string): boolean {
+	const normalizedA = normalizeEndpointUrl(a);
+	const normalizedB = normalizeEndpointUrl(b);
+	if (normalizedA === null || normalizedB === null) {
+		// If either URL is invalid, fall back to string comparison
+		return a === b;
+	}
+	return normalizedA === normalizedB;
+}
+
+/**
+ * Checks if a hostname is a loopback address that should be normalized to localhost.
+ */
+function isLoopbackAddress(hostname: string): boolean {
+	return (
+		hostname === "127.0.0.1" ||
+		hostname === "0.0.0.0" ||
+		hostname === "::1" ||
+		hostname === "[::1]"
+	);
+}

package/src/test/log.ts

@@ -0,0 +1,5 @@
+import { getLogger } from "@/common/log";
+
+export function logger() {
+	return getLogger("test");
+}

package/src/test/mod.ts

@@ -0,0 +1,99 @@
+import { serve as honoServe } from "@hono/node-server";
+import { createNodeWebSocket } from "@hono/node-ws";
+import invariant from "invariant";
+import { type TestContext, vi } from "vitest";
+import { ClientConfigSchema } from "@/client/config";
+import { type Client, createClient } from "@/client/mod";
+import { createFileSystemOrMemoryDriver } from "@/drivers/file-system/mod";
+import { createClientWithDriver, type Registry } from "@/mod";
+import { RegistryConfig, RegistryConfigSchema } from "@/registry/config";
+import { buildManagerRouter } from "@/manager/router";
+import { logger } from "./log";
+
+export interface SetupTestResult<A extends Registry<any>> {
+	client: Client<A>;
+}
+
+// Must use `TestContext` since global hooks do not work when running concurrently
+export async function setupTest<A extends Registry<any>>(
+	c: TestContext,
+	registry: A,
+): Promise<SetupTestResult<A>> {
+	// Force enable test mode
+	registry.config.test = { ...registry.config.test, enabled: true };
+
+	// Create driver
+	const driver = createFileSystemOrMemoryDriver(
+		true,
+		{ path: `/tmp/rivetkit-test-${crypto.randomUUID()}` },
+	);
+
+	// Build driver config
+	// biome-ignore lint/style/useConst: Assigned later
+	let upgradeWebSocket: any;
+	registry.config.driver = driver;
+	registry.config.inspector = {
+		enabled: true,
+		token: () => "token",
+	};
+
+	// Create router
+	const parsedConfig = registry.parseConfig();
+	const managerDriver = driver.manager?.(parsedConfig);
+	invariant(managerDriver, "missing manager driver");
+	const getUpgradeWebSocket = () => upgradeWebSocket;
+	managerDriver.setGetUpgradeWebSocket(getUpgradeWebSocket);
+	// const internalClient = createClientWithDriver(
+	// 	managerDriver,
+	// 	ClientConfigSchema.parse({}),
+	// );
+	const { router } = buildManagerRouter(
+		parsedConfig,
+		managerDriver,
+		getUpgradeWebSocket,
+	);
+
+	// Inject WebSocket
+	const nodeWebSocket = createNodeWebSocket({ app: router });
+	upgradeWebSocket = nodeWebSocket.upgradeWebSocket;
+
+	// TODO: I think this whole function is fucked, we should probably switch to calling registry.serve() directly
+	// Start server
+	const server = honoServe({
+		fetch: router.fetch,
+		hostname: "127.0.0.1",
+		port: 0,
+	});
+	if (!server.listening) {
+		await new Promise<void>((resolve) => {
+			server.once("listening", () => resolve());
+		});
+	}
+	invariant(
+		nodeWebSocket.injectWebSocket !== undefined,
+		"should have injectWebSocket",
+	);
+	nodeWebSocket.injectWebSocket(server);
+	const address = server.address();
+	invariant(address && typeof address !== "string", "missing server address");
+	const port = address.port;
+	const endpoint = `http://127.0.0.1:${port}`;
+
+	logger().info({ msg: "test server listening", port });
+
+	// Cleanup on test finish
+	c.onTestFinished(async () => {
+		await new Promise((resolve) => server.close(() => resolve(undefined)));
+	});
+
+	// Create client
+	const client = createClient<A>({
+		endpoint,
+		namespace: "default",
+		runnerName: "default",
+		disableMetadataLookup: true,
+	});
+	c.onTestFinished(async () => await client.dispose());
+
+	return { client };
+}

package/src/utils/crypto.ts

@@ -0,0 +1,24 @@
+/**
+ * Timing-safe comparison of two Uint8Arrays or strings.
+ * This prevents timing attacks by always comparing all bytes.
+ */
+export function timingSafeEqual(
+	a: Uint8Array | string,
+	b: Uint8Array | string,
+): boolean {
+	const encoder = new TextEncoder();
+	const bufferA = typeof a === "string" ? encoder.encode(a) : a;
+	const bufferB = typeof b === "string" ? encoder.encode(b) : b;
+
+	// Pad to max length to avoid leaking length information
+	const maxLength = Math.max(bufferA.byteLength, bufferB.byteLength);
+	let result = bufferA.byteLength ^ bufferB.byteLength;
+
+	for (let i = 0; i < maxLength; i++) {
+		const byteA = i < bufferA.byteLength ? bufferA[i] : 0;
+		const byteB = i < bufferB.byteLength ? bufferB[i] : 0;
+		result |= byteA ^ byteB;
+	}
+
+	return result === 0;
+}

package/src/utils/endpoint-parser.test.ts

@@ -0,0 +1,202 @@
+import { describe, expect, test } from "vitest";
+import type { z } from "zod/v4";
+import { tryParseEndpoint } from "./endpoint-parser";
+
+// Helper to create a mock Zod refinement context for testing
+function createMockCtx(): { ctx: z.RefinementCtx; issues: z.ZodIssue[] } {
+	const issues: z.ZodIssue[] = [];
+	const ctx = {
+		addIssue: (issue: z.IssueData) => {
+			issues.push(issue as z.ZodIssue);
+		},
+		path: [],
+		value: undefined,
+		issues: [],
+	} as unknown as z.RefinementCtx;
+	return { ctx, issues };
+}
+
+describe("tryParseEndpoint", () => {
+	describe("basic parsing", () => {
+		test("parses endpoint with full auth", () => {
+			const { ctx, issues } = createMockCtx();
+			const result = tryParseEndpoint(ctx, {
+				endpoint: "https://foo:bar@api.rivet.dev",
+			});
+			expect(issues).toHaveLength(0);
+			expect(result).toEqual({
+				endpoint: "https://api.rivet.dev/",
+				namespace: "foo",
+				token: "bar",
+			});
+		});
+
+		test("parses endpoint with namespace only", () => {
+			const { ctx, issues } = createMockCtx();
+			const result = tryParseEndpoint(ctx, {
+				endpoint: "https://foo@api.rivet.dev",
+			});
+			expect(issues).toHaveLength(0);
+			expect(result).toEqual({
+				endpoint: "https://api.rivet.dev/",
+				namespace: "foo",
+				token: undefined,
+			});
+		});
+
+		test("parses endpoint without auth", () => {
+			const { ctx, issues } = createMockCtx();
+			const result = tryParseEndpoint(ctx, {
+				endpoint: "https://api.rivet.dev",
+			});
+			expect(issues).toHaveLength(0);
+			expect(result).toEqual({
+				endpoint: "https://api.rivet.dev/",
+				namespace: undefined,
+				token: undefined,
+			});
+		});
+
+		test("preserves path", () => {
+			const { ctx, issues } = createMockCtx();
+			const result = tryParseEndpoint(ctx, {
+				endpoint: "https://foo:bar@api.rivet.dev/v1/actors",
+			});
+			expect(issues).toHaveLength(0);
+			expect(result).toEqual({
+				endpoint: "https://api.rivet.dev/v1/actors",
+				namespace: "foo",
+				token: "bar",
+			});
+		});
+	});
+
+	describe("validation errors", () => {
+		test("adds issue for invalid URL", () => {
+			const { ctx, issues } = createMockCtx();
+			const result = tryParseEndpoint(ctx, { endpoint: "not-a-url" });
+			expect(result).toBeUndefined();
+			expect(issues).toHaveLength(1);
+			expect(issues[0]?.message).toContain("invalid URL");
+		});
+
+		test("adds issue for query string", () => {
+			const { ctx, issues } = createMockCtx();
+			const result = tryParseEndpoint(ctx, {
+				endpoint: "https://foo:bar@api.rivet.dev?region=us",
+			});
+			expect(result).toBeUndefined();
+			expect(issues).toHaveLength(1);
+			expect(issues[0]?.message).toBe("endpoint cannot contain a query string");
+		});
+
+		test("adds issue for fragment", () => {
+			const { ctx, issues } = createMockCtx();
+			const result = tryParseEndpoint(ctx, {
+				endpoint: "https://foo:bar@api.rivet.dev#section",
+			});
+			expect(result).toBeUndefined();
+			expect(issues).toHaveLength(1);
+			expect(issues[0]?.message).toBe("endpoint cannot contain a fragment");
+		});
+
+		test("adds issue for token without namespace", () => {
+			const { ctx, issues } = createMockCtx();
+			const result = tryParseEndpoint(ctx, {
+				endpoint: "https://:token@api.rivet.dev",
+			});
+			expect(result).toBeUndefined();
+			expect(issues).toHaveLength(1);
+			expect(issues[0]?.message).toBe(
+				"endpoint cannot have a token without a namespace",
+			);
+		});
+	});
+
+	describe("duplicate credential checking", () => {
+		test("adds issue when namespace in URL and config", () => {
+			const { ctx, issues } = createMockCtx();
+			const result = tryParseEndpoint(ctx, {
+				endpoint: "https://url-ns@api.rivet.dev",
+				path: ["endpoint"],
+				namespace: "config-ns",
+			});
+			// Still returns result, but adds issue
+			expect(result).toEqual({
+				endpoint: "https://api.rivet.dev/",
+				namespace: "url-ns",
+				token: undefined,
+			});
+			expect(issues).toHaveLength(1);
+			expect(issues[0]?.message).toContain(
+				"cannot specify namespace both in endpoint URL and as a separate config option",
+			);
+			expect(issues[0]?.path).toEqual(["namespace"]);
+		});
+
+		test("adds issue when token in URL and config", () => {
+			const { ctx, issues } = createMockCtx();
+			const result = tryParseEndpoint(ctx, {
+				endpoint: "https://ns:url-token@api.rivet.dev",
+				path: ["endpoint"],
+				token: "config-token",
+			});
+			// Still returns result, but adds issue
+			expect(result).toEqual({
+				endpoint: "https://api.rivet.dev/",
+				namespace: "ns",
+				token: "url-token",
+			});
+			expect(issues).toHaveLength(1);
+			expect(issues[0]?.message).toContain(
+				"cannot specify token both in endpoint URL and as a separate config option",
+			);
+			expect(issues[0]?.path).toEqual(["token"]);
+		});
+
+		test("adds issues for both namespace and token duplicates", () => {
+			const { ctx, issues } = createMockCtx();
+			const result = tryParseEndpoint(ctx, {
+				endpoint: "https://url-ns:url-token@api.rivet.dev",
+				path: ["endpoint"],
+				namespace: "config-ns",
+				token: "config-token",
+			});
+			expect(result).toBeDefined();
+			expect(issues).toHaveLength(2);
+		});
+
+		test("no issue when namespace only in URL", () => {
+			const { ctx, issues } = createMockCtx();
+			const result = tryParseEndpoint(ctx, {
+				endpoint: "https://url-ns@api.rivet.dev",
+				path: ["endpoint"],
+				token: "config-token",
+			});
+			expect(result).toBeDefined();
+			expect(issues).toHaveLength(0);
+		});
+
+		test("no issue when namespace only in config", () => {
+			const { ctx, issues } = createMockCtx();
+			const result = tryParseEndpoint(ctx, {
+				endpoint: "https://api.rivet.dev",
+				path: ["endpoint"],
+				namespace: "config-ns",
+			});
+			expect(result).toBeDefined();
+			expect(issues).toHaveLength(0);
+		});
+	});
+
+	describe("custom path", () => {
+		test("uses custom path in error issues", () => {
+			const { ctx, issues } = createMockCtx();
+			tryParseEndpoint(ctx, {
+				endpoint: "not-a-url",
+				path: ["serverless", "publicEndpoint"],
+			});
+			expect(issues[0]?.path).toEqual(["serverless", "publicEndpoint"]);
+		});
+	});
+});

package/src/utils/endpoint-parser.ts

@@ -0,0 +1,124 @@
+import { z } from "zod/v4";
+
+export interface ParsedEndpoint {
+	endpoint: string;
+	namespace: string | undefined;
+	token: string | undefined;
+}
+
+export interface TryParseEndpointOptions {
+	/** The endpoint URL string to parse */
+	endpoint: string;
+	/** Path prefix for error messages (default: ["endpoint"]) */
+	path?: (string | number)[];
+	/** Namespace from config, to check for duplicate specification */
+	namespace?: string;
+	/** Token from config, to check for duplicate specification */
+	token?: string;
+}
+
+/**
+ * Parses an endpoint URL that may contain auth syntax for namespace and token.
+ *
+ * Uses ctx.addIssue for clean error reporting in Zod transforms.
+ *
+ * Supports formats like:
+ * - `https://namespace:token@api.rivet.dev`
+ * - `https://namespace@api.rivet.dev` (namespace only, no token)
+ * - `https://api.rivet.dev` (no auth)
+ * - `https://namespace:token@api.rivet.dev/path` (with path)
+ *
+ * Query strings and fragments are not allowed as they may conflict with
+ * runtime parameters.
+ *
+ * @param ctx - Zod refinement context for error reporting
+ * @param options - Parsing options including endpoint, path, and config values
+ * @returns ParsedEndpoint on success, undefined on error (after adding issues to ctx)
+ */
+export function tryParseEndpoint(
+	ctx: z.RefinementCtx,
+	options: TryParseEndpointOptions,
+): ParsedEndpoint | undefined {
+	const { endpoint, path = ["endpoint"], namespace: configNamespace, token: configToken } = options;
+	// Parse the URL
+	let url: URL;
+	try {
+		url = new URL(endpoint);
+	} catch {
+		ctx.addIssue({
+			code: "custom",
+			message: `invalid URL: ${endpoint}`,
+			path,
+		});
+		return undefined;
+	}
+
+	// Reject query strings
+	if (url.search) {
+		ctx.addIssue({
+			code: "custom",
+			message: "endpoint cannot contain a query string",
+			path,
+		});
+		return undefined;
+	}
+
+	// Reject fragments
+	if (url.hash) {
+		ctx.addIssue({
+			code: "custom",
+			message: "endpoint cannot contain a fragment",
+			path,
+		});
+		return undefined;
+	}
+
+	// Extract namespace and token from username and password
+	// URL stores these as percent-encoded, so we need to decode them
+	const namespace = url.username
+		? decodeURIComponent(url.username)
+		: undefined;
+	const token = url.password ? decodeURIComponent(url.password) : undefined;
+
+	// Reject token without namespace (e.g., https://:token@api.rivet.dev)
+	if (token && !namespace) {
+		ctx.addIssue({
+			code: "custom",
+			message: "endpoint cannot have a token without a namespace",
+			path,
+		});
+		return undefined;
+	}
+
+	// Check for duplicate credentials (specified both in URL and config)
+	if (namespace && configNamespace) {
+		ctx.addIssue({
+			code: "custom",
+			message:
+				"cannot specify namespace both in endpoint URL and as a separate config option",
+			path: ["namespace"],
+		});
+	}
+	if (token && configToken) {
+		ctx.addIssue({
+			code: "custom",
+			message:
+				"cannot specify token both in endpoint URL and as a separate config option",
+			path: ["token"],
+		});
+	}
+
+	// Strip auth from the URL by clearing username and password
+	url.username = "";
+	url.password = "";
+
+	// Get the cleaned endpoint without auth
+	const cleanedEndpoint = url.toString();
+
+	return {
+		endpoint: cleanedEndpoint,
+		namespace,
+		token,
+	};
+}
+