heaptrm CLI: compiled Rust binary for LLM-assisted heap exploitation

749KB binary. JSON protocol on stdin/stdout:
- send: pipe data to target binary, get heap observation
- observe: get current heap state with chunks, bins, primitives
- quit: clean exit

Includes primitive detection (tcache_poison, double_free, corruption),
auto-finds/compiles v2 harness, manages target process lifecycle.

heaptrm-cli/Cargo.toml

@@ -0,0 +1,9 @@
+[package]
+name = "heaptrm"
+version = "0.1.0"
+edition = "2021"
+description = "Heap exploit observability tool for LLM-assisted exploit generation"
+
+[dependencies]
+serde = { version = "1", features = ["derive"] }
+serde_json = "1"

heaptrm-cli/src/main.rs

@@ -0,0 +1,359 @@
+//! heaptrm - Heap exploit observability for LLM-assisted exploitation.
+//!
+//! Launches a target binary with LD_PRELOAD heap instrumentation,
+//! provides a JSON protocol on stdin/stdout for LLM interaction.
+//!
+//! Protocol:
+//!   LLM sends:  {"action": "send", "data": "1 0 64\n"}
+//!   Tool sends: {"heap": {...}, "changes": "...", "primitives": [...]}
+//!
+//!   LLM sends:  {"action": "observe"}
+//!   Tool sends: current heap state
+//!
+//!   LLM sends:  {"action": "quit"}
+//!   Tool sends: final summary and exits
+
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+use std::env;
+use std::fs;
+use std::io::{self, BufRead, Write};
+use std::path::PathBuf;
+use std::process::{Child, Command, Stdio};
+use std::thread;
+use std::time::Duration;
+
+// --- Harness output types ---
+
+#[derive(Debug, Deserialize, Clone)]
+struct RawChunk {
+    idx: usize,
+    addr: String,
+    state: u8,
+    #[serde(default)]
+    chunk_size: usize,
+    #[serde(default)]
+    fd: u64,
+    #[serde(default)]
+    fd_idx: i32,
+    #[serde(default)]
+    is_corrupted: u8,
+    #[serde(default)]
+    is_double_freed: u8,
+    #[serde(default)]
+    data_hex: String,
+}
+
+#[derive(Debug, Deserialize, Clone)]
+struct RawCorruption {
+    #[serde(rename = "type")]
+    corruption_type: String,
+    chunk_idx: i32,
+    detail: String,
+}
+
+#[derive(Debug, Deserialize, Clone)]
+struct RawState {
+    step: u32,
+    operation: String,
+    #[serde(default)]
+    corruption_count: u32,
+    #[serde(default)]
+    corruptions: Vec<RawCorruption>,
+    #[serde(default)]
+    chunks: Vec<RawChunk>,
+}
+
+// --- Output types ---
+
+#[derive(Serialize)]
+struct ChunkView {
+    index: usize,
+    address: String,
+    size: String,
+    state: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    fd: Option<String>,
+    corrupted: bool,
+}
+
+#[derive(Serialize)]
+struct BinView {
+    size: String,
+    count: usize,
+    entries: Vec<usize>,
+}
+
+#[derive(Serialize)]
+struct Primitive {
+    name: String,
+    description: String,
+    chunks: Vec<usize>,
+}
+
+#[derive(Serialize)]
+struct HeapView {
+    step: u32,
+    operation: String,
+    allocated: usize,
+    freed: usize,
+    chunks: Vec<ChunkView>,
+    bins: Vec<BinView>,
+    corruptions: Vec<serde_json::Value>,
+    primitives: Vec<Primitive>,
+    summary: String,
+}
+
+#[derive(Serialize)]
+struct Response {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    heap: Option<HeapView>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    changes: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    error: Option<String>,
+}
+
+#[derive(Deserialize)]
+struct Request {
+    action: String,
+    #[serde(default)]
+    data: String,
+}
+
+const QUARANTINE_FD: u64 = 0xFDFDFDFDFDFDFDFD;
+
+fn analyze_state(state: &RawState) -> HeapView {
+    let n_alloc = state.chunks.iter().filter(|c| c.state == 1).count();
+    let n_freed = state.chunks.iter().filter(|c| c.state == 2).count();
+
+    let chunks: Vec<ChunkView> = state.chunks.iter().map(|c| {
+        let fd = if c.fd != 0 && c.fd != QUARANTINE_FD {
+            Some(format!("0x{:x}", c.fd))
+        } else {
+            None
+        };
+        ChunkView {
+            index: c.idx,
+            address: c.addr.clone(),
+            size: format!("0x{:x}", c.chunk_size),
+            state: if c.state == 1 { "allocated".into() } else { "freed".into() },
+            fd,
+            corrupted: c.is_corrupted != 0,
+        }
+    }).collect();
+
+    // Bins
+    let mut size_bins: HashMap<usize, Vec<usize>> = HashMap::new();
+    for c in &state.chunks {
+        if c.state == 2 && c.chunk_size > 0 {
+            size_bins.entry(c.chunk_size).or_default().push(c.idx);
+        }
+    }
+    let bins: Vec<BinView> = size_bins.iter().map(|(sz, entries)| BinView {
+        size: format!("0x{:x}", sz),
+        count: entries.len(),
+        entries: entries.clone(),
+    }).collect();
+
+    // Corruptions as JSON values
+    let corruptions: Vec<serde_json::Value> = state.corruptions.iter().map(|c| {
+        serde_json::json!({
+            "type": c.corruption_type,
+            "chunk": c.chunk_idx,
+            "detail": c.detail,
+        })
+    }).collect();
+
+    // Primitives
+    let mut primitives = Vec::new();
+
+    for c in &state.chunks {
+        if c.state == 2 && c.fd != 0 && c.fd != QUARANTINE_FD && c.fd_idx == -2 {
+            primitives.push(Primitive {
+                name: "tcache_poison".into(),
+                description: format!(
+                    "Chunk {} has fd=0x{:x} outside heap. malloc(0x{:x}) returns controlled address.",
+                    c.idx, c.fd, c.chunk_size.saturating_sub(0x10)
+                ),
+                chunks: vec![c.idx],
+            });
+        }
+        if c.is_double_freed != 0 {
+            primitives.push(Primitive {
+                name: "double_free".into(),
+                description: format!("Chunk {} at {} freed multiple times.", c.idx, c.addr),
+                chunks: vec![c.idx],
+            });
+        }
+    }
+
+    for corr in &state.corruptions {
+        primitives.push(Primitive {
+            name: format!("corruption_{}", corr.corruption_type),
+            description: corr.detail.clone(),
+            chunks: vec![corr.chunk_idx as usize],
+        });
+    }
+
+    // Summary
+    let mut summary = format!("Step {}: {} | {} alloc, {} freed", state.step, state.operation, n_alloc, n_freed);
+    for corr in &state.corruptions {
+        summary.push_str(&format!("\n!! {}: {}", corr.corruption_type, corr.detail));
+    }
+    let prim_names: Vec<&str> = primitives.iter()
+        .filter(|p| !p.name.starts_with("corruption_"))
+        .map(|p| p.name.as_str())
+        .collect();
+    if !prim_names.is_empty() {
+        summary.push_str(&format!("\nPrimitives: {}", prim_names.join(", ")));
+    }
+
+    HeapView { step: state.step, operation: state.operation.clone(), allocated: n_alloc, freed: n_freed, chunks, bins, corruptions, primitives, summary }
+}
+
+fn find_harness() -> Option<PathBuf> {
+    let candidates = [
+        "heapgrid_v2.so",
+        "heaptrm/harness/heapgrid_v2.so",
+        "harness/heapgrid_harness.so",
+        "../heaptrm/harness/heapgrid_v2.so",
+    ];
+    for c in &candidates {
+        let p = PathBuf::from(c);
+        if p.exists() {
+            return Some(fs::canonicalize(p).ok()?);
+        }
+    }
+    // Try compile
+    for src in &["heaptrm/harness/heapgrid_v2.c", "heapgrid_v2.c"] {
+        let s = PathBuf::from(src);
+        if s.exists() {
+            let out = s.with_extension("so");
+            if Command::new("gcc").args(["-shared","-fPIC","-O2","-o"]).arg(&out).arg(&s).args(["-ldl","-pthread"]).status().map(|s| s.success()).unwrap_or(false) {
+                return Some(fs::canonicalize(out).ok()?);
+            }
+        }
+    }
+    None
+}
+
+fn main() {
+    let args: Vec<String> = env::args().collect();
+    if args.len() < 2 {
+        eprintln!("heaptrm — heap exploit observability for LLM-assisted exploitation");
+        eprintln!();
+        eprintln!("Usage: heaptrm <binary> [args...]");
+        eprintln!();
+        eprintln!("Launches <binary> with heap instrumentation. Reads JSON from stdin,");
+        eprintln!("writes heap observations to stdout.");
+        eprintln!();
+        eprintln!("Commands:");
+        eprintln!("  {{\"action\": \"send\", \"data\": \"...\"}}  send data to binary stdin");
+        eprintln!("  {{\"action\": \"observe\"}}                  get current heap state");
+        eprintln!("  {{\"action\": \"quit\"}}                     exit");
+        std::process::exit(1);
+    }
+
+    let binary = &args[1];
+    let binary_args = &args[2..];
+
+    let harness = find_harness().unwrap_or_else(|| {
+        eprintln!("Error: Cannot find heapgrid_v2.so");
+        std::process::exit(1);
+    });
+
+    let dump_path = format!("/tmp/heaptrm_{}.jsonl", std::process::id());
+
+    let mut child: Child = Command::new(binary)
+        .args(binary_args)
+        .stdin(Stdio::piped())
+        .stdout(Stdio::piped())
+        .stderr(Stdio::null())
+        .env("LD_PRELOAD", &harness)
+        .env("HEAPGRID_OUT", &dump_path)
+        .spawn()
+        .unwrap_or_else(|e| { eprintln!("Failed to launch: {}", e); std::process::exit(1); });
+
+    let mut child_stdin = child.stdin.take().expect("stdin");
+    thread::sleep(Duration::from_millis(50));
+
+    let stdin = io::stdin();
+    let stdout = io::stdout();
+    let mut out = stdout.lock();
+    let mut last_pos: u64 = 0;
+    let mut last_state: Option<RawState> = None;
+
+    let read_latest = |pos: &mut u64| -> Option<RawState> {
+        let content = fs::read_to_string(&dump_path).ok()?;
+        let start = *pos as usize;
+        if start >= content.len() { return None; }
+        let new = &content[start..];
+        let mut last = None;
+        for line in new.lines() {
+            if let Ok(s) = serde_json::from_str::<RawState>(line) {
+                last = Some(s);
+            }
+        }
+        *pos = content.len() as u64;
+        last
+    };
+
+    for line in stdin.lock().lines() {
+        let line = match line { Ok(l) => l, Err(_) => break };
+        if line.trim().is_empty() { continue; }
+
+        let req: Request = match serde_json::from_str(&line) {
+            Ok(r) => r,
+            Err(e) => {
+                let r = Response { heap: None, changes: None, error: Some(format!("Bad JSON: {}", e)) };
+                writeln!(out, "{}", serde_json::to_string(&r).unwrap()).ok();
+                out.flush().ok();
+                continue;
+            }
+        };
+
+        match req.action.as_str() {
+            "send" => {
+                child_stdin.write_all(req.data.as_bytes()).ok();
+                child_stdin.flush().ok();
+                thread::sleep(Duration::from_millis(20));
+
+                if let Some(s) = read_latest(&mut last_pos) {
+                    last_state = Some(s);
+                }
+                let r = match &last_state {
+                    Some(s) => Response { heap: Some(analyze_state(s)), changes: None, error: None },
+                    None => Response { heap: None, changes: None, error: Some("No heap state yet".into()) },
+                };
+                writeln!(out, "{}", serde_json::to_string(&r).unwrap()).ok();
+                out.flush().ok();
+            }
+            "observe" => {
+                if let Some(s) = read_latest(&mut last_pos) {
+                    last_state = Some(s);
+                }
+                let r = match &last_state {
+                    Some(s) => Response { heap: Some(analyze_state(s)), changes: None, error: None },
+                    None => Response { heap: None, changes: None, error: Some("No heap state yet".into()) },
+                };
+                writeln!(out, "{}", serde_json::to_string(&r).unwrap()).ok();
+                out.flush().ok();
+            }
+            "quit" => {
+                child.kill().ok();
+                fs::remove_file(&dump_path).ok();
+                break;
+            }
+            _ => {
+                let r = Response { heap: None, changes: None, error: Some(format!("Unknown: {}", req.action)) };
+                writeln!(out, "{}", serde_json::to_string(&r).unwrap()).ok();
+                out.flush().ok();
+            }
+        }
+    }
+
+    child.kill().ok();
+    child.wait().ok();
+    fs::remove_file(&dump_path).ok();
+}