git.delta.rocks / remowt / refs/commits / 6d9cf16dada2

1use std::pin::pin;23use anyhow::{anyhow, bail, Result};4use camino::Utf8PathBuf;5use futures::StreamExt as _;6use remowt_endpoints::subprocess::{ProcId, SpawnSpec, StderrSpec, StdioSpec, SubprocessClient};7use remowt_link_shared::BifConfig;8use serde::de::DeserializeOwned;9use tokio::io::{AsyncBufReadExt as _, AsyncWriteExt as _, BufReader};10use tokio::select;11use tokio_util::codec::{BytesCodec, FramedRead, LinesCodec};12use tracing::{debug, info, warn};1314use crate::forwarded::{RemowtListener, RemowtStream};15use crate::Remowt;1617#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]18pub enum StdioMode {19	#[default]20	Null,21	Pipe,22	Inherit,23}2425#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]26pub enum StderrMode {27	#[default]28	Null,29	Pipe,30	Inherit,31	MergeWithStdout,32}3334#[derive(Default)]35pub struct SpawnOptions {36	pub program: String,37	pub args: Vec<String>,38	pub env: Vec<(String, String)>,39	pub env_clear: bool,40	pub cwd: Option<Utf8PathBuf>,41	pub escalated: bool,42	pub stdin: StdioMode,43	pub stdout: StdioMode,44	pub stderr: StderrMode,45}4647pub struct RemowtChild {48	pub stdin: Option<RemowtStream>,49	pub stdout: Option<RemowtStream>,50	pub stderr: Option<RemowtStream>,51	id: ProcId,52	client: SubprocessClient<BifConfig>,53}5455impl RemowtChild {56	pub async fn wait(self) -> Result<Option<i32>> {57		let RemowtChild {58			stdin,59			stdout,60			stderr,61			id,62			client,63		} = self;64		drop(stdin);65		drop(stdout);66		drop(stderr);67		client68			.wait(id)69			.await?70			.map_err(|e| anyhow!("agent wait failed: {e}"))71	}7273	pub async fn kill(&self, signal: i32) -> Result<()> {74		self.client75			.kill(self.id, signal)76			.await?77			.map_err(|e| anyhow!("agent kill failed: {e}"))78	}79}8081fn needs_socket(m: StdioMode) -> bool {82	matches!(m, StdioMode::Pipe | StdioMode::Inherit)83}8485fn stderr_needs_socket(m: StderrMode) -> bool {86	matches!(m, StderrMode::Pipe | StderrMode::Inherit)87}8889impl Remowt {90	pub async fn spawn(&self, opts: SpawnOptions) -> Result<RemowtChild> {91		let SpawnOptions {92			program,93			args,94			env,95			env_clear,96			cwd,97			escalated,98			stdin,99			stdout,100			stderr,101		} = opts;102103		if matches!(stderr, StderrMode::MergeWithStdout) && !needs_socket(stdout) {104			bail!("stderr=MergeWithStdout requires stdout=Pipe or Inherit");105		}106107		let stdin_bound = if needs_socket(stdin) {108			Some(self.bind_runtime_unix("proc-stdin").await?)109		} else {110			None111		};112		let stdout_bound = if needs_socket(stdout) {113			Some(self.bind_runtime_unix("proc-stdout").await?)114		} else {115			None116		};117		let stderr_bound = if stderr_needs_socket(stderr) {118			Some(self.bind_runtime_unix("proc-stderr").await?)119		} else {120			None121		};122123		let stdin_spec = match &stdin_bound {124			Some((_, p)) => StdioSpec::Socket(p.clone()),125			None => StdioSpec::Null,126		};127		let stdout_spec = match &stdout_bound {128			Some((_, p)) => StdioSpec::Socket(p.clone()),129			None => StdioSpec::Null,130		};131		let stderr_spec = match (&stderr, &stderr_bound) {132			(StderrMode::Pipe | StderrMode::Inherit, Some((_, p))) => StderrSpec::Socket(p.clone()),133			(StderrMode::MergeWithStdout, _) => StderrSpec::MergeWithStdout,134			_ => StderrSpec::Null,135		};136137		let client: SubprocessClient<BifConfig> = if escalated {138			// Boxed to break the async-fn type cycle139			Box::pin(self.run0_endpoints::<SubprocessClient<BifConfig>>()).await?140		} else {141			self.endpoints()142		};143144		let spec = SpawnSpec {145			program: program.clone(),146			args,147			env,148			env_clear,149			cwd,150			stdin: stdin_spec,151			stdout: stdout_spec,152			stderr: stderr_spec,153		};154		let id = client155			.spawn(spec)156			.await?157			.map_err(|e| anyhow!("agent spawn failed: {e}"))?;158159		let (stdin_res, stdout_res, stderr_res) = tokio::join!(160			accept(stdin_bound),161			accept(stdout_bound),162			accept(stderr_bound),163		);164165		let stdin_stream = handle_stdin(stdin, stdin_res?, &program);166		let stdout_stream = handle_output(stdout, stdout_res?, &program, false);167		let stderr_stream = handle_output_err(stderr, stderr_res?, &program);168169		Ok(RemowtChild {170			stdin: stdin_stream,171			stdout: stdout_stream,172			stderr: stderr_stream,173			id,174			client,175		})176	}177178	pub fn cmd(&self, program: impl AsRef<str>) -> RemowtCommand {179		let program = program.as_ref().to_owned();180		RemowtCommand {181			program,182			args: vec![],183			env: vec![],184			remowt: self.clone(),185			escalated: false,186		}187	}188}189190async fn accept(b: Option<(RemowtListener, Utf8PathBuf)>) -> Result<Option<RemowtStream>> {191	match b {192		Some((l, _)) => Ok(Some(l.accept().await?)),193		None => Ok(None),194	}195}196197fn handle_stdin(mode: StdioMode, s: Option<RemowtStream>, program: &str) -> Option<RemowtStream> {198	match mode {199		StdioMode::Pipe => s,200		StdioMode::Inherit => {201			if let Some(s) = s {202				let program = program.to_owned();203				tokio::spawn(async move {204					let mut stdin = tokio::io::stdin();205					let mut s = s;206					if let Err(e) = tokio::io::copy(&mut stdin, &mut s).await {207						warn!(program, "stdin forward ended: {e}");208					}209					let _ = s.shutdown().await;210				});211			}212			None213		}214		StdioMode::Null => None,215	}216}217218fn handle_output(219	mode: StdioMode,220	s: Option<RemowtStream>,221	program: &str,222	is_stderr: bool,223) -> Option<RemowtStream> {224	match mode {225		StdioMode::Pipe => s,226		StdioMode::Inherit => {227			if let Some(s) = s {228				let program = program.to_owned();229				tokio::spawn(pump_to_tracing(s, program, is_stderr));230			}231			None232		}233		StdioMode::Null => None,234	}235}236237fn handle_output_err(238	mode: StderrMode,239	s: Option<RemowtStream>,240	program: &str,241) -> Option<RemowtStream> {242	match mode {243		StderrMode::Pipe => s,244		StderrMode::Inherit => {245			if let Some(s) = s {246				let program = program.to_owned();247				tokio::spawn(pump_to_tracing(s, program, true));248			}249			None250		}251		StderrMode::MergeWithStdout | StderrMode::Null => None,252	}253}254255async fn pump_to_tracing(stream: RemowtStream, program: String, is_stderr: bool) {256	let mut reader = BufReader::new(stream).lines();257	loop {258		match reader.next_line().await {259			Ok(Some(line)) => {260				if is_stderr {261					warn!(program, "{line}");262				} else {263					info!(program, "{line}");264				}265			}266			Ok(None) => break,267			Err(e) => {268				warn!(program, "child log stream error: {e}");269				break;270			}271		}272	}273}274275fn escape_bash(input: &str, out: &mut String) {276	const TO_ESCAPE: &str = "$ !\"#&'()*,;<>?[\\]^`{|}";277	if input.chars().all(|c| !TO_ESCAPE.contains(c)) {278		out.push_str(input);279		return;280	}281	out.push('\'');282	for (i, v) in input.split('\'').enumerate() {283		if i != 0 {284			out.push_str("'\"'\"'");285		}286		out.push_str(v);287	}288	out.push('\'');289}290291#[derive(Clone)]292pub struct RemowtCommand {293	program: String,294	args: Vec<String>,295	env: Vec<(String, String)>,296	remowt: Remowt,297	escalated: bool,298}299300impl RemowtCommand {301	pub fn arg(&mut self, arg: impl AsRef<str>) -> &mut Self {302		self.args.push(arg.as_ref().to_owned());303		self304	}305	pub fn args<V: AsRef<str>>(&mut self, args: impl IntoIterator<Item = V>) -> &mut Self {306		for arg in args {307			self.args.push(arg.as_ref().to_owned());308		}309		self310	}311	pub fn eqarg(&mut self, key: impl AsRef<str>, value: impl AsRef<str>) -> &mut Self {312		self.args313			.push(format!("{}={}", key.as_ref(), value.as_ref()));314		self315	}316	pub fn comparg(&mut self, key: impl AsRef<str>, value: impl AsRef<str>) -> &mut Self {317		self.args.push(key.as_ref().to_owned());318		self.args.push(value.as_ref().to_owned());319		self320	}321	pub fn env(&mut self, name: impl AsRef<str>, value: impl AsRef<str>) -> &mut Self {322		self.env323			.push((name.as_ref().to_owned(), value.as_ref().to_owned()));324		self325	}326327	pub fn sudo(mut self) -> Self {328		self.escalated = true;329		self330	}331332	/// Only for display.333	fn shell_line(&self) -> String {334		let mut out = String::new();335		if self.escalated {336			out.push_str("run0 ");337		}338		if !self.env.is_empty() {339			out.push_str("env");340			for (k, v) in &self.env {341				out.push(' ');342				assert!(!k.contains('='));343				escape_bash(k, &mut out);344				out.push('=');345				escape_bash(v, &mut out);346			}347			out.push(' ');348		}349		escape_bash(&self.program, &mut out);350		for arg in &self.args {351			out.push(' ');352			escape_bash(arg, &mut out);353		}354		out355	}356357	fn into_spawn_options(self) -> (Remowt, SpawnOptions, String) {358		let line = self.shell_line();359		let opts = SpawnOptions {360			program: self.program,361			args: self.args,362			env: self.env,363			env_clear: false,364			cwd: None,365			escalated: self.escalated,366			stdin: StdioMode::Null,367			stdout: StdioMode::Pipe,368			stderr: StderrMode::Pipe,369		};370		(self.remowt, opts, line)371	}372373	pub async fn run(self) -> Result<()> {374		run_inner(self, false).await.map(|_| ())375	}376	pub async fn run_string(self) -> Result<String> {377		let bytes = run_inner(self, true).await?.expect("want_stdout");378		Ok(String::from_utf8(bytes)?)379	}380	pub async fn run_value<T: DeserializeOwned>(self) -> Result<T> {381		let s = self.run_string().await?;382		Ok(serde_json::from_str(&s)?)383	}384}385386async fn run_inner(cmd: RemowtCommand, want_stdout: bool) -> Result<Option<Vec<u8>>> {387	let (remowt, opts, line) = cmd.into_spawn_options();388	debug!("running command {line:?} over remowt");389	let program = opts.program.clone();390	let mut child = remowt.spawn(opts).await?;391	let stderr = child.stderr.take().expect("stderr=Pipe");392	let stdout = child.stdout.take().expect("stdout=Pipe");393394	let mut err = FramedRead::new(stderr, LinesCodec::new());395	let (mut out_bytes, mut out_lines) = if want_stdout {396		(Some(FramedRead::new(stdout, BytesCodec::new())), None)397	} else {398		(None, Some(FramedRead::new(stdout, LinesCodec::new())))399	};400401	let mut buf = if want_stdout { Some(Vec::new()) } else { None };402403	let mut wait = pin!(child.wait());404	let exit = loop {405		select! {406			biased;407408			Some(e) = err.next() => {409				let e = e?;410				warn!(program = %program, "{e}");411			}412			Some(o) = async { out_bytes.as_mut()?.next().await }, if want_stdout => {413				buf.as_mut().expect("want_stdout").extend_from_slice(&o?);414			}415			Some(o) = async { out_lines.as_mut()?.next().await }, if !want_stdout => {416				let o = o?;417				info!(program = %program, "{o}");418			}419			res = &mut wait => {420				break res?;421			}422		}423	};424425	match exit {426		Some(0) => Ok(buf),427		Some(c) => bail!("command '{line}' failed with status {c}"),428		None => Err(anyhow!("command '{line}' ended without an exit status")),429	}430}