From 0ae38e1e31414c67a3be09231d208f0726c035f8 Mon Sep 17 00:00:00 2001 From: Mykola Dementiev Date: Sun, 11 Oct 2026 22:10:22 +0200 Subject: [PATCH] fix(ida): preserve non-x86 call edges and workflow uncertainty (#1906) * test(ida): expect callee edges from the call graph on non-x86 targets * fix(ida): read callee edges from the call graph on non-x86 targets * test(ida): pin address order of call graph additions * docs(ida): describe the call graph source for non-x86 callees * test(ida): pin the x86 baseline, the call graph reader and its failure modes * fix(ida): check the call graph root, scope truncation to the root and degrade without callgraph * test(ida): expect managed PE off the x86 path, possible-truncation wording and a complete ARM64 traversal * fix(ida): keep managed PE and x86 reads off the call graph path and qualify truncation * test(ida): answer only expected leaf roots and assert the requested call graph roots * test(ida): expect a call graph the producer cannot give consistently to fail the callee read * fix(ida): fail the callee read when the call graph cannot be obtained consistently * fix(ida): preserve call edge authority and workflow uncertainty * fix(ida): preserve explicit root call graph caps --------- Co-authored-by: morluto <76467478+morluto@users.noreply.github.com> --- docs/ida-provider.md | 2 +- docs/mcp-contracts.md | 7 + src/application/CallPathTracing.ts | 24 +- src/application/EnhancedTools.ts | 27 +- src/contracts/toolOutputSchemaPrimitives.ts | 1 + src/domain/callPath.test.ts | 72 ++++- src/domain/callPath.ts | 15 +- src/ida/IdaCallGraph.test.ts | 304 ++++++++++++++++++ src/ida/IdaCallGraph.ts | 153 +++++++++ src/ida/IdaCalleeRelationships.test.ts | 222 +++++++++++++ src/ida/IdaOperationRunner.ts | 48 ++- src/ida/IdaProtocolValues.ts | 41 ++- src/ida/IdaSessionClient.ts | 11 +- .../enhancedTools.callPathsAndBatches.test.ts | 65 +++- .../providers/ida/calleeRelationships.test.ts | 291 +++++++++++++++++ tests/fixtures/idaCallGraph.ts | 138 ++++++++ tests/fixtures/idaMcp.ts | 21 +- 17 files changed, 1410 insertions(+), 32 deletions(-) create mode 100644 src/ida/IdaCallGraph.test.ts create mode 100644 src/ida/IdaCallGraph.ts create mode 100644 src/ida/IdaCalleeRelationships.test.ts create mode 100644 tests/boundary/providers/ida/calleeRelationships.test.ts create mode 100644 tests/fixtures/idaCallGraph.ts diff --git a/docs/ida-provider.md b/docs/ida-provider.md index 32ba0504d..d2e939561 100644 --- a/docs/ida-provider.md +++ b/docs/ida-provider.md @@ -98,7 +98,7 @@ For guided registration, set `REA_IDA_MCP_CONFIG` before `rea setup --client Promise>; +type RelationshipCall = ( + name: AnalysisOperation, + arguments_: Readonly>, + signal?: AbortSignal, + selectorOrigin?: ExecutionOptions["selectorOrigin"], +) => Promise>; + export interface CallPathTraceInput { readonly start: string; readonly goal?: string | undefined; @@ -42,6 +50,7 @@ interface TraceState { readonly error: ReturnType; }>; readonly residual: Set; + readonly limitations: Set; queueIndex: number; traversalPath: string[]; } @@ -63,6 +72,7 @@ interface DiscoveredEdge { /** Trace direct caller/callee relationships until the graph is exhausted or the goal is reached. */ export const traceCallPath = async ( call: AnalysisCall, + callRelationship: RelationshipCall, input: CallPathTraceInput, signal?: AbortSignal, ): EnhancedResult => { @@ -88,7 +98,7 @@ export const traceCallPath = async ( if (address === undefined) break; const current = state.visited.get(address); if (current === undefined) continue; - const found = await expandAddress(call, { + const found = await expandAddress(callRelationship, { input: resolvedInput, state, address, @@ -106,12 +116,13 @@ const createTraceState = (input: CallPathTraceInput): TraceState => ({ edges: new Map(), failures: [], residual: new Set(), + limitations: new Set(), queueIndex: 0, traversalPath: input.goal === input.start ? [input.start] : [], }); const expandAddress = async ( - call: AnalysisCall, + call: RelationshipCall, expansion: AddressExpansion, ): Promise => { const { input, state, address, current, signal } = expansion; @@ -127,7 +138,7 @@ const expandAddress = async ( state.residual.add(`Call relationships were unavailable for ${address}.`); return false; } - const related = parseRelatedAddresses(result.value, relation); + const related = parseRelatedAddresses(result.value.result, relation); if (!related.ok) { state.failures.push({ address, @@ -136,6 +147,12 @@ const expandAddress = async ( state.residual.add(`Call relationships were unreadable for ${address}.`); return false; } + for (const limitation of result.value.limitations) { + state.limitations.add(limitation); + state.residual.add( + `Call relationships for ${address} are limited: ${limitation}`, + ); + } for (const relatedAddress of sortedUniqueAddresses(related.value)) { recordEdge(input, state, { address, @@ -215,6 +232,7 @@ const projectTrace = (input: CallPathTraceInput, state: TraceState) => ({ residual_unknowns: [...state.residual], limitations: [ "Direct provider call relationships may omit unresolved indirect calls.", + ...state.limitations, ], }); diff --git a/src/application/EnhancedTools.ts b/src/application/EnhancedTools.ts index 06862be0e..7e0b1addf 100644 --- a/src/application/EnhancedTools.ts +++ b/src/application/EnhancedTools.ts @@ -1,6 +1,7 @@ import { z } from "zod"; import type { + AnalysisExecution, AnalysisOperation, AnalysisOperationPort, ExecutionOptions, @@ -166,6 +167,8 @@ export class EnhancedTools { return traceCallPath( (name, arguments_, operationSignal, selectorOrigin) => this.#call(name, arguments_, operationSignal, selectorOrigin), + (name, arguments_, operationSignal, selectorOrigin) => + this.#observe(name, arguments_, operationSignal, selectorOrigin), call.input, signal, ); @@ -359,7 +362,7 @@ export class EnhancedTools { const level = String(current.depth); graph[level] ??= []; - const result = await this.#call( + const result = await this.#observe( tool, { procedure: current.address }, signal, @@ -373,7 +376,7 @@ export class EnhancedTools { }); continue; } - const related = parseRelatedAddresses(result.value, relation); + const related = parseRelatedAddresses(result.value.result, relation); if (!related.ok) { graph[level].push({ address: current.address, @@ -386,6 +389,9 @@ export class EnhancedTools { address: current.address, status: "ok", calls: [...related.value], + ...(result.value.limitations.length === 0 + ? {} + : { limitations: [...result.value.limitations] }), }); for (const address of related.value) { if (!discovered.has(address)) { @@ -523,13 +529,28 @@ export class EnhancedTools { signal?: AbortSignal, selectorOrigin?: ExecutionOptions["selectorOrigin"], ): Promise> { + const execution = await this.#observe( + name, + arguments_, + signal, + selectorOrigin, + ); + return execution.ok ? ok(execution.value.result) : execution; + } + + async #observe( + name: AnalysisOperation, + arguments_: Readonly>, + signal?: AbortSignal, + selectorOrigin?: ExecutionOptions["selectorOrigin"], + ): Promise> { if (isAborted(signal)) return err(new AnalysisCancelledError(name)); const execution = await this.analysis.execute(name, arguments_, { ...(signal === undefined ? {} : { signal }), ...(selectorOrigin === undefined ? {} : { selectorOrigin }), }); if (isAborted(signal)) return err(new AnalysisCancelledError(name)); - return execution.ok ? ok(execution.value.result) : execution; + return execution; } } diff --git a/src/contracts/toolOutputSchemaPrimitives.ts b/src/contracts/toolOutputSchemaPrimitives.ts index 51a53aa4f..1f1a8c700 100644 --- a/src/contracts/toolOutputSchemaPrimitives.ts +++ b/src/contracts/toolOutputSchemaPrimitives.ts @@ -505,6 +505,7 @@ export const graphNode = z.discriminatedUnion("status", [ address: z.string(), status: z.literal("ok"), calls: z.array(z.string()), + limitations: z.array(z.string()).optional(), }), z.object({ address: z.string(), diff --git a/src/domain/callPath.test.ts b/src/domain/callPath.test.ts index 00d9e97d4..d10995525 100644 --- a/src/domain/callPath.test.ts +++ b/src/domain/callPath.test.ts @@ -14,6 +14,10 @@ const observe = ( address: string, callees: readonly string[], provider = "rea-workflow", + options: { + readonly evidenceLimitations?: readonly string[]; + readonly dossierLimitations?: readonly string[]; + } = {}, ): Evidence => { const value = functionDossierSchema.parse({ procedure: { @@ -49,7 +53,7 @@ const observe = ( native_api: null, native_value_flow: null, basic_blocks: [], - limitations: [], + limitations: [...(options.dossierLimitations ?? [])], }); return createEvidence( { path: "/tmp/a", sha256: "a".repeat(64), format: "mach-o" }, @@ -62,6 +66,9 @@ const observe = ( result: jsonValueSchema.parse(value), confidence: "derived", authority: "shipped-artifact", + ...(options.evidenceLimitations === undefined + ? {} + : { limitations: options.evidenceLimitations }), }, ); }; @@ -126,6 +133,22 @@ describe("call path reconstruction", () => { expect(unknown.limitations).toContain( "No analyze_function Evidence covers reachable function 0x2000", ); + + const limitedEmpty = run([ + observe("0x1000", [], "rea-workflow", { + dossierLimitations: ["Call relationships are incomplete."], + }), + observe("0x4000", []), + ]); + expect(limitedEmpty).toMatchObject({ + status: "unknown", + search_scope: { exhaustive: false }, + shortest_hops: null, + paths: [], + }); + expect(limitedEmpty.limitations).toContain( + "Call relationships are incomplete.", + ); }); it("searches the complete supplied graph", () => { @@ -182,6 +205,53 @@ describe("call path topology and capacity", () => { ); }); + it("keeps reached function limitations and marks expanded coverage incomplete", () => { + const evidenceLimitation = "Call graph output was truncated."; + const dossierLimitation = "One call target could not be resolved."; + const result = run([ + observe("0x1000", ["0x4000"], "rea-workflow", { + evidenceLimitations: [evidenceLimitation], + dossierLimitations: [dossierLimitation], + }), + observe("0x4000", []), + ]); + + expect(result).toMatchObject({ + status: "found", + shortest_hops: 1, + search_scope: { exhaustive: false }, + paths: [ + { + nodes: [{ address: "0x1000" }, { address: "0x4000" }], + }, + ], + }); + expect(result.limitations).toEqual( + expect.arrayContaining([evidenceLimitation, dossierLimitation]), + ); + }); + + it("retains goal dossier limitations without expanding the goal", () => { + const goalLimitation = "The goal dossier has partial outgoing references."; + const result = run([ + observe("0x1000", ["0x4000"]), + observe("0x4000", ["0x5000"], "rea-workflow", { + dossierLimitations: [goalLimitation], + }), + ]); + + expect(result).toMatchObject({ + status: "found", + shortest_hops: 1, + search_scope: { exhaustive: true }, + }); + expect(result.limitations).toContain(goalLimitation); + expect(result.explored.nodes).toBe(2); + expect(result.limitations).not.toContain( + "No analyze_function Evidence covers reachable function 0x5000", + ); + }); + it("returns a zero-hop path when start and goal are the same function", () => { const result = run([observe("0x1000", ["0x2000"])], { goal: { address: "0x1000" }, diff --git a/src/domain/callPath.ts b/src/domain/callPath.ts index e468ba0ec..1da429625 100644 --- a/src/domain/callPath.ts +++ b/src/domain/callPath.ts @@ -229,16 +229,27 @@ const inspectSearch = ({ const reached = new Map([[start, 0]]); const queue = [start]; const limitations: string[] = []; + let incompleteExpandedCoverage = false; for (let index = 0; index < queue.length; index += 1) { const node = queue[index]; if (node === undefined) continue; const depth = reached.get(node) ?? 0; - if (node === goal) continue; const snapshot = snapshots.get(node); + if (snapshot !== undefined) { + limitations.push(...snapshot.limitations); + // The requested goal is a terminal for this search. Its outgoing edges + // cannot add a path to the goal, so its dossier limits are contextual + // but do not make the searched start-to-goal closure incomplete. + if (node === goal) continue; + if (snapshot.limitations.length > 0) incompleteExpandedCoverage = true; + } else if (node === goal) { + continue; + } if (snapshot === undefined) { limitations.push( `No analyze_function Evidence covers reachable function ${node}`, ); + incompleteExpandedCoverage = true; continue; } const neighbors = graph.outNeighbors(node).sort(compareUnicodeCodePoints); @@ -252,7 +263,7 @@ const inspectSearch = ({ graph, snapshots, reached, - exhaustive: limitations.length === 0, + exhaustive: !incompleteExpandedCoverage, limitations, }; }; diff --git a/src/ida/IdaCallGraph.test.ts b/src/ida/IdaCallGraph.test.ts new file mode 100644 index 000000000..e39c7358d --- /dev/null +++ b/src/ida/IdaCallGraph.test.ts @@ -0,0 +1,304 @@ +import { describe, expect, it } from "vitest"; +import { + AnalysisCancelledError, + AnalysisProtocolError, +} from "../domain/analysisErrorCore.js"; +import type { BinaryTarget } from "../domain/binaryTargetTypes.js"; +import type { JsonValue } from "../domain/jsonValue.js"; +import { + isX86OnlyTarget, + readCallGraphCallees, + UNRECOGNIZED_CALLEE_LIMITATION, +} from "./IdaCallGraph.js"; + +const target = ( + architecture: "x86" | "x86_64" | "arm" | "arm64", + available: readonly ("x86" | "x86_64" | "arm" | "arm64")[] = [architecture], +): BinaryTarget => ({ + path: "/fixture", + sha256: "a".repeat(64), + kind: "executable", + format: "mach-o", + architecture, + availableArchitectures: [...available], +}); + +const row = ( + edges: readonly { from: string; to: string; type: string }[], + extra: Readonly> = {}, +): JsonValue => [ + { + root: "0x1000", + nodes: [{ addr: "0x1000", name: "main" }], + edges: [...edges], + ...extra, + }, +]; +const call = (from: string, to: string) => ({ from, to, type: "call" }); +const readGraph = ( + edges: readonly { from: string; to: string; type: string }[], + extra: Readonly> = {}, + lookup: Readonly> = { + "0x20": "low", + }, +) => + readCallGraphCallees(async (tool, args) => { + if (tool === "callgraph") return row(edges, extra); + if (tool === "lookup_funcs") + return (args.queries as string[]).map((query) => { + const name = lookup[query]; + return name == null + ? { query, fn: null, error: "Not found" } + : { + query, + fn: { address: query, name, size: "0x10" }, + error: null, + }; + }); + throw new Error(`Unexpected tool: ${tool}`); + }, "0x1000"); + +describe("architecture gate for the call graph source", () => { + it.each([ + ["x86", target("x86"), true], + ["x86_64", target("x86_64"), true], + ["x86 with both x86 families", target("x86_64", ["x86", "x86_64"]), true], + ["32-bit arm", target("arm"), false], + ["arm64", target("arm64"), false], + [ + "a universal target with an arm64 slice", + target("x86_64", ["x86_64", "arm64"]), + false, + ], + ])("treats %s as x86-only: %s", (_label, candidate, expected) => { + expect( + isX86OnlyTarget(candidate), + "the gate follows every loadable slice", + ).toBe(expected); + }); + + it.each([ + ["a native PE", false, true], + ["a managed PE that IDA may load with the CLI loader", true, false], + ])("treats %s as x86-only: %s", (_label, managed, expected) => { + expect( + isX86OnlyTarget({ + path: "/fixture", + sha256: "a".repeat(64), + kind: "executable", + format: "pe", + architecture: "x86_64", + availableArchitectures: ["x86_64"], + executableRole: "application", + managed, + }), + "a managed image is not proven to be processed as x86", + ).toBe(expected); + }); + + it("never treats a non-executable target as x86-only", () => { + expect( + isX86OnlyTarget({ + path: "/fixture", + sha256: "a".repeat(64), + kind: "archive", + format: "zip", + }), + "only executables can be proven x86-only", + ).toBe(false); + }); +}); + +describe("reading root callees from the processor-neutral call graph", () => { + it("requests only the procedure as root at depth zero", async () => { + const requests: unknown[] = []; + + await readCallGraphCallees(async (tool, args) => { + requests.push([tool, args]); + return row([]); + }, "0x1000"); + + expect(requests, "one scoped request, no other tool").toEqual([ + ["callgraph", { roots: ["0x1000"], max_depth: 0 }], + ]); + }); + + it("keeps root call edges, excluding child and non-call edges", async () => { + const result = await readGraph([ + call("0x1000", "0x20"), + { from: "0x1000", to: "0x400", type: "jump" }, + call("0x20", "0x500"), + ]); + + expect( + result.callees.map(({ address }) => address), + "child edges and non-call edges are excluded", + ).toEqual(["0x20"]); + }); + + it("batches name lookup and keeps numeric order and unresolved addresses", async () => { + const requests: unknown[] = []; + const result = await readCallGraphCallees(async (tool, args) => { + requests.push([tool, args]); + if (tool === "callgraph") + return row([ + call("0x1000", "0x300"), + call("0x1000", "0x100"), + call("0x1000", "0x20"), + ]); + return [ + { + query: "0x20", + fn: { address: "0x20", name: "low", size: "0x10" }, + error: null, + }, + { + query: "0x100", + fn: { address: "0x100", name: "middle", size: "0x10" }, + error: null, + }, + { query: "0x300", fn: null, error: "Not found" }, + ]; + }, "0x1000"); + + expect( + result.callees, + "numeric order, with a real address retained when lookup misses", + ).toEqual([ + { address: "0x20", name: "low" }, + { address: "0x100", name: "middle" }, + { address: "0x300", name: "0x300" }, + ]); + expect(requests).toEqual([ + ["callgraph", { roots: ["0x1000"], max_depth: 0 }], + ["lookup_funcs", { queries: ["0x20", "0x100", "0x300"] }], + ]); + expect(result.limitations).toContain(UNRECOGNIZED_CALLEE_LIMITATION); + expect(result.limitations).toContain( + "IDA observed callee address 0x300 but could not resolve its name: Not found; the address is retained.", + ); + }); + + it("keeps observed addresses when a typed name-lookup failure prevents enrichment", async () => { + const failure = new AnalysisProtocolError("lookup reply failed"); + const result = await readCallGraphCallees(async (tool) => { + if (tool === "callgraph") return row([call("0x1000", "0x20")]); + throw failure; + }, "0x1000"); + + expect( + result.callees, + "the graph address remains a real observation", + ).toEqual([{ address: "0x20", name: "0x20" }]); + expect(result.limitations).toContain( + "IDA observed callee addresses but the name lookup failed (lookup reply failed); the addresses are retained.", + ); + }); + + it("does not hide unexpected name-lookup failures or cancellation", async () => { + const failures = [ + new Error("unexpected transport failure"), + new AnalysisCancelledError("lookup_funcs"), + ]; + for (const failure of failures) + await expect( + readCallGraphCallees(async (tool) => { + if (tool === "callgraph") return row([call("0x1000", "0x20")]); + throw failure; + }, "0x1000"), + `${failure.message} must not be converted into a name limitation`, + ).rejects.toBe(failure); + }); +}); + +describe("call graph coverage and failure handling", () => { + it.each([ + ["truncated", { truncated: true }], + [ + "per_func_capped below the edge cap", + { per_func_capped: true, max_edges_per_func: 200 }, + ], + ["limit_reason", { limit_reason: "edges" }], + ])("reports a graph flagged by %s as truncated", async (_label, extra) => { + const result = await readGraph([call("0x1000", "0x20")], extra); + + expect( + result.limitations, + "each truncation signal counts on its own", + ).toContain( + "Upstream callgraph was truncated; additional callees remain unknown.", + ); + }); + + it("keeps an explicitly uncapped root below its edge cap untruncated", async () => { + const result = await readGraph([call("0x1000", "0x20")], { + per_func_capped: false, + max_edges_per_func: 200, + }); + + expect( + result.limitations, + "an uncapped, unsaturated root has no truncation evidence", + ).toEqual([UNRECOGNIZED_CALLEE_LIMITATION]); + }); + + it("reports a graph that saturates its per-function cap with root edges", async () => { + const result = await readGraph( + [call("0x1000", "0x20"), call("0x1000", "0x100")], + { max_edges_per_func: 2 }, + ); + + expect( + result.limitations, + "reaching the cap is uncertain, so it is reported as possible truncation", + ).toContain( + "Upstream callgraph may be truncated; additional callees may remain unknown.", + ); + }); + + it("fails the read for a reply about a different root", async () => { + await expect( + readCallGraphCallees( + async () => [ + { root: "0x2000", nodes: [], edges: [call("0x1000", "0x20")] }, + ], + "0x1000", + ), + "edges of another root are never observed", + ).rejects.toThrow(/different root/u); + }); + + it("fails the read for an error row and for an unreadable reply", async () => { + await expect( + readCallGraphCallees( + async () => [ + { root: "0x1000", error: "Function not found", nodes: [], edges: [] }, + ], + "0x1000", + ), + "the upstream reason is kept", + ).rejects.toThrow("Function not found"); + await expect( + readCallGraphCallees(async () => ({ unexpected: true }), "0x1000"), + "an unreadable reply is a producer fault", + ).rejects.toThrow(/unreadable/u); + }); + + it("lets tool failures, truncated output, cancellation and transport failures surface unchanged", async () => { + const failures = [ + new AnalysisProtocolError("IDA MCP tool failed: boom"), + new AnalysisProtocolError( + "IDA MCP truncated its output; the preview cannot establish a complete observation.", + ), + new AnalysisCancelledError("callgraph"), + new Error("transport timed out"), + ]; + for (const failure of failures) + await expect( + readCallGraphCallees(async () => { + throw failure; + }, "0x1000"), + `${failure.message} is not turned into an observation`, + ).rejects.toBe(failure); + }); +}); diff --git a/src/ida/IdaCallGraph.ts b/src/ida/IdaCallGraph.ts new file mode 100644 index 000000000..53152d138 --- /dev/null +++ b/src/ida/IdaCallGraph.ts @@ -0,0 +1,153 @@ +import { + AnalysisCancelledError, + AnalysisOutputError, + AnalysisProtocolError, +} from "../domain/analysisErrorCore.js"; +import type { BinaryTarget } from "../domain/binaryTargetTypes.js"; +import type { JsonValue } from "../domain/jsonValue.js"; +import { ZodError } from "zod"; +import { + idaAddressSchema, + modernCallGraphSchema, + modernLookupBatchSchema, +} from "./IdaProtocolValues.js"; + +/** The graph only emits edges whose destinations IDA recognizes as functions. */ +export const UNRECOGNIZED_CALLEE_LIMITATION = + "This relationship read does not establish calls to targets IDA does not recognize as functions; imports and other non-function call targets remain unknown."; +const TRUNCATED_LIMITATION = + "Upstream callgraph was truncated; additional callees remain unknown."; +const POSSIBLY_TRUNCATED_LIMITATION = + "Upstream callgraph may be truncated; additional callees may remain unknown."; +const expectedReadFailure = (cause: unknown): boolean => + cause instanceof AnalysisProtocolError || + cause instanceof AnalysisOutputError || + cause instanceof ZodError; + +const errorReason = (cause: unknown): string => + cause instanceof Error ? cause.message : "Unknown upstream read failure"; + +const X86_FAMILIES: ReadonlySet = new Set(["x86", "x86_64"]); + +/** + * True only when every slice the file can load is x86-family. Upstream `callees` + * recognises x86 call instructions only, so any other or unproven processor + * must use the processor-neutral call graph for function edges. + */ +export const isX86OnlyTarget = (target: BinaryTarget): boolean => + target.kind === "executable" && + // IDA can load a managed PE with the CLI loader instead of as x86 code. + !(target.format === "pe" && target.managed) && + X86_FAMILIES.has(target.architecture) && + target.availableArchitectures.every((family) => X86_FAMILIES.has(family)); + +/** Callees observed in a root-only call graph, with what stays unknown. */ +export interface CallGraphCallees { + readonly callees: readonly { + readonly address: string; + readonly name: string; + }[]; + readonly limitations: readonly string[]; +} + +/** No function edges can be observed when the server lacks the graph tool. */ +export const unavailableCallGraph = (reason: string): CallGraphCallees => ({ + callees: [], + limitations: [ + UNRECOGNIZED_CALLEE_LIMITATION, + `Upstream callgraph was unavailable (${reason}); callee edges remain unknown.`, + ], +}); + +/** + * Read only the outgoing call edges of one procedure from upstream `callgraph`. + * `max_depth: 0` includes the root's edges without scanning child functions. + * Any failure to obtain a consistent root graph is a producer fault that + * fails the read, like a failing `callees` request: a tool error, a truncated + * transport reply, an error row, an unreadable reply and a reply for another + * root all fail the read. Cancellation still propagates. Truncation the upstream + * declares inside a valid reply stays a limitation. + */ +export const readCallGraphCallees = async ( + call: ( + tool: string, + args: Readonly>, + ) => Promise, + address: string, +): Promise => { + const raw = await call("callgraph", { roots: [address], max_depth: 0 }); + const parsed = modernCallGraphSchema.safeParse(raw); + const row = parsed.success ? parsed.data[0] : undefined; + if (row === undefined) + throw new AnalysisProtocolError("IDA callgraph reply was unreadable."); + if (row.error !== undefined) + throw new AnalysisProtocolError(`IDA callgraph failed: ${row.error}`); + const rootAddress = idaAddressSchema.safeParse(row.root); + if (!rootAddress.success || rootAddress.data !== address) + throw new AnalysisProtocolError( + `IDA callgraph answered for a different root than ${address}.`, + ); + + const edges = row.edges.filter( + ({ type, from }) => type === "call" && from === address, + ); + const targets = [...new Set(edges.map(({ to }) => to))].sort((left, right) => + BigInt(left) < BigInt(right) ? -1 : BigInt(left) > BigInt(right) ? 1 : 0, + ); + const limitations = [UNRECOGNIZED_CALLEE_LIMITATION]; + const truncated = + row.truncated === true || + row.per_func_capped === true || + (row.limit_reason ?? null) !== null; + // Reaching the per-function cap can stop the root scan without setting a + // truncation flag, so a saturated root count is possible truncation. + const saturated = + row.max_edges_per_func !== undefined && + edges.length >= row.max_edges_per_func; + if (truncated) limitations.push(TRUNCATED_LIMITATION); + else if (saturated) limitations.push(POSSIBLY_TRUNCATED_LIMITATION); + + const names = new Map(); + if (targets.length > 0) { + try { + const raw = await call("lookup_funcs", { queries: targets }); + const parsed = modernLookupBatchSchema.safeParse(raw); + if (!parsed.success) + limitations.push( + "IDA observed callee addresses but returned an unreadable name lookup result; the addresses are retained.", + ); + else { + const rows = new Map(); + for (const item of parsed.data) { + const query = idaAddressSchema.safeParse(item.query); + if (!query.success || !targets.includes(query.data)) continue; + rows.set(query.data, rows.has(query.data) ? null : item); + } + for (const target of targets) { + const item = rows.get(target); + if (item?.fn?.address === target && item.error === null) + names.set(target, item.fn.name); + else + limitations.push( + `IDA observed callee address ${target} but could not resolve its name${item?.error == null ? "" : `: ${item.error}`}; the address is retained.`, + ); + } + } + } catch (cause: unknown) { + if (cause instanceof AnalysisCancelledError) throw cause; + if (!expectedReadFailure(cause)) throw cause; + limitations.push( + `IDA observed callee addresses but the name lookup failed (${errorReason(cause)}); the addresses are retained.`, + ); + } + } + return { + callees: targets.map((to) => { + return { + address: to, + name: names.get(to) ?? to, + }; + }), + limitations, + }; +}; diff --git a/src/ida/IdaCalleeRelationships.test.ts b/src/ida/IdaCalleeRelationships.test.ts new file mode 100644 index 000000000..9bab2074a --- /dev/null +++ b/src/ida/IdaCalleeRelationships.test.ts @@ -0,0 +1,222 @@ +import { rm } from "node:fs/promises"; +import { afterEach, describe, expect, it } from "vitest"; +import { z } from "zod"; +import { + createIdaTarget, + RecordingIdaMcp, +} from "../../tests/fixtures/idaMcp.js"; +import { + callGraphSession as session, + edge, + graphRow, + removeCallGraphSessions, +} from "../../tests/fixtures/idaCallGraph.js"; +import { UNRECOGNIZED_CALLEE_LIMITATION } from "./IdaCallGraph.js"; +import { IDA_LIMITATIONS } from "./IdaProviderCapabilities.js"; +import { IdaSessionClient } from "./IdaSessionClient.js"; + +const attachedRoots: string[] = []; +afterEach(async () => { + await removeCallGraphSessions(); + for (const root of attachedRoots.splice(0)) + await rm(root, { recursive: true, force: true }); +}); + +const observedTools = (rawResult: unknown): string[] => + z + .object({ observations: z.array(z.object({ tool: z.string() })) }) + .parse(rawResult) + .observations.map(({ tool }) => tool); + +const callees = async (client: IdaSessionClient) => { + const result = await client.execute("procedure_callees", { + procedure: "main", + }); + if (!result.ok) throw result.error; + return result.value; +}; + +describe("headless IDA callee edges on non-x86 targets", () => { + it("uses the neutral call graph as the sole callee source on non-x86 targets", async () => { + const { client } = await session({ + callees: [ + { addr: "0x4000", name: "bogus_non_x86_match", type: "internal" }, + ], + callgraph: graphRow([ + edge("0x1000", "0x2000"), + edge("0x1000", "0x3000"), + edge("0x2000", "0x4000"), + ]), + }); + try { + const execution = await callees(client); + + expect( + execution.result, + "the root edges are observed and the child edge is ignored", + ).toEqual(["0x2000", "0x3000"]); + expect( + execution.limitations, + "calls to non-function targets remain unknown", + ).toContain(UNRECOGNIZED_CALLEE_LIMITATION); + expect( + execution.result, + "the unreliable callees row is neither resolved nor reported", + ).not.toContain("0x4000"); + expect( + observedTools(execution.rawResult), + "non-x86 callees is never queried", + ).not.toContain("callees"); + } finally { + expect((await client.close()).ok).toBe(true); + } + }); +}); + +describe("headless IDA call graph availability and cancellation", () => { + it("takes the call graph path for a universal target whose loaded slice is unknown", async () => { + const { client, producer } = await session( + { callees: [], callgraph: graphRow([edge("0x1000", "0x2000")]) }, + "x86_64", + ["x86_64", "arm64"], + ); + try { + const execution = await callees(client); + + expect(execution.result, "the graph edge is observed").toEqual([ + "0x2000", + ]); + expect( + producer.calls.some(({ name }) => name === "callgraph"), + "a mixed universal target cannot prove an x86 slice", + ).toBe(true); + } finally { + expect((await client.close()).ok).toBe(true); + } + }); + + it("stops at once when the caller cancels during the call graph request", async () => { + const controller = new AbortController(); + const { client, producer } = await session({ + callees: [], + callgraph: graphRow([edge("0x1000", "0x2000")]), + }); + const inner = producer.beforeCall; + producer.beforeCall = async (name, args) => { + await inner?.(name, args); + if (name === "callgraph") controller.abort(); + }; + try { + const result = await client.execute( + "procedure_callees", + { procedure: "main" }, + { signal: controller.signal }, + ); + + if (result.ok) throw new Error("expected cancellation"); + expect(result.error._tag, "cancellation is not swallowed").toBe( + "AnalysisCancelledError", + ); + } finally { + expect((await client.close()).ok).toBe(true); + } + }); +}); + +describe("headless IDA callee edges on x86 targets", () => { + it("keeps the requests, result, limitations and observations of an x86 read unchanged", async () => { + const { client, producer } = await session( + { + callees: [ + { addr: "0x2000", name: "imported_puts", type: "external" }, + { addr: "0x3000", name: "printf_stub", type: "internal" }, + ], + }, + "x86_64", + ); + try { + const execution = await callees(client); + + expect( + execution.result, + "externals come first, then resolved internals, exactly as before", + ).toEqual(["0x2000", "0x3000"]); + expect( + execution.limitations, + "exactly the provider's own limitations, nothing added on x86", + ).toEqual([...IDA_LIMITATIONS]); + expect( + producer.calls.map(({ name }) => name), + "no call graph request is made on an x86-only target", + ).toEqual([ + "idb_open", + "idb_list", + "server_health", + "lookup_funcs", + "callees", + "lookup_funcs", + "server_health", + ]); + expect( + producer.calls.find(({ name }) => name === "callees")?.args, + "the callees request is the unchanged scoped one", + ).toMatchObject({ addrs: ["0x1000"], limit: 500 }); + expect( + observedTools(execution.rawResult), + "the raw observations hold no call graph row", + ).toEqual(["lookup_funcs", "callees", "lookup_funcs"]); + } finally { + expect((await client.close()).ok).toBe(true); + } + }); + + it("keeps an empty x86 callees reply an observed empty list without any limitation", async () => { + const { client, producer } = await session({ callees: [] }, "x86_64"); + try { + const execution = await callees(client); + + expect(execution.result, "nothing is added").toEqual([]); + expect( + execution.limitations, + "no completeness warning on a trustworthy x86 target", + ).toEqual([...IDA_LIMITATIONS]); + expect( + producer.calls.some(({ name }) => name === "callgraph"), + "no call graph request even for an empty reply", + ).toBe(false); + } finally { + expect((await client.close()).ok).toBe(true); + } + }); + + it("does not use the call graph for the legacy attached profile on an ARM target", async () => { + const { root, target } = await createIdaTarget("arm64"); + attachedRoots.push(root); + const producer = new RecordingIdaMcp(target, "attached"); + const client = new IdaSessionClient( + { + command: "fixture", + args: [], + env: {}, + mode: "attached", + timeoutMs: 1000, + workspaceRoot: root, + }, + target, + producer, + ); + try { + const result = await client.execute("procedure_callees", { + procedure: "main", + }); + + expect(result.ok, "the legacy read still works").toBe(true); + expect( + producer.calls.some(({ name }) => name === "callgraph"), + "the legacy profile has no call graph source", + ).toBe(false); + } finally { + expect((await client.close()).ok).toBe(true); + } + }); +}); diff --git a/src/ida/IdaOperationRunner.ts b/src/ida/IdaOperationRunner.ts index 51de0352f..6099f379e 100644 --- a/src/ida/IdaOperationRunner.ts +++ b/src/ida/IdaOperationRunner.ts @@ -10,6 +10,7 @@ import { functionInstructionWindowSchema, } from "../domain/hopperValues.js"; import type { JsonValue } from "../domain/jsonValue.js"; +import { readCallGraphCallees, unavailableCallGraph } from "./IdaCallGraph.js"; import type { IdaMcpConnection } from "./IdaMcpConnection.js"; import type { IdaOperation } from "./IdaProviderCapabilities.js"; import { IDA_LIMITATIONS } from "./IdaProviderCapabilities.js"; @@ -68,20 +69,26 @@ export class IdaOperationRunner { private readonly connection: IdaMcpConnection, private readonly protocol: "legacy" | "modern", private readonly database?: string, - private readonly signal?: AbortSignal, - private readonly selectorOrigin?: "provider", + private readonly options: { + readonly signal?: AbortSignal | undefined; + readonly selectorOrigin?: "provider" | undefined; + /** False when the loaded processor may not be x86, so `callees` cannot be trusted alone. */ + readonly x86Only?: boolean; + /** False when the upstream server does not advertise `callgraph`. */ + readonly callGraphAvailable?: boolean; + } = {}, ) {} async call( tool: string, args: Readonly>, ): Promise { - if (this.signal?.aborted) throw new AnalysisCancelledError(tool); + if (this.options.signal?.aborted) throw new AnalysisCancelledError(tool); const scoped = this.database === undefined ? args : { ...args, database: this.database }; const result = await this.connection.call(tool, scoped); this.raw.push({ tool, arguments: scoped, result }); - if (this.signal?.aborted) throw new AnalysisCancelledError(tool); + if (this.options.signal?.aborted) throw new AnalysisCancelledError(tool); return result; } @@ -175,6 +182,33 @@ export class IdaOperationRunner { return entry; } + /** True only for modern callee reads on targets not proven x86-only; x86 reads never reach the graph code. */ + private readsCallGraph(direction: "callers" | "callees"): boolean { + return ( + this.protocol === "modern" && + direction === "callees" && + this.options.x86Only === false + ); + } + + /** Add processor-neutral function callees on targets where upstream `callees` is unsound. */ + private async addCallGraphCallees( + fn: IdaFunction, + result: Map>, + ): Promise { + const graph = + this.options.callGraphAvailable === false + ? unavailableCallGraph("the upstream server does not advertise it") + : await readCallGraphCallees( + (tool, args) => this.call(tool, args), + fn.address, + ); + this.limitations.push(...graph.limitations); + for (const callee of graph.callees) + if (!result.has(callee.address)) + result.set(callee.address, identity({ ...callee, size: "unknown" })); + } + async strings(): Promise<{ address: string; value: string }[]> { const result: { address: string; value: string }[] = []; let offset: number | null = 0; @@ -310,6 +344,10 @@ export class IdaOperationRunner { ): Promise[]> { const queries: string[] = []; const result = new Map>(); + if (this.readsCallGraph(direction)) { + await this.addCallGraphCallees(fn, result); + return [...result.values()]; + } if (this.protocol === "legacy") { const value = await this.call( direction === "callers" ? "get_callers" : "get_callees", @@ -388,7 +426,7 @@ export class IdaOperationRunner { } const fn = await this.resolve( z.string().min(1).parse(parameters.procedure), - this.selectorOrigin === "provider" ? undefined : operation, + this.options.selectorOrigin === "provider" ? undefined : operation, ); this.locations.push({ kind: "address", address: fn.address }); switch (operation) { diff --git a/src/ida/IdaProtocolValues.ts b/src/ida/IdaProtocolValues.ts index ade459a20..e786bed7c 100644 --- a/src/ida/IdaProtocolValues.ts +++ b/src/ida/IdaProtocolValues.ts @@ -85,15 +85,13 @@ export const modernPageSchema = z ) .length(1); -export const modernLookupSchema = z - .array( - z.object({ - query: z.string(), - fn: idaFunctionSchema.nullable(), - error: z.string().nullable(), - }), - ) - .length(1); +const modernLookupRowSchema = z.object({ + query: z.string(), + fn: idaFunctionSchema.nullable(), + error: z.string().nullable(), +}); +export const modernLookupBatchSchema = z.array(modernLookupRowSchema); +export const modernLookupSchema = modernLookupBatchSchema.length(1); const instructionSchema = z.union([ z.object({ @@ -182,6 +180,31 @@ export const modernCalleesSchema = z }), ) .length(1); +/** Depth-one `callgraph` reply: one row per root, with function-only call edges. */ +export const modernCallGraphSchema = z + .array( + z + .object({ + root: z.string(), + error: z.string().optional(), + nodes: z.array( + z.object({ addr: idaAddressSchema, name: z.string().nullable() }), + ), + edges: z.array( + z.object({ + from: idaAddressSchema, + to: idaAddressSchema, + type: z.string(), + }), + ), + truncated: z.boolean().optional(), + limit_reason: z.string().nullable().optional(), + max_edges_per_func: z.number().int().positive().optional(), + per_func_capped: z.boolean().optional(), + }) + .passthrough(), + ) + .length(1); export const legacyCalleesSchema = z.array( z.object({ address: idaAddressSchema, name: z.string() }).passthrough(), ); diff --git a/src/ida/IdaSessionClient.ts b/src/ida/IdaSessionClient.ts index 8de42a67d..9b502e246 100644 --- a/src/ida/IdaSessionClient.ts +++ b/src/ida/IdaSessionClient.ts @@ -21,6 +21,7 @@ import { err, ok, type Result } from "../domain/result.js"; import type { IdaConfiguration } from "./IdaConfiguration.js"; import type { IdaMcpConnection } from "./IdaMcpConnection.js"; import { readSettledHealth } from "./IdaBusyHealth.js"; +import { isX86OnlyTarget } from "./IdaCallGraph.js"; import { IdaOperationRunner } from "./IdaOperationRunner.js"; import { IDA_LIMITATIONS, @@ -103,6 +104,7 @@ export class IdaSessionClient implements AnalysisClient { #databaseCloseAcknowledged = false; #openAttempted = false; #openCompleted = false; + #callGraphAdvertised = false; #initialIdentity: string | undefined; #metadata: JsonValue = null; #lifecycle: JsonValue[] = []; @@ -227,8 +229,12 @@ export class IdaSessionClient implements AnalysisClient { this.connection, this.config.mode === "attached" ? "legacy" : "modern", this.#database, - options?.signal, - options?.selectorOrigin, + { + signal: options?.signal, + selectorOrigin: options?.selectorOrigin, + x86Only: isX86OnlyTarget(this.target), + callGraphAvailable: this.#callGraphAdvertised, + }, ); } @@ -290,6 +296,7 @@ export class IdaSessionClient implements AnalysisClient { async #initialize(): Promise { const tools = new Set(await this.connection.connect()); + this.#callGraphAdvertised = tools.has("callgraph"); const expected = this.config.mode === "attached" ? legacyTools : modernTools; const missing = expected.filter((name) => !tools.has(name)); diff --git a/tests/boundary/mcp/enhancedTools.callPathsAndBatches.test.ts b/tests/boundary/mcp/enhancedTools.callPathsAndBatches.test.ts index 5ed8e9c52..625bc27fd 100644 --- a/tests/boundary/mcp/enhancedTools.callPathsAndBatches.test.ts +++ b/tests/boundary/mcp/enhancedTools.callPathsAndBatches.test.ts @@ -1,19 +1,82 @@ import { afterEach, describe, expect, it } from "vitest"; import type { AnalysisOperationPort } from "../../../src/application/AnalysisProvider.js"; +import { createAnalysisExecution } from "../../../src/application/AnalysisProvider.js"; import { EnhancedTools } from "../../../src/application/EnhancedTools.js"; import { AnalysisCapabilityUnavailableError, AnalysisOutputError, } from "../../../src/domain/analysisErrorCore.js"; -import { err } from "../../../src/domain/result.js"; +import { err, ok as resultOk } from "../../../src/domain/result.js"; import { closeEnhancedToolResources, connect } from "./enhancedToolsHarness.js"; import { observed as ok } from "../../fixtures/analysisExecution.js"; afterEach(closeEnhancedToolResources); +const relationshipExecution = ( + result: unknown, + limitations: readonly string[] = [], +) => + resultOk( + createAnalysisExecution( + result, + { id: "fixture", name: "Fixture analysis provider", version: "1" }, + { limitations }, + ), + ); + describe("enhanced MCP tools", () => { + it("preserves per-node relationship limitations and keeps observed edges", async () => { + const limitation = "ARM64 call targets are only partially resolved."; + const client = await connect({ + execute: async (name, arguments_) => { + if (name === "procedure_address") + return relationshipExecution(String(arguments_.procedure)); + if (name !== "procedure_callees") + throw new Error(`Unexpected operation: ${name}`); + return relationshipExecution( + arguments_.procedure === "0x1" ? ["0x2"] : [], + arguments_.procedure === "0x1" ? [limitation] : [], + ); + }, + }); + const graph = await client.callTool({ + name: "get_call_graph", + arguments: { address: "0x1", direction: "forward" }, + }); + expect(graph.isError).not.toBe(true); + expect(graph.structuredContent).toMatchObject({ + normalized_result: { + "0": [ + { + address: "0x1", + status: "ok", + calls: ["0x2"], + limitations: [limitation], + }, + ], + }, + }); + + const trace = await client.callTool({ + name: "trace_call_path", + arguments: { start: "0x1", goal: "0x2" }, + }); + expect(trace.isError).not.toBe(true); + expect(trace.structuredContent).toMatchObject({ + normalized_result: { + goal_status: "reached", + traversal_path: ["0x1", "0x2"], + edges: [{ source_address: "0x1", target_address: "0x2" }], + truncated: true, + residual_unknowns: [ + `Call relationships for 0x1 are limited: ${limitation}`, + ], + }, + }); + }); + it("returns every batch item in caller order with its own result", async () => { const addresses = Array.from( { length: 37 }, diff --git a/tests/boundary/providers/ida/calleeRelationships.test.ts b/tests/boundary/providers/ida/calleeRelationships.test.ts new file mode 100644 index 000000000..ec46be03d --- /dev/null +++ b/tests/boundary/providers/ida/calleeRelationships.test.ts @@ -0,0 +1,291 @@ +import { afterEach, expect, it } from "vitest"; +import { z } from "zod"; +import { EnhancedTools } from "../../../../src/application/EnhancedTools.js"; +import { + callGraphSession, + edge, + graphRow, + removeCallGraphSessions, +} from "../../../fixtures/idaCallGraph.js"; + +afterEach(removeCallGraphSessions); + +const requestedRoots = (producer: { + calls: { name: string; args: unknown }[]; +}) => + producer.calls + .filter(({ name }) => name === "callgraph") + .map( + ({ args }) => + z.object({ roots: z.array(z.string()) }).parse(args).roots[0], + ) + .sort(); + +const armSession = () => + callGraphSession({ + callees: [], + callgraph: graphRow([edge("0x1000", "0x2000"), edge("0x1000", "0x3000")]), + leafRoots: ["0x2000", "0x3000"], + }); + +it("shows the callee edges of an ARM64 procedure in the call graph", async () => { + const { client, producer } = await armSession(); + try { + const result = await new EnhancedTools(client).execute("get_call_graph", { + address: "main", + }); + + if (!result.ok) throw result.error; + expect( + result.value, + "the root lists both observed callees instead of an empty observed list", + ).toMatchObject({ + "0": [ + { + address: "0x1000", + status: "ok", + calls: ["0x2000", "0x3000"], + limitations: expect.arrayContaining([ + "This relationship read does not establish calls to targets IDA does not recognize as functions; imports and other non-function call targets remain unknown.", + ]), + }, + ], + "1": [ + { + address: "0x2000", + status: "ok", + calls: [], + limitations: expect.arrayContaining([ + "This relationship read does not establish calls to targets IDA does not recognize as functions; imports and other non-function call targets remain unknown.", + ]), + }, + { + address: "0x3000", + status: "ok", + calls: [], + limitations: expect.arrayContaining([ + "This relationship read does not establish calls to targets IDA does not recognize as functions; imports and other non-function call targets remain unknown.", + ]), + }, + ], + }); + expect( + requestedRoots(producer), + "each reachable procedure is read once, with exactly its own root", + ).toEqual(["0x1000", "0x2000", "0x3000"]); + } finally { + expect((await client.close()).ok).toBe(true); + } +}); + +it("reaches an ARM64 callee in a call path", async () => { + const { client } = await armSession(); + try { + const result = await new EnhancedTools(client).execute("trace_call_path", { + start: "main", + goal: "0x3000", + }); + + if (!result.ok) throw result.error; + expect( + result.value, + "the goal is reached through the observed edge", + ).toMatchObject({ + goal_status: "reached", + traversal_path: ["0x1000", "0x3000"], + }); + } finally { + expect((await client.close()).ok).toBe(true); + } +}); + +it("traverses every ARM64 callee in a call path without a goal and records no failure", async () => { + const { client, producer } = await armSession(); + try { + const result = await new EnhancedTools(client).execute("trace_call_path", { + start: "main", + }); + + if (!result.ok) throw result.error; + expect( + result.value, + "the whole reachable graph is visited with exactly the observed edges and no failure", + ).toMatchObject({ + goal_status: "not_requested", + nodes: [ + { address: "0x1000", depth: 0 }, + { address: "0x2000", depth: 1 }, + { address: "0x3000", depth: 1 }, + ], + edges: [ + { source_address: "0x1000", target_address: "0x2000" }, + { source_address: "0x1000", target_address: "0x3000" }, + ], + failures: [], + truncated: true, + limitations: expect.arrayContaining([ + "This relationship read does not establish calls to targets IDA does not recognize as functions; imports and other non-function call targets remain unknown.", + ]), + residual_unknowns: expect.arrayContaining([ + expect.stringContaining("Call relationships for 0x1000 are limited:"), + ]), + }); + expect( + requestedRoots(producer), + "each reachable procedure is read once, with exactly its own root", + ).toEqual(["0x1000", "0x2000", "0x3000"]); + } finally { + expect((await client.close()).ok).toBe(true); + } +}); + +it("preserves missing callgraph coverage through direct, graph, and trace workflows", async () => { + const { client, producer } = await callGraphSession({ + tools: [ + "idb_open", + "idb_list", + "idb_close", + "server_health", + "list_funcs", + "lookup_funcs", + "find_regex", + "decompile", + "disasm", + "callees", + "xrefs_to", + ], + callees: [ + { addr: "0x3000", name: "bogus_non_x86_match", type: "internal" }, + ], + }); + try { + const direct = await client.execute("procedure_callees", { + procedure: "main", + }); + if (!direct.ok) throw direct.error; + expect( + direct.value.result, + "the x86-only producer result is ignored", + ).toEqual([]); + expect(direct.value.limitations).toEqual( + expect.arrayContaining([ + expect.stringContaining("callgraph was unavailable"), + expect.stringContaining("targets IDA does not recognize as functions"), + ]), + ); + + const application = new EnhancedTools(client); + const graph = await application.execute("get_call_graph", { + address: "main", + }); + if (!graph.ok) throw graph.error; + expect(graph.value).toMatchObject({ + "0": [ + { + status: "ok", + calls: [], + limitations: expect.arrayContaining([ + expect.stringContaining("callgraph was unavailable"), + ]), + }, + ], + }); + + const trace = await application.execute("trace_call_path", { + start: "main", + goal: "0x3000", + }); + if (!trace.ok) throw trace.error; + expect(trace.value).toMatchObject({ + goal_status: "not_reached", + edges: [], + limitations: expect.arrayContaining([ + expect.stringContaining("callgraph was unavailable"), + ]), + residual_unknowns: expect.arrayContaining([ + expect.stringContaining("Call relationships for 0x1000 are limited:"), + ]), + }); + expect( + producer.calls.some(({ name }) => name === "callees"), + "the inaccurate fallback is never queried", + ).toBe(false); + expect( + producer.calls.some(({ name }) => name === "callgraph"), + "an unadvertised tool is never requested", + ).toBe(false); + } finally { + expect((await client.close()).ok).toBe(true); + } +}); + +it("records a failed ARM64 call graph read as the node's error, not as a complete node", async () => { + const { client } = await callGraphSession({ + callees: [], + callgraphTruncatedOutput: true, + }); + try { + const result = await new EnhancedTools(client).execute("get_call_graph", { + address: "main", + }); + + if (!result.ok) throw result.error; + expect( + result.value, + "the transport-truncated graph is a recorded failure and no callee list is invented", + ).toEqual({ + "0": [ + { + address: "0x1000", + status: "error", + error: expect.objectContaining({ code: "unreadable_output" }), + }, + ], + }); + } finally { + expect((await client.close()).ok).toBe(true); + } +}); + +it("preserves an explicit root cap below the edge limit through every relationship workflow", async () => { + const { client } = await callGraphSession({ + callees: [], + callgraph: graphRow([edge("0x1000", "0x2000")], { per_func_capped: true }), + leafRoots: ["0x2000"], + }); + const limitation = + "Upstream callgraph was truncated; additional callees remain unknown."; + try { + const direct = await client.execute("procedure_callees", { + procedure: "main", + }); + if (!direct.ok) throw direct.error; + expect(direct.value.limitations).toContain(limitation); + expect(direct.value.result).toEqual(["0x2000"]); + const application = new EnhancedTools(client); + const graph = await application.execute("get_call_graph", { + address: "main", + }); + if (!graph.ok) throw graph.error; + expect(graph.value).toMatchObject({ + "0": [ + { + calls: ["0x2000"], + limitations: expect.arrayContaining([limitation]), + }, + ], + }); + const trace = await application.execute("trace_call_path", { + start: "main", + }); + if (!trace.ok) throw trace.error; + expect(trace.value).toMatchObject({ + truncated: true, + residual_unknowns: expect.arrayContaining([ + `Call relationships for 0x1000 are limited: ${limitation}`, + ]), + }); + } finally { + expect((await client.close()).ok).toBe(true); + } +}); diff --git a/tests/fixtures/idaCallGraph.ts b/tests/fixtures/idaCallGraph.ts new file mode 100644 index 000000000..b8e44d28e --- /dev/null +++ b/tests/fixtures/idaCallGraph.ts @@ -0,0 +1,138 @@ +import { rm } from "node:fs/promises"; +import { decodeIdaToolResult } from "../../src/ida/IdaMcpConnection.js"; +import { IdaSessionClient } from "../../src/ida/IdaSessionClient.js"; +import type { JsonValue } from "../../src/domain/jsonValue.js"; +import { createIdaTarget, RecordingIdaMcp } from "./idaMcp.js"; + +const roots: string[] = []; + +/** Remove the workspaces of every session opened through `callGraphSession`. */ +export const removeCallGraphSessions = async (): Promise => { + for (const root of roots.splice(0)) + await rm(root, { recursive: true, force: true }); +}; + +export const FUNCTIONS: Readonly< + Record +> = { + main: { addr: "0x1000", name: "main" }, + "0x1000": { addr: "0x1000", name: "main" }, + "0x2000": { addr: "0x2000", name: "square" }, + "0x3000": { addr: "0x3000", name: "printf_stub" }, +}; + +export const edge = (from: string, to: string, type = "call") => ({ + from, + to, + type, +}); +export const graphRow = ( + edges: readonly ReturnType[], + extra: Readonly> = {}, +): JsonValue => [ + { + root: "0x1000", + nodes: [{ addr: "0x1000", name: "main", depth: 0 }], + edges: [...edges], + max_depth: 0, + truncated: false, + limit_reason: null, + max_edges_per_func: 200, + per_func_capped: false, + ...extra, + }, +]; + +/** A realistic reply for a callee that makes no further calls. */ +const leafGraphRow = (root: string): JsonValue => [ + { + root, + nodes: [{ addr: root, name: FUNCTIONS[root]?.name ?? null, depth: 0 }], + edges: [], + max_depth: 0, + truncated: false, + limit_reason: null, + max_edges_per_func: 200, + per_func_capped: false, + }, +]; + +export interface Upstream { + /** Roots, besides the procedure itself, that the fake answers with a leaf graph; any other root is rejected. */ + readonly leafRoots?: readonly string[]; + /** Advertise the fake server's tools; omit for the default full set. */ + readonly tools?: readonly string[]; + /** Items the fake `callees` tool reports for the procedure. */ + readonly callees: readonly JsonValue[]; + /** Raw `callgraph` reply, or undefined when the tool is absent/failing. */ + readonly callgraph?: JsonValue; + readonly callgraphFailure?: string; + /** Make the fake transport report a truncated preview for the `callgraph` reply. */ + readonly callgraphTruncatedOutput?: boolean; +} + +export const callGraphSession = async ( + upstream: Upstream, + architecture: "x86_64" | "arm64" = "arm64", + available: readonly ("x86_64" | "arm64")[] = [architecture], +) => { + const { root, target } = await createIdaTarget(architecture, available); + roots.push(root); + const producer = new RecordingIdaMcp(target, "headless"); + if (upstream.tools !== undefined) producer.tools = upstream.tools; + producer.beforeCall = async (name) => { + if (name !== "callgraph") return; + if (upstream.callgraphFailure !== undefined) + decodeIdaToolResult({ + content: [{ type: "text", text: upstream.callgraphFailure }], + isError: true, + }); + if (upstream.callgraphTruncatedOutput === true) + decodeIdaToolResult({ + content: [{ type: "text", text: "preview" }], + _meta: { ida_mcp: { output_truncated: true } }, + }); + }; + producer.overrideCall = (name, args): JsonValue | undefined => { + if (name === "lookup_funcs") { + const queries = args.queries; + const addresses = Array.isArray(queries) ? queries.map(String) : []; + return addresses.map((query) => { + const fn = FUNCTIONS[query]; + return fn === undefined + ? { query, fn: null, error: "Not found" } + : { query, fn: { ...fn, size: "0x20" }, error: null }; + }); + } + if (name === "callees") + return [{ callees: [...upstream.callees], more: false }]; + if (name === "callgraph") { + const requested = Array.isArray(args.roots) ? args.roots : []; + const rootKnown = + JSON.stringify(requested) === JSON.stringify(["0x1000"]) || + (requested.length === 1 && + upstream.leafRoots?.includes(String(requested[0])) === true); + if (!rootKnown || args.max_depth !== 0) + throw new Error( + `Unexpected callgraph request: ${JSON.stringify(args)}`, + ); + return String(requested[0]) === "0x1000" + ? upstream.callgraph + : leafGraphRow(String(requested[0])); + } + return undefined; + }; + const client = new IdaSessionClient( + { + command: "fixture", + args: [], + env: {}, + mode: "headless", + timeoutMs: 1000, + workspaceRoot: root, + }, + target, + producer, + ); + return { client, producer }; +}; diff --git a/tests/fixtures/idaMcp.ts b/tests/fixtures/idaMcp.ts index ae8c1f053..00202208d 100644 --- a/tests/fixtures/idaMcp.ts +++ b/tests/fixtures/idaMcp.ts @@ -7,19 +7,27 @@ import type { JsonValue } from "../../src/domain/jsonValue.js"; import type { IdaMcpConnection } from "../../src/ida/IdaMcpConnection.js"; import { elf } from "../../src/domain/binaryTarget.fixture.js"; +/** ELF machine numbers for the architectures a fixture target can declare. */ +const ELF_MACHINE = { x86_64: 62, arm64: 183 } as const; + /** Create a test-owned executable and workspace parent without invoking an engine. */ -export const createIdaTarget = async () => { +export const createIdaTarget = async ( + architecture: keyof typeof ELF_MACHINE = "x86_64", + availableArchitectures: readonly (keyof typeof ELF_MACHINE)[] = [ + architecture, + ], +) => { const root = await mkdtemp(join(tmpdir(), "rea-ida-test-")); const path = join(root, "target.elf"); - const bytes = elf(2, 1, 62); + const bytes = elf(2, 1, ELF_MACHINE[architecture]); await writeFile(path, bytes); const target: BinaryTarget = { path, sha256: createHash("sha256").update(bytes).digest("hex"), kind: "executable", format: "elf", - architecture: "x86_64", - availableArchitectures: ["x86_64"], + architecture, + availableArchitectures: [...availableArchitectures], }; return { root, target }; }; @@ -48,6 +56,7 @@ const modernNames = [ "disasm", "callees", "xrefs_to", + "callgraph", ]; /** Stateful producer fixture at the real connection seam; no module mocking. */ @@ -82,9 +91,11 @@ export class RecordingIdaMcp implements IdaMcpConnection { this.inputSha256 = target.sha256; this.inputPath = target.path; } + /** Tool names the fake server advertises; defaults to the profile's full set. */ + tools?: readonly string[]; async connect() { this.connects += 1; - return this.mode === "attached" ? legacyNames : modernNames; + return this.tools ?? (this.mode === "attached" ? legacyNames : modernNames); } serverInfo() { return { name: "upstream-fixture", version: "sdk-version-is-not-ida" };