askar_crypto/buffer/
writer.rs

1#[cfg(feature = "alloc")]
2use alloc::vec::Vec;
3use core::{fmt, ops::Range};
4
5use super::{ResizeBuffer, WriteBuffer};
6use crate::error::Error;
7
8/// A structure wrapping a mutable pointer to a buffer
9#[derive(Debug)]
10pub struct Writer<'w, B: ?Sized> {
11    inner: &'w mut B,
12    pos: usize,
13}
14
15impl<B: ?Sized> Writer<'_, B> {
16    /// Accessor for the writer position
17    pub fn position(&self) -> usize {
18        self.pos
19    }
20}
21
22impl<'w> Writer<'w, [u8]> {
23    /// Create a new writer from a mutable byte slice
24    #[inline]
25    pub fn from_slice(slice: &'w mut [u8]) -> Self {
26        Writer {
27            inner: slice,
28            pos: 0,
29        }
30    }
31
32    /// Create a new writer from a mutable byte slice, skipping a prefix
33    #[inline]
34    pub fn from_slice_position(slice: &'w mut [u8], pos: usize) -> Self {
35        Writer { inner: slice, pos }
36    }
37}
38
39impl Writer<'_, [u8]> {
40    pub(crate) fn splice(
41        &mut self,
42        range: Range<usize>,
43        mut iter: impl ExactSizeIterator<Item = u8>,
44    ) -> Result<(), Error> {
45        assert!(range.end >= range.start);
46        let rem_len = range.len();
47        let ins_len = iter.len();
48        match ins_len {
49            _ if ins_len > rem_len => {
50                let diff = ins_len - rem_len;
51                if self.pos + diff > self.inner.len() {
52                    return Err(err_msg!(ExceededBuffer));
53                }
54                self.inner
55                    .copy_within((range.end - diff)..self.pos, range.end);
56                self.pos += diff;
57            }
58            _ if ins_len < rem_len => {
59                let diff = rem_len - ins_len;
60                self.inner
61                    .copy_within(range.end..self.pos, range.end - diff);
62                self.pos -= diff;
63            }
64            _ => {}
65        }
66        for idx in 0..ins_len {
67            self.inner[range.start + idx] = iter.next().unwrap();
68        }
69        Ok(())
70    }
71}
72
73impl AsRef<[u8]> for Writer<'_, [u8]> {
74    fn as_ref(&self) -> &[u8] {
75        &self.inner[..self.pos]
76    }
77}
78
79impl AsMut<[u8]> for Writer<'_, [u8]> {
80    fn as_mut(&mut self) -> &mut [u8] {
81        &mut self.inner[..self.pos]
82    }
83}
84
85impl WriteBuffer for Writer<'_, [u8]> {
86    fn buffer_write(&mut self, data: &[u8]) -> Result<(), Error> {
87        let total = self.inner.len();
88        let end = self.pos + data.len();
89        if end > total {
90            return Err(err_msg!(ExceededBuffer));
91        }
92        self.inner[self.pos..end].copy_from_slice(data);
93        self.pos += data.len();
94        Ok(())
95    }
96}
97
98impl ResizeBuffer for Writer<'_, [u8]> {
99    fn buffer_insert(&mut self, pos: usize, data: &[u8]) -> Result<(), Error> {
100        self.splice(pos..pos, data.iter().cloned())
101    }
102
103    fn buffer_remove(&mut self, range: Range<usize>) -> Result<(), Error> {
104        assert!(range.end >= range.start);
105        let diff = range.end - range.start;
106        self.inner.copy_within(range.end..self.pos, range.start);
107        self.pos -= diff;
108        Ok(())
109    }
110
111    fn buffer_resize(&mut self, len: usize) -> Result<(), Error> {
112        let len = self.pos + len;
113        if len > self.inner.len() {
114            return Err(err_msg!(ExceededBuffer));
115        }
116        self.pos = len;
117        Ok(())
118    }
119}
120
121#[cfg(feature = "alloc")]
122#[cfg_attr(docsrs, doc(cfg(feature = "alloc")))]
123impl<'w> Writer<'w, Vec<u8>> {
124    /// Create a new writer from a mutable Vec<u8> pointer
125    #[inline]
126    pub fn from_vec(vec: &'w mut Vec<u8>) -> Self {
127        Writer { inner: vec, pos: 0 }
128    }
129
130    /// Create a new writer from a mutable Vec<u8> pointer, skipping a prefix
131    #[inline]
132    pub fn from_vec_skip(vec: &'w mut Vec<u8>, pos: usize) -> Self {
133        Writer { inner: vec, pos }
134    }
135}
136
137impl<B: WriteBuffer + ?Sized> WriteBuffer for Writer<'_, B> {
138    fn buffer_write(&mut self, data: &[u8]) -> Result<(), Error> {
139        self.inner.buffer_write(data)
140    }
141}
142
143impl<B: ResizeBuffer + ?Sized> AsRef<[u8]> for Writer<'_, B> {
144    fn as_ref(&self) -> &[u8] {
145        &self.inner.as_ref()[self.pos..]
146    }
147}
148
149impl<B: ResizeBuffer + ?Sized> AsMut<[u8]> for Writer<'_, B> {
150    fn as_mut(&mut self) -> &mut [u8] {
151        &mut self.inner.as_mut()[self.pos..]
152    }
153}
154
155impl<B: ResizeBuffer + ?Sized> ResizeBuffer for Writer<'_, B> {
156    fn buffer_insert(&mut self, pos: usize, data: &[u8]) -> Result<(), Error> {
157        self.inner.buffer_insert(self.pos + pos, data)
158    }
159
160    fn buffer_remove(&mut self, range: Range<usize>) -> Result<(), Error> {
161        self.inner
162            .buffer_remove((self.pos + range.start)..(self.pos + range.end))
163    }
164
165    fn buffer_resize(&mut self, len: usize) -> Result<(), Error> {
166        self.inner.buffer_resize(self.pos + len)
167    }
168}
169
170impl<'b, B: ?Sized> Writer<'b, B> {
171    /// Create a new writer from a reference to a buffer implementation
172    pub fn from_buffer(buf: &'b mut B) -> Writer<'b, B> {
173        Writer { inner: buf, pos: 0 }
174    }
175}
176
177impl<'b, B: ?Sized> fmt::Write for Writer<'b, B>
178where
179    Writer<'b, B>: WriteBuffer,
180{
181    fn write_str(&mut self, s: &str) -> fmt::Result {
182        self.buffer_write(s.as_bytes()).map_err(|_| fmt::Error)
183    }
184}
185
186#[cfg(test)]
187mod tests {
188    use super::*;
189
190    #[test]
191    fn write_buffer_slice() {
192        let mut buf = [0u8; 10];
193        let mut w = Writer::from_slice(&mut buf);
194        w.buffer_write(b"he").unwrap();
195        w.buffer_write(b"y").unwrap();
196        assert_eq!(w.position(), 3);
197        assert_eq!(&buf[..3], b"hey");
198    }
199}