1use 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;
16pub 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)]
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 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
47pub 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
59pub 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
64pub 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
97const 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
116pub 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 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}