git.delta.rocks / remowt / refs/commits / 7e5126608cef

1use bytes::Bytes;2use russh::client::Msg;3use russh::{Channel, ChannelMsg};4use tokio::io::{AsyncReadExt as _, AsyncWriteExt as _, DuplexStream};5use tokio::sync::oneshot;67const BUF: usize = 64 * 1024;89pub struct RemowtChild {10	pub stdin: DuplexStream,11	pub stdout: DuplexStream,12	pub stderr: DuplexStream,13	pub exit: oneshot::Receiver<Option<u32>>,14}1516impl RemowtChild {17	/// Manage channel returned by russh exec().18	pub(crate) fn from_exec(ch: Channel<Msg>) -> Self {19		let (stdin, mut stdin_r) = tokio::io::duplex(BUF);20		let (mut out_w, stdout) = tokio::io::duplex(BUF);21		let (mut err_w, stderr) = tokio::io::duplex(BUF);22		let (exit_tx, exit) = oneshot::channel();2324		tokio::spawn(async move {25			let (mut read, write) = ch.split();2627			// Forward our stdin to the channel, signalling EOF when it closes.28			let stdin_pump = tokio::spawn(async move {29				let mut buf = vec![0u8; BUF];30				loop {31					match stdin_r.read(&mut buf).await {32						Ok(0) | Err(_) => break,33						Ok(n) => {34							if write35								.data_bytes(Bytes::copy_from_slice(&buf[..n]))36								.await37								.is_err()38							{39								return;40							}41						}42					}43				}44				let _ = write.eof().await;45			});4647			let mut code = None;48			while let Some(msg) = read.wait().await {49				match msg {50					ChannelMsg::Data { data } => {51						if out_w.write_all(&data).await.is_err() {52							break;53						}54					}55					ChannelMsg::ExtendedData { data, .. } => {56						if err_w.write_all(&data).await.is_err() {57							break;58						}59					}60					ChannelMsg::ExitStatus { exit_status } => code = Some(exit_status),61					_ => {}62				}63			}6465			// The process is gone; stop waiting on stdin we'll never forward.66			stdin_pump.abort();67			let _ = out_w.shutdown().await;68			let _ = err_w.shutdown().await;69			let _ = exit_tx.send(code);70		});7172		RemowtChild {73			stdin,74			stdout,75			stderr,76			exit,77		}78	}7980	/// Wait for the process to finish, returning its exit status.81	pub async fn wait(self) -> Option<u32> {82		self.exit.await.ok().flatten()83	}84}