diff --git a/README.md b/README.md index 2b438d05..ccbc78e1 100644 --- a/README.md +++ b/README.md @@ -213,7 +213,7 @@ The following table lists the available MCP functions for use: | `set_local_variable_type(function_address, variable_name, new_type)` | Set a local variable's type. | | `retype_variable` | Retype variable inside a given function. | | `search_functions_by_name` | Search for functions whose name contains the given substring. | -| `search_types(query, offset, count)` | Search local Types by substring (name/decl). | +| `search_types(query, offset, count, max_scan)` | Search local type names with bounded enumeration and an exact-name fast path. | | `set_comment` | Set a comment at a specific address. | | `set_function_comment` | Set a comment for a function. | | `set_function_prototype(name_or_address, prototype)` | Set a function's prototype by name or address. | @@ -243,7 +243,7 @@ These are the list of HTTP endpoints that can be called: - `/localTypes?offset=&limit=`: List local types. - `/strings?offset=&limit=`: Paginated strings. - `/strings/filter?offset=&limit=&filter=`: Filtered strings. -- `/searchTypes?query=&offset=&limit=`: Search local types by substring. +- `/searchTypes?query=&offset=&limit=&maxScan=`: Bounded type-name search. Defaults to 100 results and scanning at most 2000 names; exact names use direct lookup. - `/patch` or `/patchBytes?address=&data=&save_to_file=`: Patch raw bytes at an address (byte-level, not assembly). Can patch entire instructions by providing their bytecode. Address: hex (e.g., "0x401000") or decimal. Data: hex string (e.g., "90 90"). `save_to_file` (default True) saves to disk and re-signs on macOS. - `/renameVariables`: Batch rename locals in a function. Parameters: - Function: one of `functionAddress`, `address`, `function`, `functionName`, or `name`. diff --git a/bridge/binja_mcp_bridge.py b/bridge/binja_mcp_bridge.py index 19ee85ec..078a6dd7 100755 --- a/bridge/binja_mcp_bridge.py +++ b/bridge/binja_mcp_bridge.py @@ -514,10 +514,14 @@ def list_local_types(offset: int = 0, count: int = 200, include_libraries: bool @mcp.tool() def search_types( - query: str, offset: int = 0, count: int = 200, include_libraries: bool = False + query: str, + offset: int = 0, + count: int = 100, + max_scan: int = 2000, + include_libraries: bool = False, ) -> list: """ - Search local types whose name or declaration contains the substring. + Search local type names with bounded enumeration and an exact-name fast path. """ return safe_get( "searchTypes", @@ -525,6 +529,7 @@ def search_types( "query": query, "offset": offset, "limit": count, + "maxScan": max_scan, "includeLibraries": int(bool(include_libraries)), }, timeout=None, diff --git a/bridge/src/tools.ts b/bridge/src/tools.ts index b6ecb661..63cf8c29 100644 --- a/bridge/src/tools.ts +++ b/bridge/src/tools.ts @@ -364,18 +364,20 @@ export function registerTools(server: McpServer, client: BinjaHttpClient): void server.tool( "search_types", - "Search local types whose name or declaration contains the substring.", + "Search local type names with bounded enumeration and an exact-name fast path.", { query: z.string().describe("Search query"), offset: z.number().default(0).describe("Offset for pagination"), - count: z.number().default(200).describe("Number of results to return"), + count: z.number().int().min(1).max(1000).default(100).describe("Number of results to return"), + max_scan: z.number().int().min(1).max(20000).default(2000).describe("Maximum type names to scan"), include_libraries: z.boolean().default(false).describe("Include library types"), }, - async ({ query, offset = 0, count = 200, include_libraries = false }) => { + async ({ query, offset = 0, count = 100, max_scan = 2000, include_libraries = false }) => { const lines = await client.getLines("searchTypes", { query, offset, limit: count, + maxScan: max_scan, includeLibraries: include_libraries ? 1 : 0, }); return { content: [{ type: "text", text: lines.join("\n") }] }; diff --git a/plugin/__init__.py b/plugin/__init__.py index 2fd7b33f..72cc119d 100644 --- a/plugin/__init__.py +++ b/plugin/__init__.py @@ -603,7 +603,7 @@ class _MCPMaxUINotification(ui.UIContextNotification): def __init__(self): super().__init__() ui.UIContext.registerNotification(self) - + def _get_active_bv(self): try: ctx = ui.UIContext.activeContext() diff --git a/plugin/core/binary_operations.py b/plugin/core/binary_operations.py index 1ac91e18..765feae6 100644 --- a/plugin/core/binary_operations.py +++ b/plugin/core/binary_operations.py @@ -220,7 +220,7 @@ def list_open_binaries(self) -> list[dict[str, str]]: vb_canon = vb entries.append((canonical_id, fn, bool(vb_canon is self._current_view))) # Sort by filename for stable ordering - entries.sort(key=lambda t: (t[1] or "")) + entries.sort(key=lambda t: t[1] or "") for cid, fn, active in entries: items.append({"id": cid, "filename": fn, "active": active}) return items @@ -1647,31 +1647,147 @@ def add_type_entry(name, tobj): return results[offset : offset + limit] def search_local_types( - self, query: str, offset: int = 0, limit: int = 100, include_libraries: bool = False + self, + query: str, + offset: int = 0, + limit: int = 100, + include_libraries: bool = False, + max_scan: int = 2000, ) -> list[dict[str, Any]]: - """Search local/view types whose name or declaration contains the substring. - - Returns entries with {name, kind, type_class, decl}. - """ + """Search type names with bounded enumeration and exact-name fast path.""" if not self._current_view: raise RuntimeError("No binary loaded") if not query: return [] - ql = str(query).lower() - # Only local types by default (fast). Optionally include libraries. - all_types = self.list_local_types(0, 1_000_000, include_libraries=include_libraries) + + query_str = str(query).strip() + query_lower = query_str.lower() + offset = max(0, int(offset)) + limit = max(1, min(int(limit), 1000)) + max_scan = max(1, min(int(max_scan), 20_000)) matches: list[dict[str, Any]] = [] - for t in all_types: + seen_names: set[str] = set() + matched_count = 0 + scanned_count = 0 + + def describe_type(name: str, type_obj: Any) -> dict[str, Any]: + type_class = getattr(type_obj, "type_class", None) + kind = "unknown" try: - name = t.get("name") or "" - decl = t.get("decl") or "" - if (ql in str(name).lower()) or (ql in str(decl).lower()): - matches.append(t) + if type_class == TypeClass.StructureTypeClass: + variant = getattr(type_obj, "type", None) + if variant == StructureVariant.UnionStructureType: + kind = "union" + elif variant == StructureVariant.ClassStructureType: + kind = "class" + else: + kind = "struct" + elif type_class == TypeClass.EnumerationTypeClass: + kind = "enum" + elif type_class == TypeClass.NamedTypeReferenceClass: + kind = "typedef" + elif type_class == TypeClass.FunctionTypeClass: + kind = "function" except Exception: - continue - if isinstance(limit, int) and limit < 0: - return matches[offset:] - return matches[offset : offset + limit] + pass + + try: + declaration = str(type_obj) + except Exception: + declaration = None + return { + "name": name, + "kind": kind, + "type_class": str(type_class) if type_class is not None else None, + "decl": declaration, + } + + def add_match(name: Any, type_obj: Any) -> bool: + nonlocal matched_count + try: + name_str = str(name).strip() + except Exception: + return False + if not name_str or name_str in seen_names: + return False + seen_names.add(name_str) + if query_lower not in name_str.lower(): + return False + if matched_count >= offset: + matches.append(describe_type(name_str, type_obj)) + matched_count += 1 + return len(matches) >= limit + + # Exact lookups avoid enumerating a large type collection entirely. + get_type = getattr(self._current_view, "get_type_by_name", None) + if callable(get_type): + try: + exact_type = get_type(query_str) + if exact_type is not None: + if offset > 0: + return [] + return [describe_type(query_str, exact_type)] + except Exception: + pass + + def scan_entries(entries) -> bool: + nonlocal scanned_count + for name, type_obj in entries: + if scanned_count >= max_scan: + return True + scanned_count += 1 + if add_match(name, type_obj): + return True + return False + + try: + user_types = getattr( + getattr(self._current_view, "user_type_container", None), "types", None + ) + if user_types: + + def user_entries(): + for _type_id, entry in user_types.items(): + if isinstance(entry, (tuple, list)) and len(entry) >= 2: + yield entry[0], entry[1] + else: + yield getattr(entry, "name", None), getattr(entry, "type", entry) + + if scan_entries(user_entries()): + return matches + except Exception as e: + bn.log_warn(f"Error scanning user type names: {e}") + + try: + view_types = self._current_view.types + + def view_entries(): + for key, value in view_types.items(): + if isinstance(value, (tuple, list)) and len(value) >= 2: + yield value[0], value[1] + else: + yield getattr(value, "name", None) or key, value + + if scan_entries(view_entries()): + return matches + except Exception as e: + bn.log_warn(f"Error scanning view type names: {e}") + + if include_libraries and scanned_count < max_scan: + try: + platform = getattr(self._current_view, "platform", None) + for library in getattr(platform, "type_libraries", []) or []: + named_types = getattr(library, "named_types", None) + if not isinstance(named_types, dict): + continue + if scan_entries(named_types.items()): + return matches + except Exception as e: + bn.log_warn(f"Error scanning library type names: {e}") + + if scanned_count >= max_scan: + bn.log_warn(f"Type search stopped after {max_scan} names: query={query_str!r}") + return matches def get_type_info(self, name: str) -> dict[str, Any]: """Resolve a type by name and return detailed information. diff --git a/plugin/server/http_server.py b/plugin/server/http_server.py index 8a2df226..067b534d 100644 --- a/plugin/server/http_server.py +++ b/plugin/server/http_server.py @@ -379,33 +379,32 @@ def do_GET(self): 400, ) return - # support count=-1 to return all - eff_limit = ( - -1 - if (params.get("count") == "-1" or params.get("limit") == "-1") - else limit - ) + search_limit = max(1, min(limit, 1000)) + max_scan = parse_int_or_default(params.get("maxScan"), 2000) + max_scan = max(1, min(max_scan, 20_000)) include_libs = params.get("includeLibraries") in ( "1", "true", "True", ) - # First compute total - all_matches = self.binary_ops.search_local_types( - term, 0, -1, include_libraries=include_libs - ) - page = ( - all_matches[offset:] - if eff_limit < 0 - else all_matches[offset : offset + eff_limit] + self._set_request_stage("bounded_type_search") + page = self.binary_ops.search_local_types( + term, + offset, + search_limit, + include_libraries=include_libs, + max_scan=max_scan, ) + self._set_request_stage("format_response") self._send_json_response( { "types": page, "query": term, - "total": len(all_matches), + "total": None, "offset": offset, - "limit": eff_limit, + "limit": search_limit, + "maxScan": max_scan, + "bounded": True, "includeLibraries": include_libs, } ) diff --git a/tests/test_binary_operations.py b/tests/test_binary_operations.py new file mode 100644 index 00000000..eb76820a --- /dev/null +++ b/tests/test_binary_operations.py @@ -0,0 +1,119 @@ +import importlib.util +import sys +import types +import unittest +from pathlib import Path + + +def _load_binary_operations_module(): + info_logs = [] + binaryninja = types.ModuleType("binaryninja") + binaryninja.log_info = info_logs.append + binaryninja.log_warn = lambda _message: None + binaryninja.log_error = lambda _message: None + + def binaryninja_type(name): + value = type(name, (), {}) + setattr(binaryninja, name, value) + return value + + binaryninja.__getattr__ = binaryninja_type + sys.modules["binaryninja"] = binaryninja + + enums = types.ModuleType("binaryninja.enums") + enums.StructureVariant = type("StructureVariant", (), {}) + enums.TypeClass = type("TypeClass", (), {}) + sys.modules[enums.__name__] = enums + + for package in ("plugin", "plugin.core", "plugin.utils"): + module = types.ModuleType(package) + module.__path__ = [] + sys.modules[package] = module + + config = types.ModuleType("plugin.core.config") + config.BinaryNinjaConfig = object + sys.modules[config.__name__] = config + + string_utils = types.ModuleType("plugin.utils.string_utils") + string_utils.escape_non_ascii = lambda value: value + sys.modules[string_utils.__name__] = string_utils + + path = Path(__file__).parents[1] / "plugin" / "core" / "binary_operations.py" + spec = importlib.util.spec_from_file_location("plugin.core.binary_operations", path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module, info_logs + + +class ExactTypeSearchView: + def get_type_by_name(self, name): + if name == "logic_iobj": + return types.SimpleNamespace(type_class=None) + return None + + @property + def user_type_container(self): + raise AssertionError("exact type search must not enumerate user types") + + @property + def types(self): + raise AssertionError("exact type search must not enumerate view types") + + +class CountingTypeMap: + def __init__(self): + self.iterated = 0 + + def __bool__(self): + return True + + def items(self): + for index in range(100): + self.iterated += 1 + name = "needle_type" if index == 1 else f"type_{index}" + yield index, (name, types.SimpleNamespace(type_class=None)) + + +class BoundedTypeSearchView: + def __init__(self): + self.type_map = CountingTypeMap() + self.user_type_container = types.SimpleNamespace(types=self.type_map) + + def get_type_by_name(self, _name): + return None + + @property + def types(self): + raise AssertionError("bounded search must stop before scanning view types") + + +class BinaryOperationsTests(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.module, cls.info_logs = _load_binary_operations_module() + + def setUp(self): + self.info_logs.clear() + self.operations = self.module.BinaryOperations(object()) + + def test_exact_type_search_does_not_enumerate_types(self): + self.operations._current_view = ExactTypeSearchView() + + result = self.operations.search_local_types("logic_iobj") + + self.assertEqual(len(result), 1) + self.assertEqual(result[0]["name"], "logic_iobj") + + def test_type_search_stops_at_scan_limit(self): + view = BoundedTypeSearchView() + self.operations._current_view = view + + result = self.operations.search_local_types("needle", limit=10, max_scan=3) + + self.assertEqual([entry["name"] for entry in result], ["needle_type"]) + self.assertLessEqual(view.type_map.iterated, 4) + + +if __name__ == "__main__": + unittest.main()