use serde_json::{Map, Number, Value, json};
use std::fmt::{Display, Formatter};
use std::io;
use std::process::{ExitStatus, Stdio};
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};
use tokio::task::JoinHandle;
pub const MAX_INBOUND_LINE_BYTES: usize = 8 * 1024 * 1024;
pub const MAX_STDERR_CAPTURE_BYTES: usize = 64 * 1024;
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct RequestId(pub u64);
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum PeerId {
Number(Number),
String(String),
}
impl PeerId {
fn into_value(self) -> Value {
match self {
Self::Number(number) => Value::Number(number),
Self::String(string) => Value::String(string),
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct RpcError {
pub code: i64,
pub message: String,
pub data: Option<Value>,
}
#[derive(Clone, Debug, PartialEq)]
pub enum IncomingMessage {
Response {
id: PeerId,
result: Result<Value, RpcError>,
},
Request {
id: PeerId,
method: String,
params: Value,
},
Notification {
method: String,
params: Value,
},
}
#[derive(Clone, Debug)]
pub struct ProcessExit {
pub status: ExitStatus,
pub stderr: String,
}
#[derive(Debug)]
pub enum Error {
Io(io::Error),
Json(serde_json::Error),
InvalidMessage(String),
InboundLineTooLong,
ProcessExited(ProcessExit),
StderrTask(String),
}
impl Display for Error {
fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::Io(error) => write!(formatter, "stdio I/O failed: {error}"),
Self::Json(error) => write!(formatter, "JSON failed: {error}"),
Self::InvalidMessage(message) => {
write!(formatter, "invalid JSON-RPC message: {message}")
}
Self::InboundLineTooLong => {
write!(
formatter,
"JSON-RPC line exceeds {MAX_INBOUND_LINE_BYTES} bytes"
)
}
Self::ProcessExited(exit) => {
write!(
formatter,
"child exited with {}: {}",
exit.status, exit.stderr
)
}
Self::StderrTask(message) => write!(formatter, "stderr reader failed: {message}"),
}
}
}
impl std::error::Error for Error {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Io(error) => Some(error),
Self::Json(error) => Some(error),
_ => None,
}
}
}
impl From<io::Error> for Error {
fn from(error: io::Error) -> Self {
Self::Io(error)
}
}
impl From<serde_json::Error> for Error {
fn from(error: serde_json::Error) -> Self {
Self::Json(error)
}
}
pub struct StdioRpc {
child: Child,
stdin: Option<ChildStdin>,
stdout: BufReader<ChildStdout>,
stderr_task: Option<JoinHandle<io::Result<Vec<u8>>>>,
next_id: u64,
exit: Option<ProcessExit>,
}
impl StdioRpc {
pub fn spawn(mut command: Command) -> Result<Self, Error> {
command
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::piped())
.kill_on_drop(true);
let mut child = command.spawn()?;
let stdin = child
.stdin
.take()
.ok_or_else(|| Error::InvalidMessage("child stdin was not piped".into()))?;
let stdout = child
.stdout
.take()
.ok_or_else(|| Error::InvalidMessage("child stdout was not piped".into()))?;
let mut stderr = child
.stderr
.take()
.ok_or_else(|| Error::InvalidMessage("child stderr was not piped".into()))?;
let stderr_task = tokio::spawn(async move {
let mut captured = Vec::new();
let mut buffer = [0_u8; 8192];
loop {
let read = stderr.read(&mut buffer).await?;
if read == 0 {
return Ok(captured);
}
let remaining = MAX_STDERR_CAPTURE_BYTES.saturating_sub(captured.len());
captured.extend_from_slice(&buffer[..read.min(remaining)]);
}
});
Ok(Self {
child,
stdin: Some(stdin),
stdout: BufReader::new(stdout),
stderr_task: Some(stderr_task),
next_id: 1,
exit: None,
})
}
pub async fn send_request(
&mut self,
method: impl Into<String>,
params: Value,
) -> Result<RequestId, Error> {
let id = RequestId(self.next_id);
self.next_id = self
.next_id
.checked_add(1)
.ok_or_else(|| Error::InvalidMessage("local request ID exhausted".into()))?;
self.send(json!({"id": id.0, "method": method.into(), "params": params}))
.await?;
Ok(id)
}
pub async fn send_notification(
&mut self,
method: impl Into<String>,
params: Value,
) -> Result<(), Error> {
self.send(json!({"method": method.into(), "params": params}))
.await
}
pub async fn next(&mut self) -> Result<IncomingMessage, Error> {
let Some(line) = self.read_line().await? else {
return Err(Error::ProcessExited(self.finish_exit().await?));
};
parse_message(serde_json::from_slice(&line)?)
}
pub async fn respond(&mut self, id: PeerId, result: Value) -> Result<(), Error> {
self.send(json!({"id": id.into_value(), "result": result}))
.await
}
pub async fn respond_error(
&mut self,
id: PeerId,
code: i64,
message: impl Into<String>,
data: Option<Value>,
) -> Result<(), Error> {
let mut error = Map::new();
error.insert("code".into(), Value::Number(code.into()));
error.insert("message".into(), Value::String(message.into()));
if let Some(data) = data {
error.insert("data".into(), data);
}
self.send(json!({"id": id.into_value(), "error": error}))
.await
}
pub async fn shutdown(mut self) -> Result<ProcessExit, Error> {
self.stdin.take();
if self.child.try_wait()?.is_none()
&& let Err(error) = self.child.start_kill()
&& self.child.try_wait()?.is_none()
{
return Err(Error::Io(error));
}
self.finish_exit().await
}
async fn send(&mut self, value: Value) -> Result<(), Error> {
let mut bytes = serde_json::to_vec(&value)?;
bytes.push(b'\n');
let stdin = self
.stdin
.as_mut()
.ok_or_else(|| Error::InvalidMessage("child stdin is closed".into()))?;
stdin.write_all(&bytes).await?;
stdin.flush().await?;
Ok(())
}
async fn read_line(&mut self) -> Result<Option<Vec<u8>>, Error> {
let mut line = Vec::new();
let mut too_long = false;
let mut saw_bytes = false;
loop {
let (consumed, ended, empty) = {
let available = self.stdout.fill_buf().await?;
if available.is_empty() {
(0, false, true)
} else {
saw_bytes = true;
let newline = available.iter().position(|byte| *byte == b'\n');
let content = newline.unwrap_or(available.len());
if !too_long {
if line.len() + content > MAX_INBOUND_LINE_BYTES {
too_long = true;
} else {
line.extend_from_slice(&available[..content]);
}
}
(
content + usize::from(newline.is_some()),
newline.is_some(),
false,
)
}
};
if empty {
if too_long {
return Err(Error::InboundLineTooLong);
}
return if saw_bytes { Ok(Some(line)) } else { Ok(None) };
}
self.stdout.consume(consumed);
if ended {
return if too_long {
Err(Error::InboundLineTooLong)
} else {
Ok(Some(line))
};
}
}
}
async fn finish_exit(&mut self) -> Result<ProcessExit, Error> {
if let Some(exit) = &self.exit {
return Ok(exit.clone());
}
let status = self.child.wait().await?;
let task = self
.stderr_task
.take()
.ok_or_else(|| Error::StderrTask("stderr result was already consumed".into()))?;
let bytes = task
.await
.map_err(|error| Error::StderrTask(error.to_string()))??;
let exit = ProcessExit {
status,
stderr: String::from_utf8_lossy(&bytes).into_owned(),
};
self.exit = Some(exit.clone());
Ok(exit)
}
}
fn parse_message(value: Value) -> Result<IncomingMessage, Error> {
let Value::Object(mut fields) = value else {
return Err(Error::InvalidMessage("top level is not an object".into()));
};
if fields.remove("jsonrpc").is_some() {
return Err(Error::InvalidMessage(
"jsonrpc version member is not used on this wire".into(),
));
}
let has_method = fields.contains_key("method");
let has_result = fields.contains_key("result");
let has_error = fields.contains_key("error");
if has_method {
if has_result || has_error {
return Err(Error::InvalidMessage(
"method message also contains a response".into(),
));
}
let method = fields
.remove("method")
.and_then(|value| value.as_str().map(str::to_owned))
.ok_or_else(|| Error::InvalidMessage("method is not a string".into()))?;
let params = fields.remove("params").unwrap_or(Value::Null);
return match fields.remove("id") {
Some(id) => Ok(IncomingMessage::Request {
id: parse_peer_id(id)?,
method,
params,
}),
None => Ok(IncomingMessage::Notification { method, params }),
};
}
if fields.contains_key("params") {
return Err(Error::InvalidMessage("response contains params".into()));
}
if has_result == has_error {
return Err(Error::InvalidMessage(
"response must contain exactly one of result or error".into(),
));
}
let id = fields
.remove("id")
.ok_or_else(|| Error::InvalidMessage("response has no id".into()))?;
let result = if has_result {
Ok(fields.remove("result").expect("result presence checked"))
} else {
Err(parse_rpc_error(
fields.remove("error").expect("error presence checked"),
)?)
};
Ok(IncomingMessage::Response {
id: parse_peer_id(id)?,
result,
})
}
fn parse_peer_id(value: Value) -> Result<PeerId, Error> {
match value {
Value::Number(number) if number.as_i64().is_some() || number.as_u64().is_some() => {
Ok(PeerId::Number(number))
}
Value::String(string) => Ok(PeerId::String(string)),
_ => Err(Error::InvalidMessage(
"id is not an integer or string".into(),
)),
}
}
fn parse_rpc_error(value: Value) -> Result<RpcError, Error> {
let Value::Object(mut fields) = value else {
return Err(Error::InvalidMessage("error is not an object".into()));
};
let code = fields
.remove("code")
.and_then(|value| value.as_i64())
.ok_or_else(|| Error::InvalidMessage("error code is not an integer".into()))?;
let message = fields
.remove("message")
.and_then(|value| value.as_str().map(str::to_owned))
.ok_or_else(|| Error::InvalidMessage("error message is not a string".into()))?;
Ok(RpcError {
code,
message,
data: fields.remove("data"),
})
}