digest: Clarify panic safety of (Block)Digest::finish

src/digest.rs

@@ -24,7 +24,7 @@
 };
 use crate::{
     bits::{BitLength, FromByteLen as _},
-    cpu, debug,
+    cpu, debug, error,
     polyfill::{self, slice, sliceutil},
 };
 use core::num::Wrapping;
@@ -69,47 +69,66 @@
         leftover
     }
 
-    pub(crate) fn finish(
+    // On input, `block[..num_pending]` is the (possibly-empty) last *partial*
+    // chunk of input. It *must* be partial; that is, it is required that
+    // `num_pending < self.algorithm.block_len`.
+    //
+    // `block` may be arbitrarily overwritten.
+    pub(crate) fn try_finish(
         mut self,
-        pending: &mut [u8],
+        block: &mut [u8; MAX_BLOCK_LEN],
         num_pending: usize,
         cpu_features: cpu::Features,
-    ) -> Digest {
-        let block_len = self.algorithm.block_len();
-        assert_eq!(pending.len(), block_len);
-        assert!(num_pending < pending.len());
-        let pending = &mut pending[..block_len];
-
-        let mut padding_pos = num_pending;
-        pending[padding_pos] = 0x80;
-        padding_pos += 1;
-
-        if padding_pos > pending.len() - self.algorithm.len_len {
-            pending[padding_pos..].fill(0);
-            let (completed_bytes, leftover) = self.block_data_order(pending, cpu_features);
-            debug_assert_eq!((completed_bytes, leftover.len()), (block_len, 0));
-            // We don't increase |self.completed_bytes| because the padding
-            // isn't data, and so it isn't included in the data length.
-            padding_pos = 0;
-        }
-
-        pending[padding_pos..(block_len - 8)].fill(0);
-
-        // Output the length, in bits, in big endian order.
+    ) -> Result<Digest, FinishError> {
         let completed_bytes = self
             .completed_bytes
             .checked_add(polyfill::u64_from_usize(num_pending))
-            .unwrap();
-        let copmleted_bits = BitLength::from_byte_len(completed_bytes).unwrap();
-        pending[(block_len - 8)..].copy_from_slice(&copmleted_bits.to_be_bytes());
+            .ok_or_else(|| FinishError::too_much_input(self.completed_bytes))?;
+        let copmleted_bits = BitLength::from_byte_len(completed_bytes)
+            .map_err(|_: error::Unspecified| FinishError::too_much_input(self.completed_bytes))?;
+
+        let block_len = self.algorithm.block_len();
+        let block = &mut block[..block_len];
+
+        let padding = match block.get_mut(num_pending..) {
+            Some([separator, padding @ ..]) => {
+                *separator = 0x80;
+                padding
+            }
+            // Precondition violated.
+            unreachable => {
+                return Err(FinishError::pending_not_a_partial_block(
+                    unreachable.as_deref(),
+                ));
+            }
+        };
 
-        let (completed_bytes, leftover) = self.block_data_order(pending, cpu_features);
+        let padding = match padding
+            .len()
+            .checked_sub(self.algorithm.block_len.len_len())
+        {
+            Some(_) => padding,
+            None => {
+                padding.fill(0);
+                let (completed_bytes, leftover) = self.block_data_order(block, cpu_features);
+                debug_assert_eq!((completed_bytes, leftover.len()), (block_len, 0));
+                // We don't increase |self.completed_bytes| because the padding
+                // isn't data, and so it isn't included in the data length.
+                &mut block[..]
+            }
+        };
+
+        let (to_zero, len) = padding.split_at_mut(padding.len() - 8);
+        to_zero.fill(0);
+        len.copy_from_slice(&copmleted_bits.to_be_bytes());
+
+        let (completed_bytes, leftover) = self.block_data_order(block, cpu_features);
         debug_assert_eq!((completed_bytes, leftover.len()), (block_len, 0));
 
-        Digest {
+        Ok(Digest {
             algorithm: self.algorithm,
             value: (self.algorithm.format_output)(self.state),
-        }
+        })
     }
 
     #[must_use]
@@ -122,6 +141,37 @@
     }
 }
 
+pub(crate) enum FinishError {
+    #[allow(dead_code)]
+    TooMuchInput(u64),
+    #[allow(dead_code)]
+    PendingNotAPartialBlock(usize),
+}
+
+impl FinishError {
+    #[cold]
+    #[inline(never)]
+    fn too_much_input(completed_bytes: u64) -> Self {
+        Self::TooMuchInput(completed_bytes)
+    }
+
+    // unreachable
+    #[cold]
+    #[inline(never)]
+    fn pending_not_a_partial_block(padding: Option<&[u8]>) -> Self {
+        match padding {
+            None => Self::PendingNotAPartialBlock(0),
+            Some(padding) => Self::PendingNotAPartialBlock(padding.len()),
+        }
+    }
+}
+
+impl From<FinishError> for error::Unspecified {
+    fn from(_: FinishError) -> Self {
+        Self
+    }
+}
+
 /// A context for multi-step (Init-Update-Finish) digest calculations.
 ///
 /// # Examples
@@ -213,15 +263,15 @@
     ///
     /// `finish` consumes the context so it cannot be (mis-)used after `finish`
     /// has been called.
-    pub fn finish(mut self) -> Digest {
-        let cpu_features = cpu::features();
+    pub fn finish(self) -> Digest {
+        self.try_finish().unwrap()
+    }
 
-        let block_len = self.block.algorithm.block_len();
-        self.block.finish(
-            &mut self.pending[..block_len],
-            self.num_pending,
-            cpu_features,
-        )
+    fn try_finish(mut self) -> Result<Digest, error::Unspecified> {
+        let cpu_features = cpu::features();
+        self.block
+            .try_finish(&mut self.pending, self.num_pending, cpu_features)
+            .map_err(error::Unspecified::from)
     }
 
     /// The algorithm that this context is using.
@@ -289,9 +339,6 @@
     chaining_len: usize,
     block_len: BlockLen,
 
-    /// The length of the length in the padding.
-    len_len: usize,
-
     /// `block_data_order` processes all the full blocks of data in `data`. It
     /// returns the number of bytes processed and the unprocessed data, which
     /// is guaranteed to be less than `block_len` bytes long.
@@ -357,7 +404,6 @@
     output_len: sha1::OUTPUT_LEN,
     chaining_len: sha1::CHAINING_LEN,
     block_len: sha1::BLOCK_LEN,
-    len_len: 64 / 8,
     block_data_order: dynstate::sha1_block_data_order,
     format_output: dynstate::sha256_format_output,
     initial_state: DynState::new32([
@@ -380,7 +426,6 @@
     output_len: OutputLen::_256,
     chaining_len: SHA256_OUTPUT_LEN,
     block_len: SHA256_BLOCK_LEN,
-    len_len: 64 / 8,
     block_data_order: dynstate::sha256_block_data_order,
     format_output: dynstate::sha256_format_output,
     initial_state: DynState::new32([
@@ -403,7 +448,6 @@
     output_len: OutputLen::_384,
     chaining_len: SHA512_OUTPUT_LEN,
     block_len: SHA512_BLOCK_LEN,
-    len_len: SHA512_LEN_LEN,
     block_data_order: dynstate::sha512_block_data_order,
     format_output: dynstate::sha512_format_output,
     initial_state: DynState::new64([
@@ -426,7 +470,6 @@
     output_len: OutputLen::_512,
     chaining_len: SHA512_OUTPUT_LEN,
     block_len: SHA512_BLOCK_LEN,
-    len_len: SHA512_LEN_LEN,
     block_data_order: dynstate::sha512_block_data_order,
     format_output: dynstate::sha512_format_output,
     initial_state: DynState::new64([
@@ -453,7 +496,6 @@
     output_len: OutputLen::_256,
     chaining_len: SHA512_OUTPUT_LEN,
     block_len: SHA512_BLOCK_LEN,
-    len_len: SHA512_LEN_LEN,
     block_data_order: dynstate::sha512_block_data_order,
     format_output: dynstate::sha512_format_output,
     initial_state: DynState::new64([
@@ -516,9 +558,6 @@
 /// The length of the output of SHA-512/256, in bytes.
 pub const SHA512_256_OUTPUT_LEN: usize = OutputLen::_256.into();
 
-/// The length of the length field for SHA-512-based algorithms, in bytes.
-const SHA512_LEN_LEN: usize = 128 / 8;
-
 #[derive(Clone, Copy)]
 enum BlockLen {
     _512 = 512 / 8,
@@ -531,6 +570,21 @@
     const fn into(self) -> usize {
         self as usize
     }
+
+    #[inline(always)]
+    const fn len_len(self) -> usize {
+        let len_len = match self {
+            BlockLen::_512 => LenLen::_64,
+            BlockLen::_1024 => LenLen::_128,
+        };
+        len_len as usize
+    }
+}
+
+#[derive(Clone, Copy)]
+enum LenLen {
+    _64 = 64 / 8,
+    _128 = 128 / 8,
 }
 
 #[derive(Clone, Copy)]

src/hmac.rs

@@ -312,14 +312,19 @@ impl Context {
     /// the return value of `sign` to a tag. Use `verify` for verification
     /// instead.
     pub fn sign(self) -> Tag {
-        let cpu_features = cpu::features();
+        self.try_sign().unwrap()
+    }
 
+    fn try_sign(self) -> Result<Tag, error::Unspecified> {
+        let cpu_features = cpu::features();
         let algorithm = self.inner.algorithm();
-        let mut pending = [0u8; digest::MAX_BLOCK_LEN];
-        let pending = &mut pending[..algorithm.block_len()];
+        let buffer = &mut [0u8; digest::MAX_BLOCK_LEN];
         let num_pending = algorithm.output_len();
-        pending[..num_pending].copy_from_slice(self.inner.finish().as_ref());
-        Tag(self.outer.finish(pending, num_pending, cpu_features))
+        buffer[..num_pending].copy_from_slice(self.inner.finish().as_ref());
+        self.outer
+            .try_finish(buffer, num_pending, cpu_features)
+            .map(Tag)
+            .map_err(error::Unspecified::from)
     }
 }