From 36fc4e67d4b8c23d726234f34113ae518d4a12e7 Mon Sep 17 00:00:00 2001
From: Santo Cariotti <santo@dcariotti.me>
Date: Sun, 15 Oct 2023 18:27:15 +0200
Subject: Add multiprocessing

---
 Cargo.lock   | 10 ++++++++++
 Cargo.toml   |  1 +
 src/main.rs  | 17 ++++++++---------
 src/trace.rs | 17 ++++++++++++++---
 4 files changed, 33 insertions(+), 12 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 024a1e6..e82261e 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -114,6 +114,15 @@ version = "1.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "acbf1af155f9b9ef647e42cdc158db4b64a1b61f743629225fde6f3e0be2a7c7"
 
+[[package]]
+name = "fork"
+version = "0.1.22"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bf2ca97a59201425e7ee4d197c9c4fea282fe87a97d666a580bda889b95b8e88"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "heck"
 version = "0.4.1"
@@ -161,6 +170,7 @@ version = "0.0.1"
 dependencies = [
  "anyhow",
  "clap",
+ "fork",
  "nix",
 ]
 
diff --git a/Cargo.toml b/Cargo.toml
index 095aa18..290d613 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -9,4 +9,5 @@ authors = ["Santo Cariotti <santo@dcariotti.me>"]
 [dependencies]
 anyhow = "1.0.75"
 clap = { version = "4.4.6", features = ["derive"] }
+fork = "0.1.22"
 nix = { version = "0.27.1", features = ["ptrace"] }
diff --git a/src/main.rs b/src/main.rs
index 893b9f5..0694aa5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -2,7 +2,8 @@ mod trace;
 
 use crate::trace::{exec, trace};
 use clap::Parser;
-use std::process::Command;
+use fork::{fork, Fork};
+use nix::unistd::Pid;
 
 #[derive(Parser)]
 #[command(author, version, about, long_about = None)]
@@ -12,15 +13,13 @@ struct Args {
 
 fn main() -> anyhow::Result<()> {
     let args = Args::parse();
-    let params = args.command.split(' ').collect::<Vec<&str>>();
 
-    let mut command = Command::new(params[0]);
-    if params.len() > 1 {
-        for arg in &params[1..] {
-            command.arg(arg);
-        }
-    }
-    let pid = exec(&mut command)?;
+    let pid = match fork() {
+        Ok(Fork::Child) => return exec(&args.command),
+        Ok(Fork::Parent(child)) => Pid::from_raw(child as i32),
+        Err(err) => panic!("fork() failed: {err}"),
+    };
+
     trace(pid)?;
 
     Ok(())
diff --git a/src/trace.rs b/src/trace.rs
index 0031fcd..b00b9f9 100644
--- a/src/trace.rs
+++ b/src/trace.rs
@@ -8,16 +8,27 @@ use nix::{
 };
 use std::{os::unix::process::CommandExt, process::Command};
 
-pub fn exec(command: &mut Command) -> anyhow::Result<Pid> {
+pub fn exec(command: &String) -> anyhow::Result<()> {
+    let params: Vec<&str> = command.split(' ').collect();
+
+    let mut command = Command::new(params[0]);
+    command.args(params[1..].iter());
+
     unsafe {
         command.pre_exec(|| ptrace::traceme().map_err(|e| e.into()));
     }
-    let child = command.spawn()?;
-    Ok(Pid::from_raw(child.id() as i32))
+
+    command.exec();
+
+    Ok(())
 }
 
 pub fn trace(pid: Pid) -> anyhow::Result<()> {
     let mut have_to_print = true;
+
+    // First wait if for the parent process
+    _ = waitpid(pid, None)?;
+
     loop {
         have_to_print ^= true;
         ptrace::syscall(pid, None)?;
-- 
cgit v1.2.3-18-g5258