diff --git a/tests/test_havoc_debug_protocol.py b/tests/test_havoc_debug_protocol.py index 5a887a0..ea660b3 100644 --- a/tests/test_havoc_debug_protocol.py +++ b/tests/test_havoc_debug_protocol.py @@ -1,5 +1,13 @@ import pytest -from tools.havoc_debug_protocol import evaluate_condition +from tools.havoc_debug_protocol import decode_message, encode_message, evaluate_condition + +def test_encode_decode_message_round_trip(): + # encode_message/decode_message are the literal wire contract between the + # Windows-side debug agent and this Linux-side MCP server, built independently on + # two different machines -- prove they actually agree with each other. + op = "set_watchpoint" + params = {"addr": "0x47eb48", "size": 4, "mode": "write"} + assert decode_message(encode_message(op, params)) == {"op": op, "params": params} def test_register_condition_true(): registers = {"eax": 0x50} diff --git a/tests/test_havoc_mcp_server.py b/tests/test_havoc_mcp_server.py index dae26f3..71b1dfb 100644 --- a/tests/test_havoc_mcp_server.py +++ b/tests/test_havoc_mcp_server.py @@ -46,6 +46,86 @@ def test_client_raises_on_error_response(): client.call("set_watchpoint", {"addr": "0x5000", "size": 4, "mode": "write"}) +class _FakeCreateConnectionSocket: + """Stands in for the object socket.create_connection() returns, so tests can + capture the timeout it was called with without opening a real connection or + sleeping. Supports only what DebugAgentClient.call() actually uses.""" + + def __init__(self, response: dict): + self._unread = (json.dumps(response) + "\n").encode("utf-8") + + def __enter__(self): + return self + + def __exit__(self, *exc_info): + return False + + def sendall(self, data): + pass + + def recv(self, bufsize): + chunk, self._unread = self._unread, b"" + return chunk + + +def test_client_call_uses_a_sane_default_socket_timeout(monkeypatch): + # Regression guard for the previous hardcoded `timeout=60` on every op: normal, + # fast ops (get_status, set_breakpoint, ...) should still get some bounded default + # even though continue_execution now needs to ask for something much longer. + captured = {} + + def fake_create_connection(address, timeout=None): + captured["timeout"] = timeout + return _FakeCreateConnectionSocket({"ok": True, "result": {}}) + + monkeypatch.setattr(socket, "create_connection", fake_create_connection) + client = DebugAgentClient(host="127.0.0.1", port=1) + client.call("get_status", {}) + assert captured["timeout"] == 60 + + +def test_client_call_honors_socket_timeout_override_for_long_waits(monkeypatch): + # continue_execution can be asked to wait far longer than 60s (e.g. waiting for a + # user to alt-tab). The client must be able to use a socket timeout derived from + # that requested wait, not the fixed 60s default. + captured = {} + + def fake_create_connection(address, timeout=None): + captured["timeout"] = timeout + return _FakeCreateConnectionSocket({"ok": True, "result": {}}) + + monkeypatch.setattr(socket, "create_connection", fake_create_connection) + client = DebugAgentClient(host="127.0.0.1", port=1) + client.call("continue_execution", {"timeout": 300}, socket_timeout=330) + assert captured["timeout"] == 330 + + +class _RecordingDebugAgentClient: + """Fake DebugAgentClient that just records how it was called, so tests can verify + the continue_execution MCP tool derives the right socket_timeout without needing a + real debug agent or a real socket.""" + + def __init__(self): + self.calls = [] + + def call(self, op, params, socket_timeout=60): + self.calls.append((op, params, socket_timeout)) + return {} + + +def test_continue_execution_tool_derives_socket_timeout_from_requested_wait(): + fake_client = _RecordingDebugAgentClient() + server = build_mcp_server(fake_client) + # FastMCP's @tool() decorator returns the wrapped function unchanged, so the + # underlying callable is reachable straight off the tool manager for a direct, + # synchronous call -- no need to go through the async MCP session machinery. + continue_execution_fn = server._tool_manager.get_tool("continue_execution").fn + + continue_execution_fn(timeout=300) + + assert fake_client.calls == [("continue_execution", {"timeout": 300}, 330)] + + from tools.havoc_mcp_server import check_bearer_token def test_check_bearer_token_accepts_matching_token(): diff --git a/tools/havoc_debug_protocol.py b/tools/havoc_debug_protocol.py index e977863..8aaa486 100644 --- a/tools/havoc_debug_protocol.py +++ b/tools/havoc_debug_protocol.py @@ -71,7 +71,10 @@ class WatchpointSlots: """Tracks which of the 4 x86 hardware debug address registers (DR0-DR3) are in use and computes the DR7 control register value. Pure bookkeeping -- does not touch ctypes or an actual thread context; the debug agent applies compute_dr7()'s result - via Wow64SetThreadContext.""" + (and the DR0-DR3 addresses from addresses()) to the target thread via a native + x64 SetThreadContext call, which is what actually arms the hardware watchpoints. + Wow64SetThreadContext/Wow64GetThreadContext are used only for WOW64 (32-bit) + register readback, not for arming DR0-DR3/DR7 themselves.""" def __init__(self) -> None: self._slots: dict[int, tuple[int, int, str]] = {} # slot -> (addr, size, mode) diff --git a/tools/havoc_mcp_server.py b/tools/havoc_mcp_server.py index a9e8122..915ef35 100644 --- a/tools/havoc_mcp_server.py +++ b/tools/havoc_mcp_server.py @@ -23,8 +23,13 @@ class DebugAgentClient: self._host = host self._port = port - def call(self, op: str, params: dict) -> dict: - with socket.create_connection((self._host, self._port), timeout=60) as sock: + def call(self, op: str, params: dict, socket_timeout: float = 60) -> dict: + """`socket_timeout` bounds both connect time and every subsequent recv() on + this socket. The 60s default suits every op except continue_execution, which + can legitimately be asked to wait far longer (e.g. for a user to alt-tab) -- + callers that need to wait longer than that must pass a larger override + derived from the wait they actually requested.""" + with socket.create_connection((self._host, self._port), timeout=socket_timeout) as sock: sock.sendall(encode_message(op, params)) buf = b"" while not buf.endswith(b"\n"): @@ -107,7 +112,15 @@ def build_mcp_server(client: DebugAgentClient) -> FastMCP: @mcp.tool() def continue_execution(timeout: float) -> dict: - return client.call("continue_execution", {"timeout": timeout}) + """Resume execution until a breakpoint/watchpoint hits or `timeout` seconds + elapse. The debug agent is given `timeout + 10` seconds of headroom to + actually wait that long, so the client socket needs a correspondingly longer + timeout too -- otherwise a long wait (e.g. waiting for a user to alt-tab) + would make this call's own recv() time out while the agent is still + legitimately waiting.""" + return client.call( + "continue_execution", {"timeout": timeout}, socket_timeout=timeout + 30 + ) @mcp.tool() def read_memory(addr: str, length: int) -> dict: