schematics/src/command_runner.rs

use async_trait::async_trait;
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(String::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(&self, program: &str, args: &[&OsStr], stdin: &str) -> IoResult<Output>;

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

#[derive(Debug)]
pub struct StdCommandRunner;

#[async_trait(?Send)]
impl CommandRunner for StdCommandRunner {
  async fn run(&self, program: &str, args: &[&OsStr], input: &str) -> IoResult<Output> {
    //println!("{} {:?}", program, args);
    let mut child = Command::new(program)
      .args(args)
      .stdin(Stdio::piped())
      .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.as_bytes())
      .await
      .expect("Failed to write to stdin");
    let res = child.wait_with_output().await;
    //println!("{:?}", res);
    res
  }
}

#[derive(Debug)]
pub struct SetuidCommandRunner<'a, U: AsRef<str>, C: CommandRunner> {
  command_runner: &'a C,
  user_name: U,
}

impl<'a, U: AsRef<str>, C: CommandRunner> SetuidCommandRunner<'a, U, C> {
  pub fn new(user_name: U, command_runner: &'a C) -> Self {
    SetuidCommandRunner {
      command_runner,
      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: String) -> 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>, C: CommandRunner> CommandRunner for SetuidCommandRunner<'_, U, C> {
  async fn run(&self, program: &str, args: &[&OsStr], input: &str) -> 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.as_bytes())
      .await
      .expect("Failed to write to stdin");
    let res = child.wait_with_output().await;
    //println!("{:?}", res);
    res
  }
}

#[derive(Debug)]
pub struct SuCommandRunner<'a, C>
where
  C: CommandRunner,
{
  command_runner: &'a C,
  user_name: &'a str,
}

impl<'a, C> SuCommandRunner<'a, C>
where
  C: 'a + CommandRunner,
{
  pub fn new(user_name: &'a str, command_runner: &'a C) -> SuCommandRunner<'a, C> {
    SuCommandRunner {
      command_runner,
      user_name,
    }
  }
}

// Su doesn't set XDG_RUNTIME_DIR
// https://github.com/systemd/systemd/blob/master/src/login/pam_systemd.c#L439
#[async_trait(?Send)]
impl<'a, C> CommandRunner for SuCommandRunner<'a, C>
where
  C: 'a + CommandRunner,
{
  async fn run(&self, program: &str, args: &[&OsStr], input: &str) -> IoResult<Output> {
    let raw_new_args = [self.user_name, "-s", "/usr/bin/env", "--", program];
    let mut new_args: Vec<&OsStr> = raw_new_args.iter().map(|s| s.as_ref()).collect();
    new_args.extend_from_slice(args);
    self.command_runner.run("su", &new_args, input).await
  }
}

#[cfg(test)]
mod test {
  use crate::args;
  use crate::async_utils::run;
  use crate::command_runner::{CommandRunner, StdCommandRunner};
  use futures::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, "").fuse();
      let ps = c.run("ps", args![], "").fuse();
      futures::pin_mut!(res, ps);
      loop {
        futures::select! {
          _ = res => {},
          _ = ps => assert!((Instant::now() - start).as_millis() < 1000),
          complete => break,
        }
      }
    })
  }
}