1use anyhow::Result;
4use serde_json::{Value, json};
5use std::{
6 io::{self, BufRead, Write},
7 sync::mpsc,
8 time::Duration,
9};
10
11pub const MAX_REQUEST: usize = 4096;
12pub fn read_request(input: &mut impl BufRead) -> io::Result<Option<Vec<u8>>> {
15 let mut line = Vec::new();
16 loop {
17 let buf = input.fill_buf()?;
18 if buf.is_empty() {
19 return Ok(None);
20 }
21 let end = buf.iter().position(|b| *b == b'\n').map(|n| n + 1);
22 let n = end.unwrap_or(buf.len());
23 if line.len() + n > MAX_REQUEST {
24 return Err(io::ErrorKind::InvalidData.into());
25 }
26 line.extend_from_slice(&buf[..n]);
27 input.consume(n);
28 if end.is_some() {
29 return Ok(Some(line));
30 }
31 }
32}
33pub fn stdin_requests() -> mpsc::Receiver<Vec<u8>> {
34 let (tx, rx) = mpsc::sync_channel(16);
35 std::thread::spawn(move || {
36 let mut input = io::stdin().lock();
37 while let Ok(Some(line)) = read_request(&mut input) {
38 if tx.send(line).is_err() {
39 break;
40 }
41 }
42 });
43 rx
44}
45pub fn parse(line: &[u8]) -> Result<Value> {
46 Ok(serde_json::from_str(
47 std::str::from_utf8(line)?.trim_start_matches('\u{feff}'),
48 )?)
49}
50pub fn reply(result: Result<Value>) -> Value {
51 match result {
52 Ok(value) => json!({"ok":true,"result":value}),
53 Err(error) => json!({"ok":false,"error":format!("{error:#}")}),
54 }
55}
56pub fn write_reply(result: Result<Value>) -> io::Result<()> {
58 let response = reply(result);
59 let mut output = io::stdout().lock();
60 writeln!(output, "{response}").and_then(|_| output.flush())
61}
62pub fn serve<S>(
64 state: &mut S,
65 interval: Duration,
66 mut request: impl FnMut(&mut S, Value) -> Result<Value>,
67 mut poll: impl FnMut(&mut S),
68) {
69 let rx = stdin_requests();
70 loop {
71 match rx.recv_timeout(interval) {
72 Ok(line) => {
73 if write_reply(parse(&line).and_then(|v| request(state, v))).is_err() {
74 break;
75 }
76 }
77 Err(mpsc::RecvTimeoutError::Timeout) => poll(state),
78 Err(mpsc::RecvTimeoutError::Disconnected) => break,
79 }
80 }
81}
82#[cfg(test)]
83mod tests {
84 use super::*;
85 #[test]
86 fn framing_bounds_and_bom() {
87 let mut input = &b"\xef\xbb\xbf{\"command\":\"status\"}\n{}\n"[..];
88 assert_eq!(
89 parse(&read_request(&mut input).unwrap().unwrap()).unwrap()["command"],
90 "status"
91 );
92 assert_eq!(
93 parse(&read_request(&mut input).unwrap().unwrap()).unwrap(),
94 json!({})
95 );
96 assert!(read_request(&mut input).unwrap().is_none());
97 assert!(read_request(&mut &b"{}"[..]).unwrap().is_none());
98 assert!(read_request(&mut &vec![b'x'; MAX_REQUEST + 1][..]).is_err());
99 assert!(parse(b"\xff\n").is_err());
100 assert_eq!(reply(parse(b"bad"))["ok"], false);
101 let mut exact = vec![b' '; MAX_REQUEST];
102 exact[MAX_REQUEST - 1] = b'\n';
103 assert_eq!(
104 read_request(&mut &exact[..]).unwrap().unwrap().len(),
105 MAX_REQUEST
106 );
107 }
108}