use std::io;
use std::process::{ExitStatus, Stdio};
use kcode_jsonrpc_wire::{Message, WireError};
use serde_json::Value;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, Command};
use tokio::sync::mpsc;
use tokio::task::{JoinError, JoinHandle};
pub const MAX_INBOUND_LINE_BYTES: usize = 8 * 1024 * 1024;
pub const MAX_OUTBOUND_LINE_BYTES: usize = 8 * 1024 * 1024;
pub const MAX_STDERR_CAPTURE_BYTES: usize = 64 * 1024;
pub const MAX_INBOUND_CAPACITY: usize = 1024;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct Config {
pub inbound_capacity: usize,
}
#[derive(Clone, Debug)]
pub struct ProcessExit {
pub status: ExitStatus,
pub stderr: String,
}
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("stdio I/O failed: {0}")]
Io(#[source] io::Error),
#[error("child status query failed: {0}")]
ProcessQuery(#[source] io::Error),
#[error("child kill failed: {0}")]
ProcessKill(#[source] io::Error),
#[error("child wait failed: {0}")]
ProcessWait(#[source] io::Error),
#[error("child stderr read failed: {0}")]
StderrIo(#[source] io::Error),
#[error("JSON failed: {0}")]
Json(#[source] serde_json::Error),
#[error("inbound line is not UTF-8: {0}")]
Utf8(#[source] std::str::Utf8Error),
#[error("invalid JSON-RPC message: {0}")]
Wire(#[source] WireError),
#[error("stdout reader task failed: {0}")]
ReaderTask(#[source] JoinError),
#[error("stderr task failed: {0}")]
StderrTask(#[source] JoinError),
#[error("inbound line exceeds {MAX_INBOUND_LINE_BYTES} bytes")]
InboundLineTooLong,
#[error("stdout ended with an incomplete JSONL frame")]
IncompleteFrame,
#[error("outbound line exceeds {MAX_OUTBOUND_LINE_BYTES} bytes")]
OutboundLineTooLong,
#[error("invalid inbound capacity: {0}")]
InvalidInboundCapacity(usize),
#[error("child stdin is closed")]
StdinClosed,
#[error("stdout reader is unavailable")]
ReaderUnavailable,
#[error("child stderr reader is unavailable")]
StderrUnavailable,
#[error("child stdout closed while the child remains live")]
StdoutClosed,
#[error("child exited: {0:?}")]
ProcessExited(ProcessExit),
}
pub struct StdioRpc {
child: Child,
stdin: Option<ChildStdin>,
inbound: Option<mpsc::Receiver<Result<Message, Error>>>,
reader: Option<JoinHandle<()>>,
stderr: Option<JoinHandle<io::Result<String>>>,
exit: Option<ProcessExit>,
}
impl StdioRpc {
pub fn spawn(mut command: Command, config: Config) -> Result<Self, Error> {
if config.inbound_capacity == 0 || config.inbound_capacity > MAX_INBOUND_CAPACITY {
return Err(Error::InvalidInboundCapacity(config.inbound_capacity));
}
command
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
let mut child = command.spawn().map_err(Error::Io)?;
let stdin = child.stdin.take().ok_or(Error::StdinClosed)?;
let stdout = child.stdout.take().ok_or(Error::StdinClosed)?;
let stderr = child.stderr.take().ok_or(Error::StdinClosed)?;
let (sender, inbound) = mpsc::channel(config.inbound_capacity);
Ok(Self {
child,
stdin: Some(stdin),
inbound: Some(inbound),
reader: Some(tokio::spawn(read_stdout(BufReader::new(stdout), sender))),
stderr: Some(tokio::spawn(read_stderr(stderr))),
exit: None,
})
}
pub async fn send(&mut self, value: Value) -> Result<(), Error> {
let mut line = serde_json::to_vec(&value).map_err(Error::Json)?;
if line.len() > MAX_OUTBOUND_LINE_BYTES {
return Err(Error::OutboundLineTooLong);
}
line.push(b'\n');
let result = match self.stdin.as_mut() {
Some(stdin) => {
async {
stdin.write_all(&line).await?;
stdin.flush().await
}
.await
}
None => return Err(Error::StdinClosed),
};
if result.is_err() {
self.stdin.take();
}
result.map_err(Error::Io)
}
pub async fn next(&mut self) -> Result<Message, Error> {
let receiver = self.inbound.as_mut().ok_or(Error::ReaderUnavailable)?;
match receiver.recv().await {
Some(Ok(message)) => Ok(message),
Some(Err(error)) => {
self.inbound.take();
Err(error)
}
None => {
self.inbound.take();
self.join_reader().await?;
match self.child.try_wait().map_err(Error::ProcessQuery)? {
Some(_) => Err(Error::ProcessExited(self.reap().await?)),
None => Err(Error::StdoutClosed),
}
}
}
}
pub async fn shutdown(mut self) -> Result<ProcessExit, Error> {
self.stdin.take();
self.inbound.take();
let action = match self.child.try_wait() {
Ok(Some(_)) => Ok(()),
Ok(None) => self.child.start_kill().map_err(Error::ProcessKill),
Err(query) => match self.child.start_kill() {
Ok(()) => Err(Error::ProcessQuery(query)),
Err(kill) => Err(Error::ProcessKill(kill)),
},
};
let exit = self.reap().await;
let reader = self.abort_reader().await;
action?;
let exit = exit?;
reader?;
Ok(exit)
}
async fn reap(&mut self) -> Result<ProcessExit, Error> {
if let Some(exit) = &self.exit {
return Ok(exit.clone());
}
let status = self.child.wait().await.map_err(Error::ProcessWait)?;
self.exit = Some(ProcessExit {
status,
stderr: String::new(),
});
let stderr = self.stderr.take().ok_or(Error::StderrUnavailable)?;
let stderr = stderr
.await
.map_err(Error::StderrTask)?
.map_err(Error::StderrIo)?;
let exit = ProcessExit { status, stderr };
self.exit = Some(exit.clone());
Ok(exit)
}
async fn join_reader(&mut self) -> Result<(), Error> {
if let Some(reader) = self.reader.take() {
reader.await.map_err(Error::ReaderTask)?;
}
Ok(())
}
async fn abort_reader(&mut self) -> Result<(), Error> {
if let Some(reader) = self.reader.take() {
reader.abort();
if let Err(error) = reader.await
&& !error.is_cancelled()
{
return Err(Error::ReaderTask(error));
}
}
Ok(())
}
}
async fn read_stdout(
mut stdout: BufReader<tokio::process::ChildStdout>,
sender: mpsc::Sender<Result<Message, Error>>,
) {
loop {
let line = match read_line(&mut stdout, MAX_INBOUND_LINE_BYTES).await {
Ok(Some(line)) => line,
Ok(None) => return,
Err(error) => {
let _ = sender.send(Err(error)).await;
return;
}
};
let value = match serde_json::from_slice(&line) {
Ok(value) => value,
Err(error) => {
let _ = sender.send(Err(Error::Json(error))).await;
return;
}
};
let message = match kcode_jsonrpc_wire::parse(value) {
Ok(message) => message,
Err(error) => {
let _ = sender.send(Err(Error::Wire(error))).await;
return;
}
};
if sender.send(Ok(message)).await.is_err() {
return;
}
}
}
async fn read_line<R: tokio::io::AsyncBufRead + Unpin>(
reader: &mut R,
limit: usize,
) -> Result<Option<Vec<u8>>, Error> {
let mut line = Vec::new();
loop {
let available = reader.fill_buf().await.map_err(Error::Io)?;
if available.is_empty() {
if line.is_empty() {
return Ok(None);
}
std::str::from_utf8(&line).map_err(Error::Utf8)?;
return Err(Error::IncompleteFrame);
}
let end = available.iter().position(|byte| *byte == b'\n');
let length = end.unwrap_or(available.len());
if line.len() + length > limit {
return Err(Error::InboundLineTooLong);
}
line.extend_from_slice(&available[..length]);
reader.consume(length + usize::from(end.is_some()));
if end.is_some() {
std::str::from_utf8(&line).map_err(Error::Utf8)?;
return Ok(Some(line));
}
}
}
async fn read_stderr(mut stderr: tokio::process::ChildStderr) -> io::Result<String> {
let mut captured = Vec::new();
let mut buffer = [0; 8192];
loop {
let count = stderr.read(&mut buffer).await?;
if count == 0 {
return Ok(captured.into_iter().map(sanitize).collect());
}
let keep = (MAX_STDERR_CAPTURE_BYTES - captured.len()).min(count);
captured.extend_from_slice(&buffer[..keep]);
}
}
fn sanitize(byte: u8) -> char {
if byte.is_ascii_graphic() || matches!(byte, b' ' | b'\n' | b'\r' | b'\t') {
byte as char
} else {
'?'
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn reap_caches_exit_when_stderr_task_fails() {
let mut command = Command::new("sh");
command.arg("-c").arg("exit 0");
let mut rpc = StdioRpc::spawn(
command,
Config {
inbound_capacity: 1,
},
)
.unwrap();
rpc.stderr.as_ref().unwrap().abort();
assert!(matches!(rpc.reap().await, Err(Error::StderrTask(_))));
let exit = rpc.reap().await.unwrap();
assert_eq!(exit.status.code(), Some(0));
assert!(exit.stderr.is_empty());
rpc.join_reader().await.unwrap();
}
}