Skip to content

Commit

Permalink
rust: intercept mode commented
Browse files Browse the repository at this point in the history
  • Loading branch information
rizsotto committed Nov 17, 2024
1 parent 13e6560 commit ad0bf6d
Show file tree
Hide file tree
Showing 3 changed files with 69 additions and 12 deletions.
1 change: 1 addition & 0 deletions rust/bear/src/bin/bear.rs
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,7 @@ impl Application {
Application::Semantic(semantic) => semantic.run(),
Application::All(all) => all.run(),
};
// TODO: log the status
status.unwrap_or_else(|_| ExitCode::FAILURE)
}
}
78 changes: 66 additions & 12 deletions rust/bear/src/modes/intercept.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,13 +9,26 @@ use std::process::{Command, ExitCode};
use std::sync::Arc;
use std::{env, thread};

/// The service is responsible for collecting the events from the supervised processes.
///
/// The service is implemented as TCP server that listens on a random port on the loopback
/// interface. The address of the service can be obtained by the `address` method.
///
/// The service is started in a separate thread to dispatch the events to the consumer.
/// The consumer is a function that receives the events from the service and processes them.
/// It also runs in a separate thread. The reason for having two threads is to avoid blocking
/// the main thread of the application and decouple the collection from the processing.
pub(super) struct InterceptService {
collector: Arc<EventCollectorOnTcp>,
network_thread: Option<thread::JoinHandle<()>>,
output_thread: Option<thread::JoinHandle<()>>,
}

impl InterceptService {
/// Creates a new intercept service.
///
/// The `consumer` is a function that receives the events and processes them.
/// The function is executed in a separate thread.
pub fn new<F>(consumer: F) -> anyhow::Result<Self>
where
F: FnOnce(Receiver<Envelope>) -> anyhow::Result<()>,
Expand All @@ -27,27 +40,33 @@ impl InterceptService {

let collector_in_thread = collector_arc.clone();
let collector_thread = thread::spawn(move || {
// TODO: log failures
collector_in_thread.collect(sender).unwrap();
});
let receiver_in_thread = receiver.clone();
let output_thread = thread::spawn(move || {
// TODO: log failures
consumer(receiver_in_thread).unwrap();
});

// TODO: log the address of the service
Ok(InterceptService {
collector: collector_arc,
network_thread: Some(collector_thread),
output_thread: Some(output_thread),
})
}

/// Returns the address of the service.
pub fn address(&self) -> String {
self.collector.address()
}
}

impl Drop for InterceptService {
/// Shuts down the service.
fn drop(&mut self) {
// TODO: log the shutdown of the service and any errors
self.collector.stop().expect("Failed to stop the collector");
if let Some(thread) = self.network_thread.take() {
thread.join().expect("Failed to join the collector thread");
Expand All @@ -58,6 +77,17 @@ impl Drop for InterceptService {
}
}

/// The environment for the intercept mode.
///
/// Running the build command requires a specific environment. The environment we
/// need for intercepting the child processes is different for each intercept mode.
///
/// The `Wrapper` mode requires a temporary directory with the executables that will
/// be used to intercept the child processes. The executables are hard linked to the
/// temporary directory.
///
/// The `Preload` mode requires the path to the preload library that will be used to
/// intercept the child processes.
pub(super) enum InterceptEnvironment {
Wrapper {
bin_dir: tempfile::TempDir,
Expand All @@ -70,6 +100,11 @@ pub(super) enum InterceptEnvironment {
}

impl InterceptEnvironment {
/// Creates a new intercept environment.
///
/// The `config` is the intercept configuration that specifies the mode and the
/// required parameters for the mode. The `address` is the address of the intercept
/// service that will be used to collect the events.
pub fn new(config: &config::Intercept, address: String) -> anyhow::Result<Self> {
let result = match config {
config::Intercept::Wrapper {
Expand All @@ -92,24 +127,38 @@ impl InterceptEnvironment {
Ok(result)
}

/// Executes the build command in the intercept environment.
///
/// The method is blocking and waits for the build command to finish.
/// The method returns the exit code of the build command. Result failure
/// indicates that the build command failed to start.
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 mut child = Command::new(input.arguments[0].clone())
.args(input.arguments)
.args(input.arguments[1..].iter())
.envs(environment)
.spawn()?;

// TODO: forward signals to the child process
let result = child.wait()?;

if result.success() {
Ok(ExitCode::SUCCESS)
} else {
result
.code()
.map_or(Ok(ExitCode::FAILURE), |code| Ok(ExitCode::from(code as u8)))
}
// The exit code is not always available. When the process is killed by a signal,
// the exit code is not available. In this case, we return the `FAILURE` exit code.
let exit_code = result
.code()
.map(|code| ExitCode::from(code as u8))
.unwrap_or(ExitCode::FAILURE);

Ok(exit_code)
}

/// Returns the environment variables for the intercept environment.
///
/// The environment variables are different for each intercept mode.
/// It does not change the original environment variables, but creates
/// the environment variables that are required for the intercept mode.
fn environment(&self) -> Vec<(String, String)> {
match self {
InterceptEnvironment::Wrapper {
Expand All @@ -118,7 +167,7 @@ impl InterceptEnvironment {
let path_original = env::var("PATH").unwrap_or_else(|_| String::new());
let path_updated = InterceptEnvironment::insert_to_path(
&path_original,
Self::to_string(bin_dir.path()),
Self::path_to_string(bin_dir.path()),
);
vec![
("PATH".to_string(), path_updated),
Expand All @@ -127,8 +176,10 @@ impl InterceptEnvironment {
}
InterceptEnvironment::Preload { path, address, .. } => {
let path_original = env::var(KEY_PRELOAD_PATH).unwrap_or_else(|_| String::new());
let path_updated =
InterceptEnvironment::insert_to_path(&path_original, Self::to_string(path));
let path_updated = InterceptEnvironment::insert_to_path(
&path_original,
Self::path_to_string(path),
);
vec![
(KEY_PRELOAD_PATH.to_string(), path_updated),
(KEY_DESTINATION.to_string(), address.clone()),
Expand All @@ -137,13 +188,16 @@ impl InterceptEnvironment {
}
}

/// Manipulate a `PATH` like environment value by inserting the `first` path into
/// the original value. It removes the `first` path if it already exists in the
/// original value. And it inserts the `first` path at the beginning of the value.
fn insert_to_path(original: &str, first: String) -> String {
let mut paths: Vec<_> = original.split(':').filter(|it| it != &first).collect();
paths.insert(0, first.as_str());
paths.join(":")
}

fn to_string(path: &Path) -> String {
fn path_to_string(path: &Path) -> String {
path.to_str().unwrap_or("").to_string()
}
}
2 changes: 2 additions & 0 deletions rust/bear/src/modes/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ impl Intercept {
}
}

/// Write the envelopes into the output file.
fn write_to_file(
output_file_name: String,
envelopes: Receiver<Envelope>,
Expand All @@ -76,6 +77,7 @@ impl Intercept {

impl Mode for Intercept {
fn run(self) -> anyhow::Result<ExitCode> {
// TODO: log failures with the right context
let output_file_name = self.output.file_name.clone();
let service = InterceptService::new(move |envelopes| {
Self::write_to_file(output_file_name, envelopes)
Expand Down

0 comments on commit ad0bf6d

Please sign in to comment.