WebSocket support

packages/vite-plugin-cloudflare/package.json

@@ -37,12 +37,14 @@
 	"dependencies": {
 		"@hattip/adapter-node": "^0.0.49",
 		"miniflare": "3.20241205.0",
-		"unenv": "catalog:default"
+		"unenv": "catalog:default",
+		"ws": "^8.18.0"
 	},
 	"devDependencies": {
 		"@cloudflare/workers-shared": "^0.7.0",
 		"@cloudflare/workers-types": "catalog:default",
 		"@types/node": "catalog:default",
+		"@types/ws": "^8.5.13",
 		"@vite-plugin-cloudflare/typescript-config": "workspace:*",
 		"magic-string": "^0.30.12",
 		"tsup": "^8.3.0",

packages/vite-plugin-cloudflare/src/index.ts

@@ -21,6 +21,7 @@ import {
 } from './node-js-compat';
 import { resolvePluginConfig } from './plugin-config';
 import { getOutputDirectory, toMiniflareRequest } from './utils';
+import { handleWebSocket } from './websockets';
 import { getWarningForWorkersConfigs } from './workers-configs';
 import type { PluginConfig, ResolvedPluginConfig } from './plugin-config';
 import type { Unstable_RawConfig } from 'wrangler';
@@ -231,6 +232,8 @@ export function cloudflare(pluginConfig: PluginConfig = {}): vite.Plugin {
 			}
 		},
 		async configureServer(viteDevServer) {
+			assert(viteDevServer.httpServer, 'Unexpected error: No Vite HTTP server');
+
 			miniflare = new Miniflare(
 				getDevMiniflareOptions(resolvedPluginConfig, viteDevServer),
 			);
@@ -247,6 +250,12 @@ export function cloudflare(pluginConfig: PluginConfig = {}): vite.Plugin {
 				}) as any;
 			});
 
+			handleWebSocket(
+				viteDevServer.httpServer,
+				entryWorker.fetch,
+				viteDevServer.config.logger,
+			);
+
 			return () => {
 				viteDevServer.middlewares.use((req, res, next) => {
 					middleware(req, res, next);
@@ -267,6 +276,12 @@ export function cloudflare(pluginConfig: PluginConfig = {}): vite.Plugin {
 				}) as any;
 			});
 
+			handleWebSocket(
+				vitePreviewServer.httpServer,
+				miniflare.dispatchFetch,
+				vitePreviewServer.config.logger,
+			);
+
 			return () => {
 				vitePreviewServer.middlewares.use((req, res, next) => {
 					middleware(req, res, next);

packages/vite-plugin-cloudflare/src/utils.ts

@@ -1,6 +1,7 @@
 import * as path from 'node:path';
 import { Request as MiniflareRequest } from 'miniflare';
 import * as vite from 'vite';
+import type { IncomingHttpHeaders } from 'node:http';
 
 export function getOutputDirectory(
 	userConfig: vite.UserConfig,
@@ -23,4 +24,22 @@ export function toMiniflareRequest(request: Request): MiniflareRequest {
 	});
 }
 
+export function nodeHeadersToWebHeaders(
+	nodeHeaders: IncomingHttpHeaders,
+): Headers {
+	const headers = new Headers();
+
+	for (const [key, value] of Object.entries(nodeHeaders)) {
+		if (typeof value === 'string') {
+			headers.append(key, value);
+		} else if (Array.isArray(value)) {
+			for (const item of value) {
+				headers.append(key, item);
+			}
+		}
+	}
+
+	return headers;
+}
+
 export type Optional<T, K extends keyof T> = Omit<T, K> & Pick<Partial<T>, K>;

packages/vite-plugin-cloudflare/src/websockets.ts

@@ -0,0 +1,81 @@
+import ws from 'ws';
+import { UNKNOWN_HOST } from './shared';
+import { nodeHeadersToWebHeaders } from './utils';
+import type { Fetcher } from '@cloudflare/workers-types/experimental';
+import type { ReplaceWorkersTypes } from 'miniflare';
+import type { IncomingMessage } from 'node:http';
+import type { Duplex } from 'node:stream';
+import type * as vite from 'vite';
+
+/**
+ * This function handles 'upgrade' requests to the Vite HTTP server and forwards WebSocket events between the client and Worker environments.
+ */
+export function handleWebSocket(
+	httpServer: vite.HttpServer,
+	fetcher: ReplaceWorkersTypes<Fetcher>['fetch'],
+	logger: vite.Logger,
+) {
+	const nodeWebSocket = new ws.Server({ noServer: true });
+
+	httpServer.on(
+		'upgrade',
+		async (request: IncomingMessage, socket: Duplex, head: Buffer) => {
+			const url = new URL(request.url ?? '', UNKNOWN_HOST);
+
+			// Ignore Vite HMR WebSockets
+			if (request.headers['sec-websocket-protocol']?.startsWith('vite')) {
+				return;
+			}
+
+			const headers = nodeHeadersToWebHeaders(request.headers);
+			const response = await fetcher(url, {
+				headers,
+				method: request.method,
+			});
+			const workerWebSocket = response.webSocket;
+
+			if (!workerWebSocket) {
+				socket.destroy();
+				return;
+			}
+
+			nodeWebSocket.handleUpgrade(
+				request,
+				socket,
+				head,
+				async (clientWebSocket) => {
+					workerWebSocket.accept();
+
+					// Forward Worker events to client
+					workerWebSocket.addEventListener('message', (event) => {
+						clientWebSocket.send(event.data);
+					});
+					workerWebSocket.addEventListener('error', (event) => {
+						logger.error(
+							`WebSocket error:\n${event.error?.stack || event.error?.message}`,
+							{ error: event.error },
+						);
+					});
+					workerWebSocket.addEventListener('close', () => {
+						clientWebSocket.close();
+					});
+
+					// Forward client events to Worker
+					clientWebSocket.on('message', (event: ArrayBuffer | string) => {
+						workerWebSocket.send(event);
+					});
+					clientWebSocket.on('error', (error) => {
+						logger.error(`WebSocket error:\n${error.stack || error.message}`, {
+							error,
+						});
+					});
+					clientWebSocket.on('close', () => {
+						workerWebSocket.close();
+					});
+
+					nodeWebSocket.emit('connection', clientWebSocket, request);
+				},
+			);
+		},
+	);
+}

playground/websockets/__tests__/websockets.spec.ts

@@ -0,0 +1,41 @@
+import { expect, test, vi } from 'vitest';
+import { page } from '../../__test-utils__';
+
+async function openWebSocket() {
+	const openButton = page.getByRole('button', { name: 'Open WebSocket' });
+	const statusTextBefore = await page.textContent('h2');
+	expect(statusTextBefore).toBe('WebSocket closed');
+	await openButton.click();
+	await vi.waitFor(async () => {
+		const statusTextAfter = await page.textContent('h2');
+		expect(statusTextAfter).toBe('WebSocket open');
+	});
+}
+
+test('opens WebSocket connection', openWebSocket);
+
+test('closes WebSocket connection', async () => {
+	await openWebSocket();
+	const closeButton = page.getByRole('button', { name: 'Close WebSocket' });
+	const statusTextBefore = await page.textContent('h2');
+	expect(statusTextBefore).toBe('WebSocket open');
+	await closeButton.click();
+	await vi.waitFor(async () => {
+		const statusTextAfter = await page.textContent('h2');
+		expect(statusTextAfter).toBe('WebSocket closed');
+	});
+});
+
+test('sends and receives WebSocket messages', async () => {
+	await openWebSocket();
+	const sendButton = page.getByRole('button', { name: 'Send message' });
+	const messageTextBefore = await page.textContent('p');
+	expect(messageTextBefore).toBe('');
+	await sendButton.click();
+	await vi.waitFor(async () => {
+		const messageTextAfter = await page.textContent('p');
+		expect(messageTextAfter).toBe(
+			`Durable Object received client message: 'Client event'.`,
+		);
+	});
+});

playground/websockets/package.json

@@ -0,0 +1,19 @@
+{
+	"name": "@playground/websockets",
+	"private": true,
+	"type": "module",
+	"scripts": {
+		"build": "vite build --app",
+		"check:types": "tsc --build",
+		"dev": "vite dev",
+		"preview": "vite preview"
+	},
+	"devDependencies": {
+		"@cloudflare/workers-types": "catalog:default",
+		"@flarelabs-net/vite-plugin-cloudflare": "workspace:*",
+		"@vite-plugin-cloudflare/typescript-config": "workspace:*",
+		"typescript": "catalog:default",
+		"vite": "catalog:default",
+		"wrangler": "catalog:default"
+	}
+}

playground/websockets/src/index.html

@@ -0,0 +1,69 @@
+<!doctype html>
+<html>
+	<head>
+		<meta charset="UTF-8" />
+		<meta name="viewport" content="width=device-width, initial-scale=1.0" />
+	</head>
+	<body>
+		<main>
+			<h1>WebSockets playground</h1>
+			<button id="open" aria-label="Open WebSocket">Open</button>
+			<button id="close" aria-label="Close WebSocket">Close</button>
+			<button id="send" aria-label="Send message">Send</button>
+			<h2 id="status">WebSocket closed</h2>
+			<p id="message"></p>
+		</main>
+	</body>
+	<script type="module">
+		let ws;
+		const openButton = document.querySelector('#open');
+		const closeButton = document.querySelector('#close');
+		const sendButton = document.querySelector('#send');
+		const statusText = document.querySelector('#status');
+		const messageText = document.querySelector('#message');
+
+		openButton.addEventListener('click', () => open());
+		closeButton.addEventListener('click', () => close());
+		sendButton.addEventListener('click', () => send('Client event'));
+
+		function open() {
+			if (ws) {
+				console.log('WebSocket already open');
+
+				return;
+			}
+
+			ws = new WebSocket(`${location.origin.replace(/^http/, 'ws')}/websocket`);
+
+			ws.addEventListener('message', (event) => {
+				console.log('Received server event', event.data);
+				messageText.textContent = event.data;
+			});
+
+			ws.addEventListener('open', (event) => {
+				console.log('Client open');
+				statusText.textContent = 'WebSocket open';
+			});
+
+			ws.addEventListener('close', (event) => {
+				console.log('Client close');
+				statusText.textContent = 'WebSocket closed';
+			});
+		}
+
+		function close() {
+			if (!ws) {
+				console.log('WebSocket already closed');
+
+				return;
+			}
+
+			ws.close();
+			ws = undefined;
+		}
+
+		function send(message) {
+			ws.send(message);
+		}
+	</script>
+</html>

playground/websockets/src/index.ts

@@ -0,0 +1,44 @@
+import html from './index.html?raw';
+import { DurableObject } from 'cloudflare:workers';
+
+interface Env {
+	WEBSOCKET_SERVER: DurableObjectNamespace<WebSocketServer>;
+}
+
+export class WebSocketServer extends DurableObject {
+	override fetch() {
+		const { 0: client, 1: server } = new WebSocketPair();
+
+		this.ctx.acceptWebSocket(server);
+
+		return new Response(null, { status: 101, webSocket: client });
+	}
+
+	override async webSocketMessage(ws: WebSocket, data: string | ArrayBuffer) {
+		const decoder = new TextDecoder();
+		const message = typeof data === 'string' ? data : decoder.decode(data);
+
+		ws.send(`Durable Object received client message: '${message}'.`);
+	}
+}
+
+export default {
+	async fetch(request, env) {
+		if (request.url.endsWith('/websocket')) {
+			const upgradeHeader = request.headers.get('Upgrade');
+
+			if (!upgradeHeader || upgradeHeader !== 'websocket') {
+				return new Response('Durable Object expected Upgrade: websocket', {
+					status: 426,
+				});
+			}
+
+			const id = env.WEBSOCKET_SERVER.idFromName('id');
+			const stub = env.WEBSOCKET_SERVER.get(id);
+
+			return stub.fetch(request);
+		}
+
+		return new Response(html, { headers: { 'content-type': 'text/html' } });
+	},
+} satisfies ExportedHandler<Env>;

playground/websockets/tsconfig.json

@@ -0,0 +1,7 @@
+{
+	"files": [],
+	"references": [
+		{ "path": "./tsconfig.node.json" },
+		{ "path": "./tsconfig.worker.json" }
+	]
+}

playground/websockets/tsconfig.node.json

@@ -0,0 +1,4 @@
+{
+	"extends": ["@vite-plugin-cloudflare/typescript-config/base.json"],
+	"include": ["vite.config.ts", "__tests__"]
+}

playground/websockets/tsconfig.worker.json

@@ -0,0 +1,4 @@
+{
+	"extends": ["@vite-plugin-cloudflare/typescript-config/worker.json"],
+	"include": ["src"]
+}

playground/websockets/vite.config.ts

@@ -0,0 +1,6 @@
+import { cloudflare } from '@flarelabs-net/vite-plugin-cloudflare';
+import { defineConfig } from 'vite';
+
+export default defineConfig({
+	plugins: [cloudflare({ persistState: false })],
+});

playground/websockets/wrangler.toml

@@ -0,0 +1,10 @@
+name = "worker"
+main = "./src/index.ts"
+compatibility_date = "2024-09-09"
+
+[durable_objects]
+bindings = [{ name = "WEBSOCKET_SERVER", class_name = "WebSocketServer" }]
+
+[[migrations]]
+tag = "v1"
+new_classes = ["WebSocketServer"]