@heojeongbo/wirestage 0.0.0-stage → 0.1.1

package/src/entities/plugin/model/socket.ts

@@ -0,0 +1,113 @@
+import type { UrlMatcher } from "../../../shared/lib/match/index.js";
+
+/** One end of a routed WebSocket: the page's, or the real server's. */
+export interface SocketPeer {
+	/** Deliver a message to this end. */
+	send(data: string | Uint8Array): void;
+	/**
+	 * Close this end. The page's end takes any code and reason, as a server
+	 * may close with any code. The server's end is closed by the browser,
+	 * which sends only 1000 and 3000-4999, no reason without a code, and no
+	 * reason over 123 bytes of UTF-8: another code goes out as a normal
+	 * closure, 1000, and the reason is cut to 123 bytes, so the close still
+	 * happens.
+	 */
+	close(options?: {
+		readonly code?: number;
+		readonly reason?: string;
+	}): Promise<void>;
+	/**
+	 * Messages this end sends. One handler; a later call replaces it. Until
+	 * one is set, what this end sends is dropped, not passed on to the other.
+	 * A handler that throws, or rejects, fails the rule (see
+	 * `SocketRule.open`).
+	 */
+	onMessage(handler: (data: string | Uint8Array) => void): void;
+	/**
+	 * This end closed, by its own doing or through `close()` — and a close
+	 * the page starts is reported again once `close()` completes it, so a
+	 * handler must expect to run twice. Setting one turns off the automatic
+	 * forwarding of a close to the other end: the page's socket then hangs in
+	 * CLOSING until the handler closes the page's end itself. A handler that
+	 * throws, or rejects, fails the rule.
+	 */
+	onClose(
+		handler: (code: number | undefined, reason: string | undefined) => void,
+	): void;
+}
+
+/** A WebSocket a page opened, as the rule that takes it sees it. */
+export interface SocketView {
+	/** Links every frame of this connection. */
+	readonly correlation: string;
+	readonly url: URL;
+	readonly context: string;
+	/** The page's end. */
+	readonly page: SocketPeer;
+	/**
+	 * Open the real server's end and return it. No message is forwarded for
+	 * you: a rule that wants a proxy wires the two ends together itself, and
+	 * a direction it does not wire is dropped — a rule that sets only the
+	 * page's `onMessage` never lets the server's messages reach the page. A
+	 * close is passed on to the other end unless the end that closed has an
+	 * `onClose` handler.
+	 */
+	connectToServer(): SocketPeer;
+}
+
+export interface SocketRule {
+	/**
+	 * The `protocol` of this rule's frames: those it records, and the fault
+	 * the stage records when `open` throws.
+	 */
+	readonly protocol: string;
+	/** Names the rule in logs, diagnostics and the fault of a rule that throws. */
+	readonly label: string;
+	readonly url: UrlMatcher;
+	/** Stop taking connections after this many. */
+	readonly times?: number;
+	/**
+	 * Take a connection: called once per matching socket. The page's socket
+	 * opens when this settles — or, for a rule that connects the server,
+	 * when the server's end opens — and sooner if the rule sends to it. Set
+	 * the handlers before awaiting anything: an end drops what it sends
+	 * until it has one.
+	 *
+	 * A rule that throws, or rejects — here, or in a handler it set on
+	 * either end — is a bug in the test: the stage logs it, records a
+	 * `fault` frame (`meta.reason` `rule-error`, `meta.rule` the label,
+	 * `meta.error` the message) and closes the page's end with 1011 and a
+	 * reason naming the rule — and the server's end too, if the rule
+	 * connected it. Once per socket: what fails after that is not reported
+	 * again. The ends' `onClose` handlers hear those closes as they hear any
+	 * other.
+	 */
+	open(socket: SocketView): void | Promise<void>;
+}
+
+export interface SocketRuleHandle {
+	/** Connections this rule took. */
+	readonly hits: number;
+	readonly active: boolean;
+	/** Stop taking connections. Open ones stay as they are. */
+	remove(): Promise<void>;
+}
+
+/**
+ * Where plugins put WebSocket rules. Rules apply to every attached context,
+ * present and future, and the newest rule that matches a socket takes it; a
+ * socket no rule takes goes to the real server.
+ *
+ * Routing replaces `WebSocket` in every document of the context, so it
+ * starts with the first rule and a stage with none leaves the page alone.
+ * The replacement reaches a document as it loads: a rule added later applies
+ * to the documents loaded after it, unless the plugin called `prepare()`.
+ */
+export interface SocketRouter {
+	add(rule: SocketRule): Promise<SocketRuleHandle>;
+	/**
+	 * Start routing as soon as a context is attached, before any rule exists,
+	 * so a rule added mid-test reaches the pages already open. Call in `setup`.
+	 */
+	prepare(): void;
+}

package/src/entities/plugin/model/stream.ts

@@ -0,0 +1,46 @@
+/** How an open response ended. */
+export type StreamEnd =
+	/** The rule ended it cleanly. */
+	| "ended"
+	/** The rule reset the connection. */
+	| "cut"
+	/** The page went away: it closed the stream, navigated, or was closed. */
+	| "cancelled";
+
+/**
+ * A response that stays open, written at the test's pace: what a server
+ * stream, an event stream or a long poll is made of. A rule gets one by
+ * answering `{ type: "stream" }`.
+ */
+export interface ResponseStream {
+	/**
+	 * The request's headers as the server received them, names lower-cased.
+	 * They include what the browser's network stack adds after Playwright's
+	 * route has seen the request, such as `last-event-id`.
+	 */
+	readonly requestHeaders: Readonly<Record<string, string>>;
+	/**
+	 * Set the status and headers. They go out with the first `write`, or at
+	 * once on `flush`, and a later call replaces them until then. Without a
+	 * call the response is a bare 200. A header may carry several values
+	 * separated by a newline (`set-cookie`). `connection` is the stream
+	 * server's: it closes every connection after one response.
+	 */
+	head(status: number, headers?: Readonly<Record<string, string>>): void;
+	/** Send the head now, before any body: an EventSource opens on it. */
+	flush(): void;
+	/**
+	 * Write a chunk. Resolves once it is handed to the socket; rejects with
+	 * `CLOSED` when the response is over, the page letting go included.
+	 */
+	write(chunk: Uint8Array | string): Promise<void>;
+	/** End the response cleanly. Resolves once it is closed, as `closed` does. */
+	end(): Promise<void>;
+	/**
+	 * Reset the connection mid-response: no terminator, no trailer — what a
+	 * network drop looks like to the page.
+	 */
+	cut(): void;
+	/** Resolves when the response is over, with how, once that is on the timeline. */
+	readonly closed: Promise<StreamEnd>;
+}

package/src/features/connect/index.ts

@@ -0,0 +1,23 @@
+export { createRpcClient } from "./lib/client.js";
+export { connect } from "./lib/create-connect.js";
+export { statusOf } from "./lib/status-of.js";
+export type {
+	ClientOptions,
+	ConnectApi,
+	ConnectOptions,
+	HeldCall,
+	HoldOptions,
+	HoldStub,
+	RpcCall,
+	RpcProtocol,
+	RpcServer,
+	ServeOptions,
+	StandaloneClientOptions,
+	StreamCall,
+	StreamEnd,
+	StreamStub,
+	StubHandle,
+	StubOptions,
+	UnaryReply,
+	UnaryStub,
+} from "./model/types.js";

package/src/features/connect/lib/calls.ts

@@ -0,0 +1,58 @@
+import type { DescMessage, MessageShape } from "@bufbuild/protobuf";
+
+import type { Frame } from "../../../entities/frame/index.js";
+import type { RpcCall, RpcProtocol } from "../model/types.js";
+
+/**
+ * One method's frames as calls: one per correlation, oldest first. A frame
+ * another writer left without a correlation is a call of its own.
+ */
+export function assembleCalls<I extends DescMessage, O extends DescMessage>(
+	frames: readonly Frame[],
+): RpcCall<I, O>[] {
+	const byCall = new Map<string, Frame[]>();
+	for (const frame of frames) {
+		const key = frame.correlation ?? `#${frame.seq}`;
+		const list = byCall.get(key) ?? [];
+		list.push(frame);
+		byCall.set(key, list);
+	}
+	return [...byCall].map(([correlation, own]) => {
+		const first = own[0] as Frame;
+		const requests = own.filter(
+			(frame) => frame.direction === "out" && frame.value !== undefined,
+		);
+		// Everything that is not the caller's is the answer: messages, the
+		// status, or the fault that stood in for one.
+		const answer = own.findLast((frame) => frame.direction !== "out");
+		const status = own.findLast((frame) => frame.kind === "response");
+		const fault = own.findLast((frame) => frame.kind === "fault");
+		return {
+			correlation,
+			caller: requests[0]?.source,
+			answeredBy: answer?.source,
+			page: first.page,
+			context: first.context,
+			protocol: first.meta["rpcProtocol"] as RpcProtocol | undefined,
+			request: requests[0]?.value as MessageShape<I> | undefined,
+			requests: requests.map((frame) => frame.value as MessageShape<I>),
+			responses: own
+				.filter(
+					(frame) =>
+						frame.direction === "in" &&
+						frame.kind === "message" &&
+						frame.value !== undefined,
+				)
+				.map((frame) => frame.value as MessageShape<O>),
+			status:
+				status === undefined
+					? undefined
+					: {
+							code: Number(status.meta["code"]),
+							message: String(status.meta["message"] ?? ""),
+						},
+			fault: fault === undefined ? undefined : String(fault.meta["reason"]),
+			frames: own,
+		};
+	});
+}

package/src/features/connect/lib/client.ts

@@ -0,0 +1,203 @@
+import type { DescService } from "@bufbuild/protobuf";
+import {
+	type Client,
+	ConnectError,
+	createClient,
+	type Interceptor,
+	type Transport,
+} from "@connectrpc/connect";
+import {
+	createConnectTransport,
+	createGrpcTransport,
+	createGrpcWebTransport,
+	Http2SessionManager,
+} from "@connectrpc/connect-node";
+
+import type { CookieJarPort } from "../../../shared/api/playwright/index.js";
+import { parseSetCookie } from "../../../shared/lib/cookies/index.js";
+import type { RpcProtocol, StandaloneClientOptions } from "../model/types.js";
+import { type CallInfo, type Recorder, tapIterable } from "./recorder.js";
+
+export interface ClientDeps {
+	/** Resolved per call, so a context attached after the client exists still counts. */
+	readonly session: () => CookieJarPort | undefined;
+	/** The time cookies are stamped against; the stage's clock in a stage. */
+	readonly now?: () => number;
+	readonly recorder?: Recorder;
+	readonly correlation?: () => string;
+	/**
+	 * Handed the HTTP/2 connection a gRPC client opens, so its owner can close
+	 * it. connect-node keeps an idle one for 15 minutes, and a server in the
+	 * same process cannot finish closing while it is open (Node 22 waits).
+	 */
+	readonly onConnection?: (connection: Http2SessionManager) => void;
+}
+
+/** Sent with every call that did not set the header itself. */
+function headersInterceptor(
+	headers: Readonly<Record<string, string>>,
+): Interceptor {
+	return (next) => (request) => {
+		for (const [name, value] of Object.entries(headers))
+			if (!request.header.has(name)) request.header.set(name, value);
+		return next(request);
+	};
+}
+
+/**
+ * Shares a context's cookie jar both ways, as Playwright's own request
+ * context does — the reason hand-rolled gRPC-web helpers go through it. The
+ * jar's cookies go out on the call, read per call: a sign-in during the test
+ * is picked up, and a cookie scoped to another host or path is not sent. The
+ * cookies a response sets come back into the jar — an error response's too,
+ * which is how a sign-out or a refused sign-in clears a session.
+ *
+ * Cookies are scoped to the host the request names: its `host` header when
+ * one was set, since that is the site the server takes it for.
+ */
+function sessionInterceptor(
+	session: () => CookieJarPort | undefined,
+	now: () => number,
+): Interceptor {
+	return (next) => async (request) => {
+		const jar = session();
+		if (jar === undefined) return next(request);
+		const named = request.header.get("host");
+		const site = new URL(request.url);
+		if (named !== null) site.host = named;
+		const cookies = await jar.cookies(site.href);
+		if (cookies.length > 0 && !request.header.has("cookie"))
+			request.header.set(
+				"cookie",
+				cookies.map((c) => `${c.name}=${c.value}`).join("; "),
+			);
+		const keep = async (headers: Headers) => {
+			const set = headers
+				.getSetCookie()
+				.flatMap((line) => parseSetCookie(line, site.href, now()) ?? []);
+			if (set.length > 0) await jar.addCookies?.(set);
+		};
+		try {
+			const response = await next(request);
+			await keep(response.header);
+			return response;
+		} catch (error) {
+			// A failed call's headers ride on its error; one that never reached
+			// the server has none.
+			await keep(ConnectError.from(error).metadata);
+			throw error;
+		}
+	};
+}
+
+/** Records drive calls on the timeline: the test's requests, the backend's answers. */
+function recordingInterceptor(
+	recorder: Recorder,
+	correlation: () => string,
+	protocol: RpcProtocol,
+): Interceptor {
+	return (next) => async (request) => {
+		const call: CallInfo = {
+			correlation: correlation(),
+			from: "test",
+			answeredBy: "network",
+			rpcProtocol: protocol,
+		};
+		const method = request.method;
+		try {
+			if (!request.stream) recorder.out(call, method, request.message, true);
+			const response = await next(
+				request.stream
+					? {
+							...request,
+							message: tapIterable(request.message, (m, i) =>
+								recorder.out(call, method, m, i === 0),
+							),
+						}
+					: request,
+			);
+			if (!response.stream) {
+				recorder.in(call, method, response.message);
+				recorder.status(call, method);
+				return response;
+			}
+			return {
+				...response,
+				message: (async function* () {
+					try {
+						for await (const message of response.message) {
+							recorder.in(call, method, message);
+							yield message;
+						}
+						recorder.status(call, method);
+					} catch (error) {
+						recorder.status(call, method, ConnectError.from(error));
+						throw error;
+					}
+				})(),
+			};
+		} catch (error) {
+			recorder.status(call, method, ConnectError.from(error));
+			throw error;
+		}
+	};
+}
+
+export function createTransport(
+	options: StandaloneClientOptions,
+	deps: ClientDeps,
+): Transport {
+	const protocol = options.protocol ?? "connect";
+	// Outermost first: the defaults, then the session (which reads the `host`
+	// the defaults may have named), then the recorder, closest to the wire.
+	const interceptors: Interceptor[] = [];
+	if (options.headers !== undefined)
+		interceptors.push(headersInterceptor(options.headers));
+	interceptors.push(sessionInterceptor(deps.session, deps.now ?? Date.now));
+	if (deps.recorder !== undefined && deps.correlation !== undefined)
+		interceptors.push(
+			recordingInterceptor(deps.recorder, deps.correlation, protocol),
+		);
+	const nodeOptions = options.ignoreHttpsErrors
+		? { rejectUnauthorized: false }
+		: undefined;
+	const common = {
+		baseUrl: options.baseURL,
+		interceptors,
+		useBinaryFormat: options.binary ?? true,
+	};
+	if (protocol === "grpc") {
+		// The transport's own session manager would be out of reach; this one
+		// is handed to whoever has to close it.
+		const sessionManager = new Http2SessionManager(
+			options.baseURL,
+			undefined,
+			nodeOptions,
+		);
+		deps.onConnection?.(sessionManager);
+		return createGrpcTransport({ ...common, sessionManager });
+	}
+	const http1 = { ...common, httpVersion: "1.1" as const, nodeOptions };
+	if (protocol === "grpc-web") return createGrpcWebTransport(http1);
+	return createConnectTransport(http1);
+}
+
+/**
+ * A typed client for a real backend, outside any stage: the `rpc` helper an
+ * e2e suite otherwise writes by hand — framing, status parsing and all — per
+ * protocol. Errors arrive as `ConnectError` with the server's code.
+ *
+ * Over gRPC it holds an HTTP/2 connection, which connect-node keeps for 15
+ * minutes once idle. It does not keep the process alive, but a server in the
+ * same process waits for it before it finishes closing. A stage's client is
+ * closed with the stage.
+ */
+export function createRpcClient<S extends DescService>(
+	service: S,
+	options: StandaloneClientOptions,
+): Client<S> {
+	return createClient(
+		service,
+		createTransport(options, { session: () => options.session }),
+	);
+}

package/src/features/connect/lib/create-connect.ts

@@ -0,0 +1,164 @@
+import type { DescService } from "@bufbuild/protobuf";
+import { type Client, createClient } from "@connectrpc/connect";
+import type { Http2SessionManager } from "@connectrpc/connect-node";
+
+import {
+	definePlugin,
+	type PluginContext,
+	type WirePlugin,
+} from "../../../entities/plugin/index.js";
+import { WirestageError } from "../../../shared/lib/errors/index.js";
+import { createSequence } from "../../../shared/lib/id/index.js";
+import type {
+	ClientOptions,
+	ConnectApi,
+	ConnectOptions,
+	RpcServer,
+} from "../model/types.js";
+import { assembleCalls } from "./calls.js";
+import { createTransport } from "./client.js";
+import { stubAbort, stubHold } from "./faults.js";
+import { recordService } from "./record.js";
+import { createRecorder, endpointOf } from "./recorder.js";
+import { createRpcRules } from "./rules.js";
+import { serveRoutes } from "./serve.js";
+import { stubFailure, stubService, stubStream, stubUnary } from "./stubs.js";
+
+/**
+ * Connect, gRPC-web and gRPC: stubs answered by connect-es's own protocol
+ * handlers, recording of real traffic, drive clients, and out-of-process
+ * servers.
+ *
+ * The stubs never encode a frame by hand. Each one is a one-method Connect
+ * router whose handler runs the real server-side protocol, so the bytes the
+ * page receives are exactly what a Connect server would send — errors,
+ * trailers, JSON and binary alike.
+ */
+export function connect(
+	options: ConnectOptions = {},
+): WirePlugin<"connect", ConnectApi> {
+	return definePlugin({
+		name: "connect",
+		scope: "Connect",
+		setup: (context) => createConnect(context, options),
+	});
+}
+
+function createConnect(context: PluginContext, options: ConnectOptions) {
+	const protocol = context.name;
+	const live = (options.streaming ?? "live") === "live";
+	const recorder = createRecorder(context.timeline, protocol);
+	const rules = createRpcRules(context, options, recorder, protocol);
+	const testCorrelation = createSequence("t");
+	const serverCorrelation = createSequence("s");
+	const servers = new Set<RpcServer>();
+	/** Lets go of each call a hold still has, when the stage goes away. */
+	const heldCalls = new Set<() => void>();
+	/** The drive clients' HTTP/2 connections, closed with the stage. */
+	const connections = new Set<Http2SessionManager>();
+	/** The record rules being added, once: every context gets them. */
+	let recording: Promise<unknown> | undefined;
+
+	const api: ConnectApi = {
+		service: (service, impl, stubOptions = {}) =>
+			stubService(rules, service, impl, stubOptions),
+		unary: (method, reply, stubOptions = {}) =>
+			stubUnary(context, rules, method, reply, stubOptions),
+		stream: (method, stubOptions = {}) =>
+			stubStream(context, rules, recorder, method, stubOptions, live),
+		fail: (method, code, message, stubOptions = {}) =>
+			stubFailure(rules, method, code, message, stubOptions),
+		abort: (method, abortOptions = {}) =>
+			stubAbort(rules, recorder, method, abortOptions),
+		hold: (method, holdOptions = {}) =>
+			stubHold(context, rules, recorder, method, holdOptions, heldCalls),
+		record: (service) => recordService(rules, service),
+		async calls(method) {
+			await context.sync();
+			return assembleCalls(
+				context.timeline.frames({ protocol, endpoint: endpointOf(method) }),
+			);
+		},
+		client<S extends DescService>(
+			service: S,
+			clientOptions: ClientOptions = {},
+		): Client<S> {
+			const {
+				session,
+				baseURL: requested,
+				...transportOptions
+			} = clientOptions;
+			const base = requested ?? options.baseURL ?? "/";
+			const origin = context.attachments()[0]?.baseURL ?? context.baseURL;
+			if (!/^https?:\/\//.test(base) && origin === undefined)
+				throw new WirestageError(
+					"INVALID_ARGUMENT",
+					`client(${service.typeName}): "${base}" is relative and the stage has no baseURL; pass an absolute baseURL or give the stage one`,
+				);
+			return createClient(
+				service,
+				createTransport(
+					{
+						...transportOptions,
+						baseURL: new URL(base, origin).href.replace(/\/+$/, ""),
+						protocol:
+							transportOptions.protocol ?? options.clientProtocol ?? "connect",
+					},
+					{
+						// Resolved per call: a context attached after the client was made
+						// still lends it its session.
+						session: () =>
+							session === null
+								? undefined
+								: (session ?? context.attachments()[0]?.cookies),
+						now: () => context.clock.now(),
+						recorder,
+						correlation: testCorrelation,
+						onConnection: (connection) => connections.add(connection),
+					},
+				),
+			);
+		},
+		async serve(routes, serveOptions = {}) {
+			const server = await serveRoutes(routes, serveOptions, {
+				recorder,
+				correlation: serverCorrelation,
+			});
+			const tracked: RpcServer = {
+				url: server.url,
+				port: server.port,
+				close: async () => {
+					servers.delete(tracked);
+					await server.close();
+				},
+			};
+			servers.add(tracked);
+			return tracked;
+		},
+	};
+
+	return {
+		api,
+		async attach() {
+			const services = options.record;
+			if (services === undefined) return;
+			recording ??= (async () => {
+				for (const service of services) await api.record(service);
+			})().catch((error: unknown) => {
+				// Only the first rule can fail, installing the context's route, and
+				// the router takes it back: the next attach starts over.
+				recording = undefined;
+				throw error;
+			});
+			await recording;
+		},
+		async dispose() {
+			for (const letGo of [...heldCalls]) letGo();
+			// A drive client's connection outlives its calls; left open, a
+			// server in the test's own process would wait for it on close.
+			for (const connection of connections) connection.abort();
+			connections.clear();
+			await Promise.all([...servers].map((server) => server.close()));
+		},
+	};
+}