fix: bound incoming transport message sizes

src/acp.ts

@@ -48,11 +48,17 @@ export type Stream = {
 export function ndJsonStream(
   output: WritableStream<Uint8Array>,
   input: ReadableStream<Uint8Array>,
+  options?: NdJsonStreamOptions,
 ): Stream {
-  return createJsonStream(output, input);
+  return createJsonStream(output, input, options);
 }
 
 export { RequestError } from "./jsonrpc.js";
+export {
+  DEFAULT_MAX_MESSAGE_BYTES,
+  MessageTooLargeError,
+} from "./stream-limits.js";
+export type { NdJsonStreamOptions } from "./stream.js";
 export type {
   AnyMessage,
   AnyNotification,
@@ -65,7 +71,7 @@ export type {
   SendRequestOptions,
 } from "./jsonrpc.js";
 
-import type { WireStream } from "./stream.js";
+import type { NdJsonStreamOptions, WireStream } from "./stream.js";
 import { Connection, Handled, HandlerRegistration } from "./jsonrpc.js";
 import type {
   AnyWireMessage,

src/http-stream.ts

@@ -11,6 +11,11 @@ import {
 } from "./protocol.js";
 import { MemoryAcpCookieStore } from "./cookie-store.js";
 import { parseSseStream } from "./sse.js";
+import {
+  MessageBuffer,
+  MessageTooLargeError,
+  resolveMaxMessageBytes,
+} from "./stream-limits.js";
 
 import type { AcpCookieStore } from "./cookie-store.js";
 import type { AnyMessage } from "./jsonrpc.js";
@@ -36,10 +41,15 @@ export interface HttpStreamOptions {
    * when the stream closes/errors.
    */
   readonly cookieStore?: AcpCookieStore;
+  readonly maxMessageBytes?: number;
 }
 
 export { MemoryAcpCookieStore } from "./cookie-store.js";
 export type { AcpCookieStore } from "./cookie-store.js";
+export {
+  DEFAULT_MAX_MESSAGE_BYTES,
+  MessageTooLargeError,
+} from "./stream-limits.js";
 
 /**
  * Creates an ACP Stream over Streamable HTTP.
@@ -66,6 +76,7 @@ class HttpStreamTransport {
 
   private readonly fetchImpl: typeof globalThis.fetch;
   private readonly headers: Record<string, string>;
+  private readonly maxMessageBytes: number;
   private readonly cookiePolicy: RequestCredentials;
   private readonly cookieStore: AcpCookieStore;
   private readonly ownsCookieStore: boolean;
@@ -85,6 +96,7 @@ class HttpStreamTransport {
     private readonly serverUrl: string,
     options: HttpStreamOptions,
   ) {
+    this.maxMessageBytes = resolveMaxMessageBytes(options.maxMessageBytes);
     this.fetchImpl = resolveFetch(options.fetch);
     this.headers = options.headers ?? {};
     this.cookiePolicy = options.cookies ?? "include";
@@ -148,7 +160,11 @@ class HttpStreamTransport {
       });
 
       if (!response.ok) {
-        throw await httpError("ACP initialize failed", response);
+        throw await httpError(
+          "ACP initialize failed",
+          response,
+          this.maxMessageBytes,
+        );
       }
 
       const connectionId = response.headers.get(HEADER_CONNECTION_ID);
@@ -159,7 +175,9 @@ class HttpStreamTransport {
       cleanupConnectionId = connectionId;
       this.throwIfClosedDuringInitialize();
 
-      const body: unknown = await response.json();
+      const body: unknown = JSON.parse(
+        await readResponseText(response, this.maxMessageBytes),
+      );
       this.throwIfClosedDuringInitialize();
 
       if (!isResponseMessage(body)) {
@@ -229,8 +247,13 @@ class HttpStreamTransport {
       });
 
       if (!response.ok) {
-        throw await httpError("ACP POST failed", response);
+        throw await httpError(
+          "ACP POST failed",
+          response,
+          this.maxMessageBytes,
+        );
       }
+      void response.body?.cancel().catch(() => {});
 
       if (!("method" in message) && "id" in message) {
         const key = messageIdKey(message.id);
@@ -353,7 +376,11 @@ class HttpStreamTransport {
       });
 
       if (!response.ok) {
-        throw await httpError("ACP SSE connection failed", response);
+        throw await httpError(
+          "ACP SSE connection failed",
+          response,
+          this.maxMessageBytes,
+        );
       }
 
       if (!response.body) {
@@ -362,7 +389,10 @@ class HttpStreamTransport {
 
       lifecycle.onOpen?.();
 
-      for await (const message of parseSseStream(response.body)) {
+      for await (const message of parseSseStream(
+        response.body,
+        this.maxMessageBytes,
+      )) {
         if (this.isClosed) {
           return;
         }
@@ -509,8 +539,13 @@ class HttpStreamTransport {
     });
 
     if (!response.ok) {
-      throw await httpError("ACP DELETE failed", response);
+      throw await httpError(
+        "ACP DELETE failed",
+        response,
+        this.maxMessageBytes,
+      );
     }
+    void response.body?.cancel().catch(() => {});
   }
 
   private clearOwnedCookieStore(): void {
@@ -575,8 +610,46 @@ function resolveFetch(
   );
 }
 
-async function httpError(prefix: string, response: Response): Promise<Error> {
-  const text = await response.text().catch(() => "");
+async function readResponseText(
+  response: Response,
+  maxMessageBytes: number,
+): Promise<string> {
+  if (!response.body) {
+    return "";
+  }
+
+  const buffer = new MessageBuffer(maxMessageBytes);
+  const reader = response.body.getReader();
+  try {
+    while (true) {
+      const { value, done } = await reader.read();
+      if (done) {
+        return new TextDecoder().decode(buffer.take());
+      }
+      buffer.append(value);
+    }
+  } catch (error) {
+    void reader.cancel(error).catch(() => {});
+    throw error;
+  } finally {
+    buffer.clear();
+    reader.releaseLock();
+  }
+}
+
+async function httpError(
+  prefix: string,
+  response: Response,
+  maxMessageBytes: number,
+): Promise<Error> {
+  let text = "";
+  try {
+    text = await readResponseText(response, maxMessageBytes);
+  } catch (error) {
+    if (error instanceof MessageTooLargeError) {
+      throw error;
+    }
+  }
 
   if (text) {
     return new Error(

src/line-buffer.ts

@@ -1,3 +1,9 @@
+import {
+  MessageBuffer,
+  MessageTooLargeError,
+  resolveMaxMessageBytes,
+} from "./stream-limits.js";
+
 const newline = 0x0a;
 
 /**
@@ -6,64 +12,76 @@ const newline = 0x0a;
  * Only the newly pushed chunk is scanned for newlines, so splitting costs
  * O(total bytes) no matter how many chunks a line spans.
  *
- * Chunks passed to {@link push} must not be mutated afterwards; a chunk
- * without a newline is retained until its line completes.
+ * Chunks passed to {@link push} must not be mutated while its iterator is
+ * in use.
  */
 export class LineBuffer {
   /** Bytes of the current (incomplete) line, carried across chunks. */
-  #pending: Uint8Array[] = [];
+  readonly #pending: MessageBuffer;
+  readonly #maxMessageBytes: number;
+
+  constructor(maxMessageBytes?: number) {
+    this.#maxMessageBytes = resolveMaxMessageBytes(maxMessageBytes);
+    this.#pending = new MessageBuffer(
+      Math.min(Number.MAX_SAFE_INTEGER, this.#maxMessageBytes + 1),
+    );
+  }
 
   /**
-   * Consumes a chunk, returning each complete line without its trailing
-   * newline.
+   * Consumes a chunk, yielding each complete line without its trailing
+   * LF or CRLF.
    */
-  push(chunk: Uint8Array): Uint8Array[] {
-    const lines: Uint8Array[] = [];
+  *push(chunk: Uint8Array): Generator<Uint8Array> {
     let start = 0;
     let newlineIndex = chunk.indexOf(newline, start);
     while (newlineIndex !== -1) {
-      lines.push(this.#takeLine(chunk.subarray(start, newlineIndex)));
+      yield this.#takeLine(chunk.subarray(start, newlineIndex));
       start = newlineIndex + 1;
       newlineIndex = chunk.indexOf(newline, start);
     }
     if (start < chunk.byteLength) {
-      // Copy a partial tail so a few carried-over bytes don't pin the whole
-      // chunk's buffer. The constructor guarantees a copy, unlike slice(),
-      // which Node's Buffer subclass overrides to return a view.
-      this.#pending.push(
-        start === 0 ? chunk : new Uint8Array(chunk.subarray(start)),
-      );
+      const tail = chunk.subarray(start);
+      this.#checkLine(tail);
+      this.#pending.append(tail);
     }
-    return lines;
   }
 
   /**
    * Returns the trailing unterminated line and resets the buffer, or
    * undefined if no bytes are buffered.
    */
   flush(): Uint8Array | undefined {
-    if (this.#pending.length === 0) {
+    if (this.#pending.byteLength === 0) {
       return undefined;
     }
-    return this.#takeLine(new Uint8Array(0));
+    return stripCarriageReturn(this.#pending.take());
+  }
+
+  clear(): void {
+    this.#pending.clear();
   }
 
   #takeLine(tail: Uint8Array): Uint8Array {
-    if (this.#pending.length === 0) {
-      return tail;
+    this.#checkLine(tail);
+    if (this.#pending.byteLength === 0) {
+      return stripCarriageReturn(tail);
     }
-    let total = tail.byteLength;
-    for (const part of this.#pending) {
-      total += part.byteLength;
-    }
-    const line = new Uint8Array(total);
-    let offset = 0;
-    for (const part of this.#pending) {
-      line.set(part, offset);
-      offset += part.byteLength;
+    this.#pending.append(tail);
+    return stripCarriageReturn(this.#pending.take());
+  }
+
+  #checkLine(tail: Uint8Array): void {
+    const lastByte =
+      tail.byteLength > 0 ? tail[tail.byteLength - 1] : this.#pending.lastByte;
+    const byteLength =
+      this.#pending.byteLength + tail.byteLength - (lastByte === 0x0d ? 1 : 0);
+    if (byteLength > this.#maxMessageBytes) {
+      this.clear();
+      throw new MessageTooLargeError(this.#maxMessageBytes);
     }
-    line.set(tail, offset);
-    this.#pending = [];
-    return line;
   }
 }
+
+function stripCarriageReturn(line: Uint8Array): Uint8Array {
+  return line[line.byteLength - 1] === 0x0d ? line.subarray(0, -1) : line;
+}