Skip to content

Commit a4e7cba

Browse files
jiegecjiegec
committed
rs: add DisasmBuffer api to reuse the cs_insn allocation further
Closes #185 partially.
1 parent 85534a9 commit a4e7cba

File tree

1 file changed

+224
-21
lines changed

1 file changed

+224
-21
lines changed

‎capstone-rs/src/capstone.rs‎

Lines changed: 224 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -227,18 +227,77 @@ impl Capstone {
227227
code: &'b [u8],
228228
addr: u64,
229229
) -> CsResult<DisasmIter<'a, 'b>> {
230+
let buffer = self.create_buffer()?;
231+
Ok(DisasmIter {
232+
buffer,
233+
code: code.as_ptr(),
234+
size: code.len(),
235+
addr,
236+
_data: PhantomData,
237+
})
238+
}
239+
240+
/// Disassemble and iterate instructions from user-provided buffer `code` using `cs_disasm_iter`.
241+
/// The disassembled address of the buffer is assumed to be `addr`.
242+
/// The provided `buffer` is used to store the disassembled instruction.
243+
///
244+
/// # Examples
245+
///
246+
/// ```
247+
/// # use capstone::prelude::*;
248+
/// # let cs = Capstone::new().x86().mode(arch::x86::ArchMode::Mode32).build().unwrap();
249+
/// let mut buffer = cs.create_buffer().unwrap();
250+
/// let mut iter = cs.disasm_iter_with_buffer(b"\x90", 0x1000, &mut buffer);
251+
/// assert_eq!(iter.next().unwrap().mnemonic(), Some("nop"));
252+
/// assert!(iter.next().is_none());
253+
/// ```
254+
pub fn disasm_iter_with_buffer<'cs, 'dbuf, 'buf>(
255+
&'cs self,
256+
code: &'buf [u8],
257+
addr: u64,
258+
buffer: &'dbuf mut DisasmBuffer<'cs>,
259+
) -> DisasmIterWithBuffer<'cs, 'dbuf, 'buf> {
260+
DisasmIterWithBuffer {
261+
buffer,
262+
code: code.as_ptr(),
263+
size: code.len(),
264+
addr,
265+
_data: PhantomData,
266+
}
267+
}
268+
269+
/// Create a `DisasmBuffer`[DisasmBuffer] instance for future
270+
/// [`disasm_iter_with_buffer`](Capstone::disasm_iter_with_buffer)
271+
/// to reduce memory allocations.
272+
///
273+
/// # Examples
274+
///
275+
/// ```
276+
/// # use capstone::prelude::*;
277+
/// # let cs = Capstone::new().x86().mode(arch::x86::ArchMode::Mode32).build().unwrap();
278+
/// let mut buffer = cs.create_buffer().unwrap();
279+
/// let mut iter = cs.disasm_iter_with_buffer(b"\x90", 0x1000, &mut buffer);
280+
/// assert_eq!(iter.next().unwrap().mnemonic(), Some("nop"));
281+
/// assert!(iter.next().is_none());
282+
/// // Reuse the buffer
283+
/// let mut iter2 = cs.disasm_iter_with_buffer(b"\xc3", 0x1000, &mut buffer);
284+
/// assert_eq!(iter2.next().unwrap().mnemonic(), Some("ret"));
285+
/// assert!(iter2.next().is_none());
286+
/// ```
287+
///
288+
/// # Errors
289+
///
290+
/// If `cs_malloc` failed due to OOM, [`Err(Error::OutOfMemory)`](Error::OutOfMemory) is returned.
291+
///
292+
pub fn create_buffer<'cs>(&'cs self) -> CsResult<DisasmBuffer<'cs>> {
230293
let insn = unsafe { cs_malloc(self.csh()) };
231294
if insn.is_null() {
232295
return Err(Error::OutOfMemory);
233296
}
234-
Ok(DisasmIter {
297+
Ok(DisasmBuffer {
235298
insn,
236299
csh: self.csh,
237-
code: code.as_ptr(),
238-
size: code.len(),
239-
addr,
240-
_data1: PhantomData,
241-
_data2: PhantomData,
300+
_data: PhantomData,
242301
})
243302
}
244303

@@ -628,6 +687,27 @@ impl Drop for Capstone {
628687
}
629688
}
630689

690+
/// Structure to hold the memory buffer for a disassembled instruction.
691+
///
692+
/// Create with a [`Capstone`](Capstone) instance: [`Capstone::create_buffer()`](Capstone::create_buffer).
693+
///
694+
/// It deallocates memory via `cs_free` when dropped.
695+
///
696+
/// # Lifetimes
697+
///
698+
/// `'cs` is the lifetime of the [`Capstone`](Capstone) instance.
699+
pub struct DisasmBuffer<'cs> {
700+
insn: *mut cs_insn, // space for disassembled instruction
701+
csh: *mut c_void, // reference to the the capstone handle
702+
_data: PhantomData<&'cs ()>, // used to make sure DisasmBuffer lifetime doesn't exceed Capstone's lifetime
703+
}
704+
705+
impl<'cs> Drop for DisasmBuffer<'cs> {
706+
fn drop(&mut self) {
707+
unsafe { cs_free(self.insn, 1) };
708+
}
709+
}
710+
631711
/// Structure to handle iterative disassembly.
632712
///
633713
/// Create with a [`Capstone`](Capstone) instance: [`Capstone::disasm_iter()`](Capstone::disasm_iter).
@@ -636,24 +716,147 @@ impl Drop for Capstone {
636716
///
637717
/// `'cs` is the lifetime of the [`Capstone`](Capstone) instance.
638718
/// `'buf` is the lifetime of the user provided code buffer in [`Capstone::disasm_iter()`](Capstone::disasm_iter).
639-
///
640719
pub struct DisasmIter<'cs, 'buf> {
641-
insn: *mut cs_insn, // space for current instruction to be processed
642-
csh: *mut c_void, // reference to the the capstone handle required by disasm_iter
643-
code: *const u8, // pointer to the code buffer
644-
size: usize, // size of the code buffer
645-
addr: u64, // current address
646-
_data1: PhantomData<&'cs ()>, // used to make sure DisasmIter lifetime doesn't exceed Capstone's lifetime
647-
_data2: PhantomData<&'buf ()>, // used to make sure code lifetime doesn't exceed user provided array
720+
buffer: DisasmBuffer<'cs>, // buffer for current instruction to be processed
721+
code: *const u8, // pointer to the code buffer
722+
size: usize, // size of the code buffer
723+
addr: u64, // current address
724+
_data: PhantomData<&'buf ()>, // used to make sure code lifetime doesn't exceed user provided array
648725
}
649726

650-
impl<'cs, 'buf> Drop for DisasmIter<'cs, 'buf> {
651-
fn drop(&mut self) {
652-
unsafe { cs_free(self.insn, 1) };
727+
impl<'cs, 'buf> DisasmIter<'cs, 'buf> {
728+
/// Get next instruction if available.
729+
///
730+
/// # Examples
731+
///
732+
/// ```
733+
/// # use capstone::prelude::*;
734+
/// # let cs = Capstone::new().x86().mode(arch::x86::ArchMode::Mode32).build().unwrap();
735+
/// let code = b"\x90";
736+
/// let mut iter = cs.disasm_iter(code, 0x1000).unwrap();
737+
/// while let Some(insn) = iter.next() {
738+
/// println!("{insn}");
739+
/// }
740+
/// ```
741+
///
742+
/// At most one instruction can be accessed at the same time:
743+
///
744+
/// ```compile_fail
745+
/// # use capstone::prelude::*;
746+
/// # let cs = Capstone::new().x86().mode(arch::x86::ArchMode::Mode32).build().unwrap();
747+
/// let code = b"\x90";
748+
/// let mut iter = cs.disasm_iter(code, 0x1000).unwrap();
749+
/// let insn1 = iter.next().unwrap();
750+
/// let insn2 = iter.next().unwrap();
751+
/// // fails with: cannot borrow `iter` as mutable more than once at a time,
752+
/// // `insn1` cannot be used after calling `iter.next()` again,
753+
/// // as `iter` is mutably borrowed during the second call to `next()`.
754+
/// println!("{insn1}");
755+
/// ```
756+
pub fn next<'iter>(&'iter mut self) -> Option<Insn<'iter>> {
757+
unsafe {
758+
if cs_disasm_iter(
759+
self.buffer.csh as csh,
760+
&mut self.code,
761+
&mut self.size,
762+
&mut self.addr,
763+
self.buffer.insn,
764+
) {
765+
return Some(Insn::from_raw(self.buffer.insn));
766+
}
767+
}
768+
769+
None
770+
}
771+
772+
/// Get the slice of the code yet to be disassembled
773+
///
774+
/// ```
775+
/// # use capstone::prelude::*;
776+
/// # let cs = Capstone::new().x86().mode(arch::x86::ArchMode::Mode32).build().unwrap();
777+
/// let code = b"\x90";
778+
/// let mut iter = cs.disasm_iter(code, 0x1000).unwrap();
779+
/// assert_eq!(iter.code(), code);
780+
/// iter.next();
781+
/// assert_eq!(iter.code(), b"");
782+
/// ```
783+
pub fn code(&self) -> &[u8] {
784+
unsafe { core::slice::from_raw_parts(self.code, self.size) }
785+
}
786+
787+
/// Get the address of the next instruction to be disassembled
788+
///
789+
/// ```
790+
/// # use capstone::prelude::*;
791+
/// # let cs = Capstone::new().x86().mode(arch::x86::ArchMode::Mode32).build().unwrap();
792+
/// let code = b"\x90";
793+
/// let mut iter = cs.disasm_iter(code, 0x1000).unwrap();
794+
/// assert_eq!(iter.addr(), 0x1000);
795+
/// iter.next();
796+
/// assert_eq!(iter.addr(), 0x1001);
797+
/// ```
798+
pub fn addr(&self) -> u64 {
799+
self.addr
800+
}
801+
802+
/// Reset the iterator to disassemble in the specified code buffer
803+
///
804+
/// ```
805+
/// # use capstone::prelude::*;
806+
/// # let cs = Capstone::new().x86().mode(arch::x86::ArchMode::Mode32).build().unwrap();
807+
/// let code = b"\x90";
808+
/// let mut iter = cs.disasm_iter(code, 0x1000).unwrap();
809+
/// assert_eq!(iter.addr(), 0x1000);
810+
/// assert_eq!(iter.code(), code);
811+
/// iter.next();
812+
/// assert_eq!(iter.addr(), 0x1001);
813+
/// assert_eq!(iter.code(), b"");
814+
/// let new_code = b"\xc3";
815+
/// iter.reset(new_code, 0x2000);
816+
/// assert_eq!(iter.addr(), 0x2000);
817+
/// assert_eq!(iter.code(), new_code);
818+
/// ```
819+
pub fn reset(&mut self, code: &'buf [u8], addr: u64) {
820+
self.code = code.as_ptr();
821+
self.size = code.len();
822+
self.addr = addr;
823+
}
824+
825+
/// Extract the underlying buffer for reuse.
826+
/// ```
827+
/// # use capstone::prelude::*;
828+
/// # let cs = Capstone::new().x86().mode(arch::x86::ArchMode::Mode32).build().unwrap();
829+
/// let code = b"\x90";
830+
/// let mut iter = cs.disasm_iter(code, 0x1000).unwrap();
831+
/// assert!(iter.next().is_some());
832+
/// // reuse buffer
833+
/// let mut buffer = iter.into_buffer();
834+
/// let mut iter2 = cs.disasm_iter_with_buffer(code, 0x2000, &mut buffer);
835+
/// assert!(iter2.next().is_some());
836+
/// ```
837+
pub fn into_buffer(self) -> DisasmBuffer<'cs> {
838+
self.buffer
653839
}
654840
}
655841

656-
impl<'cs, 'buf> DisasmIter<'cs, 'buf> {
842+
/// Structure to handle iterative disassembly with provided [`DisasmBuffer`](DisasmBuffer).
843+
///
844+
/// Create with a [`Capstone`](Capstone) instance: [`Capstone::disasm_iter_with_buffer`](Capstone::disasm_iter_with_buffer).
845+
///
846+
/// # Lifetimes
847+
///
848+
/// `'cs` is the lifetime of the [`Capstone`](Capstone) instance.
849+
/// `'dbuf` is the lifetime of the user provided [`DisasmBuffer`](DisasmBuffer) in [`Capstone::disasm_iter_with_buffer`](Capstone::disasm_iter_with_buffer).
850+
/// `'buf` is the lifetime of the user provided code buffer in [`Capstone::disasm_iter_with_buffer`](Capstone::disasm_iter_with_buffer).
851+
pub struct DisasmIterWithBuffer<'cs, 'dbuf, 'buf> {
852+
buffer: &'dbuf mut DisasmBuffer<'cs>, // buffer for current instruction to be processed
853+
code: *const u8, // pointer to the code buffer
854+
size: usize, // size of the code buffer
855+
addr: u64, // current address
856+
_data: PhantomData<&'buf ()>, // used to make sure code lifetime doesn't exceed user provided array
857+
}
858+
859+
impl<'cs, 'dbuf, 'buf> DisasmIterWithBuffer<'cs, 'dbuf, 'buf> {
657860
/// Get next instruction if available.
658861
///
659862
/// # Examples
@@ -685,13 +888,13 @@ impl<'cs, 'buf> DisasmIter<'cs, 'buf> {
685888
pub fn next<'iter>(&'iter mut self) -> Option<Insn<'iter>> {
686889
unsafe {
687890
if cs_disasm_iter(
688-
self.csh as csh,
891+
self.buffer.csh as csh,
689892
&mut self.code,
690893
&mut self.size,
691894
&mut self.addr,
692-
self.insn,
895+
self.buffer.insn,
693896
) {
694-
return Some(Insn::from_raw(self.insn));
897+
return Some(Insn::from_raw(self.buffer.insn));
695898
}
696899
}
697900

0 commit comments

Comments
 (0)