Add HeapObserver: structured heap state for LLM exploit generation

Rich observations with chunk adjacency, freelist contents, corruption
events, reachable exploit primitives, and natural-language summaries.
Designed as the feedback signal for LLM-in-the-loop exploitation.

heaptrm/observe.py

@@ -0,0 +1,403 @@
+"""
+observe.py - Rich heap observability for LLM consumption.
+
+The core insight: LLMs generating exploits need structured, actionable
+feedback about heap state — not raw dumps, not ML scores. They need
+answers to:
+  - "Did my allocation land where I expected?"
+  - "Is chunk A adjacent to chunk B?"
+  - "What's in the tcache for size 0x40?"
+  - "Did my overflow corrupt the right field?"
+  - "What exploit primitives are currently reachable?"
+
+This module transforms raw harness JSONL into structured observations
+that an LLM can reason about.
+"""
+
+import json
+import os
+from dataclasses import dataclass, field
+from typing import List, Dict, Optional, Tuple
+from collections import defaultdict
+
+
+@dataclass
+class ChunkInfo:
+    index: int
+    address: str
+    size: int
+    state: str          # "allocated", "freed"
+    flags: dict
+    fd: Optional[str]   # forward pointer (freed chunks)
+    bk: Optional[str]   # backward pointer
+    fd_target: Optional[int]  # index of chunk fd points to
+    data_preview: str    # first 16 bytes hex
+    is_corrupted: bool
+    alloc_order: int
+    free_order: int
+
+
+@dataclass
+class BinInfo:
+    bin_type: str       # "tcache", "fastbin", "unsorted", "smallbin"
+    size_class: int
+    entries: List[int]  # chunk indices, head first
+
+
+@dataclass
+class Corruption:
+    step: int
+    type: str
+    chunk_index: int
+    detail: str
+
+
+@dataclass
+class Primitive:
+    name: str
+    description: str
+    ready: bool
+    requirements: List[str]
+    chunks_involved: List[int]
+
+
+@dataclass
+class HeapObservation:
+    """Complete structured observation of heap state for LLM consumption."""
+    step: int
+    operation: str
+
+    # Layout
+    chunks: List[ChunkInfo]
+    n_allocated: int
+    n_freed: int
+
+    # Bins
+    bins: List[BinInfo]
+
+    # Adjacency map: chunk_idx -> (prev_idx, next_idx)
+    adjacency: Dict[int, Tuple[Optional[int], Optional[int]]]
+
+    # Corruptions detected
+    corruptions: List[Corruption]
+    cumulative_corruptions: int
+
+    # Reachable exploit primitives
+    primitives: List[Primitive]
+
+    # Size class summary
+    size_classes: Dict[int, Dict[str, int]]  # size -> {alloc: N, freed: N}
+
+    # Human/LLM-readable summary
+    summary: str
+
+
+class HeapObserver:
+    """Transforms raw harness dumps into structured observations."""
+
+    def __init__(self):
+        self.cumulative_corruptions = 0
+        self.history: List[HeapObservation] = []
+
+    def observe(self, state: dict) -> HeapObservation:
+        """Convert a raw harness state to a structured observation."""
+        raw_chunks = state.get("chunks", [])
+        step = state.get("step", 0)
+        operation = state.get("operation", "unknown")
+
+        # Build chunk list
+        chunks = []
+        addr_to_idx = {}
+        for c in raw_chunks:
+            ci = ChunkInfo(
+                index=c.get("idx", len(chunks)),
+                address=c.get("addr", "0x0"),
+                size=c.get("chunk_size", 0),
+                state="allocated" if c.get("state") == 1 else "freed",
+                flags={"P": c.get("flag_p", 0), "M": c.get("flag_m", 0)},
+                fd=hex(c["fd"]) if c.get("fd", 0) != 0 else None,
+                bk=hex(c["bk"]) if c.get("bk", 0) != 0 else None,
+                fd_target=c.get("fd_idx") if c.get("fd_idx", -1) >= 0 else None,
+                data_preview=c.get("data_hex", "")[:32],
+                is_corrupted=c.get("is_corrupted", False),
+                alloc_order=c.get("alloc_order", 0),
+                free_order=c.get("free_order", 0),
+            )
+            chunks.append(ci)
+            addr_to_idx[c.get("addr", "")] = ci.index
+
+        n_alloc = sum(1 for c in chunks if c.state == "allocated")
+        n_freed = sum(1 for c in chunks if c.state == "freed")
+
+        # Build adjacency map (chunks sorted by address)
+        adjacency = {}
+        sorted_chunks = sorted(chunks, key=lambda c: int(c.address, 16) if c.address.startswith("0x") else 0)
+        for i, c in enumerate(sorted_chunks):
+            prev_idx = sorted_chunks[i-1].index if i > 0 else None
+            next_idx = sorted_chunks[i+1].index if i < len(sorted_chunks)-1 else None
+            adjacency[c.index] = (prev_idx, next_idx)
+
+        # Build bin info
+        bins = []
+        size_freed = defaultdict(list)
+        for c in chunks:
+            if c.state == "freed":
+                size_freed[c.size].append(c.index)
+        for size, indices in sorted(size_freed.items()):
+            bins.append(BinInfo(
+                bin_type="tcache" if size <= 0x410 else "unsorted",
+                size_class=size,
+                entries=indices,
+            ))
+
+        # Corruptions
+        corruptions = []
+        for c in state.get("corruptions", []):
+            corruptions.append(Corruption(
+                step=step,
+                type=c.get("type", "unknown"),
+                chunk_index=c.get("chunk_idx", -1),
+                detail=c.get("detail", ""),
+            ))
+        self.cumulative_corruptions += len(corruptions)
+
+        # Size class summary
+        size_classes = defaultdict(lambda: {"alloc": 0, "freed": 0})
+        for c in chunks:
+            key = c.size
+            if c.state == "allocated":
+                size_classes[key]["alloc"] += 1
+            else:
+                size_classes[key]["freed"] += 1
+
+        # Detect reachable primitives
+        primitives = self._detect_primitives(chunks, adjacency, bins, corruptions)
+
+        # Generate summary
+        summary = self._summarize(step, operation, chunks, bins, corruptions, primitives)
+
+        obs = HeapObservation(
+            step=step,
+            operation=operation,
+            chunks=chunks,
+            n_allocated=n_alloc,
+            n_freed=n_freed,
+            bins=bins,
+            adjacency=adjacency,
+            corruptions=corruptions,
+            cumulative_corruptions=self.cumulative_corruptions,
+            primitives=primitives,
+            size_classes=dict(size_classes),
+            summary=summary,
+        )
+        self.history.append(obs)
+        return obs
+
+    def _detect_primitives(self, chunks, adjacency, bins, corruptions) -> List[Primitive]:
+        """Detect which exploit primitives are currently reachable."""
+        primitives = []
+
+        # Tcache poison: freed chunk with corrupted fd
+        for c in chunks:
+            if c.state == "freed" and c.fd and c.fd_target is None and c.fd != "0x0":
+                primitives.append(Primitive(
+                    name="tcache_poison",
+                    description=f"Chunk {c.index} (freed, size {hex(c.size)}) has fd={c.fd} pointing outside heap. "
+                                f"Next malloc({c.size - 0x10}) returns attacker-controlled address.",
+                    ready=True,
+                    requirements=[],
+                    chunks_involved=[c.index],
+                ))
+
+        # Overlapping chunks: two allocated chunks at overlapping addresses
+        alloc_chunks = [(c, int(c.address, 16)) for c in chunks
+                        if c.state == "allocated" and c.address.startswith("0x")]
+        for i, (c1, a1) in enumerate(alloc_chunks):
+            for c2, a2 in alloc_chunks[i+1:]:
+                if a1 < a2 + c2.size and a2 < a1 + c1.size:
+                    primitives.append(Primitive(
+                        name="overlapping_chunks",
+                        description=f"Chunks {c1.index} and {c2.index} overlap in memory. "
+                                    f"Writing to one corrupts the other.",
+                        ready=True,
+                        requirements=[],
+                        chunks_involved=[c1.index, c2.index],
+                    ))
+
+        # Double free detected
+        seen_freed = set()
+        for c in chunks:
+            if c.state == "freed":
+                if c.address in seen_freed:
+                    primitives.append(Primitive(
+                        name="double_free",
+                        description=f"Address {c.address} freed multiple times. "
+                                    f"Tcache/fastbin contains a cycle.",
+                        ready=True,
+                        requirements=[],
+                        chunks_involved=[c.index],
+                    ))
+                seen_freed.add(c.address)
+
+        # UAF opportunity: freed chunks adjacent to allocated chunks
+        for c in chunks:
+            if c.state == "freed" and c.index in adjacency:
+                prev_idx, next_idx = adjacency[c.index]
+                for neighbor_idx in [prev_idx, next_idx]:
+                    if neighbor_idx is not None:
+                        neighbor = next((ch for ch in chunks if ch.index == neighbor_idx), None)
+                        if neighbor and neighbor.state == "allocated":
+                            primitives.append(Primitive(
+                                name="uaf_adjacent",
+                                description=f"Freed chunk {c.index} (size {hex(c.size)}) is adjacent to "
+                                            f"allocated chunk {neighbor.index} (size {hex(neighbor.size)}). "
+                                            f"UAF write to {c.index} could corrupt {neighbor.index}'s data.",
+                                ready=True,
+                                requirements=["Write to freed chunk via dangling pointer"],
+                                chunks_involved=[c.index, neighbor.index],
+                            ))
+
+        # Tcache ready: same-size chunks available for poisoning setup
+        for bin_info in bins:
+            if bin_info.bin_type == "tcache" and len(bin_info.entries) >= 1:
+                primitives.append(Primitive(
+                    name="tcache_available",
+                    description=f"Tcache bin for size {hex(bin_info.size_class)} has "
+                                f"{len(bin_info.entries)} entries. Poison fd to redirect allocation.",
+                    ready=len(bin_info.entries) >= 1,
+                    requirements=["Ability to write to freed chunk's fd pointer"],
+                    chunks_involved=bin_info.entries,
+                ))
+
+        # Coalesce opportunity: two adjacent freed chunks
+        for c in chunks:
+            if c.state == "freed" and c.index in adjacency:
+                _, next_idx = adjacency[c.index]
+                if next_idx is not None:
+                    neighbor = next((ch for ch in chunks if ch.index == next_idx), None)
+                    if neighbor and neighbor.state == "freed":
+                        primitives.append(Primitive(
+                            name="coalesce_opportunity",
+                            description=f"Freed chunks {c.index} and {neighbor.index} are adjacent. "
+                                        f"May coalesce into larger chunk on next free/malloc.",
+                            ready=True,
+                            requirements=[],
+                            chunks_involved=[c.index, neighbor.index],
+                        ))
+
+        # Metadata corruption detected
+        if corruptions:
+            for corr in corruptions:
+                primitives.append(Primitive(
+                    name=f"corruption_{corr.type}",
+                    description=f"CORRUPTION DETECTED: {corr.detail}",
+                    ready=True,
+                    requirements=[],
+                    chunks_involved=[corr.chunk_index],
+                ))
+
+        return primitives
+
+    def _summarize(self, step, operation, chunks, bins, corruptions, primitives) -> str:
+        """Generate a concise natural-language summary for LLM consumption."""
+        n_alloc = sum(1 for c in chunks if c.state == "allocated")
+        n_freed = sum(1 for c in chunks if c.state == "freed")
+
+        lines = []
+        lines.append(f"Step {step}: {operation} | {n_alloc} allocated, {n_freed} freed, {len(chunks)} total")
+
+        if bins:
+            bin_strs = [f"size {hex(b.size_class)}: {len(b.entries)} entries" for b in bins]
+            lines.append(f"Freelists: {', '.join(bin_strs)}")
+
+        if corruptions:
+            for c in corruptions:
+                lines.append(f"!! CORRUPTION: {c.type} at chunk {c.chunk_index}: {c.detail}")
+
+        ready_prims = [p for p in primitives if p.ready and "corruption" not in p.name]
+        if ready_prims:
+            prim_names = list(set(p.name for p in ready_prims))
+            lines.append(f"Primitives available: {', '.join(prim_names)}")
+
+        return "\n".join(lines)
+
+    def to_llm_context(self, obs: HeapObservation) -> str:
+        """Format observation as context for an LLM exploit generator."""
+        parts = []
+        parts.append(f"=== Heap State (step {obs.step}, after {obs.operation}) ===")
+        parts.append(f"Chunks: {obs.n_allocated} allocated, {obs.n_freed} freed")
+        parts.append("")
+
+        # Chunk table
+        parts.append("Chunks:")
+        for c in obs.chunks:
+            adj = obs.adjacency.get(c.index, (None, None))
+            adj_str = f"prev={adj[0]} next={adj[1]}" if any(adj) else ""
+            fd_str = f"fd={c.fd}" if c.fd else ""
+            corr_str = " [CORRUPTED]" if c.is_corrupted else ""
+            parts.append(f"  [{c.index}] {c.address} size={hex(c.size)} {c.state} "
+                        f"{fd_str} {adj_str}{corr_str}")
+
+        # Bins
+        if obs.bins:
+            parts.append("")
+            parts.append("Freelists:")
+            for b in obs.bins:
+                entries = " -> ".join(str(e) for e in b.entries)
+                parts.append(f"  {b.bin_type} size={hex(b.size_class)}: [{entries}]")
+
+        # Corruptions
+        if obs.corruptions:
+            parts.append("")
+            parts.append("!! CORRUPTIONS:")
+            for c in obs.corruptions:
+                parts.append(f"  {c.type}: {c.detail}")
+
+        # Primitives
+        ready = [p for p in obs.primitives if p.ready]
+        if ready:
+            parts.append("")
+            parts.append("Available primitives:")
+            for p in ready:
+                parts.append(f"  - {p.name}: {p.description}")
+
+        return "\n".join(parts)
+
+    def diff(self, prev: HeapObservation, curr: HeapObservation) -> str:
+        """Generate a diff between two observations — what changed."""
+        changes = []
+
+        # New chunks
+        prev_indices = {c.index for c in prev.chunks}
+        curr_indices = {c.index for c in curr.chunks}
+
+        for idx in curr_indices - prev_indices:
+            c = next(ch for ch in curr.chunks if ch.index == idx)
+            changes.append(f"+ Chunk {idx} allocated: size={hex(c.size)} at {c.address}")
+
+        for idx in prev_indices - curr_indices:
+            changes.append(f"- Chunk {idx} removed")
+
+        # State changes
+        for idx in prev_indices & curr_indices:
+            prev_c = next(ch for ch in prev.chunks if ch.index == idx)
+            curr_c = next(ch for ch in curr.chunks if ch.index == idx)
+            if prev_c.state != curr_c.state:
+                changes.append(f"~ Chunk {idx}: {prev_c.state} -> {curr_c.state}")
+            if prev_c.fd != curr_c.fd:
+                changes.append(f"~ Chunk {idx} fd: {prev_c.fd} -> {curr_c.fd}")
+
+        # New corruptions
+        if curr.corruptions:
+            for c in curr.corruptions:
+                changes.append(f"!! {c.type}: {c.detail}")
+
+        # New primitives
+        prev_prims = {p.name for p in prev.primitives}
+        for p in curr.primitives:
+            if p.name not in prev_prims and p.ready:
+                changes.append(f">> New primitive: {p.name} — {p.description}")
+
+        if not changes:
+            return "No significant changes."
+        return "\n".join(changes)