add mod macros; fixed (#8);

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rcheat"
-version = "0.1.0"
+version = "0.1.1"
 edition = "2018"
 authors = ["handy-sun <[email protected]>"]
 homepage = "https://github.com/handy-sun/rcheat"

src/ctrl.rs

@@ -1,3 +1,4 @@
+use crate::ceil_to_multiple;
 use crate::elf::ElfMgr;
 use crate::fmt_dump::*;
 use crate::AnyError;
@@ -106,6 +107,8 @@ pub fn trace(arg: Args) -> AnyError {
     let start = Instant::now();
     // let elf_data = match_sym_entry(&elf_bytes, &arg.keyword)?;
     let elf_mgr = ElfMgr::prase_from(&elf_bytes)?;
+    println!("[{:?}] Time of `parse elf`", start.elapsed());
+
     let entry = elf_mgr.select_sym_entry(&arg.keyword)?;
 
     let entry_addr = if elf_mgr.is_exec_elf() {
@@ -117,15 +120,17 @@ pub fn trace(arg: Args) -> AnyError {
         let file_reader = BufReader::new(file);
         let base_addr = get_base_addr(file_reader, exe_path.as_str())?;
         println!("base_addr: {:#x} ({})", base_addr, base_addr);
-        base_addr + entry.obj_addr()
+        if let Some(total_addr) = base_addr.checked_add(entry.obj_addr()) {
+            total_addr
+        } else {
+            return Err(anyhow!("Operation of base_addr add obj_addr exceeds the limit"));
+        }
     } else {
         return Err(anyhow!("Unsupport e_type:"));
     };
 
     println!("entry address: {:#x}, size: {}", entry_addr, entry.obj_size());
 
-    let duration = start.elapsed();
-    println!("[{:?}] Time of `parse elf`", duration);
     let start = Instant::now();
     pass_or_exit(&ptrace::attach(tracked_pid), "ptrace_attach")?;
 
@@ -140,31 +145,46 @@ pub fn trace(arg: Args) -> AnyError {
 
     let addr = ptrace::AddressType::from(entry_addr as ptrace::AddressType);
     let var_sz = entry.obj_size() as usize;
-    let mut peek_buf = BytesMut::with_capacity(var_sz);
-
-    let mut pos: usize = 0;
-    while pos < var_sz {
-        if pos + LONG_SIZE > var_sz {
-            let old_pos = pos;
-            pos = var_sz - LONG_SIZE;
-            peek_buf.truncate(pos);
-            println!("truncate length from: {} to {}", old_pos, pos);
-        }
-        match ptrace::read(tracked_pid, addr.wrapping_add(pos)) {
+    // It can be confirmed that this number(var_sz) must be greater than 0
+    let mut peek_buf = BytesMut::with_capacity(ceil_to_multiple!(var_sz, LONG_SIZE));
+
+    // The target address size less than c_long, only need read once, and then
+    // truncate BytesMut to real size
+    if var_sz < LONG_SIZE {
+        match ptrace::read(tracked_pid, addr) {
             Ok(long_data) => {
                 peek_buf.put(long_data.to_ne_bytes().as_ref());
+                peek_buf.truncate(var_sz);
             }
             Err(errno) => {
                 return restore_process_to_run(tracked_pid, anyhow!("peekdata at {:?}: {:?}", addr, errno));
             }
+        };
+    } else {
+        let mut pos: usize = 0;
+        while pos < var_sz {
+            if pos + LONG_SIZE > var_sz {
+                pos = var_sz - LONG_SIZE;
+                peek_buf.truncate(pos);
+            }
+            match ptrace::read(tracked_pid, addr.wrapping_add(pos)) {
+                Ok(long_data) => {
+                    peek_buf.put(long_data.to_ne_bytes().as_ref());
+                }
+                Err(errno) => {
+                    return restore_process_to_run(
+                        tracked_pid,
+                        anyhow!("peekdata at {:?}: {:?}", addr, errno),
+                    );
+                }
+            }
+            pos += LONG_SIZE;
         }
-        pos += LONG_SIZE;
     }
 
     pass_or_exit(&ptrace::detach(tracked_pid, None), "ptrace_detach")?;
 
-    let duration = start.elapsed();
-    println!("[{:?}] Time of `trace`", duration);
+    println!("[{:?}] Time of `trace and peek`", start.elapsed());
     if let Some(bytes_ref) = peek_buf.get(..) {
         let out_content = if arg.format == "dec" {
             dump_to_dec_content(bytes_ref)
@@ -183,7 +203,7 @@ mod tests {
     use super::*;
 
     #[test]
-    fn _get_base_addr() {
+    fn func_get_base_addr() {
         let exe_abs_path = "/usr/bin/test1";
         let contents = "
 00400000-004ac000 r-xp 00000000 08:02 8918620   /usr/bin/test1

src/elf/elfmgr.rs

@@ -1,5 +1,6 @@
 use std::convert::TryFrom;
 use std::io;
+use std::time::Instant;
 
 use anyhow::{anyhow, Error};
 
@@ -44,10 +45,7 @@ impl<'a> ElfMgr<'a> {
         }
 
         match Object::parse(bytes).unwrap() {
-            Object::Elf(val) => Ok(ElfMgr {
-                elf: val,
-                // satisfied_entry: SymEntry::default(),
-            }),
+            Object::Elf(val) => Ok(ElfMgr { elf: val }),
             _ => Err(anyhow!("Object format not support")),
         }
     }
@@ -61,6 +59,7 @@ impl<'a> ElfMgr<'a> {
     }
 
     pub fn select_sym_entry(&self, keyword: &String) -> Result<SymEntry, Error> {
+        let start = Instant::now();
         let strtab = &self.elf.strtab;
         let syms = self.elf.syms.to_vec();
         // let dyn_strtab = &self.elf.dynstrtab;
@@ -74,6 +73,7 @@ impl<'a> ElfMgr<'a> {
             .filter_map(|sym| self.filter_symbol(sym, strtab, keyword));
 
         let emtry_vec: Vec<SymEntry> = map_iter.collect();
+        println!("[{:?}] Time of `filter_symbol`", start.elapsed());
         println!("Matched count: {}", emtry_vec.len());
         match emtry_vec.len() {
             0 => Err(anyhow!("cannot find")),
@@ -117,7 +117,8 @@ impl<'a> ElfMgr<'a> {
                 || dem_name.contains("std::")
                 || dem_name.starts_with("__gnu_")
                 || dem_name.starts_with("__cxxabiv")
-                || dem_name.starts_with("guard variable"))
+                || dem_name.starts_with("guard variable")
+                || dem_name.ends_with(")::__func__"))
         {
             #[cfg(debug_assertions)]
             eprintln!(

src/macros.rs

@@ -0,0 +1,27 @@
+#[macro_export]
+/// Round a number up to a multiple greater than or equal to a specified numerical value of the number
+macro_rules! ceil_to_multiple {
+    ($value:expr, $multiple:expr) => {{
+        let value = $value;
+        let multiple = $multiple;
+
+        if multiple == 0 || value % multiple == 0 {
+            value
+        } else {
+            value + (multiple - (value % multiple))
+        }
+    }};
+}
+
+#[cfg(test)]
+mod tests {
+    // use super::*;
+
+    #[test]
+    fn macro_ceil_to_multiple() {
+        assert_eq!(ceil_to_multiple!(66, 10), 70);
+        assert_eq!(ceil_to_multiple!(5, 8), 8);
+        assert_eq!(ceil_to_multiple!(6, 0), 6);
+        assert_eq!(ceil_to_multiple!(0, 4), 0);
+    }
+}

src/main.rs

@@ -1,6 +1,8 @@
 mod ctrl;
 mod elf;
 mod fmt_dump;
+// #[macro_use]
+mod macros;
 
 use anyhow::{anyhow, Error};
 use clap::Parser;