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::iroh_tunnel::TunnelAddr;8use remowt_link_shared::BifConfig;9use serde::de::DeserializeOwned;10use tokio::io::AsyncWriteExt as _;11use tokio::select;12use tokio_util::codec::{BytesCodec, FramedRead, LinesCodec};13use tracing::{debug, info, warn};1415use crate::forwarded::{RemowtListener, RemowtStream};16use crate::{drain_to_tracing, Remowt};1718#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]19pub enum StdioMode {20 #[default]21 Null,22 Pipe,23 Inherit,24}2526#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]27pub enum StderrMode {28 #[default]29 Null,30 Pipe,31 Inherit,32 MergeWithStdout,33}3435#[derive(Default)]36pub struct SpawnOptions {37 pub program: String,38 pub args: Vec<String>,39 pub env: Vec<(String, String)>,40 pub env_clear: bool,41 pub cwd: Option<Utf8PathBuf>,42 pub escalated: bool,43 pub stdin: StdioMode,44 pub stdout: StdioMode,45 pub stderr: StderrMode,46}4748pub struct RemowtChild {49 pub stdin: Option<RemowtStream>,50 pub stdout: Option<RemowtStream>,51 pub stderr: Option<RemowtStream>,52 id: ProcId,53 client: SubprocessClient<BifConfig>,54}5556impl RemowtChild {57 pub async fn wait(self) -> Result<Option<i32>> {58 let RemowtChild {59 stdin,60 stdout,61 stderr,62 id,63 client,64 } = self;65 drop(stdin);66 let drain_out = async move {67 if let Some(s) = stdout {68 let _ = drain_to_tracing(s, "<child stdout>".to_owned(), false).await;69 }70 };71 let drain_err = async move {72 if let Some(s) = stderr {73 let _ = drain_to_tracing(s, "<child stderr>".to_owned(), true).await;74 }75 };76 let wait = async move {77 client78 .wait(id)79 .await?80 .map_err(|e| anyhow!("agent wait failed: {e}"))81 };82 let (code, _, _) = tokio::join!(wait, drain_out, drain_err);83 code84 }8586 pub async fn kill(&self, signal: i32) -> Result<()> {87 self.client88 .kill(self.id, signal)89 .await?90 .map_err(|e| anyhow!("agent kill failed: {e}"))91 }92}9394fn needs_socket(m: StdioMode) -> bool {95 matches!(m, StdioMode::Pipe | StdioMode::Inherit)96}9798fn stderr_needs_socket(m: StderrMode) -> bool {99 matches!(m, StderrMode::Pipe | StderrMode::Inherit)100}101102impl Remowt {103 pub async fn spawn(&self, opts: SpawnOptions) -> Result<RemowtChild> {104 let SpawnOptions {105 program,106 args,107 env,108 env_clear,109 cwd,110 escalated,111 stdin,112 stdout,113 stderr,114 } = opts;115116 if matches!(stderr, StderrMode::MergeWithStdout) && !needs_socket(stdout) {117 bail!("stderr=MergeWithStdout requires stdout=Pipe or Inherit");118 }119120 let stdin_bound = if needs_socket(stdin) {121 Some(self.bind_fast_tunnel("proc-stdin", escalated).await?)122 } else {123 None124 };125 let stdout_bound = if needs_socket(stdout) {126 Some(self.bind_fast_tunnel("proc-stdout", escalated).await?)127 } else {128 None129 };130 let stderr_bound = if stderr_needs_socket(stderr) {131 Some(self.bind_fast_tunnel("proc-stderr", escalated).await?)132 } else {133 None134 };135136 let stdin_spec = match &stdin_bound {137 Some((_, t)) => StdioSpec::Tunnel(t.clone()),138 None => StdioSpec::Null,139 };140 let stdout_spec = match &stdout_bound {141 Some((_, t)) => StdioSpec::Tunnel(t.clone()),142 None => StdioSpec::Null,143 };144 let stderr_spec = match (&stderr, &stderr_bound) {145 (StderrMode::Pipe | StderrMode::Inherit, Some((_, t))) => StderrSpec::Tunnel(t.clone()),146 (StderrMode::MergeWithStdout, _) => StderrSpec::MergeWithStdout,147 _ => StderrSpec::Null,148 };149150 let client: SubprocessClient<BifConfig> = if escalated {151 152 Box::pin(self.run0_endpoints::<SubprocessClient<BifConfig>>()).await?153 } else {154 self.endpoints()155 };156157 let spec = SpawnSpec {158 program: program.clone(),159 args,160 env,161 env_clear,162 cwd,163 stdin: stdin_spec,164 stdout: stdout_spec,165 stderr: stderr_spec,166 };167 let id = client168 .spawn(spec)169 .await?170 .map_err(|e| anyhow!("agent spawn failed: {e}"))?;171172 let (stdin_res, stdout_res, stderr_res) = tokio::join!(173 accept(stdin_bound),174 accept(stdout_bound),175 accept(stderr_bound),176 );177178 let stdin_stream = handle_stdin(stdin, stdin_res?, &program);179 let stdout_stream = handle_output(stdout, stdout_res?, &program);180 let stderr_stream = handle_output_err(stderr, stderr_res?, &program);181182 Ok(RemowtChild {183 stdin: stdin_stream,184 stdout: stdout_stream,185 stderr: stderr_stream,186 id,187 client,188 })189 }190191 pub fn cmd(&self, program: impl AsRef<str>) -> RemowtCommand {192 let program = program.as_ref().to_owned();193 RemowtCommand {194 program,195 args: vec![],196 env: vec![],197 remowt: self.clone(),198 escalated: false,199 }200 }201}202203async fn accept(b: Option<(RemowtListener, TunnelAddr)>) -> Result<Option<RemowtStream>> {204 match b {205 Some((l, _)) => Ok(Some(l.accept().await?)),206 None => Ok(None),207 }208}209210fn handle_stdin(mode: StdioMode, s: Option<RemowtStream>, program: &str) -> Option<RemowtStream> {211 match mode {212 StdioMode::Pipe => s,213 StdioMode::Inherit => {214 if let Some(s) = s {215 let program = program.to_owned();216 tokio::spawn(async move {217 let mut stdin = tokio::io::stdin();218 let mut s = s;219 if let Err(e) = tokio::io::copy(&mut stdin, &mut s).await {220 warn!(program, "stdin forward ended: {e}");221 }222 let _ = s.shutdown().await;223 });224 }225 None226 }227 StdioMode::Null => None,228 }229}230231fn handle_output(mode: StdioMode, s: Option<RemowtStream>, program: &str) -> Option<RemowtStream> {232 match mode {233 StdioMode::Pipe => s,234 StdioMode::Inherit => {235 if let Some(s) = s {236 let program = program.to_owned();237 tokio::spawn(drain_to_tracing(s, program, false));238 }239 None240 }241 StdioMode::Null => None,242 }243}244245fn handle_output_err(246 mode: StderrMode,247 s: Option<RemowtStream>,248 program: &str,249) -> Option<RemowtStream> {250 match mode {251 StderrMode::Pipe => s,252 StderrMode::Inherit => {253 if let Some(s) = s {254 let program = program.to_owned();255 tokio::spawn(drain_to_tracing(s, program, true));256 }257 None258 }259 StderrMode::MergeWithStdout | StderrMode::Null => None,260 }261}262263fn escape_bash(input: &str, out: &mut String) {264 const TO_ESCAPE: &str = "$ !\"#&'()*,;<>?[\\]^`{|}";265 if input.chars().all(|c| !TO_ESCAPE.contains(c)) {266 out.push_str(input);267 return;268 }269 out.push('\'');270 for (i, v) in input.split('\'').enumerate() {271 if i != 0 {272 out.push_str("'\"'\"'");273 }274 out.push_str(v);275 }276 out.push('\'');277}278279#[derive(Clone)]280pub struct RemowtCommand {281 program: String,282 args: Vec<String>,283 env: Vec<(String, String)>,284 remowt: Remowt,285 escalated: bool,286}287288impl RemowtCommand {289 pub fn arg(&mut self, arg: impl AsRef<str>) -> &mut Self {290 self.args.push(arg.as_ref().to_owned());291 self292 }293 pub fn args<V: AsRef<str>>(&mut self, args: impl IntoIterator<Item = V>) -> &mut Self {294 for arg in args {295 self.args.push(arg.as_ref().to_owned());296 }297 self298 }299 pub fn eqarg(&mut self, key: impl AsRef<str>, value: impl AsRef<str>) -> &mut Self {300 self.args301 .push(format!("{}={}", key.as_ref(), value.as_ref()));302 self303 }304 pub fn comparg(&mut self, key: impl AsRef<str>, value: impl AsRef<str>) -> &mut Self {305 self.args.push(key.as_ref().to_owned());306 self.args.push(value.as_ref().to_owned());307 self308 }309 pub fn env(&mut self, name: impl AsRef<str>, value: impl AsRef<str>) -> &mut Self {310 self.env311 .push((name.as_ref().to_owned(), value.as_ref().to_owned()));312 self313 }314315 pub fn sudo(mut self) -> Self {316 self.escalated = true;317 self318 }319320 321 fn shell_line(&self) -> String {322 let mut out = String::new();323 if self.escalated {324 out.push_str("run0 ");325 }326 if !self.env.is_empty() {327 out.push_str("env");328 for (k, v) in &self.env {329 out.push(' ');330 assert!(!k.contains('='));331 escape_bash(k, &mut out);332 out.push('=');333 escape_bash(v, &mut out);334 }335 out.push(' ');336 }337 escape_bash(&self.program, &mut out);338 for arg in &self.args {339 out.push(' ');340 escape_bash(arg, &mut out);341 }342 out343 }344345 fn into_spawn_options(self) -> (Remowt, SpawnOptions, String) {346 let line = self.shell_line();347 let opts = SpawnOptions {348 program: self.program,349 args: self.args,350 env: self.env,351 env_clear: false,352 cwd: None,353 escalated: self.escalated,354 stdin: StdioMode::Null,355 stdout: StdioMode::Pipe,356 stderr: StderrMode::Pipe,357 };358 (self.remowt, opts, line)359 }360361 pub async fn run(self) -> Result<()> {362 run_inner(self, false).await.map(|_| ())363 }364 pub async fn run_string(self) -> Result<String> {365 let bytes = run_inner(self, true).await?.expect("want_stdout");366 Ok(String::from_utf8(bytes)?)367 }368 pub async fn run_value<T: DeserializeOwned>(self) -> Result<T> {369 let s = self.run_string().await?;370 Ok(serde_json::from_str(&s)?)371 }372}373374async fn run_inner(cmd: RemowtCommand, want_stdout: bool) -> Result<Option<Vec<u8>>> {375 let (remowt, opts, line) = cmd.into_spawn_options();376 debug!("running command {line:?} over remowt");377 let program = opts.program.clone();378 let mut child = remowt.spawn(opts).await?;379 let stderr = child.stderr.take().expect("stderr=Pipe");380 let stdout = child.stdout.take().expect("stdout=Pipe");381382 let mut err = FramedRead::new(stderr, LinesCodec::new());383 let (mut out_bytes, mut out_lines) = if want_stdout {384 (Some(FramedRead::new(stdout, BytesCodec::new())), None)385 } else {386 (None, Some(FramedRead::new(stdout, LinesCodec::new())))387 };388389 let mut buf = if want_stdout { Some(Vec::new()) } else { None };390391 let mut wait = pin!(child.wait());392 let exit = loop {393 select! {394 biased;395396 Some(e) = err.next() => {397 let e = e?;398 warn!(program = %program, "{e}");399 }400 Some(o) = async { out_bytes.as_mut()?.next().await }, if want_stdout => {401 buf.as_mut().expect("want_stdout").extend_from_slice(&o?);402 }403 Some(o) = async { out_lines.as_mut()?.next().await }, if !want_stdout => {404 let o = o?;405 info!(program = %program, "{o}");406 }407 res = &mut wait => {408 break res?;409 }410 }411 };412413 while let Some(e) = err.next().await {414 if let Ok(line) = e {415 warn!(program = %program, "{line}");416 }417 }418 if want_stdout {419 if let Some(out_bytes) = out_bytes.as_mut() {420 while let Some(o) = out_bytes.next().await {421 if let Ok(chunk) = o {422 buf.as_mut().expect("want_stdout").extend_from_slice(&chunk);423 }424 }425 }426 } else if let Some(out_lines) = out_lines.as_mut() {427 while let Some(o) = out_lines.next().await {428 if let Ok(line) = o {429 info!(program = %program, "{line}");430 }431 }432 }433434 match exit {435 Some(0) => Ok(buf),436 Some(c) => bail!("command '{line}' failed with status {c}"),437 None => Err(anyhow!("command '{line}' ended without an exit status")),438 }439}