mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-09-19 05:12:16 +03:00
fix(mitm): bound SSE transcript retention and cancel abandoned upstream reads (#13702)
Handler-side collected strings grew without bound before the inspector clamp; abandoned streams kept the reader alive for the full upstream lifetime. createBoundedCollector caps retention at 1 MiB while keeping true responseSize; pipeSSE and server.cjs cancel on downstream close. Fixes #13395. Co-authored-by: oyi77 <oyi77@users.noreply.github.com>
This commit is contained in:
1
changelog.d/fixes/13395-mitm-pipe-bounded-collect.md
Normal file
1
changelog.d/fixes/13395-mitm-pipe-bounded-collect.md
Normal file
@@ -0,0 +1 @@
|
||||
- **fix(mitm):** bound per-request SSE transcript retention to 1 MiB and stop the upstream read when the downstream disconnects — handler-side `collected` strings grew without bound before the inspector clamp, and abandoned streams kept the reader alive for the full upstream lifetime ([#13395](https://github.com/diegosouzapw/OmniRoute/issues/13395))
|
||||
@@ -21,7 +21,7 @@
|
||||
*/
|
||||
import type { IncomingMessage, ServerResponse } from "node:http";
|
||||
import type { AgentId } from "../types";
|
||||
import { MitmHandlerBase } from "./base";
|
||||
import { MitmHandlerBase, createBoundedCollector } from "./base";
|
||||
import { TOOL_RENAME_MAP } from "@omniroute/open-sse/services/claudeCodeToolRemapper";
|
||||
|
||||
interface GeminiPart {
|
||||
@@ -171,7 +171,7 @@ export class AntigravityHandler extends MitmHandlerBase {
|
||||
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
|
||||
}
|
||||
|
||||
let collected = "";
|
||||
const sink = createBoundedCollector();
|
||||
await this.pipeSSE(upstream, res, (chunk) => {
|
||||
let chunkStr = chunk.toString();
|
||||
for (const [lower, capitalized] of Object.entries(TOOL_RENAME_MAP)) {
|
||||
@@ -180,15 +180,15 @@ export class AntigravityHandler extends MitmHandlerBase {
|
||||
`"name":"${capitalized}"`
|
||||
);
|
||||
}
|
||||
collected += chunkStr;
|
||||
sink.push(chunkStr);
|
||||
});
|
||||
|
||||
const total = this.now() - startedAt;
|
||||
this.hookBufferUpdate(intercepted, {
|
||||
status: upstream.status,
|
||||
responseHeaders: Object.fromEntries(upstream.headers.entries()),
|
||||
responseBody: collected,
|
||||
responseSize: Buffer.byteLength(collected),
|
||||
responseBody: sink.text,
|
||||
responseSize: sink.totalBytes,
|
||||
proxyLatencyMs: upstreamStart - startedAt,
|
||||
upstreamLatencyMs: total - (upstreamStart - startedAt),
|
||||
});
|
||||
|
||||
@@ -76,6 +76,51 @@ async function loadAgentBridgeHook(): Promise<{
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Ceiling for the per-request SSE text a MITM handler retains for the Traffic
|
||||
* Inspector (#13395). The inspector buffer re-clamps per body
|
||||
* (`INSPECTOR_MAX_BODY_KB`, default 1 MiB), but the handler-side `collected`
|
||||
* string grew without bound BEFORE reaching that clamp — one long-lived stream
|
||||
* kept the whole transcript in the handler closure for the request lifetime.
|
||||
* Aligned with the inspector default so the bound never hides data the UI shows.
|
||||
*/
|
||||
export const MITM_PIPE_MAX_COLLECT_BYTES = 1 * 1024 * 1024;
|
||||
|
||||
/**
|
||||
* Bounded string accumulator for piped SSE transcripts. Stops retaining past
|
||||
* `maxBytes` but keeps counting true bytes, so the inspector still reports
|
||||
* the exact `responseSize` it reported before (`Buffer.byteLength` of the
|
||||
* full transcript) and the pipe itself is unaffected — every chunk is still
|
||||
* written downstream regardless of the cap.
|
||||
*/
|
||||
export function createBoundedCollector(maxBytes: number = MITM_PIPE_MAX_COLLECT_BYTES): {
|
||||
push: (chunk: string) => void;
|
||||
text: string;
|
||||
totalBytes: number;
|
||||
truncated: boolean;
|
||||
} {
|
||||
let collected = "";
|
||||
let totalBytes = 0;
|
||||
const acc = {
|
||||
push(chunk: string): void {
|
||||
totalBytes += Buffer.byteLength(chunk);
|
||||
if (collected.length < maxBytes) {
|
||||
collected += chunk.slice(0, maxBytes - collected.length);
|
||||
}
|
||||
},
|
||||
get text(): string {
|
||||
return collected;
|
||||
},
|
||||
get totalBytes(): number {
|
||||
return totalBytes;
|
||||
},
|
||||
get truncated(): boolean {
|
||||
return totalBytes > Buffer.byteLength(collected);
|
||||
},
|
||||
};
|
||||
return acc;
|
||||
}
|
||||
|
||||
export abstract class MitmHandlerBase {
|
||||
abstract readonly agentId: AgentId;
|
||||
|
||||
@@ -151,13 +196,6 @@ export abstract class MitmHandlerBase {
|
||||
});
|
||||
}
|
||||
|
||||
/**
|
||||
* Pipe an SSE (or any chunked) upstream Response straight to the downstream
|
||||
* ServerResponse, optionally invoking `onChunk` for each received Buffer.
|
||||
*
|
||||
* Writes SSE-friendly headers before the first chunk (only if `res.headersSent`
|
||||
* is still false — handlers MAY have set custom headers first).
|
||||
*/
|
||||
protected async pipeSSE(
|
||||
upstream: Response,
|
||||
res: ServerResponse,
|
||||
@@ -179,8 +217,18 @@ export abstract class MitmHandlerBase {
|
||||
}
|
||||
|
||||
const reader = upstream.body.getReader();
|
||||
// #13395: a downstream disconnect must stop the upstream read — otherwise an
|
||||
// abandoned stream keeps the reader (and its buffers) alive for the full
|
||||
// upstream lifetime and the handler closure retains the transcript.
|
||||
let downstreamClosed = false;
|
||||
const onClose = () => {
|
||||
downstreamClosed = true;
|
||||
reader.cancel().catch(() => {});
|
||||
};
|
||||
res.once("close", onClose);
|
||||
try {
|
||||
while (true) {
|
||||
if (downstreamClosed) break;
|
||||
const { done, value } = await reader.read();
|
||||
if (done) break;
|
||||
const buf = Buffer.from(value);
|
||||
@@ -191,9 +239,16 @@ export abstract class MitmHandlerBase {
|
||||
// Inspector hook must never break the upstream pipe.
|
||||
}
|
||||
}
|
||||
if (downstreamClosed || res.closed || res.destroyed) break;
|
||||
res.write(buf);
|
||||
}
|
||||
} finally {
|
||||
res.off("close", onClose);
|
||||
try {
|
||||
reader.releaseLock();
|
||||
} catch {
|
||||
// Reader already cancelled or released.
|
||||
}
|
||||
try {
|
||||
res.end();
|
||||
} catch {
|
||||
|
||||
@@ -8,7 +8,7 @@
|
||||
*/
|
||||
import type { IncomingMessage, ServerResponse } from "node:http";
|
||||
import type { AgentId } from "../types";
|
||||
import { MitmHandlerBase } from "./base";
|
||||
import { MitmHandlerBase, createBoundedCollector } from "./base";
|
||||
|
||||
export class ClaudeCodeHandler extends MitmHandlerBase {
|
||||
readonly agentId: AgentId = "claude-code";
|
||||
@@ -51,17 +51,17 @@ export class ClaudeCodeHandler extends MitmHandlerBase {
|
||||
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
|
||||
}
|
||||
|
||||
let collected = "";
|
||||
const sink = createBoundedCollector();
|
||||
await this.pipeSSE(upstream, res, (chunk) => {
|
||||
collected += chunk.toString();
|
||||
sink.push(chunk.toString());
|
||||
});
|
||||
|
||||
const total = this.now() - startedAt;
|
||||
this.hookBufferUpdate(intercepted, {
|
||||
status: upstream.status,
|
||||
responseHeaders: Object.fromEntries(upstream.headers.entries()),
|
||||
responseBody: collected,
|
||||
responseSize: Buffer.byteLength(collected),
|
||||
responseBody: sink.text,
|
||||
responseSize: sink.totalBytes,
|
||||
proxyLatencyMs: upstreamStart - startedAt,
|
||||
upstreamLatencyMs: total - (upstreamStart - startedAt),
|
||||
});
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
*/
|
||||
import type { IncomingMessage, ServerResponse } from "node:http";
|
||||
import type { AgentId } from "../types";
|
||||
import { MitmHandlerBase } from "./base";
|
||||
import { MitmHandlerBase, createBoundedCollector } from "./base";
|
||||
|
||||
export class CodexHandler extends MitmHandlerBase {
|
||||
readonly agentId: AgentId = "codex";
|
||||
@@ -33,17 +33,17 @@ export class CodexHandler extends MitmHandlerBase {
|
||||
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
|
||||
}
|
||||
|
||||
let collected = "";
|
||||
const sink = createBoundedCollector();
|
||||
await this.pipeSSE(upstream, res, (chunk) => {
|
||||
collected += chunk.toString();
|
||||
sink.push(chunk.toString());
|
||||
});
|
||||
|
||||
const total = this.now() - startedAt;
|
||||
this.hookBufferUpdate(intercepted, {
|
||||
status: upstream.status,
|
||||
responseHeaders: Object.fromEntries(upstream.headers.entries()),
|
||||
responseBody: collected,
|
||||
responseSize: Buffer.byteLength(collected),
|
||||
responseBody: sink.text,
|
||||
responseSize: sink.totalBytes,
|
||||
proxyLatencyMs: upstreamStart - startedAt,
|
||||
upstreamLatencyMs: total - (upstreamStart - startedAt),
|
||||
});
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
*/
|
||||
import type { IncomingMessage, ServerResponse } from "node:http";
|
||||
import type { AgentId } from "../types";
|
||||
import { MitmHandlerBase } from "./base";
|
||||
import { MitmHandlerBase, createBoundedCollector } from "./base";
|
||||
|
||||
export class CopilotHandler extends MitmHandlerBase {
|
||||
readonly agentId: AgentId = "copilot";
|
||||
@@ -33,17 +33,17 @@ export class CopilotHandler extends MitmHandlerBase {
|
||||
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
|
||||
}
|
||||
|
||||
let collected = "";
|
||||
const sink = createBoundedCollector();
|
||||
await this.pipeSSE(upstream, res, (chunk) => {
|
||||
collected += chunk.toString();
|
||||
sink.push(chunk.toString());
|
||||
});
|
||||
|
||||
const total = this.now() - startedAt;
|
||||
this.hookBufferUpdate(intercepted, {
|
||||
status: upstream.status,
|
||||
responseHeaders: Object.fromEntries(upstream.headers.entries()),
|
||||
responseBody: collected,
|
||||
responseSize: Buffer.byteLength(collected),
|
||||
responseBody: sink.text,
|
||||
responseSize: sink.totalBytes,
|
||||
proxyLatencyMs: upstreamStart - startedAt,
|
||||
upstreamLatencyMs: total - (upstreamStart - startedAt),
|
||||
});
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
*/
|
||||
import type { IncomingMessage, ServerResponse } from "node:http";
|
||||
import type { AgentId } from "../types";
|
||||
import { MitmHandlerBase } from "./base";
|
||||
import { MitmHandlerBase, createBoundedCollector } from "./base";
|
||||
|
||||
export class CursorHandler extends MitmHandlerBase {
|
||||
readonly agentId: AgentId = "cursor";
|
||||
@@ -33,17 +33,17 @@ export class CursorHandler extends MitmHandlerBase {
|
||||
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
|
||||
}
|
||||
|
||||
let collected = "";
|
||||
const sink = createBoundedCollector();
|
||||
await this.pipeSSE(upstream, res, (chunk) => {
|
||||
collected += chunk.toString();
|
||||
sink.push(chunk.toString());
|
||||
});
|
||||
|
||||
const total = this.now() - startedAt;
|
||||
this.hookBufferUpdate(intercepted, {
|
||||
status: upstream.status,
|
||||
responseHeaders: Object.fromEntries(upstream.headers.entries()),
|
||||
responseBody: collected,
|
||||
responseSize: Buffer.byteLength(collected),
|
||||
responseBody: sink.text,
|
||||
responseSize: sink.totalBytes,
|
||||
proxyLatencyMs: upstreamStart - startedAt,
|
||||
upstreamLatencyMs: total - (upstreamStart - startedAt),
|
||||
});
|
||||
|
||||
@@ -10,7 +10,7 @@
|
||||
*/
|
||||
import type { IncomingMessage, ServerResponse } from "node:http";
|
||||
import type { AgentId } from "../types";
|
||||
import { MitmHandlerBase } from "./base";
|
||||
import { MitmHandlerBase, createBoundedCollector } from "./base";
|
||||
|
||||
export class KiroHandler extends MitmHandlerBase {
|
||||
readonly agentId: AgentId = "kiro";
|
||||
@@ -36,17 +36,17 @@ export class KiroHandler extends MitmHandlerBase {
|
||||
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
|
||||
}
|
||||
|
||||
let collected = "";
|
||||
const sink = createBoundedCollector();
|
||||
await this.pipeSSE(upstream, res, (chunk) => {
|
||||
collected += chunk.toString();
|
||||
sink.push(chunk.toString());
|
||||
});
|
||||
|
||||
const total = this.now() - startedAt;
|
||||
this.hookBufferUpdate(intercepted, {
|
||||
status: upstream.status,
|
||||
responseHeaders: Object.fromEntries(upstream.headers.entries()),
|
||||
responseBody: collected,
|
||||
responseSize: Buffer.byteLength(collected),
|
||||
responseBody: sink.text,
|
||||
responseSize: sink.totalBytes,
|
||||
proxyLatencyMs: upstreamStart - startedAt,
|
||||
upstreamLatencyMs: total - (upstreamStart - startedAt),
|
||||
});
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
*/
|
||||
import type { IncomingMessage, ServerResponse } from "node:http";
|
||||
import type { AgentId } from "../types";
|
||||
import { MitmHandlerBase } from "./base";
|
||||
import { MitmHandlerBase, createBoundedCollector } from "./base";
|
||||
|
||||
export class OpenCodeHandler extends MitmHandlerBase {
|
||||
readonly agentId: AgentId = "open-code";
|
||||
@@ -33,17 +33,17 @@ export class OpenCodeHandler extends MitmHandlerBase {
|
||||
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
|
||||
}
|
||||
|
||||
let collected = "";
|
||||
const sink = createBoundedCollector();
|
||||
await this.pipeSSE(upstream, res, (chunk) => {
|
||||
collected += chunk.toString();
|
||||
sink.push(chunk.toString());
|
||||
});
|
||||
|
||||
const total = this.now() - startedAt;
|
||||
this.hookBufferUpdate(intercepted, {
|
||||
status: upstream.status,
|
||||
responseHeaders: Object.fromEntries(upstream.headers.entries()),
|
||||
responseBody: collected,
|
||||
responseSize: Buffer.byteLength(collected),
|
||||
responseBody: sink.text,
|
||||
responseSize: sink.totalBytes,
|
||||
proxyLatencyMs: upstreamStart - startedAt,
|
||||
upstreamLatencyMs: total - (upstreamStart - startedAt),
|
||||
});
|
||||
|
||||
@@ -7,7 +7,7 @@
|
||||
*/
|
||||
import type { IncomingMessage, ServerResponse } from "node:http";
|
||||
import type { AgentId } from "../types";
|
||||
import { MitmHandlerBase } from "./base";
|
||||
import { MitmHandlerBase, createBoundedCollector } from "./base";
|
||||
|
||||
export class ZedHandler extends MitmHandlerBase {
|
||||
readonly agentId: AgentId = "zed";
|
||||
@@ -33,17 +33,17 @@ export class ZedHandler extends MitmHandlerBase {
|
||||
throw new Error(`OmniRoute ${upstream.status}: ${errText}`);
|
||||
}
|
||||
|
||||
let collected = "";
|
||||
const sink = createBoundedCollector();
|
||||
await this.pipeSSE(upstream, res, (chunk) => {
|
||||
collected += chunk.toString();
|
||||
sink.push(chunk.toString());
|
||||
});
|
||||
|
||||
const total = this.now() - startedAt;
|
||||
this.hookBufferUpdate(intercepted, {
|
||||
status: upstream.status,
|
||||
responseHeaders: Object.fromEntries(upstream.headers.entries()),
|
||||
responseBody: collected,
|
||||
responseSize: Buffer.byteLength(collected),
|
||||
responseBody: sink.text,
|
||||
responseSize: sink.totalBytes,
|
||||
proxyLatencyMs: upstreamStart - startedAt,
|
||||
upstreamLatencyMs: total - (upstreamStart - startedAt),
|
||||
});
|
||||
|
||||
@@ -542,47 +542,80 @@ async function intercept(req, res, bodyBuffer, override, sourceModel) {
|
||||
vlog(1, `[MITM] → forward ${forward.format} ${forward.url}`);
|
||||
|
||||
upstreamStartedAt = Date.now();
|
||||
const response = await fetch(forward.url, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
Authorization: `Bearer ${API_KEY}`,
|
||||
"x-omniroute-source": "agent-bridge",
|
||||
"x-omniroute-agent": agentId,
|
||||
},
|
||||
body: JSON.stringify(body),
|
||||
});
|
||||
// #13395: an abandoned downstream must not keep the upstream fetch + reader
|
||||
// alive. Abort the router fetch and cancel the reader on client close.
|
||||
const upstreamAbort = new AbortController();
|
||||
let downstreamClosed = false;
|
||||
const onDownstreamClose = () => {
|
||||
downstreamClosed = true;
|
||||
try {
|
||||
upstreamAbort.abort();
|
||||
} catch {
|
||||
// Abort is best-effort.
|
||||
}
|
||||
};
|
||||
res.once("close", onDownstreamClose);
|
||||
let reader = null;
|
||||
try {
|
||||
const response = await fetch(forward.url, {
|
||||
method: "POST",
|
||||
headers: {
|
||||
"Content-Type": "application/json",
|
||||
Authorization: `Bearer ${API_KEY}`,
|
||||
"x-omniroute-source": "agent-bridge",
|
||||
"x-omniroute-agent": agentId,
|
||||
},
|
||||
body: JSON.stringify(body),
|
||||
signal: upstreamAbort.signal,
|
||||
});
|
||||
|
||||
captureStatus = response.status;
|
||||
respHeaders = headersToObject(response.headers);
|
||||
captureStatus = response.status;
|
||||
respHeaders = headersToObject(response.headers);
|
||||
|
||||
if (!response.ok) {
|
||||
const errText = await response.text().catch(() => "");
|
||||
respBody = errText.slice(0, INGEST_MAX_BODY);
|
||||
respSize = Buffer.byteLength(errText);
|
||||
throw new Error(`OmniRoute ${response.status}: ${errText}`);
|
||||
}
|
||||
if (!response.ok) {
|
||||
const errText = await response.text().catch(() => "");
|
||||
respBody = errText.slice(0, INGEST_MAX_BODY);
|
||||
respSize = Buffer.byteLength(errText);
|
||||
throw new Error(`OmniRoute ${response.status}: ${errText}`);
|
||||
}
|
||||
|
||||
res.writeHead(200, {
|
||||
"Content-Type": "text/event-stream",
|
||||
"Cache-Control": "no-cache",
|
||||
Connection: "keep-alive",
|
||||
"X-Accel-Buffering": "no",
|
||||
});
|
||||
res.writeHead(200, {
|
||||
"Content-Type": "text/event-stream",
|
||||
"Cache-Control": "no-cache",
|
||||
Connection: "keep-alive",
|
||||
"X-Accel-Buffering": "no",
|
||||
});
|
||||
|
||||
const reader = response.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
reader = response.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
|
||||
while (true) {
|
||||
const { done, value } = await reader.read();
|
||||
if (done) {
|
||||
res.end();
|
||||
break;
|
||||
while (true) {
|
||||
if (downstreamClosed) break;
|
||||
const { done, value } = await reader.read();
|
||||
if (done) {
|
||||
res.end();
|
||||
break;
|
||||
}
|
||||
const text = decoder.decode(value, { stream: true });
|
||||
if (respBody.length < INGEST_MAX_BODY) respBody += text;
|
||||
respSize += value ? value.length : 0;
|
||||
if (downstreamClosed || res.closed || res.destroyed) break;
|
||||
res.write(text);
|
||||
}
|
||||
} finally {
|
||||
res.off("close", onDownstreamClose);
|
||||
if (reader) {
|
||||
try {
|
||||
await reader.cancel();
|
||||
} catch {
|
||||
// Reader already done or released.
|
||||
}
|
||||
try {
|
||||
reader.releaseLock();
|
||||
} catch {
|
||||
// Already released.
|
||||
}
|
||||
}
|
||||
const text = decoder.decode(value, { stream: true });
|
||||
if (respBody.length < INGEST_MAX_BODY) respBody += text;
|
||||
respSize += value ? value.length : 0;
|
||||
res.write(text);
|
||||
}
|
||||
} catch (error) {
|
||||
// Log the raw message locally (server console only) but never expose it
|
||||
|
||||
@@ -44,10 +44,29 @@ function fakeRes(): { res: ServerResponse; out: HarnessResult } {
|
||||
responseChunks: [],
|
||||
};
|
||||
let headersSent = false;
|
||||
const listeners = new Map<string, Set<(...args: unknown[]) => void>>();
|
||||
const res = {
|
||||
get headersSent() {
|
||||
return headersSent;
|
||||
},
|
||||
closed: false,
|
||||
destroyed: false,
|
||||
once(event: string, fn: (...args: unknown[]) => void) {
|
||||
let set = listeners.get(event);
|
||||
if (!set) {
|
||||
set = new Set();
|
||||
listeners.set(event, set);
|
||||
}
|
||||
set.add(fn);
|
||||
return res;
|
||||
},
|
||||
off(event: string, fn: (...args: unknown[]) => void) {
|
||||
listeners.get(event)?.delete(fn);
|
||||
return res;
|
||||
},
|
||||
emitClose() {
|
||||
for (const fn of [...(listeners.get("close") ?? [])]) fn();
|
||||
},
|
||||
writeHead(s: number) {
|
||||
out.status = s;
|
||||
headersSent = true;
|
||||
@@ -59,7 +78,7 @@ function fakeRes(): { res: ServerResponse; out: HarnessResult } {
|
||||
end(c?: Buffer | string) {
|
||||
if (c) out.responseChunks.push(typeof c === "string" ? c : c.toString());
|
||||
},
|
||||
} as unknown as ServerResponse;
|
||||
} as unknown as ServerResponse & { emitClose: () => void };
|
||||
return { res, out };
|
||||
}
|
||||
|
||||
|
||||
143
tests/unit/mitm-pipe-bounded-collect-13395.test.ts
Normal file
143
tests/unit/mitm-pipe-bounded-collect-13395.test.ts
Normal file
@@ -0,0 +1,143 @@
|
||||
// #13395: MITM pipe paths must not retain unbounded transcripts, and an
|
||||
// abandoned downstream must stop the upstream read.
|
||||
import test from "node:test";
|
||||
import assert from "node:assert/strict";
|
||||
import fs from "node:fs";
|
||||
import path from "node:path";
|
||||
import url from "node:url";
|
||||
import type { ServerResponse } from "node:http";
|
||||
import { Readable } from "node:stream";
|
||||
import {
|
||||
MitmHandlerBase,
|
||||
createBoundedCollector,
|
||||
MITM_PIPE_MAX_COLLECT_BYTES,
|
||||
} from "../../src/mitm/handlers/base.ts";
|
||||
import type { AgentId } from "../../src/mitm/types.ts";
|
||||
|
||||
class ExposedHandler extends MitmHandlerBase {
|
||||
readonly agentId: AgentId = "antigravity";
|
||||
async intercept(): Promise<void> {
|
||||
throw new Error("not used");
|
||||
}
|
||||
pipe(
|
||||
upstream: Response,
|
||||
res: ServerResponse,
|
||||
onChunk?: (c: Buffer) => void
|
||||
): Promise<void> {
|
||||
return this.pipeSSE(upstream, res, onChunk);
|
||||
}
|
||||
}
|
||||
|
||||
function trackingRes() {
|
||||
const listeners = new Map<string, Set<(...a: unknown[]) => void>>();
|
||||
const written: string[] = [];
|
||||
let offCalls = 0;
|
||||
const res = {
|
||||
headersSent: false,
|
||||
closed: false,
|
||||
destroyed: false,
|
||||
once(event: string, fn: (...a: unknown[]) => void) {
|
||||
let s = listeners.get(event);
|
||||
if (!s) {
|
||||
s = new Set();
|
||||
listeners.set(event, s);
|
||||
}
|
||||
s.add(fn);
|
||||
return res;
|
||||
},
|
||||
off(event: string, fn: (...a: unknown[]) => void) {
|
||||
offCalls += 1;
|
||||
listeners.get(event)?.delete(fn);
|
||||
return res;
|
||||
},
|
||||
emitClose() {
|
||||
for (const fn of [...(listeners.get("close") ?? [])]) fn();
|
||||
},
|
||||
writeHead() {
|
||||
(res as { headersSent: boolean }).headersSent = true;
|
||||
},
|
||||
write(c: Buffer | string) {
|
||||
written.push(typeof c === "string" ? c : c.toString());
|
||||
return true;
|
||||
},
|
||||
end() {},
|
||||
} as unknown as ServerResponse;
|
||||
return { res, written, listeners, offCalls: () => offCalls };
|
||||
}
|
||||
|
||||
function chunkedUpstream(chunks: string[], delayMs = 0): Response {
|
||||
const iterable = (async function* () {
|
||||
for (const c of chunks) {
|
||||
if (delayMs > 0) await new Promise((r) => setTimeout(r, delayMs));
|
||||
yield Buffer.from(c);
|
||||
}
|
||||
})();
|
||||
const stream = Readable.toWeb(Readable.from(iterable)) as unknown as ReadableStream<Uint8Array>;
|
||||
return new Response(stream, { status: 200 });
|
||||
}
|
||||
|
||||
test("bounded collector retains at most the cap but counts the true total", () => {
|
||||
const sink = createBoundedCollector(100);
|
||||
sink.push("a".repeat(60));
|
||||
sink.push("b".repeat(60));
|
||||
assert.equal(sink.text.length, 100);
|
||||
assert.equal(sink.totalBytes, 120);
|
||||
assert.equal(sink.truncated, true);
|
||||
});
|
||||
|
||||
test("bounded collector passes small transcripts through untouched", () => {
|
||||
const sink = createBoundedCollector(100);
|
||||
sink.push("hello");
|
||||
assert.equal(sink.text, "hello");
|
||||
assert.equal(sink.totalBytes, 5);
|
||||
assert.equal(sink.truncated, false);
|
||||
});
|
||||
|
||||
test("default collect ceiling is 1 MiB", () => {
|
||||
assert.equal(MITM_PIPE_MAX_COLLECT_BYTES, 1 * 1024 * 1024);
|
||||
});
|
||||
|
||||
test("pipeSSE delivers every chunk when the downstream stays open", async () => {
|
||||
const h = new ExposedHandler();
|
||||
const { res, written } = trackingRes();
|
||||
await h.pipe(chunkedUpstream(["x".repeat(10), "y".repeat(10), "z".repeat(10)]), res);
|
||||
assert.equal(written.join(""), "x".repeat(10) + "y".repeat(10) + "z".repeat(10));
|
||||
});
|
||||
|
||||
test("pipeSSE stops the upstream read after downstream close", async () => {
|
||||
const h = new ExposedHandler();
|
||||
const { res, written } = trackingRes();
|
||||
const seen: string[] = [];
|
||||
const pipe = h.pipe(
|
||||
chunkedUpstream(Array.from({ length: 50 }, (_, i) => `c${i};`), 5),
|
||||
res,
|
||||
(c) => seen.push(c.toString())
|
||||
);
|
||||
// Let a few chunks flow, then abandon the downstream.
|
||||
await new Promise((r) => setTimeout(r, 25));
|
||||
(res as unknown as { emitClose: () => void }).emitClose();
|
||||
await pipe;
|
||||
const totalUpstream = Array.from({ length: 50 }, (_, i) => `c${i};`).join("");
|
||||
assert.ok(
|
||||
seen.join("").length < totalUpstream.length,
|
||||
`abandoned pipe must stop early (read ${seen.join("").length} of ${totalUpstream.length})`
|
||||
);
|
||||
assert.ok(written.join("").length <= seen.join("").length, "no writes after close");
|
||||
});
|
||||
|
||||
test("pipeSSE detaches its close listener when the stream completes", async () => {
|
||||
const h = new ExposedHandler();
|
||||
const t = trackingRes();
|
||||
await h.pipe(chunkedUpstream(["done"]), t.res);
|
||||
assert.equal(t.listeners.get("close")?.size ?? 0, 0, "close listener must be removed");
|
||||
assert.ok(t.offCalls() >= 1, "off(close) must run");
|
||||
});
|
||||
|
||||
test("server.cjs aborts the router fetch and cancels the reader on downstream close (#13395)", async () => {
|
||||
const here = path.dirname(url.fileURLToPath(import.meta.url));
|
||||
const src = fs.readFileSync(path.resolve(here, "../../src/mitm/server.cjs"), "utf8");
|
||||
assert.match(src, /new AbortController\(\)/, "must create an abort controller per intercept");
|
||||
assert.match(src, /signal:\s*upstreamAbort\.signal/, "router fetch must take the abort signal");
|
||||
assert.match(src, /res\.once\(\s*"close"/, "must listen for downstream close");
|
||||
assert.match(src, /reader\.cancel\(\)/, "must cancel the upstream reader on close");
|
||||
});
|
||||
Reference in New Issue
Block a user