tor_cell/
slicewriter.rs

Help
1//! A Writer that can put its results into an buffer of known byte size without
2//! changing that size.
3//!
4//! TODO: This API is crate-private in tor-cell because tor-cell is the only one
5//! to use it -- and only uses it in one place.  Its existence may be an argument
6//! for refactoring the Writer API entirely.
7//!
8//! NOTE: This will likely change or go away in the future.
9
10use thiserror::Error;
11
12use tor_bytes::Writer;
13
14/// An error that occurred while trying to unwrap a SliceWriter.
15#[non_exhaustive]
16#[derive(Clone, Debug, Error)]
17pub(crate) enum SliceWriterError {
18    /// We've tried to write more than would fit into a fixed-size slice
19    #[error("Tried to write more than would fit into a fixed-size slice.")]
20    Truncated,
21}
22
23/// An object that supports writing into a byte-slice of fixed size.
24///
25/// Since the writer API does not allow all `write_*` functions to report errors,
26/// this type defers any truncated-data errors until you try to retrieve the
27/// inner data.
28///
29//
30// TODO: in theory we could have a version of this that used MaybeUninit, but I
31// don't think that would be worth it.
32pub(crate) struct SliceWriter<T> {
33    /// The object we're writing into.  Must have fewer than usize::LEN bytes.
34    data: T,
35    /// Our current write position within that object.
36    offset: usize,
37}
38
39impl<T> Writer for SliceWriter<T>
40where
41    T: AsMut<[u8]>,
42{
43    fn write_all(&mut self, b: &[u8]) {
44        let new_offset = self.offset.saturating_add(b.len());
45        if new_offset <= self.data.as_mut().len() {
46            // Note that if we reach this case, the addition was not saturating.
47            self.data.as_mut()[self.offset..new_offset].copy_from_slice(b);
48            self.offset = new_offset;
49        } else {
50            self.offset = usize::MAX;
51        }
52    }
53}
54impl<T> SliceWriter<T> {
55    /// Construct a new SliceWriter
56    ///
57    /// Typically, you would want to use this on a type that implements
58    /// `AsMut<[u8]>`, or else it won't be very useful.
59    ///
60    /// Preexisting bytes in the `data` object will be unchanged, unless you use
61    /// the [`Writer`] API to write to them.
62    pub(crate) fn new(data: T) -> Self {
63        Self { data, offset: 0 }
64    }
65
66    /// Try to extract the data from this `SliceWriter`.
67    ///
68    /// On success (if we did not write "off the end" of the underlying object),
69    /// return the object and the number of bytes we wrote into it.  (Bytes
70    /// after that position are unchanged.)
71    ///
72    /// On failure (if we tried to write too much), return an error.
73    pub(crate) fn try_unwrap(self) -> Result<(T, usize), SliceWriterError> {
74        let offset = self.offset()?;
75        Ok((self.data, offset))
76    }
77
78    /// Return the number of bytes written into this `SliceWriter` so far,
79    /// or an error if it has overflowed.
80    pub(crate) fn offset(&self) -> Result<usize, SliceWriterError> {
81        if self.offset != usize::MAX {
82            Ok(self.offset)
83        } else {
84            Err(SliceWriterError::Truncated)
85        }
86    }
87}
88
89#[cfg(test)]
90mod test {
91    // @@ begin test lint list maintained by maint/add_warning @@
92    #![allow(clippy::bool_assert_comparison)]
93    #![allow(clippy::clone_on_copy)]
94    #![allow(clippy::dbg_macro)]
95    #![allow(clippy::mixed_attributes_style)]
96    #![allow(clippy::print_stderr)]
97    #![allow(clippy::print_stdout)]
98    #![allow(clippy::single_char_pattern)]
99    #![allow(clippy::unwrap_used)]
100    #![allow(clippy::unchecked_duration_subtraction)]
101    #![allow(clippy::useless_vec)]
102    #![allow(clippy::needless_pass_by_value)]
103    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
104    use super::*;
105
106    #[test]
107    fn basics() {
108        let mut w = SliceWriter::new([0_u8; 16]);
109        w.write_u8(b'h');
110        w.write_u16(0x656c);
111        w.write_u32(0x6c6f2077);
112        w.write_all(b"orld!");
113        let (a, len) = w.try_unwrap().unwrap();
114
115        assert_eq!(a.as_ref(), b"hello world!\0\0\0\0");
116        assert_eq!(len, 12);
117    }
118
119    #[test]
120    fn full_is_ok() {
121        let mut w = SliceWriter::new([0_u8; 4]);
122        w.write_u8(1);
123        w.write_u16(0x0203);
124        w.write_u8(4);
125        let (a, len) = w.try_unwrap().unwrap();
126
127        assert_eq!(a.as_ref(), [1, 2, 3, 4]);
128        assert_eq!(len, 4);
129    }
130
131    #[test]
132    fn too_full_is_not_ok() {
133        let mut w = SliceWriter::new([0_u8; 5]);
134        w.write_u32(12);
135        w.write_u32(12);
136        assert!(matches!(w.try_unwrap(), Err(SliceWriterError::Truncated)));
137    }
138}
tor_cell/slicewriter.rs

tor_cell/
slicewriter.rs
