netget/server/msgpack_rpc/
wire.rs

1//! MessagePack and the MessagePack-RPC envelope (request `[0, msgid, method, params]`,
2//! response `[1, msgid, error, result]`, notification `[2, method, params]`), shared by the
3//! server and the client.
4//!
5//! The stream has no framing, so bytes are buffered (at most `MAX_MESSAGE`) and a message is
6//! decoded only from a complete in-memory buffer: no length the peer announces can allocate
7//! more than the bytes already received, and nesting stops at `MAX_DEPTH`.
8//!
9//! Values map to JSON one-to-one where JSON can say them; binary becomes `{"$bin": hex}`, an
10//! extension `{"$ext": type, "hex": hex}`, and a map whose keys are not strings uses each key's
11//! JSON text as its name. Those objects encode back to the same MessagePack.
12use anyhow::{bail, ensure, Context, Result};
13use serde_json::{json, Map, Number, Value};
14use std::time::Duration;
15use tokio::io::{AsyncRead, AsyncReadExt};
16
17/// Largest single message either side accepts.
18pub const MAX_MESSAGE: usize = 1024 * 1024;
19/// Deepest nesting decoded.
20pub const MAX_DEPTH: usize = 32;
21/// Deadline for the rest of a message once its first byte has arrived.
22pub const IO_TIMEOUT: Duration = Duration::from_secs(30);
23
24pub const REQUEST: u64 = 0;
25pub const RESPONSE: u64 = 1;
26pub const NOTIFICATION: u64 = 2;
27
28/// Decoding outcome for a buffer that may hold only part of a message.
29enum Need {
30    More,
31}
32
33struct Cursor<'a> {
34    buf: &'a [u8],
35    pos: usize,
36}
37
38type Step<T> = std::result::Result<Result<T>, Need>;
39
40macro_rules! need {
41    ($e:expr) => {
42        match $e {
43            Ok(Ok(v)) => v,
44            Ok(Err(e)) => return Ok(Err(e)),
45            Err(n) => return Err(n),
46        }
47    };
48}
49
50impl Cursor<'_> {
51    fn take(&mut self, n: usize) -> Step<&[u8]> {
52        if self.buf.len() - self.pos < n {
53            return Err(Need::More);
54        }
55        let s = &self.buf[self.pos..self.pos + n];
56        self.pos += n;
57        Ok(Ok(s))
58    }
59    fn uint(&mut self, n: usize) -> Step<u64> {
60        let b = need!(self.take(n));
61        Ok(Ok(b.iter().fold(0u64, |acc, x| (acc << 8) | u64::from(*x))))
62    }
63    /// A count of `n` elements needs at least `n` more bytes; refuse it before looping.
64    fn count(&mut self, n: u64, per_item: usize) -> Step<usize> {
65        let remaining = (self.buf.len() - self.pos) as u64;
66        if n.saturating_mul(per_item as u64) > MAX_MESSAGE as u64 {
67            return Ok(Err(anyhow::anyhow!(
68                "MessagePack count {n} exceeds the message bound"
69            )));
70        }
71        if n.saturating_mul(per_item as u64) > remaining {
72            return Err(Need::More);
73        }
74        Ok(Ok(n as usize))
75    }
76    fn hex(b: &[u8]) -> String {
77        b.iter().map(|x| format!("{x:02x}")).collect()
78    }
79    fn value(&mut self, depth: usize) -> Step<Value> {
80        if depth > MAX_DEPTH {
81            return Ok(Err(anyhow::anyhow!(
82                "MessagePack nested deeper than {MAX_DEPTH}"
83            )));
84        }
85        let tag = need!(self.take(1))[0];
86        let v = match tag {
87            0x00..=0x7f => json!(tag),
88            0xe0..=0xff => json!(tag as i8),
89            0xc0 => Value::Null,
90            0xc2 => json!(false),
91            0xc3 => json!(true),
92            0xcc => json!(need!(self.uint(1))),
93            0xcd => json!(need!(self.uint(2))),
94            0xce => json!(need!(self.uint(4))),
95            0xcf => json!(need!(self.uint(8))),
96            0xd0 => json!(need!(self.uint(1)) as u8 as i8),
97            0xd1 => json!(need!(self.uint(2)) as u16 as i16),
98            0xd2 => json!(need!(self.uint(4)) as u32 as i32),
99            0xd3 => json!(need!(self.uint(8)) as i64),
100            0xca => {
101                let f = f32::from_bits(need!(self.uint(4)) as u32);
102                Number::from_f64(f64::from(f))
103                    .map(Value::Number)
104                    .unwrap_or(Value::Null)
105            }
106            0xcb => Number::from_f64(f64::from_bits(need!(self.uint(8))))
107                .map(Value::Number)
108                .unwrap_or(Value::Null),
109            0xa0..=0xbf | 0xd9 | 0xda | 0xdb => {
110                let n = match tag {
111                    0xd9 => need!(self.uint(1)),
112                    0xda => need!(self.uint(2)),
113                    0xdb => need!(self.uint(4)),
114                    t => u64::from(t & 0x1f),
115                };
116                let n = need!(self.count(n, 1));
117                let b = need!(self.take(n));
118                match std::str::from_utf8(b) {
119                    Ok(s) => json!(s),
120                    Err(_) => json!({"$bin": Self::hex(b)}),
121                }
122            }
123            0xc4..=0xc6 => {
124                let n = need!(self.uint(1 << (tag - 0xc4)));
125                let n = need!(self.count(n, 1));
126                json!({"$bin": Self::hex(need!(self.take(n)))})
127            }
128            0x90..=0x9f | 0xdc | 0xdd => {
129                let n = match tag {
130                    0xdc => need!(self.uint(2)),
131                    0xdd => need!(self.uint(4)),
132                    t => u64::from(t & 0x0f),
133                };
134                let n = need!(self.count(n, 1));
135                let mut items = Vec::with_capacity(n);
136                for _ in 0..n {
137                    items.push(need!(self.value(depth + 1)));
138                }
139                Value::Array(items)
140            }
141            0x80..=0x8f | 0xde | 0xdf => {
142                let n = match tag {
143                    0xde => need!(self.uint(2)),
144                    0xdf => need!(self.uint(4)),
145                    t => u64::from(t & 0x0f),
146                };
147                let n = need!(self.count(n, 2));
148                let mut map = Map::new();
149                for _ in 0..n {
150                    let k = need!(self.value(depth + 1));
151                    let v = need!(self.value(depth + 1));
152                    let key = match k {
153                        Value::String(s) => s,
154                        other => other.to_string(),
155                    };
156                    map.insert(key, v);
157                }
158                Value::Object(map)
159            }
160            0xd4..=0xd8 | 0xc7..=0xc9 => {
161                let n = match tag {
162                    0xd4..=0xd8 => 1usize << (tag - 0xd4),
163                    0xc7 => need!(self.uint(1)) as usize,
164                    0xc8 => need!(self.uint(2)) as usize,
165                    _ => need!(self.uint(4)) as usize,
166                };
167                let ext = need!(self.take(1))[0] as i8;
168                let n = need!(self.count(n as u64, 1));
169                json!({"$ext": ext, "hex": Self::hex(need!(self.take(n)))})
170            }
171            0xc1 => return Ok(Err(anyhow::anyhow!("MessagePack byte 0xc1 is never used"))),
172        };
173        Ok(Ok(v))
174    }
175}
176
177/// Decode one value from the front of `buf`: `Ok(Some((value, used)))`, `Ok(None)` when the
178/// buffer holds only part of it, `Err` when it is malformed.
179pub fn decode(buf: &[u8]) -> Result<Option<(Value, usize)>> {
180    let mut c = Cursor { buf, pos: 0 };
181    match c.value(0) {
182        Ok(Ok(v)) => Ok(Some((v, c.pos))),
183        Ok(Err(e)) => Err(e),
184        Err(Need::More) => Ok(None),
185    }
186}
187
188fn unhex(s: &str) -> Result<Vec<u8>> {
189    ensure!(
190        s.is_ascii() && s.len().is_multiple_of(2),
191        "hex must be an even number of hex digits"
192    );
193    (0..s.len())
194        .step_by(2)
195        .map(|i| u8::from_str_radix(&s[i..i + 2], 16).context("invalid hex"))
196        .collect()
197}
198
199fn put_len(out: &mut Vec<u8>, n: usize, fix: Option<(u8, usize)>, tags: [u8; 3]) {
200    match fix {
201        Some((base, max)) if n <= max => out.push(base | n as u8),
202        _ if tags[0] != 0 && n <= 0xff => {
203            out.push(tags[0]);
204            out.push(n as u8);
205        }
206        _ if n <= 0xffff => {
207            out.push(tags[1]);
208            out.extend((n as u16).to_be_bytes());
209        }
210        _ => {
211            out.push(tags[2]);
212            out.extend((n as u32).to_be_bytes());
213        }
214    }
215}
216
217/// Encode a JSON value as MessagePack (the inverse of `decode`'s mapping).
218pub fn encode(v: &Value, out: &mut Vec<u8>, depth: usize) -> Result<()> {
219    ensure!(depth <= MAX_DEPTH, "value nested deeper than {MAX_DEPTH}");
220    match v {
221        Value::Null => out.push(0xc0),
222        Value::Bool(b) => out.push(if *b { 0xc3 } else { 0xc2 }),
223        Value::Number(n) => {
224            if let Some(u) = n.as_u64() {
225                if u <= 0x7f {
226                    out.push(u as u8);
227                } else {
228                    out.push(0xcf);
229                    out.extend(u.to_be_bytes());
230                }
231            } else if let Some(i) = n.as_i64() {
232                if i >= -32 {
233                    out.push(i as i8 as u8);
234                } else {
235                    out.push(0xd3);
236                    out.extend(i.to_be_bytes());
237                }
238            } else {
239                out.push(0xcb);
240                out.extend(n.as_f64().unwrap_or(0.0).to_be_bytes());
241            }
242        }
243        Value::String(s) => {
244            put_len(out, s.len(), Some((0xa0, 31)), [0xd9, 0xda, 0xdb]);
245            out.extend_from_slice(s.as_bytes());
246        }
247        Value::Array(items) => {
248            put_len(out, items.len(), Some((0x90, 15)), [0, 0xdc, 0xdd]);
249            for i in items {
250                encode(i, out, depth + 1)?;
251            }
252        }
253        Value::Object(map) => {
254            if let (1, Some(Value::String(h))) = (map.len(), map.get("$bin")) {
255                let b = unhex(h)?;
256                put_len(out, b.len(), None, [0xc4, 0xc5, 0xc6]);
257                out.extend(b);
258                return Ok(());
259            }
260            if let (2, Some(t), Some(Value::String(h))) = (
261                map.len(),
262                map.get("$ext").and_then(Value::as_i64),
263                map.get("hex"),
264            ) {
265                let b = unhex(h)?;
266                ensure!((-128..=127).contains(&t), "extension type out of range");
267                match b.len() {
268                    1 => out.push(0xd4),
269                    2 => out.push(0xd5),
270                    4 => out.push(0xd6),
271                    8 => out.push(0xd7),
272                    16 => out.push(0xd8),
273                    n => put_len(out, n, None, [0xc7, 0xc8, 0xc9]),
274                }
275                out.push(t as i8 as u8);
276                out.extend(b);
277                return Ok(());
278            }
279            put_len(out, map.len(), Some((0x80, 15)), [0, 0xde, 0xdf]);
280            for (k, v) in map {
281                encode(&Value::String(k.clone()), out, depth + 1)?;
282                encode(v, out, depth + 1)?;
283            }
284        }
285    }
286    Ok(())
287}
288
289pub fn encode_message(v: &Value) -> Result<Vec<u8>> {
290    let mut out = Vec::new();
291    encode(v, &mut out, 0)?;
292    ensure!(
293        out.len() <= MAX_MESSAGE,
294        "message exceeds {MAX_MESSAGE} bytes"
295    );
296    Ok(out)
297}
298
299/// A stream of messages: buffered bytes and the reader they come from.
300pub struct Stream<R> {
301    reader: R,
302    buf: Vec<u8>,
303}
304
305impl<R: AsyncRead + Unpin> Stream<R> {
306    pub fn new(reader: R) -> Self {
307        Self {
308            reader,
309            buf: Vec::new(),
310        }
311    }
312
313    /// The next whole message. `None` on a clean EOF between messages; the first byte may
314    /// wait `idle`, the rest must follow within `IO_TIMEOUT`.
315    pub async fn next(&mut self, idle: Duration) -> Result<Option<Value>> {
316        let mut deadline = None;
317        loop {
318            if let Some((v, used)) = decode(&self.buf)? {
319                self.buf.drain(..used);
320                return Ok(Some(v));
321            }
322            ensure!(
323                self.buf.len() < MAX_MESSAGE,
324                "message exceeds {MAX_MESSAGE} bytes"
325            );
326            let wait = if self.buf.is_empty() {
327                idle
328            } else {
329                deadline
330                    .get_or_insert_with(|| tokio::time::Instant::now() + IO_TIMEOUT)
331                    .saturating_duration_since(tokio::time::Instant::now())
332            };
333            let mut chunk = vec![0u8; 64 * 1024];
334            let n = match tokio::time::timeout(wait, self.reader.read(&mut chunk)).await {
335                Ok(r) => r?,
336                Err(_) if self.buf.is_empty() => {
337                    bail!("MessagePack-RPC peer idle for {}s", idle.as_secs())
338                }
339                Err(_) => bail!(
340                    "MessagePack-RPC message incomplete after {}s",
341                    IO_TIMEOUT.as_secs()
342                ),
343            };
344            if n == 0 {
345                ensure!(self.buf.is_empty(), "connection closed inside a message");
346                return Ok(None);
347            }
348            self.buf.extend_from_slice(&chunk[..n]);
349        }
350    }
351}
352
353/// An RPC message, checked against the envelope.
354#[derive(Debug, Clone, PartialEq)]
355pub enum Rpc {
356    Request {
357        msgid: u64,
358        method: String,
359        params: Value,
360    },
361    Response {
362        msgid: u64,
363        error: Value,
364        result: Value,
365    },
366    Notification {
367        method: String,
368        params: Value,
369    },
370}
371
372pub fn parse(v: Value) -> Result<Rpc> {
373    let Value::Array(items) = v else {
374        bail!("an RPC message is an array")
375    };
376    let kind = items
377        .first()
378        .and_then(Value::as_u64)
379        .context("an RPC message starts with its type")?;
380    let msgid = |v: &Value| {
381        v.as_u64()
382            .filter(|m| *m <= u64::from(u32::MAX))
383            .context("msgid must be a 32-bit unsigned integer")
384    };
385    let params = |v: &Value| -> Result<Value> {
386        ensure!(v.is_array(), "params must be an array");
387        Ok(v.clone())
388    };
389    Ok(match (kind, items.len()) {
390        (REQUEST, 4) => Rpc::Request {
391            msgid: msgid(&items[1])?,
392            method: items[2]
393                .as_str()
394                .context("method must be a string")?
395                .to_string(),
396            params: params(&items[3])?,
397        },
398        (RESPONSE, 4) => Rpc::Response {
399            msgid: msgid(&items[1])?,
400            error: items[2].clone(),
401            result: items[3].clone(),
402        },
403        (NOTIFICATION, 3) => Rpc::Notification {
404            method: items[1]
405                .as_str()
406                .context("method must be a string")?
407                .to_string(),
408            params: params(&items[2])?,
409        },
410        (k, n) => bail!("RPC message of type {k} with {n} elements"),
411    })
412}
413
414pub fn request(msgid: u64, method: &str, params: &Value) -> Result<Vec<u8>> {
415    encode_message(&json!([REQUEST, msgid, method, params]))
416}
417
418pub fn response(msgid: u64, error: &Value, result: &Value) -> Result<Vec<u8>> {
419    encode_message(&json!([RESPONSE, msgid, error, result]))
420}
421
422pub fn notification(method: &str, params: &Value) -> Result<Vec<u8>> {
423    encode_message(&json!([NOTIFICATION, method, params]))
424}
netget/server/msgpack_rpc/wire.rs

netget/server/msgpack_rpc/
wire.rs
