diff --git a/packages/mcp/src/lib/drawio-layout.ts b/packages/mcp/src/lib/drawio-layout.ts index 15117a97..e3ff71e4 100644 --- a/packages/mcp/src/lib/drawio-layout.ts +++ b/packages/mcp/src/lib/drawio-layout.ts @@ -18,6 +18,23 @@ import { normalizeInput, parseCells, type DrawioCell } from "./drawio-xml.js"; const DEFAULT_W = 140; const DEFAULT_H = 60; +// DoS bounds for the in-process ELK layout. The mxGraph XML is LLM-supplied +// (layout:"elk" in drawio_create/drawio_update) and elkjs runs synchronously on +// the MCP server's event loop, so an unbounded graph would block it for +// seconds-to-minutes. A ~1MB XML (well under the stage-1 16MB cap) can carry +// thousands of nodes. We cap the graph size and race the layout against a +// wall-clock timeout; on either bound we fall back to the ORIGINAL model, the +// same best-effort contract the catch already honours. +// - 500 nodes lays out in well under a second; beyond that ELK cost climbs +// steeply, so refuse and leave the (already-valid) model untouched. +// - Edges dominate the layered-crossing cost, so allow a bit more headroom +// (1000) than nodes but still bound them. +// - 5s is generous for any graph within the caps yet short enough that a +// pathological input can never wedge the server. +const ELK_MAX_NODES = 500; +const ELK_MAX_EDGES = 1000; +const ELK_TIMEOUT_MS = 5000; + // Spacing is set >=150px on purpose so an ELK layout never trips the linter's // "gap between adjacent shapes < 150px" quality warning (acceptance #3). const LAYOUT_OPTIONS: Record = { @@ -118,6 +135,13 @@ export async function applyElkLayout(inputXml: string): Promise { edges.push({ id: c.id || `e${edges.length}`, sources: [c.source], targets: [c.target] }); } + // DoS guard: refuse to lay out an oversized LLM-supplied graph. elkjs runs + // in-process on the event loop, so bound the work before we ever call it and + // return the original model unchanged (best-effort, same as the catch below). + if (vertices.length > ELK_MAX_NODES || edges.length > ELK_MAX_EDGES) { + return modelXml; + } + const graph: ElkGraph = { id: "root", layoutOptions: LAYOUT_OPTIONS, @@ -126,15 +150,26 @@ export async function applyElkLayout(inputXml: string): Promise { }; let laid: ElkGraph; + let timer: ReturnType | undefined; try { // elkjs ships a CJS default export whose interop shape varies across // module systems; resolve the real constructor at runtime, then cast (the // runtime call is verified — see the layout unit test). const Ctor: any = (ELK as any).default ?? ELK; const elk = new Ctor(); - laid = (await elk.layout(graph as any)) as ElkGraph; + // Race the layout against a wall-clock timeout so a graph that is under the + // node/edge caps but still pathologically slow can never wedge the server. + const timeout = new Promise((_, reject) => { + timer = setTimeout( + () => reject(new Error("ELK layout timed out")), + ELK_TIMEOUT_MS, + ); + }); + laid = (await Promise.race([elk.layout(graph as any), timeout])) as ElkGraph; } catch { - return modelXml; // best-effort: keep the model as-is on any ELK failure + return modelXml; // best-effort: keep the model as-is on timeout or ELK failure + } finally { + if (timer) clearTimeout(timer); } // Collect computed geometry per node id (coords are parent-relative already). diff --git a/packages/mcp/test/unit/drawio-layout.test.mjs b/packages/mcp/test/unit/drawio-layout.test.mjs index e88f3c41..a6998d17 100644 --- a/packages/mcp/test/unit/drawio-layout.test.mjs +++ b/packages/mcp/test/unit/drawio-layout.test.mjs @@ -78,6 +78,29 @@ test("edges and cell count are preserved by layout", async () => { assert.equal(cells.filter((c) => c.vertex).length, 4); }); +test("DoS guard: a graph over the node cap is returned unchanged, quickly", async () => { + // 600 vertices > ELK_MAX_NODES (500): the layout must be SKIPPED and the + // input returned verbatim, without ever handing the graph to elkjs. This + // exercises the cap path that bounds the in-process, event-loop-blocking + // layout on LLM-supplied XML. + const model = stackedGraph(600, []); + const t0 = Date.now(); + const laid = await applyElkLayout(model); + const dt = Date.now() - t0; + // normalizeInput may reserialize, but geometry must be untouched: every + // vertex is still stacked at (10,10), i.e. no ELK coordinates were applied. + const cells = parseCells(laid); + const verts = cells.filter((c) => c.vertex); + assert.equal(verts.length, 600, "all vertices survived"); + for (const v of verts) { + assert.equal(v.geometry.x, 10, "x untouched -> layout was skipped"); + assert.equal(v.geometry.y, 10, "y untouched -> layout was skipped"); + } + // Returning the input without an ELK pass is essentially instant; assert it + // did not hang. Generous bound to stay non-flaky on a loaded CI box. + assert.ok(dt < 2000, `cap path should be fast, took ${dt}ms`); +}); + test("layout is best-effort: an empty/degenerate model is returned intact", async () => { const model = '';