use std::collections::HashMap; use anyhow::{Context, Result}; use command_group::CommandGroup; use nix::sys::signal; use nix::sys::wait::{self, WaitStatus}; use nix::unistd::Pid; use std::fs::File as StdFile; use std::os::unix::io::FromRawFd; use std::os::unix::io::IntoRawFd; use std::process::Command; use std::process::Stdio; use std::sync::Arc; use tokio::fs::File; use tokio::io::AsyncBufReadExt; use tokio::io::BufReader; use tokio::signal::unix; use tokio::sync::oneshot; use tokio::sync::Mutex; mod buffer; pub use buffer::Logs; pub struct Process { cmd: String, env: HashMap, cwd: String, } type WaitChannel = oneshot::Receiver; pub struct Child { pub pid: Pid, ch: WaitChannel, } impl Child { pub fn new(pid: Pid, ch: WaitChannel) -> Child { Child { pid, ch } } pub async fn wait(self) -> Result { Ok(self.ch.await?) } } type Handler = oneshot::Sender; impl Process { pub fn new>(cmd: S, cwd: S, env: Option>) -> Process { let env = env.unwrap_or_default(); Process { env, cmd: cmd.into(), cwd: cwd.into(), } } } #[derive(Clone)] pub enum Log { None, Stdout, Ring(String), } #[derive(Clone)] pub struct ProcessManager { table: Arc>>, ring: buffer::Ring, env: Environ, } impl ProcessManager { pub fn new(cap: usize) -> ProcessManager { ProcessManager { table: Arc::new(Mutex::new(HashMap::new())), ring: buffer::Ring::new(cap), env: Environ::new(), } } fn wait_process() -> Vec { let mut statuses: Vec = Vec::new(); loop { let status = match wait::waitpid(Option::None, Some(wait::WaitPidFlag::WNOHANG)) { Ok(status) => status, Err(_) => { return statuses; } }; match status { WaitStatus::StillAlive => break, _ => statuses.push(status), } } statuses } pub fn start(&self) { let table = Arc::clone(&self.table); let mut signals = match unix::signal(unix::SignalKind::child()) { Ok(s) => s, Err(err) => { panic!("failed to bind to signals: {}", err); } }; tokio::spawn(async move { loop { signals.recv().await; let mut table = table.lock().await; for exited in Self::wait_process() { if let Some(pid) = exited.pid() { if let Some(sender) = table.remove(&pid) { if sender.send(exited).is_err() { debug!("failed to send exit state to process: {}", pid); } } } } } }); } fn sink(&self, file: File, prefix: String) { let ring = self.ring.clone(); let reader = BufReader::new(file); tokio::spawn(async move { let mut lines = reader.lines(); while let Ok(line) = lines.next_line().await { let _ = match line { Some(line) => ring.push(format!("{}: {}", prefix, line)).await, None => break, }; } }); } pub async fn stream(&self, existing_logs: bool, follow: bool) -> Logs { self.ring.stream(existing_logs, follow).await } pub fn signal(&self, pid: Pid, sig: signal::Signal) -> Result<()> { Ok(signal::killpg(pid, sig)?) } pub async fn run(&self, cmd: Process, log: Log) -> Result { let args = shlex::split(&cmd.cmd).context("failed to parse command")?; if args.is_empty() { bail!("invalid command"); } let mut child = Command::new(&args[0]); let child = if !cmd.cwd.is_empty() { child.current_dir(&cmd.cwd) } else { child.current_dir("/") }; let child = child.args(&args[1..]).envs(&self.env.0).envs(cmd.env); let child = match log { Log::None => child.stdout(Stdio::null()).stderr(Stdio::null()), Log::Ring(_) => child.stdout(Stdio::piped()).stderr(Stdio::piped()), _ => child, // default to inherit }; let mut table = self.table.lock().await; let mut child = child .group_spawn() .context("failed to spawn command")? .into_inner(); if let Log::Ring(prefix) = log { let _ = self .ring .push(format!("[-] {}: ------------ [start] ------------", prefix)) .await; if let Some(out) = child.stdout.take() { let out = File::from_std(unsafe { StdFile::from_raw_fd(out.into_raw_fd()) }); self.sink(out, format!("[+] {}", prefix)) } if let Some(out) = child.stderr.take() { let out = File::from_std(unsafe { StdFile::from_raw_fd(out.into_raw_fd()) }); self.sink(out, format!("[-] {}", prefix)) } } let (tx, rx) = oneshot::channel(); let id = child.id(); let pid = Pid::from_raw(id as i32); table.insert(pid, tx); Ok(Child::new(pid, rx)) } } #[derive(Clone)] struct Environ(HashMap); impl Environ { fn new() -> Environ { let env = match Environ::parse("/etc/environment") { Ok(r) => r, Err(err) => { error!("failed to load /etc/environment file: {}", err); HashMap::new() } }; Environ(env) } fn parse

(p: P) -> Result, std::io::Error> where P: AsRef, { let mut m = HashMap::new(); let txt = match std::fs::read_to_string(p) { Ok(txt) => txt, Err(err) if err.kind() == std::io::ErrorKind::NotFound => { info!("skipping /etc/environment file because it does not exist"); "".into() } Err(err) => return Err(err), }; for line in txt.lines() { let line = line.trim(); if line.starts_with('#') { continue; } let parts: Vec<&str> = line.splitn(2, '=').collect(); let key = String::from(parts[0]); let value = match parts.len() { 2 => String::from(parts[1]), _ => String::default(), }; //m.into_iter() m.insert(key, value); } Ok(m) } }