feat: add onAuth lifecycle hooks

Author: NathanFlurry

docs/workers/quickstart.mdx

@@ -7,10 +7,11 @@ description: Start building awesome documentation in under 5 minutes
 ```ts registry.ts
 import { setup } from "rivetkit";
 import { worker } from "rivetkit/worker";
-import { workflow } from "rivetkit/workflow";
-import { realtime } from "rivetkit/realtime";
 
 const counter = worker({
+	onAuth: async () => {
+		// Allow public access
+	},
 	state: {
 		count: 0,
 	},
@@ -36,18 +37,24 @@ export type Registry = typeof registry;
 ```
 
 ```ts server.ts
+// With router
 import { registry } from "./registry";
 
-const registry = new Hono();
-app.route("/registry", registry.handler);
+const app = new Hono();
+app.route("/registry", c => registry.handler(c.req.raw));
+serve(app);
+
+// Without router
+import { serve } from "@rivetkit/node";
+
 serve(registry);
 ```
 
 ```ts client.ts
 import { createClient } from "rivetkit/client";
 import type { Registry } from "./registry";
 
-const client = createClient<Registry>("http://localhost:8080/registry");
+const client = createClient<Registry>("http://localhost:8080");
 ```
 
 <Steps>

packages/core/src/manager/auth.ts

@@ -0,0 +1,106 @@
+import * as errors from "@/worker/errors";
+import { Hono, type Context as HonoContext, type Next } from "hono";
+import type { WorkerQuery } from "./protocol/query";
+import type { AuthIntent } from "@/worker/config";
+import type { AnyWorkerDefinition } from "@/worker/definition";
+import type { RegistryConfig } from "@/registry/config";
+
+/**
+ * Get authentication intents from a worker query
+ */
+export function getIntentsFromQuery(query: WorkerQuery): Set<AuthIntent> {
+	const intents = new Set<AuthIntent>();
+
+	if ("getForId" in query) {
+		intents.add("get");
+	} else if ("getForKey" in query) {
+		intents.add("get");
+	} else if ("getOrCreateForKey" in query) {
+		intents.add("get");
+		intents.add("create");
+	} else if ("create" in query) {
+		intents.add("create");
+	}
+
+	return intents;
+}
+
+/**
+ * Get worker name from a worker query
+ */
+export function getWorkerNameFromQuery(query: WorkerQuery): string {
+	if ("getForId" in query) {
+		throw new errors.InvalidRequest(
+			"Cannot determine worker name from getForId query",
+		);
+	} else if ("getForKey" in query) {
+		return query.getForKey.name;
+	} else if ("getOrCreateForKey" in query) {
+		return query.getOrCreateForKey.name;
+	} else if ("create" in query) {
+		return query.create.name;
+	} else {
+		throw new errors.InvalidRequest("Invalid query format");
+	}
+}
+
+/**
+ * Authenticate a request using the worker's onAuth function
+ */
+export async function authenticateRequest(
+	c: HonoContext,
+	workerDefinition: AnyWorkerDefinition,
+	intents: Set<AuthIntent>,
+	params: unknown,
+): Promise<unknown> {
+	const { onAuth } = workerDefinition.config;
+
+	if (!onAuth) {
+		throw new errors.Forbidden(
+			"Worker requires authentication but no onAuth handler is defined",
+		);
+	}
+
+	try {
+		return await onAuth({
+			req: c.req.raw,
+			intents,
+			params,
+		});
+	} catch (error) {
+		if (errors.WorkerError.isWorkerError(error)) {
+			throw error;
+		}
+		throw new errors.Forbidden("Authentication failed");
+	}
+}
+
+/**
+ * Simplified authentication for endpoints that combines all auth steps
+ */
+export async function authenticateEndpoint(
+	c: HonoContext,
+	registryConfig: RegistryConfig,
+	query: WorkerQuery,
+	additionalIntents: AuthIntent[],
+	params?: unknown,
+): Promise<unknown> {
+	// Get base intents from query
+	const intents = getIntentsFromQuery(query);
+
+	// Add endpoint-specific intents
+	for (const intent of additionalIntents) {
+		intents.add(intent);
+	}
+
+	// Get worker definition
+	const workerName = getWorkerNameFromQuery(query);
+	const workerDefinition = registryConfig.workers[workerName];
+	if (!workerDefinition) {
+		throw new errors.WorkerNotFound(workerName);
+	}
+
+	// Authenticate
+	return await authenticateRequest(c, workerDefinition, intents, params);
+}
+

packages/core/src/manager/router.ts

@@ -55,8 +55,9 @@ import {
 import type { WorkerQuery } from "./protocol/query";
 import { VERSION } from "@/utils";
 import { ConnRoutingHandler } from "@/worker/conn-routing-handler";
-import { ClientDriver, createClientWithDriver } from "@/client/client";
-import { Transport, TransportSchema } from "@/worker/protocol/message/mod";
+import { ClientDriver } from "@/client/client";
+import { Transport } from "@/worker/protocol/message/mod";
+import { authenticateEndpoint, authenticateMessage } from "./auth";
 
 type ManagerRouterHandler = {
 	onConnectInspector?: ManagerInspectorConnHandler;
@@ -141,7 +142,10 @@ export function createManagerRouter(
 
 			return cors({
 				...corsConfig,
-				allowHeaders: [...(registryConfig.cors?.allowHeaders ?? []), ...ALL_HEADERS],
+				allowHeaders: [
+					...(registryConfig.cors?.allowHeaders ?? []),
+					...ALL_HEADERS,
+				],
 			})(c, next);
 		});
 	}
@@ -194,7 +198,9 @@ export function createManagerRouter(
 			responses: buildOpenApiResponses(ResolveResponseSchema),
 		});
 
-		router.openapi(resolveRoute, (c) => handleResolveRequest(c, driver));
+		router.openapi(resolveRoute, (c) =>
+			handleResolveRequest(c, registryConfig, driver),
+		);
 	}
 
 	// GET /workers/connect/websocket
@@ -667,6 +673,20 @@ async function handleSseConnectRequest(
 
 		const query = params.data.query;
 
+		// Parse connection parameters for authentication
+		const connParams = params.data.connParams
+			? JSON.parse(params.data.connParams)
+			: undefined;
+
+		// Authenticate the request
+		const authData = await authenticateEndpoint(
+			c,
+			registryConfig,
+			query,
+			["connect"],
+			connParams,
+		);
+
 		// Get the worker ID and meta
 		const { workerId, meta } = await queryWorker(c, query, driver);
 		invariant(workerId, "Missing worker ID");
@@ -682,6 +702,7 @@ async function handleSseConnectRequest(
 				driverConfig,
 				handler.routingHandler.inline.handlers.onConnectSse,
 				workerId,
+				authData,
 			);
 		} else if ("custom" in handler.routingHandler) {
 			logger().debug("using custom proxy mode for sse connection");
@@ -790,6 +811,14 @@ async function handleWebSocketConnectRequest(
 			throw new errors.InvalidRequest(params.error);
 		}
 
+		// For WebSocket, params come later over the socket, so we auth without params for now
+		const authData = await authenticateEndpoint(
+			c,
+			registryConfig,
+			params.data.query,
+			["connect"],
+		);
+
 		// Get the worker ID and meta
 		const { workerId, meta } = await queryWorker(c, params.data.query, driver);
 		logger().debug("found worker for websocket connection", { workerId, meta });
@@ -811,6 +840,7 @@ async function handleWebSocketConnectRequest(
 					driverConfig,
 					onConnectWebSocket,
 					workerId,
+					authData,
 				)();
 			})(c, noopNext());
 		} else if ("custom" in handler.routingHandler) {
@@ -878,6 +908,9 @@ async function handleWebSocketConnectRequest(
 
 /**
  * Handle a connection message request to a worker
+ *
+ * There is no authentication handler on this request since the connection
+ * token is used to authenticate the message.
  */
 async function handleMessageRequest(
 	c: HonoContext,
@@ -900,6 +933,22 @@ async function handleMessageRequest(
 		}
 		const { workerId, connId, encoding, connToken } = params.data;
 
+		// TODO: This endpoint can be used to exhause resources (DoS attack) on an worker if you know the worker ID:
+		// 1. Get the worker ID (usually this is reasonably secure, but we don't assume worker ID is sensitive)
+		// 2. Spam messages to the worker (the conn token can be invalid)
+		// 3. The worker will be exhausted processing messages — even if the token is invalid
+		//
+		// The solution is we need to move the authorization of the connection token to this request handler
+		// AND include the worker ID in the connection token so we can verify that it has permission to send
+		// a message to that worker. This would require changing the token to a JWT so we can include a secure
+		// payload, but this requires managing a private key & managing key rotations.
+		//
+		// All other solutions (e.g. include the worker name as a header or include the worker name in the worker ID)
+		// have exploits that allow the caller to send messages to arbitrary workers.
+		//
+		// Currently, we assume this is not a critical problem because requests will likely get rate
+		// limited before enough messages are passed to the worker to exhaust resources.
+
 		// Handle based on mode
 		if ("inline" in handler.routingHandler) {
 			logger().debug("using inline proxy mode for connection message");
@@ -972,6 +1021,20 @@ async function handleActionRequest(
 			throw new errors.InvalidRequest(params.error);
 		}
 
+		// Parse connection parameters for authentication
+		const connParams = params.data.connParams
+			? JSON.parse(params.data.connParams)
+			: undefined;
+
+		// Authenticate the request
+		const authData = await authenticateEndpoint(
+			c,
+			registryConfig,
+			params.data.query,
+			["action"],
+			connParams,
+		);
+
 		// Get the worker ID and meta
 		const { workerId, meta } = await queryWorker(c, params.data.query, driver);
 		logger().debug("found worker for action", { workerId, meta });
@@ -988,6 +1051,7 @@ async function handleActionRequest(
 				handler.routingHandler.inline.handlers.onAction,
 				actionName,
 				workerId,
+				authData,
 			);
 		} else if ("custom" in handler.routingHandler) {
 			logger().debug("using custom proxy mode for action call");
@@ -1031,6 +1095,7 @@ async function handleActionRequest(
  */
 async function handleResolveRequest(
 	c: HonoContext,
+	registryConfig: RegistryConfig,
 	driver: ManagerDriver,
 ): Promise<Response> {
 	const encoding = getRequestEncoding(c.req, false);
@@ -1046,8 +1111,13 @@ async function handleResolveRequest(
 		throw new errors.InvalidRequest(params.error);
 	}
 
+	const query = params.data.query;
+
+	// Authenticate the request
+	await authenticateEndpoint(c, registryConfig, query, []);
+
 	// Get the worker ID and meta
-	const { workerId, meta } = await queryWorker(c, params.data.query, driver);
+	const { workerId, meta } = await queryWorker(c, query, driver);
 	logger().debug("resolved worker", { workerId, meta });
 	invariant(workerId, "Missing worker ID");
 

packages/core/src/topologies/standalone/topology.ts

@@ -157,6 +157,7 @@ export class StandaloneTopology {
 							connState,
 							CONN_DRIVER_GENERIC_WEBSOCKET,
 							{ encoding: opts.encoding } satisfies GenericWebSocketDriverState,
+							opts.authData,
 						);
 					},
 					onMessage: async (message) => {
@@ -199,6 +200,7 @@ export class StandaloneTopology {
 							connState,
 							CONN_DRIVER_GENERIC_SSE,
 							{ encoding: opts.encoding } satisfies GenericSseDriverState,
+							opts.authData,
 						);
 					},
 					onClose: async () => {
@@ -224,6 +226,7 @@ export class StandaloneTopology {
 						connState,
 						CONN_DRIVER_GENERIC_HTTP,
 						{} satisfies GenericHttpDriverState,
+						opts.authData,
 					);
 
 					// Call action