Return Result from Expander::fill_bytes()

Author: daxpedda

ed448-goldilocks/src/field/element.rs

@@ -464,15 +464,15 @@ mod tests {
             )
             .unwrap();
             let mut data = Array::<u8, U84>::default();
-            expander.fill_bytes(&mut data);
+            expander.fill_bytes(&mut data).unwrap();
             // TODO: This should be `Curve448FieldElement`.
             let u0 = Ed448FieldElement::from_okm(&data).0;
             let mut e_u0 = *expected_u0;
             e_u0.reverse();
             let mut e_u1 = *expected_u1;
             e_u1.reverse();
             assert_eq!(u0.to_bytes(), e_u0);
-            expander.fill_bytes(&mut data);
+            expander.fill_bytes(&mut data).unwrap();
             // TODO: This should be `Curve448FieldElement`.
             let u1 = Ed448FieldElement::from_okm(&data).0;
             assert_eq!(u1.to_bytes(), e_u1);
@@ -498,14 +498,14 @@ mod tests {
             )
             .unwrap();
             let mut data = Array::<u8, U84>::default();
-            expander.fill_bytes(&mut data);
+            expander.fill_bytes(&mut data).unwrap();
             let u0 = Ed448FieldElement::from_okm(&data).0;
             let mut e_u0 = *expected_u0;
             e_u0.reverse();
             let mut e_u1 = *expected_u1;
             e_u1.reverse();
             assert_eq!(u0.to_bytes(), e_u0);
-            expander.fill_bytes(&mut data);
+            expander.fill_bytes(&mut data).unwrap();
             let u1 = Ed448FieldElement::from_okm(&data).0;
             assert_eq!(u1.to_bytes(), e_u1);
         }

hash2curve/src/hash2field.rs

@@ -51,7 +51,9 @@ where
     let mut tmp = Array::<u8, <T as FromOkm>::Length>::default();
     let mut expander = E::expand_message(data, domain, len_in_bytes)?;
     Ok(core::array::from_fn(|_| {
-        expander.fill_bytes(&mut tmp);
+        expander
+            .fill_bytes(&mut tmp)
+            .expect("never exceeds `len_in_bytes`");
         T::from_okm(&tmp)
     }))
 }

hash2curve/src/hash2field/expand_msg.rs

@@ -39,8 +39,12 @@ pub trait ExpandMsg<K> {
 
 /// Expander that, call `read` until enough bytes have been consumed.
 pub trait Expander {
-    /// Fill the array with the expanded bytes
-    fn fill_bytes(&mut self, okm: &mut [u8]);
+    /// Fill the array with the expanded bytes, returning how many bytes were read.
+    ///
+    /// # Errors
+    ///
+    /// If no bytes are left.
+    fn fill_bytes(&mut self, okm: &mut [u8]) -> Result<usize>;
 }
 
 /// The domain separation tag

hash2curve/src/hash2field/expand_msg/xmd.rs

@@ -108,10 +108,16 @@ where
     HashT: BlockSizeUser + Default + FixedOutput + HashMarker,
     HashT::OutputSize: IsLessOrEqual<HashT::BlockSize, Output = True>,
 {
-    fn fill_bytes(&mut self, okm: &mut [u8]) {
+    fn fill_bytes(&mut self, okm: &mut [u8]) -> Result<usize> {
+        let mut read_bytes = 0;
+
         for b in okm {
             if self.remaining == 0 {
-                return;
+                if read_bytes == 0 {
+                    return Err(Error);
+                } else {
+                    return Ok(read_bytes);
+                }
             }
 
             if self.offset == self.b_vals.len() {
@@ -135,7 +141,10 @@ where
             *b = self.b_vals[self.offset];
             self.offset += 1;
             self.remaining -= 1;
+            read_bytes += 1;
         }
+
+        Ok(read_bytes)
     }
 }
 
@@ -210,7 +219,7 @@ mod test {
             )?;
 
             let mut uniform_bytes = Array::<u8, L>::default();
-            expander.fill_bytes(&mut uniform_bytes);
+            expander.fill_bytes(&mut uniform_bytes).unwrap();
 
             assert_eq!(uniform_bytes.as_slice(), self.uniform_bytes);
             Ok(())

hash2curve/src/hash2field/expand_msg/xof.rs

@@ -5,11 +5,11 @@ use core::{fmt, num::NonZero, ops::Mul};
 use digest::{
     CollisionResistance, ExtendableOutput, HashMarker, Update, XofReader, typenum::IsGreaterOrEqual,
 };
-use elliptic_curve::Result;
 use elliptic_curve::array::{
     ArraySize,
     typenum::{Prod, True, U2},
 };
+use elliptic_curve::{Error, Result};
 
 /// Implements `expand_message_xof` via the [`ExpandMsg`] trait:
 /// <https://www.rfc-editor.org/rfc/rfc9380.html#name-expand_message_xof>
@@ -78,14 +78,15 @@ impl<HashT> Expander for ExpandMsgXof<HashT>
 where
     HashT: Default + ExtendableOutput + Update + HashMarker,
 {
-    fn fill_bytes(&mut self, okm: &mut [u8]) {
+    fn fill_bytes(&mut self, okm: &mut [u8]) -> Result<usize> {
         if self.remaining == 0 {
-            return;
+            return Err(Error);
         }
 
         let bytes_to_read = self.remaining.min(okm.len().try_into().unwrap_or(u16::MAX));
         self.reader.read(&mut okm[..bytes_to_read.into()]);
         self.remaining -= bytes_to_read;
+        Ok(bytes_to_read.into())
     }
 }
 
@@ -147,7 +148,7 @@ mod test {
             )?;
 
             let mut uniform_bytes = Array::<u8, L>::default();
-            expander.fill_bytes(&mut uniform_bytes);
+            expander.fill_bytes(&mut uniform_bytes).unwrap();
 
             assert_eq!(uniform_bytes.as_slice(), self.uniform_bytes);
             Ok(())