block-buffer: add block_mode_processing methods (#315)

Author: newpavlov

block-buffer/src/buffer.rs

@@ -3,8 +3,9 @@ use block_padding::Padding;
 
 use crate::{
     utils::{to_blocks, to_blocks_mut},
-    Block, DigestBuffer, ParBlock,
+    Block, DigestBuffer, InvalidLength, ParBlock,
 };
+use core::slice;
 use generic_array::{typenum::U1, ArrayLength};
 
 /// Buffer for block processing of data.
@@ -93,6 +94,65 @@ impl<BlockSize: ArrayLength<u8>> BlockBuffer<BlockSize> {
         self.process_data(data, &mut gen_block, set, |f| f(), unreachable);
     }
 
+    /// Process `data` in blocks and write result to `out_buf`, storing
+    /// leftovers for future use.
+    #[inline]
+    pub fn block_mode_processing<'a>(
+        &mut self,
+        mut data: &[u8],
+        buf: &'a mut [u8],
+        mut process: impl FnMut(&mut [Block<BlockSize>]),
+    ) -> Result<&'a [u8], InvalidLength> {
+        let pos = self.get_pos();
+        let rem = self.remaining();
+        let mut blocks_processed = 0;
+        let (_, mut buf_blocks, _) = to_blocks_mut::<BlockSize, U1>(buf);
+        if pos != 0 {
+            let n = data.len();
+            if n < rem {
+                // double slicing allows to remove panic branches
+                self.buffer[pos..][..n].copy_from_slice(data);
+                self.set_pos_unchecked(pos + n);
+                return Ok(&buf[..0]);
+            }
+            if buf_blocks.is_empty() {
+                return Err(InvalidLength);
+            }
+
+            let (l, r) = buf_blocks.split_at_mut(1);
+            let buf_block = &mut l[0];
+            buf_blocks = r;
+            let (l, r) = data.split_at(rem);
+            data = r;
+
+            buf_block[..pos].copy_from_slice(&self.buffer[..pos]);
+            buf_block[pos..].copy_from_slice(l);
+
+            process(slice::from_mut(buf_block));
+            blocks_processed += 1;
+        }
+
+        let (data_blocks, leftover) = to_blocks(data);
+        let buf_blocks = buf_blocks
+            .get_mut(..data_blocks.len())
+            .ok_or(InvalidLength)?;
+        buf_blocks.clone_from_slice(data_blocks);
+        process(buf_blocks);
+        blocks_processed += buf_blocks.len();
+
+        let n = leftover.len();
+        self.buffer[..n].copy_from_slice(leftover);
+        self.set_pos_unchecked(n);
+
+        let res = unsafe {
+            let res_len = BlockSize::USIZE * blocks_processed;
+            // SAFETY: number of processed blocks never exceeds capacity of `buf`
+            debug_assert!(buf.len() >= res_len);
+            buf.get_unchecked(..res_len)
+        };
+        Ok(res)
+    }
+
     /// Compress remaining data after padding it with `delim`, zeros and
     /// the `suffix` bytes. If there is not enough unused space, `compress`
     /// will be called twice.
@@ -215,7 +275,7 @@ impl<B: ArrayLength<u8>> DigestBuffer<B> for BlockBuffer<B> {
             let (left, right) = input.split_at(r);
             input = right;
             self.buffer[pos..].copy_from_slice(left);
-            compress(core::slice::from_ref(&self.buffer));
+            compress(slice::from_ref(&self.buffer));
         }
 
         let (blocks, leftover) = to_blocks(input);

block-buffer/src/lazy.rs

@@ -1,5 +1,9 @@
-use crate::{utils::to_blocks, Block, DigestBuffer};
-use generic_array::ArrayLength;
+use crate::{
+    utils::{to_blocks, to_blocks_mut},
+    Block, DigestBuffer, InvalidLength,
+};
+use core::slice;
+use generic_array::{typenum::U1, ArrayLength};
 
 /// Buffer for lazy block processing of data.
 #[derive(Clone, Default)]
@@ -9,6 +13,65 @@ pub struct LazyBlockBuffer<BlockSize: ArrayLength<u8>> {
 }
 
 impl<BlockSize: ArrayLength<u8>> LazyBlockBuffer<BlockSize> {
+    /// Process `data` in blocks and write result to `out_buf`, storing
+    /// leftovers for future use.
+    #[inline]
+    pub fn block_mode_processing<'a>(
+        &mut self,
+        mut data: &[u8],
+        buf: &'a mut [u8],
+        mut process: impl FnMut(&mut [Block<BlockSize>]),
+    ) -> Result<&'a [u8], InvalidLength> {
+        let pos = self.get_pos();
+        let rem = self.remaining();
+        let mut blocks_processed = 0;
+        let (_, mut buf_blocks, _) = to_blocks_mut::<BlockSize, U1>(buf);
+        if pos != 0 {
+            let n = data.len();
+            if n <= rem {
+                // double slicing allows to remove panic branches
+                self.buffer[pos..][..n].copy_from_slice(data);
+                self.set_pos_unchecked(pos + n);
+                return Ok(&buf[..0]);
+            }
+            if buf_blocks.is_empty() {
+                return Err(InvalidLength);
+            }
+
+            let (l, r) = buf_blocks.split_at_mut(1);
+            let buf_block = &mut l[0];
+            buf_blocks = r;
+            let (l, r) = data.split_at(rem);
+            data = r;
+
+            buf_block[..pos].copy_from_slice(&self.buffer[..pos]);
+            buf_block[pos..].copy_from_slice(l);
+
+            process(slice::from_mut(buf_block));
+            blocks_processed += 1;
+        }
+
+        let (data_blocks, leftover) = to_blocks_lazy(data);
+        let buf_blocks = buf_blocks
+            .get_mut(..data_blocks.len())
+            .ok_or(InvalidLength)?;
+        buf_blocks.clone_from_slice(data_blocks);
+        process(buf_blocks);
+        blocks_processed += buf_blocks.len();
+
+        let n = leftover.len();
+        self.buffer[..n].copy_from_slice(leftover);
+        self.set_pos_unchecked(n);
+
+        let res = unsafe {
+            let res_len = BlockSize::USIZE * blocks_processed;
+            // SAFETY: number of processed blocks never exceeds capacity of `buf`
+            debug_assert!(buf.len() >= res_len);
+            buf.get_unchecked(..res_len)
+        };
+        Ok(res)
+    }
+
     /// Pad remaining data with zeros and call `compress` with resulting block.
     pub fn pad_zeros(&mut self) -> &mut Block<BlockSize> {
         let pos = self.get_pos();
@@ -17,6 +80,15 @@ impl<BlockSize: ArrayLength<u8>> LazyBlockBuffer<BlockSize> {
         &mut self.buffer
     }
 
+    /// Return block if buffer is full, otherwise returns `None`.
+    #[inline]
+    pub fn get_full_block(&mut self) -> Option<&mut Block<BlockSize>> {
+        match self.remaining() {
+            0 => Some(&mut self.buffer),
+            _ => None,
+        }
+    }
+
     /// Return size of the internall buffer in bytes.
     #[inline]
     pub fn size(&self) -> usize {
@@ -82,17 +154,7 @@ impl<B: ArrayLength<u8>> DigestBuffer<B> for LazyBlockBuffer<B> {
             compress(core::slice::from_ref(&self.buffer));
         }
 
-        let (mut blocks, mut leftover) = to_blocks(input);
-        if leftover.is_empty() {
-            debug_assert!(!blocks.is_empty());
-            let m = blocks.len() - 1;
-            // SAFETY: at this stage `input` always contains at least one byte,
-            // so either `leftover` is not empty or we have at least one block
-            unsafe {
-                leftover = blocks.get_unchecked(m);
-                blocks = blocks.get_unchecked(..m);
-            }
-        }
+        let (blocks, leftover) = to_blocks_lazy(input);
         compress(blocks);
 
         let n = leftover.len();
@@ -105,3 +167,18 @@ impl<B: ArrayLength<u8>> DigestBuffer<B> for LazyBlockBuffer<B> {
         self.pos = 0;
     }
 }
+
+fn to_blocks_lazy<N: ArrayLength<u8>>(data: &[u8]) -> (&[Block<N>], &[u8]) {
+    let (mut blocks, mut leftover) = to_blocks(data);
+    if leftover.is_empty() {
+        debug_assert!(!blocks.is_empty());
+        let m = blocks.len() - 1;
+        // SAFETY: at this stage `input` always contains at least one byte,
+        // so either `leftover` is not empty or we have at least one block
+        unsafe {
+            leftover = blocks.get_unchecked(m);
+            blocks = blocks.get_unchecked(..m);
+        }
+    }
+    (blocks, leftover)
+}

block-buffer/src/lib.rs

@@ -11,6 +11,7 @@
 pub use block_padding;
 pub use generic_array;
 
+use core::fmt;
 use generic_array::{ArrayLength, GenericArray};
 
 mod buffer;
@@ -34,3 +35,14 @@ pub trait DigestBuffer<BlockSize: ArrayLength<u8>>: Default {
     /// Reset buffer by setting cursor position to zero.
     fn reset(&mut self);
 }
+
+/// Error type used by the [`BlockBuffer::block_mode_processing`] and
+/// [`LazyBlockBuffer::block_mode_processing`] methods.
+#[derive(Copy, Clone, Debug)]
+pub struct InvalidLength;
+
+impl fmt::Display for InvalidLength {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
+        f.write_str("Invalid Length")
+    }
+}