use std::process::Stdio; use std::time::Duration; use tokio::io::AsyncWriteExt; use tokio::process::Command; const COMMAND_TIMEOUT: Duration = Duration::from_secs(60); #[derive(Clone, Debug)] pub struct Runner { ssh: Option, sudo: bool, kubectl: String, } impl Runner { pub fn new(ssh: Option, sudo: bool, kubectl: Option) -> Self { Self { ssh, sudo, kubectl: kubectl.unwrap_or_else(|| "kubectl".to_string()), } } pub fn kubectl_path(&self) -> &str { &self.kubectl } pub fn ssh_host(&self) -> Option<&str> { self.ssh.as_deref() } pub async fn kubectl(&self, args: &[String], stdin: Option<&str>) -> Result { let mut command = Vec::with_capacity(args.len() + 1); command.push(self.kubectl.clone()); command.extend(args.iter().cloned()); self.run(&command, stdin).await } pub async fn shell(&self, script: &str) -> Result { self.shell_with_stdin(script, None).await } pub async fn shell_with_stdin( &self, script: &str, stdin: Option<&str>, ) -> Result { let command = vec!["sh".to_string(), "-c".to_string(), script.to_string()]; self.run(&command, stdin).await } /// Spawns `kubectl ` with stdout piped for streaming (watches). /// The child is killed when dropped. pub fn spawn_kubectl_stream(&self, args: &[String]) -> Result { let mut argv = Vec::with_capacity(args.len() + 1); argv.push(self.kubectl.clone()); argv.extend(args.iter().cloned()); let mut command = self.command(&argv); command .stdin(Stdio::null()) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .kill_on_drop(true); command .spawn() .map_err(|error| format!("failed to spawn {:?}: {error}", argv[0])) } /// The local or `ssh` command for `args`, with `sudo -n` when enabled. fn command(&self, args: &[String]) -> Command { let command = match &self.ssh { Some(host) => { let mut remote = Vec::with_capacity(args.len() + 2); if self.sudo { remote.push("sudo".to_string()); remote.push("-n".to_string()); } remote.extend(args.iter().cloned()); let script = remote .iter() .map(|argument| shell_quote(argument)) .collect::>() .join(" "); let mut command = Command::new("ssh"); command .arg("-o") .arg("BatchMode=yes") .arg("-o") .arg("ConnectTimeout=10") .arg(host) .arg(script); command } None => { let (program, rest) = args.split_first().expect("checked non-empty"); let mut command = if self.sudo { let mut command = Command::new("sudo"); command.arg("-n").arg(program); command } else { Command::new(program) }; command.args(rest); command } }; command } pub async fn run(&self, args: &[String], stdin: Option<&str>) -> Result { if args.is_empty() { return Err("refusing to run an empty command".to_string()); } let mut command = self.command(args); command .stdin(if stdin.is_some() { Stdio::piped() } else { Stdio::null() }) .stdout(Stdio::piped()) .stderr(Stdio::piped()); let mut child = command .spawn() .map_err(|error| format!("failed to spawn {:?}: {error}", args[0]))?; if let Some(input) = stdin { let mut pipe = child .stdin .take() .ok_or_else(|| "failed to open command stdin".to_string())?; pipe.write_all(input.as_bytes()) .await .map_err(|error| format!("failed to write command stdin: {error}"))?; pipe.shutdown() .await .map_err(|error| format!("failed to close command stdin: {error}"))?; } let output = tokio::time::timeout(COMMAND_TIMEOUT, child.wait_with_output()) .await .map_err(|_| { format!( "command timed out after {}s: {}", COMMAND_TIMEOUT.as_secs(), args.join(" ") ) })? .map_err(|error| format!("command failed: {error}"))?; let stdout = String::from_utf8_lossy(&output.stdout).to_string(); if output.status.success() { return Ok(stdout); } let stderr = String::from_utf8_lossy(&output.stderr); let detail = if stderr.trim().is_empty() { stdout } else { stderr.to_string() }; Err(format!( "command exited with {}: {}", output.status, detail.trim() )) } } pub fn shell_quote(value: &str) -> String { if value .chars() .all(|character| character.is_ascii_alphanumeric() || "-_./:=@".contains(character)) { return value.to_string(); } format!("'{}'", value.replace('\'', "'\\''")) } #[cfg(test)] mod tests { use super::shell_quote; #[test] fn quotes_only_what_the_remote_shell_needs() { assert_eq!(shell_quote("kubectl"), "kubectl"); assert_eq!( shell_quote("/var/lib/kcp/admin.kubeconfig"), "/var/lib/kcp/admin.kubeconfig" ); assert_eq!( shell_quote("*kcp/admin.kubeconfig"), "'*kcp/admin.kubeconfig'" ); assert_eq!(shell_quote("a b"), "'a b'"); assert_eq!(shell_quote("it's"), "'it'\\''s'"); } }