From 86bb1f9cdaafdff4d974f5690d72707e93302bb5 Mon Sep 17 00:00:00 2001 From: MrBBot Date: Sat, 13 Nov 2021 10:11:00 +0000 Subject: [PATCH] Wait for worker response before opening WebSocket in client, closes #88 --- packages/http-server/src/index.ts | 72 +++++++++++-------------- packages/http-server/test/index.spec.ts | 28 ++++++---- 2 files changed, 49 insertions(+), 51 deletions(-) diff --git a/packages/http-server/src/index.ts b/packages/http-server/src/index.ts index 218d6e767..302904166 100644 --- a/packages/http-server/src/index.ts +++ b/packages/http-server/src/index.ts @@ -18,7 +18,7 @@ import { import { randomHex } from "@miniflare/shared"; import { coupleWebSocket } from "@miniflare/web-sockets"; import { BodyInit, Headers } from "undici"; -import StandardWebSocket, { WebSocketServer } from "ws"; +import { WebSocketServer } from "ws"; import { getAccessibleHosts } from "./helpers"; import { HTTPPlugin, RequestMeta } from "./plugin"; @@ -305,38 +305,6 @@ export function createRequestListener( }; } -export type WebSocketUpgradeListener = ( - ws: StandardWebSocket, - req: http.IncomingMessage -) => void; - -export function createWebSocketUpgradeListener< - Plugins extends CorePluginSignatures ->( - mf: MiniflareCore, - listener: RequestListener -): WebSocketUpgradeListener { - return async (ws, req) => { - // Handle request in worker - const response = await listener(req); - - // Check web socket response was returned - const webSocket = response?.webSocket; - if (response?.status !== 101 || !webSocket) { - ws.close(1002, "Protocol Error"); - mf.log.error( - new TypeError( - "Web Socket request did not return status 101 Switching Protocols response with Web Socket" - ) - ); - return; - } - - // Couple the web socket here - await coupleWebSocket(ws, webSocket); - }; -} - export async function createServer( mf: MiniflareCore, options?: http.ServerOptions & https.ServerOptions @@ -355,17 +323,39 @@ export async function createServer( } // Setup WebSocket servers - const upgrader = createWebSocketUpgradeListener(mf, listener); const webSocketServer = new WebSocketServer({ noServer: true }); - webSocketServer.on("connection", upgrader); const liveReloadServer = new WebSocketServer({ noServer: true }); - server.on("upgrade", (request, socket, head) => { + server.on("upgrade", async (request, socket, head) => { + // Only interested in pathname so base URL doesn't matter const { pathname } = new URL(request.url ?? "", "http://localhost"); - const server = - pathname === "/cdn-cgi/mf/reload" ? liveReloadServer : webSocketServer; - server.handleUpgrade(request, socket as any, head, (ws) => - server.emit("connection", ws, request) - ); + if (pathname === "/cdn-cgi/mf/reload") { + // If this is the for live-reload, handle the request ourselves + liveReloadServer.handleUpgrade(request, socket as any, head, (ws) => { + liveReloadServer.emit("connection", ws, request); + }); + } else { + // Otherwise, handle the request in the worker + const response = await listener(request); + + // Check web socket response was returned + const webSocket = response?.webSocket; + if (response?.status !== 101 || !webSocket) { + socket.write("HTTP/1.1 500 Internal Server Error\r\n\r\n"); + socket.destroy(); + mf.log.error( + new TypeError( + "Web Socket request did not return status 101 Switching Protocols response with Web Socket" + ) + ); + return; + } + + // Accept and couple the Web Socket + webSocketServer.handleUpgrade(request, socket as any, head, (ws) => { + void coupleWebSocket(ws, webSocket); + webSocketServer.emit("connection", ws, request); + }); + } }); const reloadListener = () => { // Reload all connected live reload clients diff --git a/packages/http-server/test/index.spec.ts b/packages/http-server/test/index.spec.ts index 53444ab6e..d4f29f361 100644 --- a/packages/http-server/test/index.spec.ts +++ b/packages/http-server/test/index.spec.ts @@ -32,7 +32,12 @@ import { } from "@miniflare/shared-test"; import { MessageEvent, WebSocketPlugin } from "@miniflare/web-sockets"; import test, { ExecutionContext, Macro } from "ava"; -import StandardWebSocket, { Data, Event as WebSocketEvent } from "ws"; +import StandardWebSocket, { + Data, + CloseEvent as WebSocketCloseEvent, + ErrorEvent as WebSocketErrorEvent, + Event as WebSocketEvent, +} from "ws"; function listen( t: ExecutionContext, @@ -487,7 +492,10 @@ test("createServer: handles web socket upgrades", async (t) => { const mf = useMiniflareWithHandler( { HTTPPlugin, WebSocketPlugin }, {}, - (globals) => { + async (globals) => { + // Simulate slow response, WebSocket must not open until worker responds + await new Promise((resolve) => globals.setTimeout(resolve, 1000)); + const [client, worker] = Object.values(new globals.WebSocketPair()); worker.accept(); worker.addEventListener("message", (e: MessageEvent) => { @@ -523,18 +531,18 @@ test("createServer: expects status 101 and web socket response for upgrades", as const port = await listen(t, await createServer(mf)); const ws = new StandardWebSocket(`ws://localhost:${port}`); - const [eventTrigger, eventPromise] = triggerPromise<{ - code: number; - reason: string; - }>(); - ws.addEventListener("close", eventTrigger); - const event = await eventPromise; + const [closeTrigger, closePromise] = triggerPromise(); + const [errorTrigger, errorPromise] = triggerPromise(); + ws.addEventListener("close", closeTrigger); + ws.addEventListener("error", errorTrigger); + const closeEvent = await closePromise; + const errorEvent = await errorPromise; t.deepEqual(log.logsAtLevel(LogLevel.ERROR), [ "TypeError: Web Socket request did not return status 101 Switching Protocols response with Web Socket", ]); - t.is(event.code, 1002); - t.is(event.reason, "Protocol Error"); + t.is(closeEvent.code, 1006); + t.is(errorEvent.message, "Unexpected server response: 500"); }); test("createServer: notifies connected live reload clients on reload", async (t) => { const mf = useMiniflareWithHandler({ HTTPPlugin }, {}, (globals) => {