block_buffer/

read.rs

use super::{Array, BlockSizes, Error};
use core::fmt;

/// Buffer for reading block-generated data.
pub struct ReadBuffer<BS: BlockSizes> {
    /// The first byte of the block is used as cursor position.
    /// `&buffer[usize::from(buffer[0])..]` is interpreted as unread bytes.
    /// The cursor position is always bigger than zero and smaller than or equal to block size.
    buffer: Array<u8, BS>,
}

impl<BS: BlockSizes> fmt::Debug for ReadBuffer<BS> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ReadBuffer")
            .field("remaining_data", &self.remaining())
            .finish_non_exhaustive()
    }
}

impl<BS: BlockSizes> Default for ReadBuffer<BS> {
    #[inline]
    fn default() -> Self {
        let buffer = Default::default();
        let mut res = Self { buffer };
        // SAFETY: `BS::USIZE` satisfies the `set_pos_unchecked` safety contract
        unsafe { res.set_pos_unchecked(BS::USIZE) };
        res
    }
}

impl<BS: BlockSizes> Clone for ReadBuffer<BS> {
    #[inline]
    fn clone(&self) -> Self {
        let buffer = self.buffer.clone();
        Self { buffer }
    }
}

impl<BS: BlockSizes> ReadBuffer<BS> {
    /// Return current cursor position, i.e. how many bytes were read from the buffer.
    #[inline(always)]
    pub fn get_pos(&self) -> usize {
        let pos = self.buffer[0];
        if pos == 0 || pos > BS::U8 {
            debug_assert!(false);
            // SAFETY: `pos` never breaks the invariant
            unsafe {
                core::hint::unreachable_unchecked();
            }
        }
        pos as usize
    }

    /// Return size of the internal buffer in bytes.
    #[inline(always)]
    pub fn size(&self) -> usize {
        BS::USIZE
    }

    /// Return number of remaining bytes in the internal buffer.
    #[inline(always)]
    pub fn remaining(&self) -> usize {
        self.size() - self.get_pos()
    }

    /// Set cursor position.
    ///
    /// # Safety
    /// `pos` must be smaller than or equal to the buffer block size and be bigger than zero.
    #[inline(always)]
    #[allow(clippy::cast_possible_truncation)]
    unsafe fn set_pos_unchecked(&mut self, pos: usize) {
        debug_assert!(pos != 0 && pos <= BS::USIZE);
        self.buffer[0] = pos as u8;
    }

    /// Read up to `len` bytes of remaining data in the buffer.
    ///
    /// Returns slice with length of `ret_len = min(len, buffer.remaining())` bytes
    /// and sets the cursor position to `buffer.get_pos() + ret_len`.
    #[inline(always)]
    pub fn read_cached(&mut self, len: usize) -> &[u8] {
        let rem = self.remaining();
        let new_len = core::cmp::min(rem, len);
        let pos = self.get_pos();

        // SAFETY: `pos + new_len` is not equal to zero and not bigger than block size
        unsafe { self.set_pos_unchecked(pos + new_len) };
        &self.buffer[pos..][..new_len]
    }

    /// Write new block and consume `read_len` bytes from it.
    ///
    /// If `read_len` is equal to zero, immediately returns without calling the closures.
    /// Otherwise, the method calls `gen_block` to fill the internal buffer,
    /// passes to `read_fn` slice with first `read_len` bytes of the block,
    /// and sets the cursor position to `read_len`.
    ///
    /// # Panics
    /// If `read_len` is bigger than block size.
    #[inline(always)]
    pub fn write_block(
        &mut self,
        read_len: usize,
        gen_block: impl FnOnce(&mut Array<u8, BS>),
        read_fn: impl FnOnce(&[u8]),
    ) {
        if read_len == 0 {
            return;
        }
        assert!(read_len < BS::USIZE);

        let g = ResetGuard(self);
        let buf = &mut g.0.buffer;

        // Note that generated block is likely to break the `ReadBuffer` invariant.
        // We restore it using `set_pos_unchecked` below and in case if one of the closures
        // panic the buffer gets reset by the guard.
        gen_block(buf);
        read_fn(&buf[..read_len]);

        core::mem::forget(g);

        // We checked that `read_len` satisfies the `set_pos_unchecked` safety contract
        unsafe { self.set_pos_unchecked(read_len) };
    }

    /// Reset buffer into exhausted state.
    pub fn reset(&mut self) {
        self.buffer[0] = BS::U8;
    }

    /// Write remaining data inside buffer into `buf`, fill remaining space
    /// in `buf` with blocks generated by `gen_block`, and save leftover data
    /// from the last generated block into the buffer for future use.
    #[inline]
    pub fn read(&mut self, buf: &mut [u8], mut gen_block: impl FnMut(&mut Array<u8, BS>)) {
        let head_ks = self.read_cached(buf.len());
        let (head, buf) = buf.split_at_mut(head_ks.len());
        let (blocks, tail) = Array::slice_as_chunks_mut(buf);

        head.copy_from_slice(head_ks);
        for block in blocks {
            gen_block(block);
        }

        self.write_block(tail.len(), gen_block, |tail_ks| {
            tail.copy_from_slice(tail_ks);
        });
    }

    /// Serialize buffer into a byte array.
    #[inline]
    pub fn serialize(&self) -> Array<u8, BS> {
        let pos = self.get_pos();
        let mut res = self.buffer.clone();
        // zeroize "garbage" data
        for b in &mut res[1..pos] {
            *b = 0;
        }
        res
    }

    /// Deserialize buffer from a byte array.
    ///
    /// # Errors
    /// - If the first byte is `0`.
    /// - If the first byte is bigger than `BS`.
    #[inline]
    pub fn deserialize(buffer: &Array<u8, BS>) -> Result<Self, Error> {
        let pos = usize::from(buffer[0]);
        if pos == 0 || pos > BS::USIZE || buffer[1..pos].iter().any(|&b| b != 0) {
            Err(Error)
        } else {
            let buffer = buffer.clone();
            Ok(Self { buffer })
        }
    }
}

#[cfg(feature = "zeroize")]
impl<BS: BlockSizes> Drop for ReadBuffer<BS> {
    fn drop(&mut self) {
        use zeroize::Zeroize;
        self.buffer.zeroize();
    }
}

#[cfg(feature = "zeroize")]
impl<BS: BlockSizes> zeroize::ZeroizeOnDrop for ReadBuffer<BS> {}

/// Resets the referenced buffer on drop.
struct ResetGuard<'a, BS: BlockSizes>(&'a mut ReadBuffer<BS>);

impl<BS: BlockSizes> Drop for ResetGuard<'_, BS> {
    fn drop(&mut self) {
        self.0.reset();
    }
}