Skip to main content

ozzy_proto/
layout.rs

1//! Bounds-checked big-endian fields for wire and file layouts.
2//!
3//! [`Reader`] and [`Writer`] walk counted layouts and fail with a
4//! [`LayoutError`] where the bytes run out; each codec converts that into its
5//! own error. [`Fixed`] reads and writes the fields of a fixed-size record at
6//! offsets checked when the code is built, so a record whose length was checked
7//! once needs no further checks.
8
9/// A layout that does not fit its bytes.
10#[derive(Debug, Clone, Copy, PartialEq, Eq, thiserror::Error)]
11pub enum LayoutError {
12    /// The input ends before a field.
13    #[error("layout is truncated: need {needed} bytes, have {available}")]
14    Truncated {
15        /// Input length the field needs.
16        needed: usize,
17        /// Input length.
18        available: usize,
19    },
20    /// Bytes remain after the last field.
21    #[error("layout has {0} trailing bytes")]
22    Trailing(usize),
23    /// The output cannot hold a field.
24    #[error("layout overflows its output: need {needed} bytes, have {capacity}")]
25    Overflow {
26        /// Output length the field needs.
27        needed: usize,
28        /// Output length.
29        capacity: usize,
30    },
31}
32
33/// Sequential reader over borrowed bytes.
34#[derive(Debug, Clone, Copy)]
35pub struct Reader<'a> {
36    rest: &'a [u8],
37    len: usize,
38}
39
40impl<'a> Reader<'a> {
41    /// Read `input` from its start.
42    #[must_use]
43    pub const fn new(input: &'a [u8]) -> Self {
44        Self {
45            rest: input,
46            len: input.len(),
47        }
48    }
49
50    /// Bytes consumed so far.
51    #[must_use]
52    #[expect(
53        clippy::arithmetic_side_effects,
54        reason = "rest is a suffix of the original input"
55    )]
56    pub const fn offset(&self) -> usize {
57        self.len - self.rest.len()
58    }
59
60    /// Bytes not yet consumed.
61    #[must_use]
62    pub const fn rest(&self) -> &'a [u8] {
63        self.rest
64    }
65
66    /// The next `count` bytes.
67    pub fn take(&mut self, count: usize) -> Result<&'a [u8], LayoutError> {
68        let (bytes, rest) = self
69            .rest
70            .split_at_checked(count)
71            .ok_or_else(|| self.truncated(count))?;
72        self.rest = rest;
73        Ok(bytes)
74    }
75
76    /// The next `N` bytes, borrowed.
77    pub fn chunk<const N: usize>(&mut self) -> Result<&'a [u8; N], LayoutError> {
78        let (bytes, rest) = self
79            .rest
80            .split_first_chunk()
81            .ok_or_else(|| self.truncated(N))?;
82        self.rest = rest;
83        Ok(bytes)
84    }
85
86    /// The next `N` bytes, copied.
87    pub fn array<const N: usize>(&mut self) -> Result<[u8; N], LayoutError> {
88        self.chunk().copied()
89    }
90
91    /// The next byte.
92    pub fn u8(&mut self) -> Result<u8, LayoutError> {
93        let [byte] = self.array()?;
94        Ok(byte)
95    }
96
97    /// The next big-endian `u16`.
98    pub fn u16(&mut self) -> Result<u16, LayoutError> {
99        self.array().map(u16::from_be_bytes)
100    }
101
102    /// The next big-endian `u32`.
103    pub fn u32(&mut self) -> Result<u32, LayoutError> {
104        self.array().map(u32::from_be_bytes)
105    }
106
107    /// The next big-endian `u64`.
108    pub fn u64(&mut self) -> Result<u64, LayoutError> {
109        self.array().map(u64::from_be_bytes)
110    }
111
112    /// Require that every byte was consumed.
113    pub const fn finish(self) -> Result<(), LayoutError> {
114        if self.rest.is_empty() {
115            Ok(())
116        } else {
117            Err(LayoutError::Trailing(self.rest.len()))
118        }
119    }
120
121    fn truncated(&self, count: usize) -> LayoutError {
122        LayoutError::Truncated {
123            needed: self.offset().saturating_add(count),
124            available: self.len,
125        }
126    }
127}
128
129/// Sequential writer into a caller-sized buffer.
130#[derive(Debug)]
131pub struct Writer<'a> {
132    output: &'a mut [u8],
133    offset: usize,
134}
135
136impl<'a> Writer<'a> {
137    /// Write `output` from its start.
138    pub const fn new(output: &'a mut [u8]) -> Self {
139        Self { output, offset: 0 }
140    }
141
142    /// Write the first `len` bytes of `output`, or fail when it is shorter.
143    pub fn with_len(output: &'a mut [u8], len: usize) -> Result<Self, LayoutError> {
144        let capacity = output.len();
145        let output = output.get_mut(..len).ok_or(LayoutError::Overflow {
146            needed: len,
147            capacity,
148        })?;
149        Ok(Self::new(output))
150    }
151
152    /// Bytes written so far.
153    #[must_use]
154    pub const fn offset(&self) -> usize {
155        self.offset
156    }
157
158    /// Write `bytes` next, or nothing when they do not fit.
159    pub fn bytes(&mut self, bytes: &[u8]) -> Result<(), LayoutError> {
160        let overflow = LayoutError::Overflow {
161            needed: self.offset.saturating_add(bytes.len()),
162            capacity: self.output.len(),
163        };
164        let end = self.offset.checked_add(bytes.len()).ok_or(overflow)?;
165        let field = self.output.get_mut(self.offset..end).ok_or(overflow)?;
166        field.copy_from_slice(bytes);
167        self.offset = end;
168        Ok(())
169    }
170
171    /// Write one byte.
172    pub fn u8(&mut self, value: u8) -> Result<(), LayoutError> {
173        self.bytes(&[value])
174    }
175
176    /// Write a big-endian `u16`.
177    pub fn u16(&mut self, value: u16) -> Result<(), LayoutError> {
178        self.bytes(&value.to_be_bytes())
179    }
180
181    /// Write a big-endian `u32`.
182    pub fn u32(&mut self, value: u32) -> Result<(), LayoutError> {
183        self.bytes(&value.to_be_bytes())
184    }
185
186    /// Write a big-endian `u64`.
187    pub fn u64(&mut self, value: u64) -> Result<(), LayoutError> {
188        self.bytes(&value.to_be_bytes())
189    }
190}
191
192/// The `N`-byte record at `at`.
193pub fn record<const N: usize>(input: &[u8], at: usize) -> Result<&[u8; N], LayoutError> {
194    input
195        .get(at..)
196        .and_then(<[u8]>::first_chunk)
197        .ok_or(LayoutError::Truncated {
198            needed: at.saturating_add(N),
199            available: input.len(),
200        })
201}
202
203/// The mutable `N`-byte record at `at`.
204pub fn record_mut<const N: usize>(
205    output: &mut [u8],
206    at: usize,
207) -> Result<&mut [u8; N], LayoutError> {
208    let capacity = output.len();
209    output
210        .get_mut(at..)
211        .and_then(<[u8]>::first_chunk_mut)
212        .ok_or(LayoutError::Overflow {
213            needed: at.saturating_add(N),
214            capacity,
215        })
216}
217
218/// Overwrite the `N` bytes at `at`, as when filling a placeholder.
219pub fn patch<const N: usize>(
220    output: &mut [u8],
221    at: usize,
222    bytes: [u8; N],
223) -> Result<(), LayoutError> {
224    *record_mut(output, at)? = bytes;
225    Ok(())
226}
227
228/// Fields of a fixed-size record at offsets checked when the code is built.
229///
230/// A field that does not fit the record fails the build, so the accessors
231/// cannot fail. Get the record from a slice with `first_chunk`,
232/// `split_first_chunk`, or `as_chunks`.
233pub trait Fixed {
234    /// The `M` bytes at `AT`, copied.
235    fn array<const AT: usize, const M: usize>(&self) -> [u8; M];
236
237    /// The `M` bytes at `AT`, borrowed.
238    fn slice<const AT: usize, const M: usize>(&self) -> &[u8];
239
240    /// Overwrite the `M` bytes at `AT`.
241    fn put<const AT: usize, const M: usize>(&mut self, bytes: [u8; M]);
242
243    /// The byte at `AT`.
244    fn u8_at<const AT: usize>(&self) -> u8 {
245        let [byte] = self.array::<AT, 1>();
246        byte
247    }
248
249    /// The big-endian `u16` at `AT`.
250    fn u16_at<const AT: usize>(&self) -> u16 {
251        u16::from_be_bytes(self.array::<AT, 2>())
252    }
253
254    /// The big-endian `u32` at `AT`.
255    fn u32_at<const AT: usize>(&self) -> u32 {
256        u32::from_be_bytes(self.array::<AT, 4>())
257    }
258
259    /// The big-endian `u64` at `AT`.
260    fn u64_at<const AT: usize>(&self) -> u64 {
261        u64::from_be_bytes(self.array::<AT, 8>())
262    }
263
264    /// Write the byte at `AT`.
265    fn put_u8<const AT: usize>(&mut self, value: u8) {
266        self.put::<AT, 1>([value]);
267    }
268
269    /// Write a big-endian `u16` at `AT`.
270    fn put_u16<const AT: usize>(&mut self, value: u16) {
271        self.put::<AT, 2>(value.to_be_bytes());
272    }
273
274    /// Write a big-endian `u32` at `AT`.
275    fn put_u32<const AT: usize>(&mut self, value: u32) {
276        self.put::<AT, 4>(value.to_be_bytes());
277    }
278
279    /// Write a big-endian `u64` at `AT`.
280    fn put_u64<const AT: usize>(&mut self, value: u64) {
281        self.put::<AT, 8>(value.to_be_bytes());
282    }
283}
284
285// Each body asserts its bounds at build time; the `None` arms are dead.
286impl<const N: usize> Fixed for [u8; N] {
287    fn array<const AT: usize, const M: usize>(&self) -> [u8; M] {
288        const { assert!(AT + M <= N, "field exceeds the record") };
289        self.get(AT..)
290            .and_then(<[u8]>::first_chunk)
291            .copied()
292            .unwrap_or([0; M])
293    }
294
295    fn slice<const AT: usize, const M: usize>(&self) -> &[u8] {
296        const { assert!(AT + M <= N, "field exceeds the record") };
297        self.get(AT..)
298            .and_then(|tail| tail.get(..M))
299            .unwrap_or_default()
300    }
301
302    fn put<const AT: usize, const M: usize>(&mut self, bytes: [u8; M]) {
303        const { assert!(AT + M <= N, "field exceeds the record") };
304        if let Some(field) = self.get_mut(AT..).and_then(<[u8]>::first_chunk_mut) {
305            *field = bytes;
306        }
307    }
308}
309
310#[cfg(test)]
311mod tests {
312    use super::*;
313
314    #[test]
315    fn reader_takes_fields_in_order_and_stops_at_the_end() {
316        let input = [1, 0, 2, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 4, 5, 6];
317        let short = |needed| LayoutError::Truncated {
318            needed,
319            available: 17,
320        };
321        let mut reader = Reader::new(&input);
322        assert_eq!(reader.u8(), Ok(1));
323        assert_eq!(reader.u16(), Ok(2));
324        assert_eq!(reader.u32(), Ok(3));
325        assert_eq!(reader.u64(), Ok(4));
326        assert_eq!(reader.offset(), 15);
327        assert_eq!(reader.finish(), Err(LayoutError::Trailing(2)));
328        assert_eq!(reader.array::<3>(), Err(short(18)));
329        assert_eq!(reader.rest(), &[5, 6]);
330        assert_eq!(reader.take(2), Ok(&[5, 6][..]));
331        assert_eq!(reader.u8(), Err(short(18)));
332        assert_eq!(reader.take(usize::MAX), Err(short(usize::MAX)));
333        assert_eq!(reader.finish(), Ok(()));
334    }
335
336    #[test]
337    fn writer_fills_its_buffer_and_refuses_overflow() {
338        let mut output = [0; 7];
339        let mut writer = Writer::new(&mut output);
340        assert_eq!(writer.u8(1), Ok(()));
341        assert_eq!(writer.u16(2), Ok(()));
342        assert_eq!(
343            writer.u64(3),
344            Err(LayoutError::Overflow {
345                needed: 11,
346                capacity: 7
347            })
348        );
349        assert_eq!(writer.u32(4), Ok(()));
350        assert_eq!(writer.offset(), 7);
351        assert!(writer.u8(5).is_err());
352        assert_eq!(output, [1, 0, 2, 0, 0, 0, 4]);
353        assert!(Writer::with_len(&mut output, 8).is_err());
354        let mut writer = Writer::with_len(&mut output, 1).unwrap();
355        assert!(writer.u16(9).is_err());
356    }
357
358    #[test]
359    fn fixed_fields_round_trip_at_their_offsets() {
360        let mut record = [0; 16];
361        record.put_u8::<0>(1);
362        record.put_u16::<1>(2);
363        record.put_u32::<3>(3);
364        record.put_u64::<7>(4);
365        record.put::<15, 1>([9]);
366        assert_eq!(record, [1, 0, 2, 0, 0, 0, 3, 0, 0, 0, 0, 0, 0, 0, 4, 9]);
367        assert_eq!(record.u8_at::<0>(), 1);
368        assert_eq!(record.u16_at::<1>(), 2);
369        assert_eq!(record.u32_at::<3>(), 3);
370        assert_eq!(record.u64_at::<7>(), 4);
371        assert_eq!(record.array::<14, 2>(), [4, 9]);
372        assert_eq!(record.slice::<13, 3>(), &[0, 4, 9]);
373    }
374
375    #[test]
376    fn records_are_bounded_by_their_input() {
377        let mut bytes = [1, 2, 3, 4, 5];
378        assert_eq!(record::<2>(&bytes, 3), Ok(&[4, 5]));
379        assert_eq!(
380            record::<2>(&bytes, 4),
381            Err(LayoutError::Truncated {
382                needed: 6,
383                available: 5
384            })
385        );
386        assert!(record::<1>(&bytes, usize::MAX).is_err());
387        record_mut::<2>(&mut bytes, 0).unwrap().put_u16::<0>(0x0909);
388        assert_eq!(bytes, [9, 9, 3, 4, 5]);
389        assert!(record_mut::<6>(&mut bytes, 0).is_err());
390    }
391
392    #[test]
393    fn patch_fills_a_placeholder_inside_its_output() {
394        let mut output = [0; 6];
395        assert_eq!(patch(&mut output, 2, [1, 2, 3, 4]), Ok(()));
396        assert_eq!(output, [0, 0, 1, 2, 3, 4]);
397        assert_eq!(
398            patch(&mut output, 3, [5, 6, 7, 8]),
399            Err(LayoutError::Overflow {
400                needed: 7,
401                capacity: 6
402            })
403        );
404        assert!(patch(&mut output, usize::MAX, [9]).is_err());
405        assert_eq!(output, [0, 0, 1, 2, 3, 4]);
406    }
407}