From 3256f4b607e846ea6199ab70ec5adcc3b4672bbc Mon Sep 17 00:00:00 2001
From: pyr0ball <pyroballpcs@gmail.com>
Date: Sat, 4 Jul 2026 06:58:46 -0700
Subject: [PATCH] feat: add watchpoint slot/DR7 bookkeeping to MCP protocol
 module

---
 tests/test_havoc_debug_protocol.py | 47 ++++++++++++++++++++++++++++++
 tools/havoc_debug_protocol.py      | 47 ++++++++++++++++++++++++++++++
 2 files changed, 94 insertions(+)

diff --git a/tests/test_havoc_debug_protocol.py b/tests/test_havoc_debug_protocol.py
index 4410ef7..5a887a0 100644
--- a/tests/test_havoc_debug_protocol.py
+++ b/tests/test_havoc_debug_protocol.py
@@ -41,3 +41,50 @@ def test_comparison_operators():
 def test_unknown_register_raises():
     with pytest.raises(KeyError):
         evaluate_condition({"register": "bogus", "op": "==", "value": "0x0"}, {}, None)
+
+from tools.havoc_debug_protocol import WatchpointSlots, WatchpointLimitError
+
+def test_allocate_returns_free_slots_in_order():
+    slots = WatchpointSlots()
+    assert slots.allocate(0x47c48c, 4, "write") == 0
+    assert slots.allocate(0x47da88, 4, "readwrite") == 1
+
+def test_allocate_raises_past_four():
+    slots = WatchpointSlots()
+    slots.allocate(0x1000, 4, "write")
+    slots.allocate(0x2000, 4, "write")
+    slots.allocate(0x3000, 4, "write")
+    slots.allocate(0x4000, 4, "write")
+    with pytest.raises(WatchpointLimitError):
+        slots.allocate(0x5000, 4, "write")
+
+def test_release_frees_slot_for_reuse():
+    slots = WatchpointSlots()
+    slot0 = slots.allocate(0x1000, 4, "write")
+    slots.release(slot0)
+    assert slots.allocate(0x2000, 4, "write") == 0
+
+def test_compute_dr7_enables_local_bit_per_active_slot():
+    slots = WatchpointSlots()
+    slots.allocate(0x1000, 4, "write")  # slot 0 -> L0 = bit 0
+    dr7 = slots.compute_dr7()
+    assert dr7 & 0b1 == 0b1        # L0 enabled
+    assert dr7 & 0b100 == 0        # L1 not enabled
+
+def test_compute_dr7_encodes_write_mode_and_4byte_len():
+    slots = WatchpointSlots()
+    slots.allocate(0x1000, 4, "write")
+    dr7 = slots.compute_dr7()
+    rw0 = (dr7 >> 16) & 0b11
+    len0 = (dr7 >> 18) & 0b11
+    assert rw0 == 0b01   # write-only
+    assert len0 == 0b11  # 4 bytes
+
+def test_compute_dr7_encodes_readwrite_and_1byte_len():
+    slots = WatchpointSlots()
+    slots.allocate(0x2000, 1, "readwrite")
+    dr7 = slots.compute_dr7()
+    rw0 = (dr7 >> 16) & 0b11
+    len0 = (dr7 >> 18) & 0b11
+    assert rw0 == 0b11   # read/write
+    assert len0 == 0b00  # 1 byte
diff --git a/tools/havoc_debug_protocol.py b/tools/havoc_debug_protocol.py
index c0fa025..e977863 100644
--- a/tools/havoc_debug_protocol.py
+++ b/tools/havoc_debug_protocol.py
@@ -57,3 +57,50 @@ def encode_message(op: str, params: dict) -> bytes:
 def decode_message(line: bytes) -> dict:
     """Decode a single NDJSON line (request or response) back into a dict."""
     return json.loads(line.decode("utf-8"))
+
+
+class WatchpointLimitError(Exception):
+    """Raised when all 4 hardware debug address registers are already in use."""
+
+
+_MODE_RW_BITS = {"write": 0b01, "readwrite": 0b11}
+_SIZE_LEN_BITS = {1: 0b00, 2: 0b01, 4: 0b11}
+
+
+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."""
+
+    def __init__(self) -> None:
+        self._slots: dict[int, tuple[int, int, str]] = {}  # slot -> (addr, size, mode)
+
+    def allocate(self, addr: int, size: int, mode: str) -> int:
+        if mode not in _MODE_RW_BITS:
+            raise ValueError(f"mode must be 'write' or 'readwrite', got {mode!r}")
+        if size not in _SIZE_LEN_BITS:
+            raise ValueError(f"size must be 1, 2, or 4 bytes, got {size}")
+        for slot in range(4):
+            if slot not in self._slots:
+                self._slots[slot] = (addr, size, mode)
+                return slot
+        raise WatchpointLimitError("all 4 hardware debug registers (DR0-DR3) are in use")
+
+    def release(self, slot: int) -> None:
+        self._slots.pop(slot, None)
+
+    def addresses(self) -> dict[int, int]:
+        """slot -> address, for setting DR0-DR3 themselves."""
+        return {slot: addr for slot, (addr, _size, _mode) in self._slots.items()}
+
+    def compute_dr7(self) -> int:
+        dr7 = 0
+        for slot, (_addr, size, mode) in self._slots.items():
+            dr7 |= 1 << (slot * 2)  # local enable bit (L0=bit0, L1=bit2, L2=bit4, L3=bit6)
+            rw_bits = _MODE_RW_BITS[mode]
+            len_bits = _SIZE_LEN_BITS[size]
+            shift = 16 + slot * 4
+            dr7 |= rw_bits << shift
+            dr7 |= len_bits << (shift + 2)
+        return dr7