diff --git a/README.md b/README.md index 6f74391..96f76cc 100644 --- a/README.md +++ b/README.md @@ -142,6 +142,15 @@ For a binary that is already imported into the project, use `-process` instead: -preScript GAMCPStartServerScript.java "host=127.0.0.1" "port=8080" ``` +To keep a headless MCP session open after analysis completes, run the server as a post-script with wait mode: + +```bash +"$GHIDRA_INSTALL_DIR/support/analyzeHeadless" /tmp/ghidra-projects McpHeadless \ + -process binary_name \ + -scriptPath "$GHIDRASSISTMCP_EXT/ghidra_scripts" \ + -postScript GAMCPStartServerScript.java "host=127.0.0.1" "port=8080" "wait=true" +``` + MCP clients can connect to: ```text @@ -150,7 +159,7 @@ SSE messages: http://127.0.0.1:8080/message Streamable HTTP: http://127.0.0.1:8080/mcp ``` -The headless MCP server runs inside the `analyzeHeadless` JVM and uses the loaded `currentProgram`. Keep that process alive while clients are connected; when `analyzeHeadless` exits, the MCP server exits with it. +The headless MCP server runs inside the `analyzeHeadless` JVM and uses the loaded `currentProgram`. The server holds a program consumer while it is running so MCP requests do not race against program database closure. Use `wait=true` when you want `analyzeHeadless` to stay open for interactive MCP clients; cancel the script or terminate the process to stop the server. ## Available Tools @@ -179,6 +188,8 @@ GhidrAssistMCP provides 38 tools organized into categories. Several tools use an | `get_current_function` | Get function at current cursor position | | `get_function_stack_layout` | Get stack frame layout with variable offsets | | `get_basic_blocks` | Get basic block information for a function | +| `create_function` | Create/define a function at an address | +| `disassemble_at` | Disassemble code at an address | ### Binary Information @@ -275,7 +286,7 @@ Rename multiple symbols in one operation. | ------ | ----------- | | `list` | List all available data types | | `get_info` | Get detailed data type information and structure definitions | -| `set` | Set data type at a specific address | +| `set` | Set data type at a specific address, including arrays with `array_count` or suffix syntax like `int[16]` | | `delete` | Delete a data type by name (optionally scoped by `category`) | #### `bookmarks` - Bookmark Management Tool @@ -454,6 +465,54 @@ If multiple types share the same name across categories, pass `category` (or pas } ``` +### Set an Array Data Type + +```json +{ + "method": "tools/call", + "params": { + "name": "types", + "arguments": { + "action": "set", + "address": "0x00402000", + "data_type": "int[16]" + } + } +} +``` + +Equivalent form: + +```json +{ + "method": "tools/call", + "params": { + "name": "types", + "arguments": { + "action": "set", + "address": "0x00402000", + "data_type": "int", + "array_count": 16 + } + } +} +``` + +### Create a Function + +```json +{ + "method": "tools/call", + "params": { + "name": "create_function", + "arguments": { + "address": "0x00401000", + "name": "mainWndProc" + } + } +} +``` + ### Rename Function (Action-Based) ```json @@ -592,7 +651,7 @@ GhidrAssistMCP/ - `rename_symbol`: `target_type: function|data|variable` - `comments`: `action: get|set|list|remove` - `variables`: `action: list|rename|set_type|set_prototype` with `scope: auto|local|global` for rename -- `types`: `action: list|get_info|set|delete` +- `types`: `action: list|get|set|create_struct|create_enum|create_typedef|delete` - `bookmarks`: `action: list|set|remove` - `xrefs`: `address|function` with `include_calls` parameter diff --git a/src/main/java/ghidrassistmcp/GhidrAssistMCPHeadlessServer.java b/src/main/java/ghidrassistmcp/GhidrAssistMCPHeadlessServer.java index 7f17785..0201ba9 100644 --- a/src/main/java/ghidrassistmcp/GhidrAssistMCPHeadlessServer.java +++ b/src/main/java/ghidrassistmcp/GhidrAssistMCPHeadlessServer.java @@ -81,6 +81,10 @@ public class GhidrAssistMCPHeadlessServer { Msg.info(this, "Headless MCP server stopped"); } catch (Exception e) { Msg.error(this, "Error stopping headless MCP server", e); + } finally { + if (headlessBackend != null) { + headlessBackend.shutdown(); + } } server = null; headlessBackend = null; @@ -111,35 +115,64 @@ public class GhidrAssistMCPHeadlessServer { private static class HeadlessBackend extends GhidrAssistMCPBackend { private volatile Program currentProgram; + private final Object programConsumer = new Object(); HeadlessBackend(Program program) { super(); - this.currentProgram = program; - // Notify backend of the initial program - onProgramActivated(program); + setProgram(program); } - void setProgram(Program program) { - if (this.currentProgram != null) { - onProgramDeactivated(this.currentProgram); + synchronized void setProgram(Program program) { + Program oldProgram = this.currentProgram; + if (oldProgram != null) { + onProgramDeactivated(oldProgram); + releaseProgram(oldProgram); } this.currentProgram = program; if (program != null) { + if (!program.isClosed()) { + program.addConsumer(programConsumer); + } onProgramActivated(program); } } @Override public Program getCurrentProgram() { + if (currentProgram != null && currentProgram.isClosed()) { + Msg.warn(this, "Headless current program is closed; clearing program reference"); + currentProgram = null; + } return currentProgram; } @Override public List getAllOpenPrograms() { - if (currentProgram != null) { - return Collections.singletonList(currentProgram); + Program program = getCurrentProgram(); + if (program != null) { + return Collections.singletonList(program); } return Collections.emptyList(); } + + synchronized void shutdown() { + Program program = currentProgram; + currentProgram = null; + if (program != null) { + onProgramDeactivated(program); + releaseProgram(program); + } + getTaskManager().shutdown(); + } + + private void releaseProgram(Program program) { + try { + if (!program.isClosed() && program.isUsedBy(programConsumer)) { + program.release(programConsumer); + } + } catch (Exception e) { + Msg.warn(this, "Failed to release headless program consumer: " + e.getMessage(), e); + } + } } } diff --git a/src/main/java/ghidrassistmcp/scripts/GAMCPStartServerScript.java b/src/main/java/ghidrassistmcp/scripts/GAMCPStartServerScript.java index b35c178..c532c4d 100644 --- a/src/main/java/ghidrassistmcp/scripts/GAMCPStartServerScript.java +++ b/src/main/java/ghidrassistmcp/scripts/GAMCPStartServerScript.java @@ -24,8 +24,9 @@ public class GAMCPStartServerScript extends GhidraScript { String host = "localhost"; int port = 8080; + boolean waitForClients = false; - // Parse optional arguments: host=... port=... + // Parse optional arguments: host=... port=... wait=true|false String[] args = getScriptArgs(); if (args != null) { for (String arg : args) { @@ -37,6 +38,8 @@ public class GAMCPStartServerScript extends GhidraScript { } catch (NumberFormatException e) { Msg.warn(this, "Invalid port argument, using default 8080"); } + } else if (arg.startsWith("wait=")) { + waitForClients = Boolean.parseBoolean(arg.substring(5)); } } } @@ -46,11 +49,30 @@ public class GAMCPStartServerScript extends GhidraScript { if (mcpServer.isRunning()) { Msg.info(this, "MCP server already running, updating program reference"); mcpServer.setProgram(currentProgram); + if (waitForClients) { + waitUntilCancelled(mcpServer); + } return; } Msg.info(this, "Starting headless MCP server for: " + currentProgram.getName()); mcpServer.start(currentProgram, host, port); Msg.info(this, "Headless MCP server ready on " + host + ":" + port); + if (waitForClients) { + waitUntilCancelled(mcpServer); + } + } + + private void waitUntilCancelled(GhidrAssistMCPHeadlessServer mcpServer) { + Msg.info(this, "Headless MCP server wait mode enabled; cancel the script or terminate analyzeHeadless to stop"); + try { + while (!monitor.isCancelled() && mcpServer.isRunning()) { + Thread.sleep(1000); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + mcpServer.stop(); + } } } diff --git a/src/main/java/ghidrassistmcp/tools/DataTypeResolver.java b/src/main/java/ghidrassistmcp/tools/DataTypeResolver.java new file mode 100644 index 0000000..34a3bb7 --- /dev/null +++ b/src/main/java/ghidrassistmcp/tools/DataTypeResolver.java @@ -0,0 +1,125 @@ +package ghidrassistmcp.tools; + +import java.util.ArrayList; +import java.util.List; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +import ghidra.program.model.data.ArrayDataType; +import ghidra.program.model.data.DataType; +import ghidra.program.model.data.DataTypeManager; + +final class DataTypeResolver { + + private static final Pattern ARRAY_SUFFIX = Pattern.compile("^(.+)\\[(\\d+)]\\s*$"); + + private DataTypeResolver() { + } + + static Result resolve(DataTypeManager dtm, String dataTypeName, Object arrayCountValue) { + if (dataTypeName == null || dataTypeName.isBlank()) { + return Result.error("Data type name is required"); + } + + String baseTypeName = dataTypeName.trim(); + Integer arrayCount = parseArrayCount(arrayCountValue); + if (arrayCountValue != null && arrayCount == null) { + return Result.error("array_count must be an integer"); + } + + Matcher matcher = ARRAY_SUFFIX.matcher(baseTypeName); + if (matcher.matches()) { + if (arrayCount != null) { + return Result.error("Specify array size either with data_type suffix syntax or array_count, not both"); + } + baseTypeName = matcher.group(1).trim(); + arrayCount = Integer.parseInt(matcher.group(2)); + } + + if (arrayCount != null && arrayCount < 1) { + return Result.error("array_count must be >= 1"); + } + + DataType baseType = resolveBaseType(dtm, baseTypeName); + if (baseType == null) { + return Result.error("Data type not found: " + baseTypeName); + } + + if (arrayCount == null) { + return Result.ok(baseType); + } + + int elementLength = baseType.getLength(); + if (elementLength <= 0) { + return Result.error("Cannot build array of variable-length type '" + baseType.getName() + "'"); + } + + return Result.ok(new ArrayDataType(baseType, arrayCount, elementLength)); + } + + private static Integer parseArrayCount(Object value) { + if (value == null) { + return null; + } + if (value instanceof Number) { + double doubleValue = ((Number) value).doubleValue(); + int intValue = ((Number) value).intValue(); + return doubleValue == intValue ? intValue : null; + } + if (value instanceof String) { + try { + return Integer.parseInt(((String) value).trim()); + } catch (NumberFormatException e) { + return null; + } + } + return null; + } + + private static DataType resolveBaseType(DataTypeManager dtm, String dataTypeName) { + DataType dataType = null; + if (dataTypeName.startsWith("/")) { + dataType = dtm.getDataType(dataTypeName); + } + if (dataType == null) { + dataType = dtm.getDataType("/" + dataTypeName); + } + if (dataType == null) { + dataType = dtm.getDataType(dataTypeName); + } + if (dataType != null) { + return dataType; + } + + List allTypes = new ArrayList<>(); + dtm.getAllDataTypes(allTypes); + for (DataType candidate : allTypes) { + if (candidate.getName().equals(dataTypeName)) { + return candidate; + } + } + return null; + } + + static final class Result { + final DataType dataType; + final String errorMessage; + + private Result(DataType dataType, String errorMessage) { + this.dataType = dataType; + this.errorMessage = errorMessage; + } + + static Result ok(DataType dataType) { + return new Result(dataType, null); + } + + static Result error(String errorMessage) { + return new Result(null, errorMessage); + } + + boolean isError() { + return errorMessage != null; + } + } +} diff --git a/src/main/java/ghidrassistmcp/tools/SetDataTypeTool.java b/src/main/java/ghidrassistmcp/tools/SetDataTypeTool.java index bb8efa5..f554927 100644 --- a/src/main/java/ghidrassistmcp/tools/SetDataTypeTool.java +++ b/src/main/java/ghidrassistmcp/tools/SetDataTypeTool.java @@ -8,7 +8,6 @@ import java.util.Map; import ghidra.program.model.address.Address; import ghidra.program.model.data.DataType; -import ghidra.program.model.data.DataTypeManager; import ghidra.program.model.listing.Program; import ghidra.program.model.util.CodeUnitInsertionException; import ghidrassistmcp.McpTool; @@ -44,7 +43,11 @@ public class SetDataTypeTool implements McpTool { return new McpSchema.JsonSchema("object", Map.of( "address", new McpSchema.JsonSchema("string", null, null, null, null, null), - "data_type", new McpSchema.JsonSchema("string", null, null, null, null, null) + "data_type", new McpSchema.JsonSchema("string", null, null, null, null, null), + "array_count", Map.of( + "type", "integer", + "description", "Optional: wraps data_type in an array of this many elements. Alternatively use suffix syntax like 'int[16]'." + ) ), List.of("address", "data_type"), null, null, null); } @@ -76,32 +79,14 @@ public class SetDataTypeTool implements McpTool { .build(); } - // Find the data type - DataTypeManager dtm = currentProgram.getDataTypeManager(); - DataType dataType = dtm.getDataType("/" + dataTypeName); - if (dataType == null) { - // Try finding by name without path - dataType = dtm.getDataType(dataTypeName); - } - - if (dataType == null) { - // Search for type in all categories if simple lookup failed - var allTypes = new java.util.ArrayList(); - dtm.getAllDataTypes(allTypes); - for (DataType dt : allTypes) { - if (dt.getName().equals(dataTypeName)) { - dataType = dt; - break; - } - } - } - - if (dataType == null) { + DataTypeResolver.Result resolvedType = + DataTypeResolver.resolve(currentProgram.getDataTypeManager(), dataTypeName, arguments.get("array_count")); + if (resolvedType.isError()) { return McpSchema.CallToolResult.builder() - .addTextContent("Data type not found: " + dataTypeName + - ". Use built-in types like 'int', 'char', 'void*', etc.") + .addTextContent(resolvedType.errorMessage) .build(); } + DataType dataType = resolvedType.dataType; // Start transaction int transactionID = currentProgram.startTransaction("Set Data Type"); diff --git a/src/main/java/ghidrassistmcp/tools/TypesTool.java b/src/main/java/ghidrassistmcp/tools/TypesTool.java index f77b287..7443608 100644 --- a/src/main/java/ghidrassistmcp/tools/TypesTool.java +++ b/src/main/java/ghidrassistmcp/tools/TypesTool.java @@ -58,7 +58,8 @@ public class TypesTool implements McpTool { Map.entry("offset", Map.of("type", "integer", "description", "Pagination offset (for list)")), Map.entry("limit", Map.of("type", "integer", "description", "Pagination limit (for list, default 100)")), Map.entry("address", Map.of("type", "string", "description", "Address (for set action)")), - Map.entry("data_type", Map.of("type", "string", "description", "Data type name (for set action)")), + Map.entry("data_type", Map.of("type", "string", "description", "Data type name (for set action). Supports array suffix syntax like 'int[16]'.")), + Map.entry("array_count", Map.of("type", "integer", "description", "For set action: wraps data_type in an array of this many elements")), Map.entry("size", Map.of("type", "integer", "description", "Size in bytes (for create_struct/create_enum)")), Map.entry("packed", Map.of("type", "boolean", "description", "Enable packing (for create_struct)")), Map.entry("values", Map.of("type", "object", "description", "Enum values as name:value pairs (for create_enum)")), @@ -192,10 +193,10 @@ public class TypesTool implements McpTool { try { address = program.getAddressFactory().getAddress(addressStr); } catch (Exception e) { return result("Invalid address: " + addressStr); } - DataTypeManager dtm = program.getDataTypeManager(); - DataType dataType = dtm.getDataType("/" + dataTypeName); - if (dataType == null) dataType = dtm.getDataType(dataTypeName); - if (dataType == null) return result("Data type not found: " + dataTypeName); + DataTypeResolver.Result resolvedType = + DataTypeResolver.resolve(program.getDataTypeManager(), dataTypeName, arguments.get("array_count")); + if (resolvedType.isError()) return result(resolvedType.errorMessage); + DataType dataType = resolvedType.dataType; int txId = program.startTransaction("Set Data Type"); try {