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

            
10
use thiserror::Error;
11

            
12
use tor_bytes::Writer;
13

            
14
/// An error that occurred while trying to unwrap a SliceWriter.
15
#[non_exhaustive]
16
#[derive(Clone, Debug, Error)]
17
pub(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.
32
pub(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

            
39
impl<T> Writer for SliceWriter<T>
40
where
41
    T: AsMut<[u8]>,
42
{
43
33775
    fn write_all(&mut self, b: &[u8]) {
44
33775
        let new_offset = self.offset.saturating_add(b.len());
45
33775
        if new_offset <= self.data.as_mut().len() {
46
33773
            // Note that if we reach this case, the addition was not saturating.
47
33773
            self.data.as_mut()[self.offset..new_offset].copy_from_slice(b);
48
33773
            self.offset = new_offset;
49
33773
        } else {
50
2
            self.offset = usize::MAX;
51
2
        }
52
33775
    }
53
}
54
impl<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
5404
    pub(crate) fn new(data: T) -> Self {
63
5404
        Self { data, offset: 0 }
64
5404
    }
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
5404
    pub(crate) fn try_unwrap(self) -> Result<(T, usize), SliceWriterError> {
74
5404
        let offset = self.offset()?;
75
5402
        Ok((self.data, offset))
76
5404
    }
77

            
78
    /// Return the number of bytes written into this `SliceWriter` so far,
79
    /// or an error if it has overflowed.
80
21598
    pub(crate) fn offset(&self) -> Result<usize, SliceWriterError> {
81
21598
        if self.offset != usize::MAX {
82
21596
            Ok(self.offset)
83
        } else {
84
2
            Err(SliceWriterError::Truncated)
85
        }
86
21598
    }
87
}
88

            
89
#[cfg(test)]
90
mod 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::print_stderr)]
96
    #![allow(clippy::print_stdout)]
97
    #![allow(clippy::single_char_pattern)]
98
    #![allow(clippy::unwrap_used)]
99
    #![allow(clippy::unchecked_duration_subtraction)]
100
    #![allow(clippy::useless_vec)]
101
    #![allow(clippy::needless_pass_by_value)]
102
    //! <!-- @@ end test lint list maintained by maint/add_warning @@ -->
103
    use super::*;
104

            
105
    #[test]
106
    fn basics() {
107
        let mut w = SliceWriter::new([0_u8; 16]);
108
        w.write_u8(b'h');
109
        w.write_u16(0x656c);
110
        w.write_u32(0x6c6f2077);
111
        w.write_all(b"orld!");
112
        let (a, len) = w.try_unwrap().unwrap();
113

            
114
        assert_eq!(a.as_ref(), b"hello world!\0\0\0\0");
115
        assert_eq!(len, 12);
116
    }
117

            
118
    #[test]
119
    fn full_is_ok() {
120
        let mut w = SliceWriter::new([0_u8; 4]);
121
        w.write_u8(1);
122
        w.write_u16(0x0203);
123
        w.write_u8(4);
124
        let (a, len) = w.try_unwrap().unwrap();
125

            
126
        assert_eq!(a.as_ref(), [1, 2, 3, 4]);
127
        assert_eq!(len, 4);
128
    }
129

            
130
    #[test]
131
    fn too_full_is_not_ok() {
132
        let mut w = SliceWriter::new([0_u8; 5]);
133
        w.write_u32(12);
134
        w.write_u32(12);
135
        assert!(matches!(w.try_unwrap(), Err(SliceWriterError::Truncated)));
136
    }
137
}