block_buffer/
read.rs

1use super::{Array, ArraySize, Error};
2
3use core::{fmt, slice};
4#[cfg(feature = "zeroize")]
5use zeroize::Zeroize;
6
7/// Buffer for reading block-generated data.
8pub struct ReadBuffer<BS: ArraySize> {
9    // The first byte of the block is used as position.
10    buffer: Array<u8, BS>,
11}
12
13impl<BS: ArraySize> fmt::Debug for ReadBuffer<BS> {
14    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
15        f.debug_struct("ReadBuffer")
16            .field("remaining_data", &self.get_pos())
17            .finish()
18    }
19}
20
21impl<BS: ArraySize> Default for ReadBuffer<BS> {
22    #[inline]
23    fn default() -> Self {
24        let mut buffer = Array::<u8, BS>::default();
25        buffer[0] = BS::U8;
26        Self { buffer }
27    }
28}
29
30impl<BS: ArraySize> Clone for ReadBuffer<BS> {
31    #[inline]
32    fn clone(&self) -> Self {
33        Self {
34            buffer: self.buffer.clone(),
35        }
36    }
37}
38
39impl<BS: ArraySize> ReadBuffer<BS> {
40    /// Return current cursor position.
41    #[inline(always)]
42    pub fn get_pos(&self) -> usize {
43        let pos = self.buffer[0];
44        if pos == 0 || pos > BS::U8 {
45            debug_assert!(false);
46            // SAFETY: `pos` never breaks the invariant
47            unsafe {
48                core::hint::unreachable_unchecked();
49            }
50        }
51        pos as usize
52    }
53
54    /// Return size of the internal buffer in bytes.
55    #[inline(always)]
56    pub fn size(&self) -> usize {
57        BS::USIZE
58    }
59
60    /// Return number of remaining bytes in the internal buffer.
61    #[inline(always)]
62    pub fn remaining(&self) -> usize {
63        self.size() - self.get_pos()
64    }
65
66    #[inline(always)]
67    fn set_pos_unchecked(&mut self, pos: usize) {
68        debug_assert!(pos <= BS::USIZE);
69        self.buffer[0] = pos as u8;
70    }
71
72    /// Write remaining data inside buffer into `data`, fill remaining space
73    /// in `data` with blocks generated by `gen_block`, and save leftover data
74    /// from the last generated block into buffer for future use.
75    #[inline]
76    pub fn read(&mut self, mut data: &mut [u8], mut gen_block: impl FnMut(&mut Array<u8, BS>)) {
77        let pos = self.get_pos();
78        let r = self.remaining();
79        let n = data.len();
80
81        if r != 0 {
82            if n < r {
83                // double slicing allows to remove panic branches
84                data.copy_from_slice(&self.buffer[pos..][..n]);
85                self.set_pos_unchecked(pos + n);
86                return;
87            }
88            let (left, right) = data.split_at_mut(r);
89            data = right;
90            left.copy_from_slice(&self.buffer[pos..]);
91        }
92
93        let (blocks, leftover) = Self::to_blocks_mut(data);
94        for block in blocks {
95            gen_block(block);
96        }
97
98        let n = leftover.len();
99        if n != 0 {
100            let mut block = Default::default();
101            gen_block(&mut block);
102            leftover.copy_from_slice(&block[..n]);
103            self.buffer = block;
104            self.set_pos_unchecked(n);
105        } else {
106            self.set_pos_unchecked(BS::USIZE);
107        }
108    }
109
110    /// Serialize buffer into a byte array.
111    #[inline]
112    pub fn serialize(&self) -> Array<u8, BS> {
113        let mut res = self.buffer.clone();
114        let pos = self.get_pos();
115        // zeroize "garbage" data
116        for b in res[1..pos].iter_mut() {
117            *b = 0;
118        }
119        res
120    }
121
122    /// Deserialize buffer from a byte array.
123    #[inline]
124    pub fn deserialize(buffer: &Array<u8, BS>) -> Result<Self, Error> {
125        let pos = buffer[0];
126        if pos == 0 || pos > BS::U8 || buffer[1..pos as usize].iter().any(|&b| b != 0) {
127            Err(Error)
128        } else {
129            Ok(Self {
130                buffer: buffer.clone(),
131            })
132        }
133    }
134
135    /// Split message into mutable slice of parallel blocks, blocks, and leftover bytes.
136    #[inline(always)]
137    fn to_blocks_mut(data: &mut [u8]) -> (&mut [Array<u8, BS>], &mut [u8]) {
138        let nb = data.len() / BS::USIZE;
139        let (left, right) = data.split_at_mut(nb * BS::USIZE);
140        let p = left.as_mut_ptr() as *mut Array<u8, BS>;
141        // SAFETY: we guarantee that `blocks` does not point outside of `data`, and `p` is valid for
142        // mutation
143        let blocks = unsafe { slice::from_raw_parts_mut(p, nb) };
144        (blocks, right)
145    }
146}
147
148#[cfg(feature = "zeroize")]
149impl<BS: ArraySize> Zeroize for ReadBuffer<BS> {
150    #[inline]
151    fn zeroize(&mut self) {
152        self.buffer.zeroize();
153    }
154}