lambda-microvm-hook-server 0.1.0

use crate::MicroVmError;
use nix::sys::signal::{Signal, killpg};
use nix::unistd::Pid;
use std::collections::BTreeMap;
use std::future::Future;
use std::os::unix::process::CommandExt;
use std::pin::Pin;
use std::process::Stdio;
use tokio::process::Command;
use tokio_util::sync::CancellationToken;

pub(crate) type CommandResult = Pin<Box<dyn Future<Output = Result<(), MicroVmError>> + Send>>;

pub(crate) fn spawn_command(
    command: String,
    args: Vec<String>,
    environment: BTreeMap<String, String>,
    cancellation: CancellationToken,
) -> Result<CommandResult, MicroVmError> {
    let mut child = {
        let mut process = Command::new(command);
        process.args(args).envs(environment).stdin(Stdio::null()).stdout(Stdio::inherit()).stderr(Stdio::inherit());
        process.as_std_mut().process_group(0);
        process.spawn().map_err(MicroVmError::CommandSpawn)?
    };

    let process_group = child.id().map(|id| Pid::from_raw(id.cast_signed()));
    Ok(Box::pin(async move {
        tokio::select! {
            status = child.wait() => {
                let status = status.map_err(MicroVmError::CommandWait)?;
                if status.success() { Ok(()) } else { Err(MicroVmError::CommandFailed) }
            }
            () = cancellation.cancelled() => {
                if let Some(process_group) = process_group {
                    let _ = killpg(process_group, Signal::SIGKILL);
                } else {
                    let _ = child.kill().await;
                }
                child.wait().await.map_err(MicroVmError::CommandWait)?;
                Ok(())
            }
        }
    }))
}