Skip to main content

mountos_native/
frame.rs

1//! Wire frame codec. Mirrors the Go reference in go/frame/frame.go.
2//!
3//! Layout: 24-byte big-endian header (magic, version, flags, method, request id,
4//! meta length, payload length), then Protobuf meta, then raw payload bytes.
5
6use crate::error::{Error, Result};
7use std::io::{self, IoSlice, Read, Write};
8
9pub const MAGIC: [u8; 4] = *b"MFS1";
10pub const VERSION: u8 = 1;
11pub const HEADER_SIZE: usize = 24;
12
13pub const FLAG_RESPONSE: u8 = 1;
14pub const FLAG_ERROR: u8 = 2;
15pub const FLAG_STREAM_END: u8 = 4;
16/// Server-initiated frame (a notice). The request id is the id of the call it answers.
17pub const FLAG_NOTIFY: u8 = 8;
18
19pub const MAX_META_BYTES: u32 = 16 << 20;
20pub const MAX_PAYLOAD_BYTES: u32 = 8 << 20;
21
22/// One decoded message. Build it with `Frame::new`. New fields can be
23/// added in a minor version.
24#[derive(Debug, Default, Clone, PartialEq, Eq)]
25#[non_exhaustive]
26pub struct Frame {
27  pub flags: u8,
28  pub method: u16,
29  pub request_id: u64,
30  pub meta: Vec<u8>,
31  pub payload: Vec<u8>,
32}
33
34impl Frame {
35  /// Builds a frame from its five parts.
36  pub fn new(flags: u8, method: u16, request_id: u64, meta: Vec<u8>, payload: Vec<u8>) -> Frame {
37    Frame {
38      flags,
39      method,
40      request_id,
41      meta,
42      payload,
43    }
44  }
45}
46
47/// Rejects sections over the limits. Callers run this before they send a byte, so an oversize
48/// request fails alone and leaves the connection usable.
49pub fn check_limits(meta_len: usize, payload_len: usize) -> Result<()> {
50  if meta_len as u64 > MAX_META_BYTES as u64 {
51    return Err(Error::TooLarge(format!("meta {meta_len} > {MAX_META_BYTES}")));
52  }
53  if payload_len as u64 > MAX_PAYLOAD_BYTES as u64 {
54    return Err(Error::TooLarge(format!("payload {payload_len} > {MAX_PAYLOAD_BYTES}")));
55  }
56  Ok(())
57}
58
59/// Writes one frame. Rejects sections over the limits before sending anything.
60pub fn write<W: Write>(w: &mut W, f: &Frame) -> Result<()> {
61  write_parts(w, f.flags, f.method, f.request_id, &f.meta, &f.payload)
62}
63
64/// Writes one frame from borrowed sections with a single vectored write, so the payload is not
65/// copied and the header, meta and payload leave in one call.
66pub fn write_parts<W: Write>(
67  w: &mut W,
68  flags: u8,
69  method: u16,
70  request_id: u64,
71  meta: &[u8],
72  payload: &[u8],
73) -> Result<()> {
74  check_limits(meta.len(), payload.len())?;
75  let mut hdr = [0u8; HEADER_SIZE];
76  hdr[0..4].copy_from_slice(&MAGIC);
77  hdr[4] = VERSION;
78  hdr[5] = flags;
79  hdr[6..8].copy_from_slice(&method.to_be_bytes());
80  hdr[8..16].copy_from_slice(&request_id.to_be_bytes());
81  hdr[16..20].copy_from_slice(&(meta.len() as u32).to_be_bytes());
82  hdr[20..24].copy_from_slice(&(payload.len() as u32).to_be_bytes());
83  let mut parts = [IoSlice::new(&hdr), IoSlice::new(meta), IoSlice::new(payload)];
84  let mut rest: &mut [IoSlice<'_>] = &mut parts;
85  IoSlice::advance_slices(&mut rest, 0);
86  while !rest.is_empty() {
87    match w.write_vectored(rest) {
88      Ok(0) => return Err(io::Error::from(io::ErrorKind::WriteZero).into()),
89      Ok(n) => IoSlice::advance_slices(&mut rest, n),
90      Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
91      Err(e) => return Err(e.into()),
92    }
93  }
94  Ok(())
95}
96
97/// Step a large section grows by. The buffer follows the bytes that arrived, never the length
98/// the header claims, so a peer that claims a large section and stalls holds at most one step.
99const READ_STEP: usize = 64 << 10;
100
101fn read_section<R: Read>(r: &mut R, n: usize) -> io::Result<Vec<u8>> {
102  if n <= READ_STEP {
103    let mut b = vec![0u8; n];
104    r.read_exact(&mut b)?;
105    return Ok(b);
106  }
107  let mut buf = Vec::with_capacity(READ_STEP);
108  while buf.len() < n {
109    let got = buf.len();
110    buf.resize(got + (n - got).min(READ_STEP), 0);
111    r.read_exact(&mut buf[got..])?;
112  }
113  Ok(buf)
114}
115
116/// Reads one frame. Returns `Ok(None)` only when the stream ends cleanly at a
117/// frame boundary. A mid-frame end is an `UnexpectedEof` I/O error.
118pub fn read<R: Read>(r: &mut R) -> Result<Option<Frame>> {
119  let mut hdr = [0u8; HEADER_SIZE];
120  let mut got = 0;
121  while got < HEADER_SIZE {
122    match r.read(&mut hdr[got..]) {
123      Ok(0) if got == 0 => return Ok(None),
124      Ok(0) => return Err(io::Error::from(io::ErrorKind::UnexpectedEof).into()),
125      Ok(n) => got += n,
126      Err(e) if e.kind() == io::ErrorKind::Interrupted => {}
127      Err(e) => return Err(e.into()),
128    }
129  }
130  if hdr[0..4] != MAGIC {
131    return Err(Error::Protocol("bad magic".into()));
132  }
133  if hdr[4] != VERSION {
134    return Err(Error::Protocol(format!("unsupported frame version {}", hdr[4])));
135  }
136  let meta_len = u32::from_be_bytes(hdr[16..20].try_into().expect("4-byte header field"));
137  let payload_len = u32::from_be_bytes(hdr[20..24].try_into().expect("4-byte header field"));
138  if meta_len > MAX_META_BYTES {
139    return Err(Error::Protocol(format!("meta {meta_len} > {MAX_META_BYTES}")));
140  }
141  if payload_len > MAX_PAYLOAD_BYTES {
142    return Err(Error::Protocol(format!("payload {payload_len} > {MAX_PAYLOAD_BYTES}")));
143  }
144  let meta = read_section(r, meta_len as usize)?;
145  let payload = read_section(r, payload_len as usize)?;
146  Ok(Some(Frame {
147    flags: hdr[5],
148    method: u16::from_be_bytes(hdr[6..8].try_into().expect("2-byte header field")),
149    request_id: u64::from_be_bytes(hdr[8..16].try_into().expect("8-byte header field")),
150    meta,
151    payload,
152  }))
153}
154
155#[cfg(test)]
156mod tests {
157  use super::*;
158
159  #[test]
160  fn roundtrip_and_layout() {
161    let f = Frame {
162      flags: FLAG_RESPONSE,
163      method: 14,
164      request_id: 7,
165      meta: vec![1, 2],
166      payload: vec![9; 5],
167    };
168    let mut buf = Vec::new();
169    write(&mut buf, &f).unwrap();
170    assert_eq!(&buf[0..6], b"MFS1\x01\x01");
171    assert_eq!(&buf[16..24], &[0, 0, 0, 2, 0, 0, 0, 5]);
172    assert_eq!(read(&mut &buf[..]).unwrap(), Some(f));
173  }
174
175  #[test]
176  fn clean_eof_truncation_and_bad_header() {
177    assert_eq!(read(&mut &[][..]).unwrap(), None);
178    let mut buf = Vec::new();
179    write(
180      &mut buf,
181      &Frame {
182        meta: vec![1; 4],
183        ..Default::default()
184      },
185    )
186    .unwrap();
187    assert!(matches!(read(&mut &buf[..10]), Err(Error::Io(_))));
188    assert!(matches!(read(&mut &buf[..HEADER_SIZE + 2]), Err(Error::Io(_))));
189    let mut bad = buf.clone();
190    bad[0] = b'X';
191    assert!(matches!(read(&mut &bad[..]), Err(Error::Protocol(_))));
192    let mut bad = buf.clone();
193    bad[4] = 2;
194    assert!(matches!(read(&mut &bad[..]), Err(Error::Protocol(_))));
195    let mut big = buf.clone();
196    big[20..24].copy_from_slice(&(MAX_PAYLOAD_BYTES + 1).to_be_bytes());
197    assert!(matches!(read(&mut &big[..]), Err(Error::Protocol(_))));
198  }
199
200  #[test]
201  fn write_rejects_oversize_payload() {
202    let f = Frame {
203      payload: vec![0; MAX_PAYLOAD_BYTES as usize + 1],
204      ..Default::default()
205    };
206    assert!(matches!(write(&mut Vec::new(), &f), Err(Error::TooLarge(_))));
207  }
208
209  /// Serves a header, then zeros, then ends. Records the largest buffer a read was asked to fill.
210  struct Probe {
211    head: Vec<u8>,
212    zeros: usize,
213    max_ask: usize,
214  }
215
216  impl Read for Probe {
217    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
218      if !self.head.is_empty() {
219        let n = buf.len().min(self.head.len());
220        buf[..n].copy_from_slice(&self.head[..n]);
221        self.head.drain(..n);
222        return Ok(n);
223      }
224      self.max_ask = self.max_ask.max(buf.len());
225      let n = buf.len().min(self.zeros);
226      buf[..n].fill(0);
227      self.zeros -= n;
228      Ok(n)
229    }
230  }
231
232  #[test]
233  fn large_claim_is_read_in_steps() {
234    let mut hdr = Vec::new();
235    write(&mut hdr, &Frame::default()).unwrap();
236    hdr[20..24].copy_from_slice(&MAX_PAYLOAD_BYTES.to_be_bytes());
237    let mut p = Probe {
238      head: hdr,
239      zeros: 100_000,
240      max_ask: 0,
241    };
242    assert!(matches!(read(&mut p), Err(Error::Io(e)) if e.kind() == io::ErrorKind::UnexpectedEof));
243    assert!(p.max_ask <= READ_STEP, "asked for {}", p.max_ask);
244  }
245
246  #[test]
247  fn write_uses_one_call_per_frame() {
248    struct Count(usize, Vec<u8>);
249    impl Write for Count {
250      fn write(&mut self, b: &[u8]) -> io::Result<usize> {
251        self.0 += 1;
252        self.1.extend_from_slice(b);
253        Ok(b.len())
254      }
255      fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> {
256        self.0 += 1;
257        let mut n = 0;
258        for b in bufs {
259          self.1.extend_from_slice(b);
260          n += b.len();
261        }
262        Ok(n)
263      }
264      fn flush(&mut self) -> io::Result<()> {
265        Ok(())
266      }
267    }
268    let f = Frame {
269      meta: vec![1; 3],
270      payload: vec![2; 9],
271      ..Default::default()
272    };
273    let mut c = Count(0, Vec::new());
274    write(&mut c, &f).unwrap();
275    assert_eq!(c.0, 1);
276    assert_eq!(read(&mut &c.1[..]).unwrap(), Some(f));
277  }
278}