Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 32 additions & 1 deletion packages/fresh/src/context.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<Request, UpgradeSource> =
globalAny[UPGRADE_SOURCE_KEY] ??
(globalAny[UPGRADE_SOURCE_KEY] = new WeakMap<Request, UpgradeSource>());

/**
* Duck-type check: treats the argument as managed-mode handlers when at least
* one of the handler keys (`open`, `message`, `close`, `error`) is a
Expand Down Expand Up @@ -595,7 +617,16 @@ export class Context<State> {
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 };
Expand Down
1 change: 1 addition & 0 deletions packages/fresh/src/internals.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
11 changes: 11 additions & 0 deletions packages/plugin-vite/demo/routes/tests/ws_echo.ts
Original file line number Diff line number Diff line change
@@ -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);
},
});
},
});
71 changes: 70 additions & 1 deletion packages/plugin-vite/src/plugins/dev_server.ts
Original file line number Diff line number Diff line change
@@ -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";
Expand All @@ -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;

Expand Down
24 changes: 24 additions & 0 deletions packages/plugin-vite/tests/dev_server_test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<string>((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();
Expand Down
Loading