diff --git a/packages/vite-plugin-cloudflare/package.json b/packages/vite-plugin-cloudflare/package.json index 1f5d190..70ac030 100644 --- a/packages/vite-plugin-cloudflare/package.json +++ b/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", diff --git a/packages/vite-plugin-cloudflare/src/index.ts b/packages/vite-plugin-cloudflare/src/index.ts index 2e7faa0..37ce7de 100644 --- a/packages/vite-plugin-cloudflare/src/index.ts +++ b/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); diff --git a/packages/vite-plugin-cloudflare/src/utils.ts b/packages/vite-plugin-cloudflare/src/utils.ts index db43576..b2c0544 100644 --- a/packages/vite-plugin-cloudflare/src/utils.ts +++ b/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 = Omit & Pick, K>; diff --git a/packages/vite-plugin-cloudflare/src/websockets.ts b/packages/vite-plugin-cloudflare/src/websockets.ts new file mode 100644 index 0000000..e4b53b5 --- /dev/null +++ b/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['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); + }, + ); + }, + ); +} diff --git a/playground/websockets/__tests__/websockets.spec.ts b/playground/websockets/__tests__/websockets.spec.ts new file mode 100644 index 0000000..e5275c1 --- /dev/null +++ b/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'.`, + ); + }); +}); diff --git a/playground/websockets/package.json b/playground/websockets/package.json new file mode 100644 index 0000000..6f04372 --- /dev/null +++ b/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" + } +} diff --git a/playground/websockets/src/index.html b/playground/websockets/src/index.html new file mode 100644 index 0000000..9c70cf7 --- /dev/null +++ b/playground/websockets/src/index.html @@ -0,0 +1,69 @@ + + + + + + + +
+

WebSockets playground

+ + + +

WebSocket closed

+

+
+ + + diff --git a/playground/websockets/src/index.ts b/playground/websockets/src/index.ts new file mode 100644 index 0000000..4426dd7 --- /dev/null +++ b/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; +} + +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; diff --git a/playground/websockets/tsconfig.json b/playground/websockets/tsconfig.json new file mode 100644 index 0000000..b52af70 --- /dev/null +++ b/playground/websockets/tsconfig.json @@ -0,0 +1,7 @@ +{ + "files": [], + "references": [ + { "path": "./tsconfig.node.json" }, + { "path": "./tsconfig.worker.json" } + ] +} diff --git a/playground/websockets/tsconfig.node.json b/playground/websockets/tsconfig.node.json new file mode 100644 index 0000000..c24133d --- /dev/null +++ b/playground/websockets/tsconfig.node.json @@ -0,0 +1,4 @@ +{ + "extends": ["@vite-plugin-cloudflare/typescript-config/base.json"], + "include": ["vite.config.ts", "__tests__"] +} diff --git a/playground/websockets/tsconfig.worker.json b/playground/websockets/tsconfig.worker.json new file mode 100644 index 0000000..b2ab8f7 --- /dev/null +++ b/playground/websockets/tsconfig.worker.json @@ -0,0 +1,4 @@ +{ + "extends": ["@vite-plugin-cloudflare/typescript-config/worker.json"], + "include": ["src"] +} diff --git a/playground/websockets/vite.config.ts b/playground/websockets/vite.config.ts new file mode 100644 index 0000000..f7c1b36 --- /dev/null +++ b/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 })], +}); diff --git a/playground/websockets/wrangler.toml b/playground/websockets/wrangler.toml new file mode 100644 index 0000000..d3d793e --- /dev/null +++ b/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"] diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 16e9adc..ecd3e80 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -74,6 +74,9 @@ importers: unenv: specifier: catalog:default version: unenv-nightly@2.0.0-20241218-183400-5d6aec3 + ws: + specifier: ^8.18.0 + version: 8.18.0 devDependencies: '@cloudflare/workers-shared': specifier: ^0.7.0 @@ -84,6 +87,9 @@ importers: '@types/node': specifier: catalog:default version: 22.10.1 + '@types/ws': + specifier: ^8.5.13 + version: 8.5.13 '@vite-plugin-cloudflare/typescript-config': specifier: workspace:* version: link:../typescript-config @@ -411,6 +417,27 @@ importers: specifier: catalog:default version: 3.94.0(@cloudflare/workers-types@4.20241205.0) + playground/websockets: + devDependencies: + '@cloudflare/workers-types': + specifier: catalog:default + version: 4.20241205.0 + '@flarelabs-net/vite-plugin-cloudflare': + specifier: workspace:* + version: link:../../packages/vite-plugin-cloudflare + '@vite-plugin-cloudflare/typescript-config': + specifier: workspace:* + version: link:../../packages/typescript-config + typescript: + specifier: catalog:default + version: 5.7.2 + vite: + specifier: catalog:default + version: 6.0.7(@types/node@22.10.1) + wrangler: + specifier: catalog:default + version: 3.94.0(@cloudflare/workers-types@4.20241205.0) + playground/worker: devDependencies: '@cloudflare/workers-types': @@ -1603,6 +1630,9 @@ packages: '@types/react@19.0.1': resolution: {integrity: sha512-YW6614BDhqbpR5KtUYzTA+zlA7nayzJRA9ljz9CQoxthR0sDisYZLuvSMsil36t4EH/uAt8T52Xb4sVw17G+SQ==} + '@types/ws@8.5.13': + resolution: {integrity: sha512-osM/gWBTPKgHV8XkTunnegTRIsvF6owmf5w+JtAfOw472dptdm0dlGv4xCt6GwQRcC2XVOvvRE/0bAoQcL2QkA==} + '@vitejs/plugin-react@4.3.4': resolution: {integrity: sha512-SCCPBJtYLdE8PX/7ZQAs1QAZ8Jqwih+0VBLum1EGqmCCQal+MIUqLCzj3ZUy8ufbC0cAM4LRlSTm7IQJwWT4ug==} engines: {node: ^14.18.0 || >=16.0.0} @@ -3847,6 +3877,10 @@ snapshots: dependencies: csstype: 3.1.3 + '@types/ws@8.5.13': + dependencies: + '@types/node': 22.10.1 + '@vitejs/plugin-react@4.3.4(vite@6.0.7(@types/node@22.10.1))': dependencies: '@babel/core': 7.26.0