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 139 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 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}