diff --git a/packages/fresh/src/context.ts b/packages/fresh/src/context.ts index 3c2d9c07815..da835177623 100644 --- a/packages/fresh/src/context.ts +++ b/packages/fresh/src/context.ts @@ -58,6 +58,28 @@ export interface WebSocketUpgradeOptions { protocol?: string; } +/** + * Side channel used by `@fresh/plugin-vite` (and any other adapter sitting on + * top of `node:http`) to hand a raw socket + buffered head to `ctx.upgrade()`. + * `Deno.upgradeWebSocket()` normally pulls these off the request handled by + * `Deno.serve`, but a request synthesized from `node:http` carries neither, so + * the adapter stashes them here and `ctx.upgrade()` forwards them to Deno. + * + * Stored on `globalThis` so a single WeakMap is shared even when this module + * is evaluated more than once (e.g. once by Deno for the Vite plugin and once + * by Vite's SSR runner for the user's server code). + */ +const UPGRADE_SOURCE_KEY: unique symbol = Symbol.for( + "fresh.upgradeSourceMap", +) as typeof UPGRADE_SOURCE_KEY; +// deno-lint-ignore no-explicit-any +type UpgradeSource = { socket: any; head: any }; +// deno-lint-ignore no-explicit-any +const globalAny = globalThis as any; +export const upgradeSourceMap: WeakMap = + globalAny[UPGRADE_SOURCE_KEY] ?? + (globalAny[UPGRADE_SOURCE_KEY] = new WeakMap()); + /** * Duck-type check: treats the argument as managed-mode handlers when at least * one of the handler keys (`open`, `message`, `close`, `error`) is a @@ -595,7 +617,16 @@ export class Context { throw new HttpError(400, "Expected a WebSocket upgrade request"); } - const { socket, response } = Deno.upgradeWebSocket(this.req, options); + const source = upgradeSourceMap.get(this.req); + if (source !== undefined) upgradeSourceMap.delete(this.req); + const upgradeOptions = source + ? { ...options, socket: source.socket, head: source.head } + : options; + const { socket, response } = Deno.upgradeWebSocket( + this.req, + // deno-lint-ignore no-explicit-any + upgradeOptions as any, + ); if (handlers === undefined) { return { socket, response }; diff --git a/packages/fresh/src/internals.ts b/packages/fresh/src/internals.ts index 69455f77b7f..a0dec9dcbc5 100644 --- a/packages/fresh/src/internals.ts +++ b/packages/fresh/src/internals.ts @@ -4,3 +4,4 @@ export { setBuildCache, setErrorInterceptor } from "./app.ts"; export { IslandPreparer, ProdBuildCache } from "./build_cache.ts"; export { path }; export { ASSET_CACHE_BUST_KEY } from "./constants.ts"; +export { upgradeSourceMap } from "./context.ts"; diff --git a/packages/plugin-vite/demo/routes/tests/ws_echo.ts b/packages/plugin-vite/demo/routes/tests/ws_echo.ts new file mode 100644 index 00000000000..770e5d4fa6b --- /dev/null +++ b/packages/plugin-vite/demo/routes/tests/ws_echo.ts @@ -0,0 +1,11 @@ +import { define } from "../../utils.ts"; + +export const handler = define.handlers({ + GET(ctx) { + return ctx.upgrade({ + message(socket, e) { + socket.send("echo: " + e.data); + }, + }); + }, +}); diff --git a/packages/plugin-vite/src/plugins/dev_server.ts b/packages/plugin-vite/src/plugins/dev_server.ts index c1ca76d6c1f..a69a4dbc57d 100644 --- a/packages/plugin-vite/src/plugins/dev_server.ts +++ b/packages/plugin-vite/src/plugins/dev_server.ts @@ -1,7 +1,7 @@ import type { DevEnvironment, Plugin } from "vite"; import * as path from "@std/path"; import { contentType as getStdContentType } from "@std/media-types/content-type"; -import { ASSET_CACHE_BUST_KEY } from "fresh/internal"; +import { ASSET_CACHE_BUST_KEY, upgradeSourceMap } from "fresh/internal"; import { createRequest, sendResponse } from "@remix-run/node-fetch-server"; import { hashCode } from "../shared.ts"; import type { ResolvedFreshViteConfig } from "../utils.ts"; @@ -28,6 +28,75 @@ export function devServer(freshConfig: ResolvedFreshViteConfig): Plugin[] { `^(${base})?/(@(vite|fs|id)|\\.vite)/`, ); + // WebSocket upgrade requests don't flow through Connect middleware — + // node:http dispatches them as `upgrade` events on the http.Server. + // Vite's own HMR server identifies its connections via the + // `sec-websocket-protocol` header (`vite-hmr`/`vite-ping`), so we + // skip those and let Vite handle them. Everything else we route into + // the Fresh handler with the raw socket + head stashed in + // `upgradeSourceMap` so `ctx.upgrade()` can finish the handshake via + // `Deno.upgradeWebSocket(req, { socket, head })`. + server.httpServer?.on("upgrade", (nodeReq, nodeSocket, head) => { + const protocolHeader = nodeReq.headers["sec-websocket-protocol"]; + if ( + typeof protocolHeader === "string" && + /\b(vite-hmr|vite-ping)\b/.test(protocolHeader) + ) { + return; + } + + (async () => { + try { + const serverCfg = server.config.server; + const protocol = serverCfg.https ? "https" : "http"; + const host = nodeReq.headers.host ?? "localhost"; + const url = new URL( + nodeReq.url ?? "/", + `${protocol}://${host}`, + ); + + const headers = new Headers(); + for (const [key, value] of Object.entries(nodeReq.headers)) { + if (value === undefined || value === null) continue; + if (Array.isArray(value)) { + for (const v of value) { + if (typeof v === "string") headers.append(key, v); + } + } else if (typeof value === "string") { + headers.set(key, value); + } + } + + const req = new Request(url, { + method: nodeReq.method ?? "GET", + headers, + }); + + upgradeSourceMap.set(req, { socket: nodeSocket, head }); + + const mod = await server.ssrLoadModule("fresh:server_entry"); + mod.setErrorInterceptor((err: unknown) => { + if (err instanceof Error) server.ssrFixStacktrace(err); + }); + + await mod.default.fetch(req); + // If the handler called ctx.upgrade(), the 101 has already been + // written to nodeSocket and the WebSocket is live. If it didn't, + // we have no usable channel left to report an error on — close + // the socket so the client doesn't hang. + if (!nodeSocket.destroyed && upgradeSourceMap.has(req)) { + nodeSocket.destroy(); + } + upgradeSourceMap.delete(req); + } catch (err) { + if (err instanceof Error) server.ssrFixStacktrace(err); + // deno-lint-ignore no-console + console.error("[fresh] WebSocket upgrade failed:", err); + if (!nodeSocket.destroyed) nodeSocket.destroy(); + } + })(); + }); + server.middlewares.use(async (nodeReq, nodeRes, next) => { const serverCfg = server.config.server; diff --git a/packages/plugin-vite/tests/dev_server_test.ts b/packages/plugin-vite/tests/dev_server_test.ts index ffc30f92ee2..abb381fe242 100644 --- a/packages/plugin-vite/tests/dev_server_test.ts +++ b/packages/plugin-vite/tests/dev_server_test.ts @@ -27,6 +27,30 @@ integrationTest("vite dev - launches", async () => { expect(text).toContain("it works"); }); +integrationTest("vite dev - upgrades WebSocket connections", async () => { + const address = demoServer.address(); + const wsUrl = address.replace(/^http/, "ws") + "/tests/ws_echo"; + + const ws = new WebSocket(wsUrl); + const reply = await new Promise((resolve, reject) => { + const timer = setTimeout( + () => reject(new Error("timed out waiting for echo")), + 5000, + ); + ws.onopen = () => ws.send("hello"); + ws.onmessage = (e) => { + clearTimeout(timer); + resolve(typeof e.data === "string" ? e.data : ""); + }; + ws.onerror = () => { + clearTimeout(timer); + reject(new Error("websocket errored")); + }; + }); + ws.close(); + expect(reply).toBe("echo: hello"); +}); + integrationTest("vite dev - serves static files", async () => { const res = await fetch(`${demoServer.address()}/test_static/foo.txt`); const text = await res.text();