Trace attaching a pid

5 files changed, 201 insertions, 111 deletions

diff --git a/src/cli.rs b/src/cli.rs
new file mode 100644
index 0000000..685b26f
--- /dev/null
+++ b/src/cli.rs
@@ -0,0 +1,22 @@
+use clap::Parser;
+
+#[derive(Parser)]
+#[command(author, version, about, long_about = None)]
+pub struct Args {
+    /// Command to execute from ptrace
+    #[arg(short, long)]
+    pub command: Option<String>,
+
+    /// Attach the tracing to an existing process ID. We're using the `-p` short flag because
+    /// strace uses it
+    #[arg(short = 'p', long)]
+    pub attach: Option<i32>,
+
+    /// Write the output to a file instead of the standard output
+    #[arg(short = 'f', long = "file")]
+    pub file_to_print: Option<String>,
+
+    /// If defined, it hides the TUI
+    #[arg(long = "no-tui", default_value_t = false)]
+    pub no_tui: bool,
+}
diff --git a/src/main.rs b/src/main.rs
index 19a9533..136c8c6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,50 +1,46 @@
+mod cli;
 mod registers;
 mod trace;
 mod ui;
-use std::{
-    io::{self, Write},
-    str,
-};
 
+use crate::cli::Args;
 use crate::trace::{exec, trace};
-use crate::ui::run_tui;
+use crate::ui::UI;
+
 use clap::Parser;
 use fork::{fork, Fork};
 use nix::unistd::Pid;
-
-#[derive(Parser)]
-#[command(author, version, about, long_about = None)]
-struct Args {
-    /// Command to execute from ptrace
-    command: String,
-
-    /// Write the output to a file instead of the standard output
-    #[arg(short = 'f', long = "file")]
-    file_to_print: Option<String>,
-
-    /// If defined, it hides the TUI
-    #[arg(long = "no-tui", default_value_t = false)]
-    no_tui: bool,
-}
+use trace::attach;
 
 /// Create a fork of the program and execute the process in the child. Parent gets the pid
 /// value and trace it.
 fn main() -> anyhow::Result<()> {
     let args = Args::parse();
 
-    let pid = match fork() {
-        Ok(Fork::Child) => return exec(&args.command),
-        Ok(Fork::Parent(child)) => Pid::from_raw(child),
-        Err(err) => panic!("fork() failed: {err}"),
+    let pid = if args.command.is_some() {
+        match fork() {
+            Ok(Fork::Child) => return exec(&args.command.unwrap()),
+            Ok(Fork::Parent(child)) => Pid::from_raw(child),
+            Err(err) => panic!("fork() failed: {err}"),
+        }
+    } else if args.attach.is_some() {
+        let pid = Pid::from_raw(args.attach.unwrap());
+
+        if attach(pid).is_ok() {
+            pid
+        } else {
+            panic!("Unable to attach to process {pid}");
+        }
+    } else {
+        panic!("You must define a command or a PID to attach");
     };
-    let registers = trace(pid, args.file_to_print)?;
 
     if !args.no_tui {
-        run_tui(pid, &registers)?;
+        let mut ui = UI::new();
+
+        ui.start(pid, &args)?;
     } else {
-        for line in registers {
-            writeln!(io::stdout(), "{}", line.output())?;
-        }
+        trace(pid, &args, true)?;
     }
 
     Ok(())
diff --git a/src/registers.rs b/src/registers.rs
index d73424d..cf2aa98 100644
--- a/src/registers.rs
+++ b/src/registers.rs
@@ -6,6 +6,7 @@ use ratatui::{
 };
 
 /// Struct used to manipulate registers data from https://docs.rs/libc/0.2.147/libc/struct.user_regs_struct.html
+#[derive(Debug)]
 pub struct RegistersData {
     orig_rax: u64,
     rdi: u64,
diff --git a/src/trace.rs b/src/trace.rs
index 3612519..c097fba 100644
--- a/src/trace.rs
+++ b/src/trace.rs
@@ -1,3 +1,4 @@
+use crate::cli::Args;
 use crate::registers::RegistersData;
 use nix::{
     sys::{
@@ -7,7 +8,14 @@ use nix::{
     },
     unistd::Pid,
 };
-use std::{fs::File, io::Write, os::unix::process::CommandExt, process::Command};
+use std::{
+    fs::File,
+    // fs::File,
+    io::{self, Write},
+    os::unix::process::CommandExt,
+    process::Command,
+    str,
+};
 
 /// Exec the `command` value tracing it with `ptrace` lib
 pub fn exec(command: &str) -> anyhow::Result<()> {
@@ -25,54 +33,82 @@ pub fn exec(command: &str) -> anyhow::Result<()> {
     Ok(())
 }
 
-/// Trace a process with `pid` ID and returns a list of `RegistersData`
-pub fn trace(pid: Pid, file_to_print: Option<String>) -> anyhow::Result<Vec<RegistersData>> {
-    // Since you have to do 2 syscalls (start and end) you have to alternate the print value,
-    // because it could be equals except for the `rax` register.
-    let mut have_to_print = true;
+/// Attach a ptrace status to a `pid`
+pub fn attach(pid: Pid) -> anyhow::Result<()> {
+    ptrace::attach(pid)?;
+
+    Ok(())
+}
 
+/// Trace a process with `pid` ID and returns a list of `RegistersData`
+pub fn trace(pid: Pid, args: &Args, run_loop: bool) -> anyhow::Result<Vec<RegistersData>> {
     // First wait for the parent process
     _ = waitpid(pid, None)?;
 
-    // If `fiole_to_print` is not None, create a new file with that value for redirecting all the
+    // FIXME: file writing on attachment
+    // If `file_to_print` is not None, create a new file with that value for redirecting all the
     // output (also in stdout)
     let mut f = None;
-    if let Some(filename) = file_to_print {
+    if let Some(filename) = &args.file_to_print {
         f = Some(File::create(filename)?);
     }
 
     let mut lines: Vec<RegistersData> = Vec::new();
 
-    loop {
-        have_to_print ^= true;
-        ptrace::syscall(pid, None)?;
-        let status = waitpid(pid, None)?;
+    if run_loop {
+        // Since you have to do 2 syscalls (start and end) you have to alternate the print value,
+        // because it could be equals except for the `rax` register.
+        let mut have_to_print = true;
 
-        match status {
-            // Break the loop if the process exists
-            WaitStatus::Exited(_pid, _) => {
-                break;
-            }
-            // Match the stopped value for a process
-            WaitStatus::Stopped(pid, signal) => {
-                match signal {
-                    Signal::SIGTRAP => {
-                        if have_to_print {
-                            let reg = RegistersData::new(ptrace::getregs(pid)?);
-
-                            if let Some(ref mut f) = f {
-                                writeln!(f, "{}", reg.output())?;
-                            }
-
-                            lines.push(reg);
+        loop {
+            match trace_next(pid)? {
+                Some(reg) => {
+                    have_to_print ^= true;
+                    if have_to_print {
+                        if let Some(ref mut f) = f {
+                            writeln!(f, "{}", reg.output())?;
+                        }
+
+                        if args.no_tui {
+                            writeln!(io::stdout(), "{}", reg.output())?;
                         }
+
+                        lines.push(reg);
                     }
-                    _ => {}
-                };
+                }
+                None => {
+                    break;
+                }
             }
-            _ => {}
-        };
+        }
     }
-
     Ok(lines)
 }
+
+/// Get the next step for a ptrace process
+pub fn trace_next(pid: Pid) -> anyhow::Result<Option<RegistersData>> {
+    ptrace::syscall(pid, None)?;
+    let status = waitpid(pid, None).unwrap();
+
+    match status {
+        // Match the stopped value for a process
+        WaitStatus::Stopped(pid, signal) => {
+            match signal {
+                Signal::SIGTRAP => {
+                    let reg = RegistersData::new(ptrace::getregs(pid)?);
+                    return Ok(Some(reg));
+                }
+                _ => {}
+            };
+        }
+        _ => {}
+    };
+
+    Ok(None)
+}
+
+/// Kill a process traced by ptrace
+pub fn trace_kill(pid: Pid) -> anyhow::Result<()> {
+    let _ = ptrace::kill(pid);
+    Ok(())
+}
diff --git a/src/ui.rs b/src/ui.rs
index f881c41..8fd44b1 100644
--- a/src/ui.rs
+++ b/src/ui.rs
@@ -1,4 +1,8 @@
-use crate::registers::RegistersData;
+use crate::{
+    cli::Args,
+    registers::RegistersData,
+    trace::{trace, trace_kill, trace_next},
+};
 use crossterm::{
     event::{self, Event, KeyCode},
     terminal::{disable_raw_mode, enable_raw_mode, EnterAlternateScreen, LeaveAlternateScreen},
@@ -8,19 +12,95 @@ use nix::unistd::Pid;
 use ratatui::{prelude::*, widgets::*};
 use std::io::{self, stdout};
 
-struct UI {
+pub struct UI {
     height: usize,
     max_lines: usize,
     scroll: usize,
+    lines: Vec<RegistersData>,
 }
 
 impl UI {
-    fn new() -> UI {
+    pub fn new() -> UI {
         UI {
             height: 0,
             max_lines: 0,
             scroll: 0,
+            lines: vec![],
+        }
+    }
+
+    pub fn add_line(&mut self, registers: RegistersData) {
+        self.lines.push(registers);
+        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()).collect();
+        let paragraph = Paragraph::new(lines)
+            .block(
+                Block::default()
+                    .border_style(Style::default().fg(Color::Yellow))
+                    .title(format!("[{pid}]"))
+                    .title(
+                        block::Title::from(format!(
+                            "[lines {}-{}]",
+                            self.scroll,
+                            self.scroll + self.height
+                        ))
+                        .position(block::Position::Bottom)
+                        .alignment(Alignment::Right),
+                    )
+                    .borders(Borders::ALL),
+            )
+            .scroll((self.scroll as u16, 0));
+
+        paragraph
+    }
+
+    pub fn start(&mut self, pid: Pid, args: &Args) -> anyhow::Result<()> {
+        enable_raw_mode()?;
+        stdout().execute(EnterAlternateScreen)?;
+        let mut terminal = Terminal::new(CrosstermBackend::new(stdout()))?;
+
+        let mut have_to_print = true;
+        let mut have_to_trace = !args.command.is_some();
+        let mut should_quit = false;
+
+        let lines = trace(pid, &args, args.command.is_some())?;
+        if lines.len() > 1 {
+            for line in lines {
+                self.add_line(line);
+            }
+        }
+        while !should_quit {
+            if have_to_trace {
+                if let Some(reg) = trace_next(pid)? {
+                    have_to_print ^= true;
+                    if have_to_print {
+                        self.add_line(reg);
+                    }
+                } else {
+                    have_to_trace = false;
+                }
+            }
+
+            self.height = terminal.get_frame().size().height as usize;
+            terminal.draw(|frame| {
+                let size = frame.size();
+
+                frame.render_widget(self.get_paragraph(pid), size);
+            })?;
+
+            should_quit = handle_events(self)?;
         }
+
+        // FIXME: avoid this kill without Rust errors
+        let _ = trace_kill(pid);
+
+        disable_raw_mode()?;
+        stdout().execute(LeaveAlternateScreen)?;
+
+        Ok(())
     }
 }
 
@@ -55,48 +135,3 @@ fn handle_events(ui: &mut UI) -> io::Result<bool> {
     }
     Ok(false)
 }
-
-pub fn run_tui(pid: Pid, registers: &Vec<RegistersData>) -> anyhow::Result<()> {
-    enable_raw_mode()?;
-    stdout().execute(EnterAlternateScreen)?;
-    let mut terminal = Terminal::new(CrosstermBackend::new(stdout()))?;
-
-    let mut ui = UI::new();
-    ui.max_lines = registers.len() + 1;
-
-    let mut should_quit = false;
-    while !should_quit {
-        ui.height = terminal.get_frame().size().height as usize;
-        terminal.draw(move |frame| {
-            let size = frame.size();
-            let lines: Vec<Line> = registers.iter().map(|x| x.output_ui()).collect();
-
-            frame.render_widget(
-                Paragraph::new(lines)
-                    .block(
-                        Block::default()
-                            .border_style(Style::default().fg(Color::Yellow))
-                            .title(format!("[{pid}]"))
-                            .title(
-                                block::Title::from(format!(
-                                    "[lines {}-{}]",
-                                    ui.scroll,
-                                    ui.scroll + ui.height
-                                ))
-                                .position(block::Position::Bottom)
-                                .alignment(Alignment::Right),
-                            )
-                            .borders(Borders::ALL),
-                    )
-                    .scroll((ui.scroll as u16, 0)),
-                size,
-            );
-        })?;
-        should_quit = handle_events(&mut ui)?;
-    }
-
-    disable_raw_mode()?;
-    stdout().execute(LeaveAlternateScreen)?;
-
-    Ok(())
-}