kcode_jsonrpc_stdio/
lib.rs

1use serde_json::{Map, Number, Value, json};
2use std::fmt::{Display, Formatter};
3use std::io;
4use std::process::{ExitStatus, Stdio};
5use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader};
6use tokio::process::{Child, ChildStdin, ChildStdout, Command};
7use tokio::task::JoinHandle;
8
9pub const MAX_INBOUND_LINE_BYTES: usize = 8 * 1024 * 1024;
10pub const MAX_STDERR_CAPTURE_BYTES: usize = 64 * 1024;
11
12#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
13pub struct RequestId(pub u64);
14
15#[derive(Clone, Debug, Eq, PartialEq)]
16pub enum PeerId {
17    Number(Number),
18    String(String),
19}
20
21impl PeerId {
22    fn into_value(self) -> Value {
23        match self {
24            Self::Number(number) => Value::Number(number),
25            Self::String(string) => Value::String(string),
26        }
27    }
28}
29
30#[derive(Clone, Debug, PartialEq)]
31pub struct RpcError {
32    pub code: i64,
33    pub message: String,
34    pub data: Option<Value>,
35}
36
37#[derive(Clone, Debug, PartialEq)]
38pub enum IncomingMessage {
39    Response {
40        id: PeerId,
41        result: Result<Value, RpcError>,
42    },
43    Request {
44        id: PeerId,
45        method: String,
46        params: Value,
47    },
48    Notification {
49        method: String,
50        params: Value,
51    },
52}
53
54#[derive(Clone, Debug)]
55pub struct ProcessExit {
56    pub status: ExitStatus,
57    pub stderr: String,
58}
59
60#[derive(Debug)]
61pub enum Error {
62    Io(io::Error),
63    Json(serde_json::Error),
64    InvalidMessage(String),
65    InboundLineTooLong,
66    ProcessExited(ProcessExit),
67    StderrTask(String),
68}
69
70impl Display for Error {
71    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
72        match self {
73            Self::Io(error) => write!(formatter, "stdio I/O failed: {error}"),
74            Self::Json(error) => write!(formatter, "JSON failed: {error}"),
75            Self::InvalidMessage(message) => {
76                write!(formatter, "invalid JSON-RPC message: {message}")
77            }
78            Self::InboundLineTooLong => {
79                write!(
80                    formatter,
81                    "JSON-RPC line exceeds {MAX_INBOUND_LINE_BYTES} bytes"
82                )
83            }
84            Self::ProcessExited(exit) => {
85                write!(
86                    formatter,
87                    "child exited with {}: {}",
88                    exit.status, exit.stderr
89                )
90            }
91            Self::StderrTask(message) => write!(formatter, "stderr reader failed: {message}"),
92        }
93    }
94}
95
96impl std::error::Error for Error {
97    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
98        match self {
99            Self::Io(error) => Some(error),
100            Self::Json(error) => Some(error),
101            _ => None,
102        }
103    }
104}
105
106impl From<io::Error> for Error {
107    fn from(error: io::Error) -> Self {
108        Self::Io(error)
109    }
110}
111
112impl From<serde_json::Error> for Error {
113    fn from(error: serde_json::Error) -> Self {
114        Self::Json(error)
115    }
116}
117
118pub struct StdioRpc {
119    child: Child,
120    stdin: Option<ChildStdin>,
121    stdout: BufReader<ChildStdout>,
122    stderr_task: Option<JoinHandle<io::Result<Vec<u8>>>>,
123    next_id: u64,
124    exit: Option<ProcessExit>,
125}
126
127impl StdioRpc {
128    pub fn spawn(mut command: Command) -> Result<Self, Error> {
129        command
130            .stdin(Stdio::piped())
131            .stdout(Stdio::piped())
132            .stderr(Stdio::piped())
133            .kill_on_drop(true);
134        let mut child = command.spawn()?;
135        let stdin = child
136            .stdin
137            .take()
138            .ok_or_else(|| Error::InvalidMessage("child stdin was not piped".into()))?;
139        let stdout = child
140            .stdout
141            .take()
142            .ok_or_else(|| Error::InvalidMessage("child stdout was not piped".into()))?;
143        let mut stderr = child
144            .stderr
145            .take()
146            .ok_or_else(|| Error::InvalidMessage("child stderr was not piped".into()))?;
147        let stderr_task = tokio::spawn(async move {
148            let mut captured = Vec::new();
149            let mut buffer = [0_u8; 8192];
150            loop {
151                let read = stderr.read(&mut buffer).await?;
152                if read == 0 {
153                    return Ok(captured);
154                }
155                let remaining = MAX_STDERR_CAPTURE_BYTES.saturating_sub(captured.len());
156                captured.extend_from_slice(&buffer[..read.min(remaining)]);
157            }
158        });
159        Ok(Self {
160            child,
161            stdin: Some(stdin),
162            stdout: BufReader::new(stdout),
163            stderr_task: Some(stderr_task),
164            next_id: 1,
165            exit: None,
166        })
167    }
168
169    pub async fn send_request(
170        &mut self,
171        method: impl Into<String>,
172        params: Value,
173    ) -> Result<RequestId, Error> {
174        let id = RequestId(self.next_id);
175        self.next_id = self
176            .next_id
177            .checked_add(1)
178            .ok_or_else(|| Error::InvalidMessage("local request ID exhausted".into()))?;
179        self.send(json!({"id": id.0, "method": method.into(), "params": params}))
180            .await?;
181        Ok(id)
182    }
183
184    pub async fn send_notification(
185        &mut self,
186        method: impl Into<String>,
187        params: Value,
188    ) -> Result<(), Error> {
189        self.send(json!({"method": method.into(), "params": params}))
190            .await
191    }
192
193    pub async fn next(&mut self) -> Result<IncomingMessage, Error> {
194        let Some(line) = self.read_line().await? else {
195            return Err(Error::ProcessExited(self.finish_exit().await?));
196        };
197        parse_message(serde_json::from_slice(&line)?)
198    }
199
200    pub async fn respond(&mut self, id: PeerId, result: Value) -> Result<(), Error> {
201        self.send(json!({"id": id.into_value(), "result": result}))
202            .await
203    }
204
205    pub async fn respond_error(
206        &mut self,
207        id: PeerId,
208        code: i64,
209        message: impl Into<String>,
210        data: Option<Value>,
211    ) -> Result<(), Error> {
212        let mut error = Map::new();
213        error.insert("code".into(), Value::Number(code.into()));
214        error.insert("message".into(), Value::String(message.into()));
215        if let Some(data) = data {
216            error.insert("data".into(), data);
217        }
218        self.send(json!({"id": id.into_value(), "error": error}))
219            .await
220    }
221
222    pub async fn shutdown(mut self) -> Result<ProcessExit, Error> {
223        self.stdin.take();
224        if self.child.try_wait()?.is_none()
225            && let Err(error) = self.child.start_kill()
226            && self.child.try_wait()?.is_none()
227        {
228            return Err(Error::Io(error));
229        }
230        self.finish_exit().await
231    }
232
233    async fn send(&mut self, value: Value) -> Result<(), Error> {
234        let mut bytes = serde_json::to_vec(&value)?;
235        bytes.push(b'\n');
236        let stdin = self
237            .stdin
238            .as_mut()
239            .ok_or_else(|| Error::InvalidMessage("child stdin is closed".into()))?;
240        stdin.write_all(&bytes).await?;
241        stdin.flush().await?;
242        Ok(())
243    }
244
245    async fn read_line(&mut self) -> Result<Option<Vec<u8>>, Error> {
246        let mut line = Vec::new();
247        let mut too_long = false;
248        let mut saw_bytes = false;
249        loop {
250            let (consumed, ended, empty) = {
251                let available = self.stdout.fill_buf().await?;
252                if available.is_empty() {
253                    (0, false, true)
254                } else {
255                    saw_bytes = true;
256                    let newline = available.iter().position(|byte| *byte == b'\n');
257                    let content = newline.unwrap_or(available.len());
258                    if !too_long {
259                        if line.len() + content > MAX_INBOUND_LINE_BYTES {
260                            too_long = true;
261                        } else {
262                            line.extend_from_slice(&available[..content]);
263                        }
264                    }
265                    (
266                        content + usize::from(newline.is_some()),
267                        newline.is_some(),
268                        false,
269                    )
270                }
271            };
272            if empty {
273                if too_long {
274                    return Err(Error::InboundLineTooLong);
275                }
276                return if saw_bytes { Ok(Some(line)) } else { Ok(None) };
277            }
278            self.stdout.consume(consumed);
279            if ended {
280                return if too_long {
281                    Err(Error::InboundLineTooLong)
282                } else {
283                    Ok(Some(line))
284                };
285            }
286        }
287    }
288
289    async fn finish_exit(&mut self) -> Result<ProcessExit, Error> {
290        if let Some(exit) = &self.exit {
291            return Ok(exit.clone());
292        }
293        let status = self.child.wait().await?;
294        let task = self
295            .stderr_task
296            .take()
297            .ok_or_else(|| Error::StderrTask("stderr result was already consumed".into()))?;
298        let bytes = task
299            .await
300            .map_err(|error| Error::StderrTask(error.to_string()))??;
301        let exit = ProcessExit {
302            status,
303            stderr: String::from_utf8_lossy(&bytes).into_owned(),
304        };
305        self.exit = Some(exit.clone());
306        Ok(exit)
307    }
308}
309
310fn parse_message(value: Value) -> Result<IncomingMessage, Error> {
311    let Value::Object(mut fields) = value else {
312        return Err(Error::InvalidMessage("top level is not an object".into()));
313    };
314    if fields.remove("jsonrpc").is_some() {
315        return Err(Error::InvalidMessage(
316            "jsonrpc version member is not used on this wire".into(),
317        ));
318    }
319    let has_method = fields.contains_key("method");
320    let has_result = fields.contains_key("result");
321    let has_error = fields.contains_key("error");
322    if has_method {
323        if has_result || has_error {
324            return Err(Error::InvalidMessage(
325                "method message also contains a response".into(),
326            ));
327        }
328        let method = fields
329            .remove("method")
330            .and_then(|value| value.as_str().map(str::to_owned))
331            .ok_or_else(|| Error::InvalidMessage("method is not a string".into()))?;
332        let params = fields.remove("params").unwrap_or(Value::Null);
333        return match fields.remove("id") {
334            Some(id) => Ok(IncomingMessage::Request {
335                id: parse_peer_id(id)?,
336                method,
337                params,
338            }),
339            None => Ok(IncomingMessage::Notification { method, params }),
340        };
341    }
342    if fields.contains_key("params") {
343        return Err(Error::InvalidMessage("response contains params".into()));
344    }
345    if has_result == has_error {
346        return Err(Error::InvalidMessage(
347            "response must contain exactly one of result or error".into(),
348        ));
349    }
350    let id = fields
351        .remove("id")
352        .ok_or_else(|| Error::InvalidMessage("response has no id".into()))?;
353    let result = if has_result {
354        Ok(fields.remove("result").expect("result presence checked"))
355    } else {
356        Err(parse_rpc_error(
357            fields.remove("error").expect("error presence checked"),
358        )?)
359    };
360    Ok(IncomingMessage::Response {
361        id: parse_peer_id(id)?,
362        result,
363    })
364}
365
366fn parse_peer_id(value: Value) -> Result<PeerId, Error> {
367    match value {
368        Value::Number(number) if number.as_i64().is_some() || number.as_u64().is_some() => {
369            Ok(PeerId::Number(number))
370        }
371        Value::String(string) => Ok(PeerId::String(string)),
372        _ => Err(Error::InvalidMessage(
373            "id is not an integer or string".into(),
374        )),
375    }
376}
377
378fn parse_rpc_error(value: Value) -> Result<RpcError, Error> {
379    let Value::Object(mut fields) = value else {
380        return Err(Error::InvalidMessage("error is not an object".into()));
381    };
382    let code = fields
383        .remove("code")
384        .and_then(|value| value.as_i64())
385        .ok_or_else(|| Error::InvalidMessage("error code is not an integer".into()))?;
386    let message = fields
387        .remove("message")
388        .and_then(|value| value.as_str().map(str::to_owned))
389        .ok_or_else(|| Error::InvalidMessage("error message is not a string".into()))?;
390    Ok(RpcError {
391        code,
392        message,
393        data: fields.remove("data"),
394    })
395}
kcode_jsonrpc_stdio/lib.rs

kcode_jsonrpc_stdio/
lib.rs
