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

mountos_native/
frame.rs
