Add structure to monitor arg:value for a register

3 files changed, 89 insertions, 59 deletions

diff --git a/src/registers.rs b/src/registers.rs
index 44d7271..9491bcf 100644
--- a/src/registers.rs
+++ b/src/registers.rs
@@ -12,33 +12,61 @@ use crate::arch::linux::x86_64::*;
 use crate::arch::syscall_name;
 use crate::trace::read_memory;
 
+#[derive(Clone, Debug)]
+/// Structure use to monitor what a register has for (argument: value)
+struct RegisterOutput {
+    /// Value for a register, by default is a number which could be a real value or a memory
+    /// address
+    value: String,
+    /// Argument for a register, eg: "const char *buf"
+    argument: &'static str,
+}
+
+impl RegisterOutput {
+    fn new(address: u64, argument: &'static str) -> Self {
+        Self {
+            value: address.to_string(),
+            argument,
+        }
+    }
+}
+
 /// Struct used to manipulate registers data from https://docs.rs/libc/0.2.147/libc/struct.user_regs_struct.html
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug)]
 pub struct RegistersData {
     timestamp: DateTime<Local>,
     orig_rax: u64,
-    rdi: u64,
-    rsi: u64,
-    rdx: u64,
-    r10: u64,
-    r8: u64,
-    r9: u64,
+    rdi: RegisterOutput,
+    rsi: RegisterOutput,
+    rdx: RegisterOutput,
+    r10: RegisterOutput,
+    r8: RegisterOutput,
+    r9: RegisterOutput,
     rax: u64,
 }
 
 impl RegistersData {
     /// Create new `RegistersData` from an `user_regs_struct`'C structure
     pub fn new(registers: user_regs_struct) -> RegistersData {
+        let (rdi, rsi, rdx, r10, r8, r9) = (
+            RegisterOutput::new(registers.rdi, rdi(registers.orig_rax)),
+            RegisterOutput::new(registers.rsi, rsi(registers.orig_rax)),
+            RegisterOutput::new(registers.rdx, rdx(registers.orig_rax)),
+            RegisterOutput::new(registers.r10, r10(registers.orig_rax)),
+            RegisterOutput::new(registers.r8, r8(registers.orig_rax)),
+            RegisterOutput::new(registers.r9, r9(registers.orig_rax)),
+        );
+
         RegistersData {
             timestamp: Local::now(),
             orig_rax: registers.orig_rax,
-            rdi: registers.rdi,
-            rsi: registers.rsi,
-            rdx: registers.rdx,
-            r10: registers.r10,
-            r8: registers.r8,
-            r9: registers.r9,
             rax: registers.rax,
+            rdi,
+            rsi,
+            rdx,
+            r10,
+            r8,
+            r9,
         }
     }
 
@@ -53,7 +81,7 @@ impl RegistersData {
     }
 
     /// Returns a good string which shows the output for a line
-    pub fn output(&self, pid: Pid) -> String {
+    pub fn output(&mut self, pid: Pid) -> String {
         let mut output = format!("[{}]: ", self.date());
 
         if !self.name().is_empty() {
@@ -62,33 +90,34 @@ impl RegistersData {
             output.push_str(&format!("{}(", self.orig_rax.yellow().bold()));
         }
 
-        let mut has_param = false;
+        let mut has_reg = false;
 
-        let params = [
-            (self.rdi, rdi(self.orig_rax)),
-            (self.rsi, rsi(self.orig_rax)),
-            (self.rdx, rdx(self.orig_rax)),
-            (self.r10, r10(self.orig_rax)),
-            (self.r8, r8(self.orig_rax)),
-            (self.r9, r9(self.orig_rax)),
+        let mut regs = [
+            &mut self.rdi,
+            &mut self.rsi,
+            &mut self.rdx,
+            &mut self.r10,
+            &mut self.r8,
+            &mut self.r9,
         ];
 
-        for param in params {
-            if !param.1.is_empty() {
-                let output_param = param.1.to_owned() + ":";
-                let output_value = if output_param.starts_with("const char *")
-                    || output_param.starts_with("char *")
+        for reg in &mut regs {
+            if !reg.argument.is_empty() {
+                let output_reg = reg.argument.to_owned() + ":";
+                reg.value = if (output_reg.starts_with("const char *")
+                    || output_reg.starts_with("char *"))
+                    && !reg.value.starts_with("\"")
                 {
-                    read_memory(pid, param.0)
+                    read_memory(pid, reg.value.parse::<u64>().unwrap())
                 } else {
-                    param.0.to_string()
+                    reg.value.to_string()
                 };
-                output.push_str(&format!("{} {}, ", output_param.blue(), output_value));
-                has_param = true;
+                output.push_str(&format!("{} {}, ", output_reg.blue(), reg.value));
+                has_reg = true;
             }
         }
 
-        if has_param {
+        if has_reg {
             output.remove(output.len() - 1);
             output.remove(output.len() - 1);
         }
@@ -98,7 +127,7 @@ impl RegistersData {
     }
 
     /// Returns a good line for TUI
-    pub fn output_ui(&self, pid: Pid) -> Line {
+    pub fn output_ui(&mut self, _pid: Pid) -> Line {
         let mut spans: Vec<Span> = vec![];
         spans.push(Span::raw(format!("[{}]: ", self.date())));
         if !self.name().is_empty() {
@@ -115,35 +144,33 @@ impl RegistersData {
             ));
         }
 
-        let params = [
-            (self.rdi, rdi(self.orig_rax)),
-            (self.rsi, rsi(self.orig_rax)),
-            (self.rdx, rdx(self.orig_rax)),
-            (self.r10, r10(self.orig_rax)),
-            (self.r8, r8(self.orig_rax)),
-            (self.r9, r9(self.orig_rax)),
+        let mut regs = [
+            &mut self.rdi,
+            &mut self.rsi,
+            &mut self.rdx,
+            &mut self.r10,
+            &mut self.r8,
+            &mut self.r9,
         ];
 
-        for param in params {
-            if !param.1.is_empty() {
-                let output_param = param.1.to_owned() + ":";
+        for reg in &mut regs {
+            if !reg.argument.is_empty() {
+                let output_reg = reg.argument.to_owned() + ":";
                 spans.push(Span::styled(
-                    format!("{} ", output_param),
+                    format!("{} ", output_reg),
                     Style::default().fg(Color::Blue),
                 ));
 
-                // FIXME: read memory does not work
-                let output_value = if output_param.starts_with("const char *")
-                    || output_param.starts_with("char *")
-                {
-                    read_memory(pid, param.0)
-                } else {
-                    param.0.to_string()
-                };
-                spans.push(Span::styled(
-                    format!("{}, ", output_value),
-                    Style::default(),
-                ));
+                //  FIXME: read memory does not work
+                // reg.value = if (output_reg.starts_with("const char *")
+                //     || output_reg.starts_with("char *"))
+                //     && !reg.value.starts_with("\"")
+                // {
+                //     read_memory(pid, reg.value.parse::<u64>().unwrap())
+                // } else {
+                //     reg.value.to_string()
+                // };
+                spans.push(Span::styled(format!("{}, ", reg.value), Style::default()));
             }
         }
 
diff --git a/src/trace.rs b/src/trace.rs
index e424d09..8ce01cf 100644
--- a/src/trace.rs
+++ b/src/trace.rs
@@ -64,7 +64,7 @@ pub fn trace(pid: Pid, args: &Args) -> anyhow::Result<Vec<RegistersData>> {
         Some(filter) => filter.split(',').collect::<Vec<&str>>(),
         None => vec![],
     };
-    while let Some(reg) = trace_next(pid)? {
+    while let Some(mut reg) = trace_next(pid)? {
         have_to_print ^= true;
         if have_to_print {
             if !filters.is_empty() && !filters.contains(&reg.name()) {
diff --git a/src/ui.rs b/src/ui.rs
index f846349..92e2a04 100644
--- a/src/ui.rs
+++ b/src/ui.rs
@@ -34,8 +34,11 @@ impl UI {
         self.max_lines = self.lines.len() + 1;
     }
 
-    pub fn get_paragraph(&self, pid: Pid) -> Paragraph {
-        let lines: Vec<Line> = self.lines.iter().map(|x| x.output_ui(pid)).collect();
+    pub fn get_paragraph(&mut self, pid: Pid) -> Paragraph {
+        let mut lines: Vec<Line> = vec![];
+        for line in &mut self.lines {
+            lines.push(line.output_ui(pid));
+        }
         let paragraph = Paragraph::new(lines)
             .block(
                 Block::default()