use std::io::{Read, Write};
use std::process::{Child, ExitStatus};
use std::sync::mpsc::{Receiver, Sender, SyncSender};
use std::thread;
use crate::broker::protocol_v2::{session_frame, SessionExit, SessionFrame};
pub trait FrameSink: Clone + Send + 'static {
fn send(&self, frame: SessionFrame) -> Result<(), SessionFrame>;
}
impl FrameSink for Sender<SessionFrame> {
fn send(&self, frame: SessionFrame) -> Result<(), SessionFrame> {
Sender::send(self, frame).map_err(|e| e.0)
}
}
impl FrameSink for SyncSender<SessionFrame> {
fn send(&self, frame: SessionFrame) -> Result<(), SessionFrame> {
SyncSender::send(self, frame).map_err(|e| e.0)
}
}
pub trait SessionChild {
type Stdin: Write + Send + 'static;
type Stdout: Read + Send + 'static;
type Stderr: Read + Send + 'static;
fn take_stdin(&mut self) -> Option<Self::Stdin>;
fn take_stdout(&mut self) -> Option<Self::Stdout>;
fn take_stderr(&mut self) -> Option<Self::Stderr>;
fn wait_session(&mut self) -> std::io::Result<SessionExit>;
}
impl SessionChild for Child {
type Stdin = std::process::ChildStdin;
type Stdout = std::process::ChildStdout;
type Stderr = std::process::ChildStderr;
fn take_stdin(&mut self) -> Option<Self::Stdin> {
self.stdin.take()
}
fn take_stdout(&mut self) -> Option<Self::Stdout> {
self.stdout.take()
}
fn take_stderr(&mut self) -> Option<Self::Stderr> {
self.stderr.take()
}
fn wait_session(&mut self) -> std::io::Result<SessionExit> {
Ok(session_exit_from_status(&self.wait()?))
}
}
fn pump_output_stream<S, F, K>(stream: Option<S>, wrap: F, out: &K)
where
S: Read,
F: Fn(Vec<u8>) -> session_frame::Kind,
K: FrameSink,
{
let Some(mut stream) = stream else {
return;
};
let mut buf = [0u8; 8192];
loop {
match stream.read(&mut buf) {
Ok(0) => break,
Ok(n) => {
let frame = SessionFrame {
kind: Some(wrap(buf[..n].to_vec())),
};
if out.send(frame).is_err() {
break;
}
}
Err(err) if err.kind() == std::io::ErrorKind::Interrupted => continue,
Err(_) => break,
}
}
}
fn session_exit_from_status(status: &ExitStatus) -> SessionExit {
#[cfg(unix)]
{
use std::os::unix::process::ExitStatusExt;
SessionExit {
code: status.code().unwrap_or(-1),
signal: status.signal().unwrap_or(0),
metadata: Default::default(),
}
}
#[cfg(windows)]
{
SessionExit {
code: status.code().unwrap_or(-1),
signal: 0,
metadata: Default::default(),
}
}
}
pub fn run_child_session<C: SessionChild, K: FrameSink>(
mut child: C,
out: K,
stdin_rx: Receiver<SessionFrame>,
) -> std::io::Result<SessionExit> {
let child_stdout = child.take_stdout();
let child_stderr = child.take_stderr();
let mut child_stdin = child.take_stdin();
let stdin_handle = thread::spawn(move || {
for frame in stdin_rx {
match frame.kind {
Some(session_frame::Kind::Stdin(bytes)) => {
if let Some(writer) = child_stdin.as_mut() {
if writer
.write_all(&bytes)
.and_then(|()| writer.flush())
.is_err()
{
child_stdin = None;
}
}
}
Some(session_frame::Kind::StdinEof(_)) => {
child_stdin = None;
}
_ => {}
}
}
drop(child_stdin);
});
let out_for_stdout = out.clone();
let stdout_handle = thread::spawn(move || {
pump_output_stream(child_stdout, session_frame::Kind::Stdout, &out_for_stdout);
});
let out_for_stderr = out.clone();
let stderr_handle = thread::spawn(move || {
pump_output_stream(child_stderr, session_frame::Kind::Stderr, &out_for_stderr);
});
let _ = stdout_handle.join();
let _ = stderr_handle.join();
let exit = child.wait_session()?;
let _ = stdin_handle.join();
let _ = out.send(SessionFrame {
kind: Some(session_frame::Kind::Exit(exit.clone())),
});
Ok(exit)
}
#[cfg(test)]
mod tests;