@yourrealm/sdk 0.0.0-alpha.1378.ce03c8dc

package/src/worker-protocol/uplink/rpc.ts

@@ -0,0 +1,309 @@
+// Request/response over one mux stream. The OPEN header is UTF-8 JSON
+// `{method, headers}`; the request body is that stream's DATA until END.
+// The response is `[u16 headLen][head JSON][body…]` then END — the length
+// prefix is what makes the head robust to fragment boundaries.
+//
+// The method names themselves are not defined here: the step that rehosts a
+// call owns its schema.
+
+import { MAX_PAYLOAD_BYTES, MAX_RPC_HEAD_BYTES } from "./constants.ts";
+import { UplinkProtocolError } from "./errors.ts";
+import { concatBytes, decodeUtf8, encodeUtf8 } from "./frame.ts";
+import type { Mux, Stream } from "./mux.ts";
+
+export interface RpcRequest {
+  method: string;
+  headers: Record<string, string>;
+  body: ReadableStream<Uint8Array>;
+}
+
+export interface RpcResponse {
+  status: number;
+  headers: Record<string, string>;
+  body: ReadableStream<Uint8Array>;
+}
+
+export type RpcHandler = (req: RpcRequest) => RpcResponse | Promise<RpcResponse>;
+
+export class Rpc {
+  readonly #mux: Mux;
+  readonly #handlers: Record<string, RpcHandler>;
+
+  constructor(mux: Mux, handlers: Record<string, RpcHandler>) {
+    this.#mux = mux;
+    this.#handlers = handlers;
+    mux.onOpen((stream, header) => {
+      void this.#serve(stream, header);
+    });
+  }
+
+  /** Resolves as soon as the response head is parsed; the body streams after. */
+  async call(
+    method: string,
+    headers: Record<string, string>,
+    body?: Uint8Array | ReadableStream<Uint8Array>,
+    opts?: { signal?: AbortSignal },
+  ): Promise<RpcResponse> {
+    const header = encodeUtf8(JSON.stringify({ method, headers }));
+    if (header.length > MAX_PAYLOAD_BYTES) {
+      throw new RangeError(`OPEN header ${header.length} exceeds ${MAX_PAYLOAD_BYTES} bytes`);
+    }
+    const signal = opts?.signal;
+    signal?.throwIfAborted();
+    const stream = await this.#mux.open(header);
+    let aborted = false;
+    let abortReason: unknown;
+    const onAbort = () => {
+      aborted = true;
+      abortReason = signal?.reason;
+      stream.reset(reasonText(abortReason));
+    };
+    signal?.addEventListener("abort", onAbort, { once: true });
+    try {
+      // Not awaited: a handler may answer before it has drained the request
+      // body, and the two directions are independent. A write failure shows
+      // up on the read side.
+      void sendBody(stream, body).catch(() => {});
+      const { head, body: rest } = await readHead(stream.readable);
+      return { status: head.status, headers: head.headers, body: rest };
+    } catch (e) {
+      if (aborted) throw abortReason;
+      throw e;
+    } finally {
+      signal?.removeEventListener("abort", onAbort);
+    }
+  }
+
+  async #serve(stream: Stream, header: Uint8Array): Promise<void> {
+    let open: { method: string; headers: Record<string, string> };
+    try {
+      open = parseOpenHeader(header);
+    } catch (e) {
+      stream.reset(asError(e).message);
+      return;
+    }
+    const handler = this.#handlers[open.method];
+    let res: RpcResponse;
+    if (handler === undefined) res = jsonResponse(404, { error: "unknown method" });
+    else {
+      try {
+        res = await handler({ method: open.method, headers: open.headers, body: stream.readable });
+      } catch (e) {
+        res = textResponse(500, asError(e).message);
+      }
+    }
+    try {
+      await writeResponse(stream, res);
+      // A handler that ignored the body (404, or a throw) would otherwise
+      // leave the caller stalled on a window that never reopens.
+      if (!stream.readable.locked) await drain(stream.readable);
+    } catch {
+      stream.reset("response failed");
+    }
+  }
+}
+
+/** JSON body helper so a method handler stays one line. */
+export function jsonResponse(
+  status: number,
+  value: unknown,
+  headers?: Record<string, string>,
+): RpcResponse {
+  return {
+    status,
+    headers: { "content-type": "application/json", ...headers },
+    body: bytesBody(encodeUtf8(JSON.stringify(value))),
+  };
+}
+
+/** Drains a request or response body and parses it as JSON. */
+export async function readJson<T>(body: ReadableStream<Uint8Array>): Promise<T> {
+  return JSON.parse(decodeUtf8(await readAll(body))) as T;
+}
+
+function textResponse(status: number, text: string): RpcResponse {
+  return {
+    status,
+    headers: { "content-type": "text/plain; charset=utf-8" },
+    body: bytesBody(encodeUtf8(text)),
+  };
+}
+
+function bytesBody(bytes: Uint8Array): ReadableStream<Uint8Array> {
+  return new ReadableStream<Uint8Array>(
+    {
+      pull: (controller) => {
+        if (bytes.length > 0) controller.enqueue(bytes);
+        controller.close();
+      },
+    },
+    { highWaterMark: 0 },
+  );
+}
+
+async function sendBody(
+  stream: Stream,
+  body: Uint8Array | ReadableStream<Uint8Array> | undefined,
+): Promise<void> {
+  if (body === undefined) {
+    await stream.writable.close();
+    return;
+  }
+  if (body instanceof Uint8Array) {
+    const writer = stream.writable.getWriter();
+    await writer.write(body);
+    await writer.close();
+    return;
+  }
+  await body.pipeTo(stream.writable);
+}
+
+async function writeResponse(stream: Stream, res: RpcResponse): Promise<void> {
+  const head = encodeUtf8(JSON.stringify({ status: res.status, headers: res.headers }));
+  if (head.length > MAX_RPC_HEAD_BYTES) {
+    throw new RangeError(`response head ${head.length} exceeds ${MAX_RPC_HEAD_BYTES} bytes`);
+  }
+  const prefix = new Uint8Array(2);
+  new DataView(prefix.buffer).setUint16(0, head.length);
+  const writer = stream.writable.getWriter();
+  await writer.write(concatBytes(prefix, head));
+  writer.releaseLock();
+  await res.body.pipeTo(stream.writable);
+}
+
+interface ResponseHead {
+  status: number;
+  headers: Record<string, string>;
+}
+
+async function readHead(
+  readable: ReadableStream<Uint8Array>,
+): Promise<{ head: ResponseHead; body: ReadableStream<Uint8Array> }> {
+  const reader = readable.getReader();
+  let buf: Uint8Array = new Uint8Array(0);
+  for (;;) {
+    if (buf.length >= 2) {
+      const headLen = new DataView(buf.buffer, buf.byteOffset, buf.byteLength).getUint16(0);
+      if (buf.length >= 2 + headLen) {
+        const head = parseResponseHead(buf.subarray(2, 2 + headLen));
+        return { head, body: continuation(reader, buf.slice(2 + headLen)) };
+      }
+    }
+    const { value, done } = await reader.read();
+    if (done) throw new UplinkProtocolError("response ended before its head");
+    buf = concatBytes(buf, value);
+  }
+}
+
+// The rest of the stream, after the bytes already pulled with the head.
+function continuation(
+  reader: ReadableStreamDefaultReader<Uint8Array>,
+  rest: Uint8Array,
+): ReadableStream<Uint8Array> {
+  let leftover: Uint8Array | undefined = rest.length > 0 ? rest : undefined;
+  return new ReadableStream<Uint8Array>(
+    {
+      pull: async (controller) => {
+        if (leftover !== undefined) {
+          const chunk = leftover;
+          leftover = undefined;
+          controller.enqueue(chunk);
+          return;
+        }
+        const { value, done } = await reader.read();
+        if (done) controller.close();
+        else controller.enqueue(value);
+      },
+      cancel: async (reason) => {
+        await reader.cancel(reason);
+      },
+    },
+    { highWaterMark: 0 },
+  );
+}
+
+function parseOpenHeader(header: Uint8Array): { method: string; headers: Record<string, string> } {
+  const value = parseJson(header, "OPEN header");
+  const method = (value as { method?: unknown }).method;
+  if (typeof method !== "string" || method === "") {
+    throw new UplinkProtocolError("OPEN header has no method");
+  }
+  return { method, headers: readHeaders(value, "OPEN header") };
+}
+
+function parseResponseHead(head: Uint8Array): ResponseHead {
+  const value = parseJson(head, "response head");
+  const status = (value as { status?: unknown }).status;
+  if (typeof status !== "number" || !Number.isInteger(status)) {
+    throw new UplinkProtocolError("response head has no status");
+  }
+  return { status, headers: readHeaders(value, "response head") };
+}
+
+function parseJson(bytes: Uint8Array, what: string): Record<string, unknown> {
+  let value: unknown;
+  try {
+    value = JSON.parse(decodeUtf8(bytes));
+  } catch {
+    throw new UplinkProtocolError(`${what} is not JSON`);
+  }
+  if (typeof value !== "object" || value === null || Array.isArray(value)) {
+    throw new UplinkProtocolError(`${what} is not an object`);
+  }
+  return value as Record<string, unknown>;
+}
+
+function readHeaders(value: Record<string, unknown>, what: string): Record<string, string> {
+  const headers = value.headers;
+  if (headers === undefined) return {};
+  if (typeof headers !== "object" || headers === null || Array.isArray(headers)) {
+    throw new UplinkProtocolError(`${what} has non-object headers`);
+  }
+  const out: Record<string, string> = {};
+  for (const [key, entry] of Object.entries(headers)) {
+    if (typeof entry !== "string") {
+      throw new UplinkProtocolError(`${what} header ${key} is not a string`);
+    }
+    out[key] = entry;
+  }
+  return out;
+}
+
+async function readAll(body: ReadableStream<Uint8Array>): Promise<Uint8Array> {
+  const reader = body.getReader();
+  const chunks: Uint8Array[] = [];
+  let size = 0;
+  for (;;) {
+    const { value, done } = await reader.read();
+    if (done) break;
+    chunks.push(value);
+    size += value.length;
+  }
+  const out = new Uint8Array(size);
+  let offset = 0;
+  for (const chunk of chunks) {
+    out.set(chunk, offset);
+    offset += chunk.length;
+  }
+  return out;
+}
+
+async function drain(body: ReadableStream<Uint8Array>): Promise<void> {
+  const reader = body.getReader();
+  try {
+    while (!(await reader.read()).done) {
+      /* discard */
+    }
+  } finally {
+    reader.releaseLock();
+  }
+}
+
+function asError(value: unknown): Error {
+  return value instanceof Error ? value : new Error(String(value));
+}
+
+function reasonText(reason: unknown): string {
+  if (reason === undefined || reason === null) return "aborted";
+  return reason instanceof Error ? reason.message : String(reason);
+}