Merge pull request #294 from cloudflare/feat/ws-ticketing-auth

Feat: Auth via ticketing for Websocket

Author: AshishKumar4

worker/agents/core/codingAgent.ts

@@ -29,6 +29,8 @@ import { ProjectObjective } from "./objectives/base";
 import { FileOutputType } from "../schemas";
 import { SecretsClient, type UserSecretsStoreStub } from '../../services/secrets/SecretsClient';
 import { StateMigration } from './stateMigration';
+import { PendingWsTicket, TicketConsumptionResult } from '../../types/auth-types';
+import { WsTicketManager } from '../../utils/wsTicketManager';
 
 const DEFAULT_CONVERSATION_SESSION_ID = 'default';
 
@@ -43,6 +45,10 @@ export class CodeGeneratorAgent extends Agent<Env, AgentState> implements AgentI
     private objective!: ProjectObjective<BaseProjectState>;
     private secretsClient: SecretsClient | null = null;
     protected static readonly PROJECT_NAME_PREFIX_MAX_LENGTH = 20;
+    
+    /** Ticket manager for WebSocket authentication */
+    private ticketManager = new WsTicketManager();
+    
     // Services
     readonly fileManager: FileManager;
     readonly deploymentManager: DeploymentManager;
@@ -778,4 +784,34 @@ export class CodeGeneratorAgent extends Agent<Env, AgentState> implements AgentI
     clearGitHubToken(): void {
         this.objective.clearGitHubToken();
     }
+
+    // ==========================================
+    // WebSocket Ticket Management
+    // ==========================================
+
+    /**
+     * Store a WebSocket ticket for later consumption
+     * Called by controller after ownership is verified via JWT
+     */
+    storeWsTicket(ticket: PendingWsTicket): void {
+        this.ticketManager.store(ticket);
+        this.logger().info('WebSocket ticket stored', {
+            userId: ticket.user.id,
+            expiresIn: Math.floor((ticket.expiresAt - Date.now()) / 1000),
+        });
+    }
+
+    /**
+     * Consume a WebSocket ticket (one-time use)
+     * Returns user session if valid, null otherwise
+     */
+    consumeWsTicket(token: string): TicketConsumptionResult | null {
+        const result = this.ticketManager.consume(token);
+        if (result) {
+            this.logger().info('Ticket consumed successfully', { userId: result.user.id });
+        } else {
+            this.logger().warn('Ticket consumption failed', { tokenPrefix: token.slice(0, 10) });
+        }
+        return result;
+    }
 }

worker/api/controllers/agent/controller.ts

@@ -22,6 +22,7 @@ import { getPreviewDomain } from 'worker/utils/urls';
 import { ImageType, uploadImage } from 'worker/utils/images';
 import { ProcessedImageAttachment } from 'worker/types/image-attachment';
 import { getTemplateImportantFiles } from 'worker/services/sandbox/utils';
+import { hasTicketParam } from '../../../middleware/auth/ticketAuth';
 
 const defaultCodeGenArgs: Partial<CodeGenArgs> = {
     language: 'typescript',
@@ -209,6 +210,10 @@ export class CodingAgentController extends BaseController {
     /**
      * Handle WebSocket connections for code generation
      * This routes the WebSocket connection directly to the Agent
+     * 
+     * Supports two authentication methods:
+     * 1. Ticket-based auth (SDK): ?ticket=tk_xxx in URL
+     * 2. JWT-based auth (Browser): Cookie/Header with origin validation
      */
     static async handleWebSocketConnection(
         request: Request,
@@ -217,52 +222,43 @@ export class CodingAgentController extends BaseController {
         context: RouteContext
     ): Promise<Response> {
         try {
-            const chatId = context.pathParams.agentId; // URL param is still agentId for backward compatibility
-            if (!chatId) {
+            const agentId = context.pathParams.agentId;
+            if (!agentId) {
                 return CodingAgentController.createErrorResponse('Missing agent ID parameter', 400);
             }
 
             // Ensure the request is a WebSocket upgrade request
             if (request.headers.get('Upgrade') !== 'websocket') {
                 return new Response('Expected WebSocket upgrade', { status: 426 });
             }
-            
-            // Validate WebSocket origin
-            if (!validateWebSocketOrigin(request, env)) {
-                return new Response('Forbidden: Invalid origin', { status: 403 });
-            }
 
-            // Extract user for rate limiting
-            const user = context.user!;
+            // User already authenticated via ticket OR JWT by middleware
+            const user = context.user;
             if (!user) {
-                return CodingAgentController.createErrorResponse('Missing user', 401);
+                return CodingAgentController.createErrorResponse('Authentication required', 401);
             }
 
-            this.logger.info(`WebSocket connection request for chat: ${chatId}`);
-            
-            // Log request details for debugging
-            const headers: Record<string, string> = {};
-            request.headers.forEach((value, key) => {
-                headers[key] = value;
-            });
-            this.logger.info('WebSocket request details', {
-                headers,
-                url: request.url,
-                chatId
+            // Origin validation only for non-ticket auth (ticket auth is origin-agnostic)
+            const isTicketAuth = hasTicketParam(request);
+            if (!isTicketAuth && !validateWebSocketOrigin(request, env)) {
+                return new Response('Forbidden: Invalid origin', { status: 403 });
+            }
+
+            this.logger.info('WebSocket connection authorized', {
+                agentId,
+                userId: user.id,
+                authMethod: isTicketAuth ? 'ticket' : 'jwt',
             });
 
             try {
                 // Get the agent instance to handle the WebSocket connection
-                const agentInstance = await getAgentStub(env, chatId);
-                
-                this.logger.info(`Successfully got agent instance for chat: ${chatId}`);
+                const agentInstance = await getAgentStub(env, agentId);
 
                 // Let the agent handle the WebSocket connection directly
                 return agentInstance.fetch(request);
             } catch (error) {
-                this.logger.error(`Failed to get agent instance with ID ${chatId}:`, error);
+                this.logger.error(`Failed to get agent instance with ID ${agentId}:`, error);
                 // Return an appropriate WebSocket error response
-                // We need to emulate a WebSocket response even for errors
                 const { 0: client, 1: server } = new WebSocketPair();
 
                 server.accept();

worker/api/controllers/ticket/controller.ts

@@ -0,0 +1,127 @@
+/**
+ * Ticket Controller
+ *
+ * Creates one-time-use tickets for WebSocket authentication.
+ * Tickets are stored in the appropriate DO and consumed on connection.
+ */
+
+import { BaseController } from '../baseController';
+import { RouteContext } from '../../types/route-context';
+import { createLogger } from '../../../logger';
+import { checkAppOwnership } from '../../../middleware/auth/routeAuth';
+import { generateTicketToken, getResourceStub } from '../../../middleware/auth/ticketAuth';
+import type { TicketResourceType } from '../../../middleware/auth/routeAuth';
+import type { PendingWsTicket, AuthUser } from '../../../types/auth-types';
+
+const TICKET_TTL_MS = 15_000;
+
+interface CreateTicketRequest {
+	resourceType: TicketResourceType;
+	resourceId?: string;
+}
+
+// ============================================================================
+// Helpers
+// ============================================================================
+
+async function verifyOwnership(
+	user: AuthUser,
+	resourceType: TicketResourceType,
+	resourceId: string,
+	env: Env
+): Promise<boolean> {
+	switch (resourceType) {
+		case 'agent':
+			return checkAppOwnership(user, { agentId: resourceId }, env);
+		case 'vault':
+			return resourceId === user.id;
+	}
+}
+
+function resolveResourceId(
+	body: CreateTicketRequest,
+	userId: string
+): { resourceId: string } | { error: string } {
+	if (body.resourceType === 'vault') {
+		return { resourceId: userId };
+	}
+	if (!body.resourceId) {
+		return { error: 'resourceId is required for agent tickets' };
+	}
+	return { resourceId: body.resourceId };
+}
+
+// ============================================================================
+// Controller
+// ============================================================================
+
+export class TicketController extends BaseController {
+	static readonly logger = createLogger('TicketController');
+
+	/**
+	 * Create a WebSocket ticket
+	 * POST /api/ws-ticket
+	 */
+	static async createTicket(
+		request: Request,
+		env: Env,
+		_ctx: ExecutionContext,
+		context: RouteContext
+	): Promise<Response> {
+		const user = context.user;
+		if (!user) {
+			return this.createErrorResponse('Authentication required', 401);
+		}
+
+		// Parse and validate request
+		let body: CreateTicketRequest;
+		try {
+			body = await request.json() as CreateTicketRequest;
+		} catch {
+			return this.createErrorResponse('Invalid JSON body', 400);
+		}
+
+		if (!body.resourceType || !['agent', 'vault'].includes(body.resourceType)) {
+			return this.createErrorResponse('Invalid resourceType', 400);
+		}
+
+		// Resolve resource ID
+		const resolved = resolveResourceId(body, user.id);
+		if ('error' in resolved) {
+			return this.createErrorResponse(resolved.error, 400);
+		}
+		const { resourceId } = resolved;
+
+		// Verify ownership
+		if (!await verifyOwnership(user, body.resourceType, resourceId, env)) {
+			this.logger.warn('Ticket creation denied', { userId: user.id, resourceType: body.resourceType, resourceId });
+			return this.createErrorResponse('Access denied', 403);
+		}
+
+		// Create and store ticket
+		const now = Date.now();
+		const ticket: PendingWsTicket = {
+			token: generateTicketToken(body.resourceType, resourceId),
+			user,
+			sessionId: context.sessionId ?? `ticket:${body.resourceType}:${resourceId}`,
+			createdAt: now,
+			expiresAt: now + TICKET_TTL_MS,
+		};
+
+		try {
+			const stub = await getResourceStub(env, body.resourceType, resourceId);
+			await stub.storeWsTicket(ticket);
+
+			this.logger.info('Ticket created', { resourceType: body.resourceType, resourceId, userId: user.id });
+
+			return this.createSuccessResponse({
+				ticket: ticket.token,
+				expiresIn: Math.floor(TICKET_TTL_MS / 1000),
+				expiresAt: new Date(ticket.expiresAt).toISOString(),
+			});
+		} catch (error) {
+			this.logger.error('Failed to create ticket', { resourceType: body.resourceType, resourceId, error });
+			return this.createErrorResponse('Failed to create ticket', 500);
+		}
+	}
+}

worker/api/routes/codegenRoutes.ts

@@ -19,9 +19,11 @@ export function setupCodegenRoutes(app: Hono<AppEnv>): void {
     // APP EDITING ROUTES (/chat/:id frontend)
     // ========================================
 
-    // WebSocket for app editing - OWNER ONLY (for /chat/:id route)
-    // Only the app owner should be able to connect and modify via WebSocket
-    app.get('/api/agent/:agentId/ws', setAuthLevel(AuthConfig.ownerOnly), adaptController(CodingAgentController, CodingAgentController.handleWebSocketConnection));
+    // WebSocket for app editing - OWNER ONLY with ticket support
+    // Supports ticket-based auth (SDK) or JWT-based auth (browser)
+    app.get('/api/agent/:agentId/ws', setAuthLevel(AuthConfig.ownerOnly, { 
+        ticketAuth: { resourceType: 'agent', paramName: 'agentId' } 
+    }), adaptController(CodingAgentController, CodingAgentController.handleWebSocketConnection));
 
     // Connect to existing agent for editing - OWNER ONLY
     // Only the app owner should be able to connect for editing purposes

worker/api/routes/index.ts

@@ -11,6 +11,7 @@ import { setupCodegenRoutes } from './codegenRoutes';
 import { setupScreenshotRoutes } from './imagesRoutes';
 import { setupSentryRoutes } from './sentryRoutes';
 import { setupCapabilitiesRoutes } from './capabilitiesRoutes';
+import { setupTicketRoutes } from './ticketRoutes';
 import { Hono } from "hono";
 import { AppEnv } from "../../types/appenv";
 import { setupStatusRoutes } from './statusRoutes';
@@ -33,6 +34,9 @@ export function setupRoutes(app: Hono<AppEnv>): void {
     // Authentication and user management routes
     setupAuthRoutes(app);
 
+    // WebSocket ticket routes
+    setupTicketRoutes(app);
+    
     // Codegen routes
     setupCodegenRoutes(app);
 

worker/api/routes/ticketRoutes.ts

@@ -0,0 +1,19 @@
+/**
+ * WebSocket Ticket Routes
+ */
+
+import { Hono } from 'hono';
+import { AppEnv } from '../../types/appenv';
+import { AuthConfig, setAuthLevel } from '../../middleware/auth/routeAuth';
+import { adaptController } from '../honoAdapter';
+import { TicketController } from '../controllers/ticket/controller';
+
+export function setupTicketRoutes(app: Hono<AppEnv>): void {
+	// Create WebSocket ticket - requires authentication
+	// Ownership check is done in the controller based on resourceType
+	app.post(
+		'/api/ws-ticket',
+		setAuthLevel(AuthConfig.authenticated),
+		adaptController(TicketController, TicketController.createTicket)
+	);
+}

worker/api/routes/userSecretsRoutes.ts

@@ -13,9 +13,12 @@ export function setupUserSecretsRoutes(app: Hono<AppEnv>): void {
 	const vaultRouter = new Hono<AppEnv>();
 
 	// WebSocket connection to vault DO - must be before other routes
+	// Supports ticket-based auth (SDK) or JWT-based auth (browser)
 	vaultRouter.get(
 		'/ws',
-		setAuthLevel(AuthConfig.authenticated),
+		setAuthLevel(AuthConfig.authenticated, {
+			ticketAuth: { resourceType: 'vault' }
+		}),
 		adaptController(UserSecretsController, UserSecretsController.handleWebSocketConnection)
 	);