Support cs_disasm_iter

capstone-rs/Cargo.toml

@@ -22,6 +22,7 @@ libc = { version = "0.2", default-features = false }
 macho = "0.*"
 criterion = "0.3"
 rayon = "1.1"
+object = "0.26.2"
 
 [[bench]]
 name = "my_benchmark"

capstone-rs/examples/recursive.rs

@@ -0,0 +1,87 @@
+use std::collections::{HashMap, VecDeque};
+use std::env;
+use std::fs;
+use std::process;
+
+use object::{Object, ObjectSection, SectionKind};
+
+use capstone;
+use capstone::arch::x86::X86Insn;
+use capstone::prelude::*;
+
+fn main() {
+    let args: Vec<_> = env::args().collect();
+    if args.len() != 2 {
+        eprintln!("Usage: {} <file>", args[0]);
+        process::exit(-1);
+    }
+
+    let buf = if let Ok(bd) = fs::read(&args[1]) {
+        bd
+    } else {
+        eprintln!("cannot read file");
+        process::exit(-1);
+    };
+
+    let obj = if let Ok(od) = object::File::parse(&*buf) {
+        od
+    } else {
+        eprintln!("cannot parse file");
+        process::exit(-1);
+    };
+
+    let mut addr_queue: VecDeque<u64> = VecDeque::new();
+    let mut addr_seen: HashMap<u64, bool> = HashMap::new();
+
+    for section in obj.sections() {
+        if section.kind() == SectionKind::Text {
+            println!("{:x?} ", section);
+            addr_queue.push_back(section.address());
+        }
+    }
+
+    let cs = if let Ok(cs) = Capstone::new()
+        .x86()
+        .mode(arch::x86::ArchMode::Mode64)
+        .detail(true)
+        .build()
+    {
+        cs
+    } else {
+        eprintln!("failed to create capstone handle");
+        process::exit(-1);
+    };
+
+    let mut disasm = cs.get_disasm_iter();
+
+    while !addr_queue.is_empty() {
+        let addr = addr_queue.pop_front().unwrap();
+        if let Some(_) = addr_seen.get(&addr) {
+            continue;
+        }
+        addr_seen.insert(addr, true);
+
+        println!("---> addr: {:#02x?}", addr);
+
+        let offset = addr as usize;
+        let mut cur_insn = disasm.disasm_iter(&buf, offset, addr);
+        while let Ok(insn) = cur_insn {
+            if insn.id() == InsnId(X86Insn::X86_INS_INVALID as u32) {
+                break;
+            }
+            println!("{}", insn);
+            match X86Insn::from(insn.id().0) {
+                X86Insn::X86_INS_HLT => break,
+                X86Insn::X86_INS_CALL => break,
+                X86Insn::X86_INS_JMP => break,
+                X86Insn::X86_INS_RET => break,
+                _ => {}
+            }
+
+            // add logic here to add more targets to the addr_queue
+            // ...
+
+            cur_insn = disasm.disasm_iter_continue(&buf);
+        }
+    }
+}

capstone-rs/src/capstone.rs

@@ -165,6 +165,18 @@ impl Capstone {
         }
     }
 
+    /// Creates an instance of DisasmIter structure
+    ///
+    pub fn get_disasm_iter<'a>(&'a self) -> DisasmIter<'a> {
+        DisasmIter {
+            insn: unsafe { cs_malloc(self.csh()) },
+            csh: self.csh,
+            _covariant: PhantomData,
+            offset: 0,
+            addr: 0,
+        }
+    }
+
     /// Disassemble all instructions in buffer
     ///
     /// ```
@@ -423,3 +435,64 @@ impl Drop for Capstone {
         unsafe { cs_close(&mut self.csh()) };
     }
 }
+
+/// Structure to handle iterative disassembly
+///
+/// Create with a capstone instance `get_disasm_iter()`
+///
+pub struct DisasmIter<'a> {
+    insn: *mut cs_insn, // space for current instruction to be processed
+    csh: *mut c_void,   // reference to the the capstone handle required by disasm_iter
+    offset: usize,
+    addr: u64,
+    _covariant: PhantomData<&'a ()>, // used to make sure DIasmIter lifetime doesn't exceed Capstone's lifetime
+}
+
+impl<'a> Drop for DisasmIter<'a> {
+    fn drop(&mut self) {
+        unsafe { cs_free(self.insn, 1) };
+    }
+}
+
+impl<'a> DisasmIter<'a> {
+    /// Used to continue to the next instruction without jumping
+    ///
+    /// usage shown in examples/recursive.rs
+    pub fn disasm_iter_continue(&mut self, code: &[u8]) -> CsResult<Insn> {
+        self.disasm_iter(code, self.offset, self.addr)
+    }
+
+    /// Used to start the iterative disassembly
+    ///
+    /// usage shown in examples/recursive.rs
+    pub fn disasm_iter(&mut self, code: &[u8], offset: usize, addr: u64) -> CsResult<Insn> {
+        let code_len = code.len();
+        let code_ptr = &mut code[offset..].as_ptr();
+        let mut c = code_len - offset;
+        let mut a = addr;
+        let ret = unsafe {
+            cs_disasm_iter(
+                self.csh as csh, // capstone handle
+                code_ptr,        // double pointer to code to disassemble; automatically incremented
+                &mut c,          // number of bytes left to disassemble; automatically decremented
+                &mut a,          // automatically incremented address
+                self.insn,       // pointer to cs_insn object
+            )
+        };
+        if ret {
+            self.offset = code_len - c;
+            self.addr = a;
+            let insn = unsafe { Insn::from_raw(self.insn) };
+            Ok(insn)
+        } else {
+            Err(Error::CustomError("not disasm"))
+        }
+    }
+
+    pub fn offset(&self) -> usize {
+        self.offset
+    }
+    pub fn addr(&self) -> u64 {
+        self.addr
+    }
+}