rust: intercept wrapper simplified

Author: rizsotto

rust/bear/src/bin/wrapper.rs

@@ -20,8 +20,7 @@ extern crate core;
 use anyhow::{Context, Result};
 use bear::intercept::create_reporter;
 use bear::intercept::supervise::supervise;
-use bear::intercept::{Event, Execution, ProcessId};
-use std::path::{Path, PathBuf};
+use bear::intercept::{Event, Execution};
 
 /// Implementation of the wrapper process.
 ///
@@ -31,57 +30,43 @@ use std::path::{Path, PathBuf};
 /// the execution.
 fn main() -> Result<()> {
     env_logger::init();
-    // Find out what is the executable name the execution was started with
-    let executable = file_name_from_arguments()?;
-    log::info!("Executable as called: {:?}", executable);
+    // Capture the current process execution details
+    let execution = Execution::capture().with_context(|| "Failed to capture the execution")?;
+    log::info!("Execution captured: {:?}", execution);
     // Read the PATH variable and find the next executable with the same name
-    let real_executable = next_in_path(&executable)?;
-    log::info!("Executable to call: {:?}", real_executable);
+    let real_executable = next_in_path(&execution.executable)?;
+    let real_execution = execution.with_executable(&real_executable);
+    log::info!("Execution to call: {:?}", real_execution);
 
     // Reporting failures shall not fail the execution.
-    match into_execution(&real_executable).and_then(report) {
+    match report(real_execution.clone()) {
         Ok(_) => log::info!("Execution reported"),
         Err(e) => log::error!("Execution reporting failed: {}", e),
     }
 
     // Execute the real executable with the same arguments
-    let mut command = std::process::Command::new(real_executable);
-    let exit_status = supervise(command.args(std::env::args().skip(1)))?;
+    let exit_status = supervise(real_execution)?;
     log::info!("Execution finished with status: {:?}", exit_status);
     // Return the child process status code
     std::process::exit(exit_status.code().unwrap_or(1));
 }
 
-/// Get the file name of the executable from the arguments.
-///
-/// Since the executable will be called via soft link, the first argument
-/// will be the name of the soft link. This function returns the file name
-/// of the soft link.
-fn file_name_from_arguments() -> Result<PathBuf> {
-    std::env::args()
-        .next()
-        .ok_or_else(|| anyhow::anyhow!("Cannot get first argument"))
-        .and_then(|arg| match PathBuf::from(arg).file_name() {
-            Some(file_name) => Ok(PathBuf::from(file_name)),
-            None => Err(anyhow::anyhow!(
-                "Cannot get the file name from the argument"
-            )),
-        })
-}
-
 /// Find the next executable in the PATH variable.
 ///
 /// The function reads the PATH variable and tries to find the next executable
 /// with the same name as the given executable. It returns the path to the
 /// executable.
-fn next_in_path(target: &Path) -> Result<PathBuf> {
-    let path = std::env::var("PATH")?;
+fn next_in_path(current_exe: &std::path::Path) -> Result<std::path::PathBuf> {
+    let target = current_exe
+        .file_name()
+        .with_context(|| "Cannot get the file name of the executable")?;
+    let path =
+        std::env::var("PATH").with_context(|| "Cannot get the PATH variable from environment")?;
+
     log::debug!("PATH: {}", path);
-    // The `current_exe` is a canonical path to the current executable.
-    let current_exe = std::env::current_exe()?;
 
-    path.split(':')
-        .map(|dir| Path::new(dir).join(target))
+    std::env::split_paths(&path)
+        .map(|dir| dir.join(target))
         .filter(|path| path.is_file())
         .find(|path| {
             // We need to compare it with the real path of the candidate executable to avoid
@@ -96,22 +81,10 @@ fn next_in_path(target: &Path) -> Result<PathBuf> {
 }
 
 fn report(execution: Execution) -> Result<()> {
-    let event = Event {
-        pid: ProcessId(std::process::id()),
-        execution,
-    };
+    let event = Event::new(execution);
 
     create_reporter()
         .with_context(|| "Cannot create execution reporter")?
         .report(event)
         .with_context(|| "Sending execution failed")
 }
-
-fn into_execution(path_buf: &Path) -> Result<Execution> {
-    Ok(Execution {
-        executable: path_buf.to_path_buf(),
-        arguments: std::env::args().collect(),
-        working_dir: std::env::current_dir()?,
-        environment: std::env::vars().collect(),
-    })
-}

rust/bear/src/intercept/mod.rs

@@ -17,7 +17,7 @@ use anyhow::Context;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::path::{Path, PathBuf};
-use std::process::{Command, ExitCode};
+use std::process::ExitCode;
 use std::sync::mpsc::{channel, Receiver, Sender};
 use std::sync::Arc;
 use std::{fmt, thread};
@@ -73,24 +73,28 @@ pub trait Collector {
     fn stop(&self) -> Result<(), anyhow::Error>;
 }
 
-/// Process id is an OS identifier for a process.
-#[derive(Serialize, Deserialize, Debug, PartialEq, Clone)]
-pub struct ProcessId(pub u32);
-
 /// Represent a relevant life cycle event of a process.
 ///
 /// In the current implementation, we only have one event, the `Started` event.
 /// This event is sent when a process is started. It contains the process id
 /// and the execution information.
 #[derive(Serialize, Deserialize, Debug, PartialEq, Clone)]
 pub struct Event {
-    pub pid: ProcessId,
+    pub pid: u32,
     pub execution: Execution,
 }
 
+impl Event {
+    /// Creates a new event that is originated from the current process.
+    pub fn new(execution: Execution) -> Self {
+        let pid = std::process::id();
+        Event { pid, execution }
+    }
+}
+
 impl fmt::Display for Event {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "Event pid={}, execution={}", self.pid.0, self.execution)
+        write!(f, "Event pid={}, execution={}", self.pid, self.execution)
     }
 }
 
@@ -107,6 +111,69 @@ pub struct Execution {
     pub environment: HashMap<String, String>,
 }
 
+impl Execution {
+    /// Capture the execution information of the current process.
+    pub fn capture() -> anyhow::Result<Self> {
+        let executable = std::env::current_exe()?;
+        let arguments = std::env::args().collect();
+        let working_dir = std::env::current_dir()?;
+        let environment = std::env::vars().collect();
+
+        Ok(Self {
+            executable,
+            arguments,
+            working_dir,
+            environment,
+        })
+    }
+
+    pub fn with_executable(&self, executable: &Path) -> Self {
+        let mut updated = self.clone();
+        updated.executable = executable.to_path_buf();
+        updated
+    }
+
+    pub fn with_environment(&self, environment: HashMap<String, String>) -> Self {
+        let mut updated = self.clone();
+        updated.environment = environment;
+        updated
+    }
+}
+
+impl TryFrom<args::BuildCommand> for Execution {
+    type Error = anyhow::Error;
+
+    /// Converts the `BuildCommand` to an `Execution` object.
+    fn try_from(value: args::BuildCommand) -> Result<Self, Self::Error> {
+        let executable = value
+            .arguments
+            .first()
+            .ok_or_else(|| anyhow::anyhow!("No executable found"))?
+            .clone()
+            .into();
+        let arguments = value.arguments.to_vec();
+        let working_dir = std::env::current_dir()?;
+        let environment = std::env::vars().collect();
+
+        Ok(Self {
+            executable,
+            arguments,
+            working_dir,
+            environment,
+        })
+    }
+}
+
+impl From<Execution> for std::process::Command {
+    fn from(val: Execution) -> Self {
+        let mut command = std::process::Command::new(val.executable);
+        command.args(val.arguments.iter().skip(1));
+        command.envs(val.environment);
+        command.current_dir(val.working_dir);
+        command
+    }
+}
+
 impl fmt::Display for Execution {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
         write!(
@@ -145,7 +212,7 @@ pub fn event(
     environment: HashMap<&str, &str>,
 ) -> Event {
     Event {
-        pid: ProcessId(pid),
+        pid,
         execution: execution(executable, arguments, working_dir, environment),
     }
 }
@@ -286,13 +353,9 @@ impl InterceptEnvironment {
     pub fn execute_build_command(&self, input: args::BuildCommand) -> anyhow::Result<ExitCode> {
         // TODO: record the execution of the build command
 
-        let environment = self.environment();
-        let process = input.arguments[0].clone();
-        let arguments = input.arguments[1..].to_vec();
-
-        let mut child = Command::new(process);
-
-        let exit_status = supervise(child.args(arguments).envs(environment))?;
+        let child: Execution = TryInto::<Execution>::try_into(input)?
+            .with_environment(self.environment().into_iter().collect());
+        let exit_status = supervise(child)?;
         log::info!("Execution finished with status: {:?}", exit_status);
 
         // The exit code is not always available. When the process is killed by a signal,

rust/bear/src/intercept/supervise.rs

@@ -1,7 +1,8 @@
 // SPDX-License-Identifier: GPL-3.0-or-later
 
+use crate::intercept::Execution;
 use anyhow::Result;
-use std::process::{Command, ExitStatus};
+use std::process::ExitStatus;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::thread;
@@ -12,13 +13,13 @@ use std::time;
 /// It starts the command and waits for its completion. It also forwards
 /// signals to the child process. The method returns the exit status of the
 /// child process.
-pub fn supervise(command: &mut Command) -> Result<ExitStatus> {
+pub fn supervise(execution: Execution) -> Result<ExitStatus> {
     let signaled = Arc::new(AtomicUsize::new(0));
     for signal in signal_hook::consts::TERM_SIGNALS {
         signal_hook::flag::register_usize(*signal, Arc::clone(&signaled), *signal as usize)?;
     }
 
-    let mut child = command.spawn()?;
+    let mut child = Into::<std::process::Command>::into(execution).spawn()?;
     loop {
         // Forward signals to the child process, but don't exit the loop while it is running
         if signaled.swap(0usize, Ordering::SeqCst) != 0 {