schematics/src/command_runner.rs

use async_trait::async_trait;
#[cfg(test)]
use mockall::mock;
use std::error::Error;
use std::ffi::OsStr;
use std::io::Result as IoResult;
use std::process::Output;
use std::process::Stdio;
use tokio::io::AsyncWriteExt;
use tokio::process::Command;

#[macro_export]
macro_rules! args {
    ($($x:expr),*) => (
        &[$($x.as_ref()),*]
    );
    ($($x:expr,)*) => (args![$($x),*]) // handle trailing commas
}

fn check_success(output: Output) -> Result<Output, Box<dyn Error>> {
  if output.status.success() {
    Ok(output)
  } else {
    Err(std::str::from_utf8(&output.stderr)?.into())
  }
}

pub fn is_success(res: Result<Output, impl Error + 'static>) -> Result<Output, Box<dyn Error>> {
  check_success(res?)
}

pub fn get_output(output: Output) -> Result<Vec<u8>, Box<dyn Error>> {
  Ok(check_success(output)?.stdout)
}

#[async_trait(?Send)]
pub trait CommandRunner {
  async fn run<'a>(&self, program: &str, args: &'a [&'a OsStr], input: &[u8]) -> IoResult<Output>;

  async fn run_with_args<'a>(&self, program: &str, args: &'a [&'a OsStr]) -> IoResult<Output> {
    self.run(program, args, b"").await
  }
  async fn get_output<'a>(
    &self,
    program: &str,
    args: &'a [&'a OsStr],
  ) -> Result<Vec<u8>, Box<dyn Error>> {
    let output = self.run_with_args(program, args).await?;
    get_output(output)
  }
  async fn run_successfully<'a>(
    &self,
    program: &str,
    args: &'a [&'a OsStr],
  ) -> Result<(), Box<dyn Error>> {
    is_success(self.run(program, args, b"").await)?;
    Ok(())
  }
  async fn get_stderr<'a>(
    &self,
    program: &str,
    args: &'a [&'a OsStr],
  ) -> Result<Vec<u8>, Box<dyn Error>> {
    Ok(is_success(self.run_with_args(program, args).await)?.stderr)
  }
}

#[cfg(test)]
mock! {
  pub CommandRunner {
  }
  #[async_trait(?Send)]
  impl CommandRunner for CommandRunner {
    async fn run<'a>(&self, program: &str, args: &'a [&'a OsStr], input: &[u8]) -> IoResult<Output>;

    async fn run_with_args<'a>(&self, program: &str, args: &'a [&'a OsStr]) -> IoResult<Output>;
    async fn get_output<'a>(&self, program: &str, args: &'a [&'a OsStr]) -> Result<Vec<u8>, Box<dyn Error>>;
    async fn run_successfully<'a>(&self, program: &str, args: &'a [&'a OsStr]) -> Result<(), Box<dyn Error>>;
    async fn get_stderr<'a>(&self, program: &str, args: &'a [&'a OsStr]) -> Result<Vec<u8>, Box<dyn Error>>;
  }
}

#[derive(Debug)]
pub struct StdCommandRunner;

#[async_trait(?Send)]
impl CommandRunner for StdCommandRunner {
  async fn run<'a>(&self, program: &str, args: &'a [&'a OsStr], input: &[u8]) -> IoResult<Output> {
    //println!("{} {:?}", program, args);
    let mut child = Command::new(program)
      .args(args)
      .stdin(Stdio::piped())
      .stdout(Stdio::piped())
      .stderr(Stdio::piped())
      .spawn()?;
    let stdin = child.stdin.as_mut().expect("Failed to open stdin");
    stdin.write_all(input).await?;
    let res = child.wait_with_output().await;
    //println!("{:?}", res);
    #[allow(clippy::let_and_return)]
    res
  }
}

#[derive(Debug)]
pub struct SetuidCommandRunner<U: AsRef<str>> {
  user_name: U,
}

impl<U: AsRef<str>> SetuidCommandRunner<U> {
  pub const fn new(user_name: U) -> Self {
    Self { user_name }
  }
}

use std::env;
use users::get_user_by_name;

struct TempSetEnv<'a> {
  name: &'a str,
  old_value: Option<String>,
}

impl<'a> TempSetEnv<'a> {
  fn new(name: &'a str, new_value: impl AsRef<OsStr>) -> TempSetEnv<'a> {
    let old_value = env::var(name);
    env::set_var(name, new_value);
    TempSetEnv {
      name,
      old_value: old_value.ok(),
    }
  }
}

impl Drop for TempSetEnv<'_> {
  fn drop(&mut self) {
    match self.old_value {
      Some(ref val) => env::set_var(self.name, val),
      None => env::remove_var(self.name),
    }
  }
}

#[async_trait(?Send)]
impl<U: AsRef<str>> CommandRunner for SetuidCommandRunner<U> {
  async fn run<'a>(&self, program: &str, args: &'a [&'a OsStr], input: &[u8]) -> IoResult<Output> {
    let uid = get_user_by_name(self.user_name.as_ref())
      .expect("User does not exist")
      .uid();
    let set_home = TempSetEnv::new("HOME", format!("/home/{}", self.user_name.as_ref()));
    let set_dbus = TempSetEnv::new("XDG_RUNTIME_DIR", format!("/run/user/{uid}"));
    //println!("{} {:?}", program, args);
    let mut child = Command::new(program)
      .args(args)
      .stdin(Stdio::piped())
      .uid(uid)
      .gid(uid)
      .stdout(Stdio::piped())
      .stderr(Stdio::piped())
      .spawn()
      .expect("Failed to spawn child process");
    let stdin = child.stdin.as_mut().expect("Failed to open stdin");
    stdin
      .write_all(input)
      .await
      .expect("Failed to write to stdin");
    let res = child.wait_with_output().await;
    drop(set_home);
    drop(set_dbus);
    //println!("{:?}", res);
    res
  }
}

#[cfg(test)]
mod test {
  use crate::args;
  use crate::async_utils::run;
  use crate::command_runner::{CommandRunner, StdCommandRunner};
  use futures_util::future::FutureExt;
  use std::time::Instant;

  #[test]
  fn test() {
    let c = StdCommandRunner;
    run(async {
      let args = args!["1"];
      let start = Instant::now();
      let res = c.run("sleep", args, b"").fuse();
      let ps = c.run("ps", args![], b"").fuse();
      futures_util::pin_mut!(res, ps);
      loop {
        futures_util::select! {
          _ = res => {},
          _ = ps => assert!(start.elapsed().as_millis() < 1000),
          complete => break,
        }
      }
    });
  }
}