Return Result from Expander::fill_bytes()

Author: daxpedda

ed448-goldilocks/src/field/element.rs

@@ -465,15 +465,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);
@@ -499,14 +499,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

@@ -6,6 +6,7 @@ pub(super) mod xof;
 use core::num::NonZero;
 
 use digest::{Digest, ExtendableOutput, Update, XofReader};
+use elliptic_curve::Error;
 use elliptic_curve::array::{Array, ArraySize};
 use xmd::ExpandMsgXmdError;
 use xof::ExpandMsgXofError;
@@ -42,8 +43,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, Error>;
 }
 
 /// The domain separation tag

hash2curve/src/hash2field/expand_msg/xmd.rs

@@ -11,6 +11,7 @@ use digest::{
     },
     block_api::BlockSizeUser,
 };
+use elliptic_curve::Error;
 
 /// Implements `expand_message_xof` via the [`ExpandMsg`] trait:
 /// <https://www.rfc-editor.org/rfc/rfc9380.html#name-expand_message_xmd>
@@ -107,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, Error> {
+        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() {
@@ -134,7 +141,10 @@ where
             *b = self.b_vals[self.offset];
             self.offset += 1;
             self.remaining -= 1;
+            read_bytes += 1;
         }
+
+        Ok(read_bytes)
     }
 }
 
@@ -232,7 +242,7 @@ mod test {
             .unwrap();
 
             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);
         }

hash2curve/src/hash2field/expand_msg/xof.rs

@@ -3,6 +3,7 @@
 use super::{Domain, ExpandMsg, Expander};
 use core::{fmt, num::NonZero, ops::Mul};
 use digest::{CollisionResistance, ExtendableOutput, HashMarker, Update, XofReader};
+use elliptic_curve::Error;
 use elliptic_curve::array::{
     ArraySize,
     typenum::{IsGreaterOrEqual, Prod, True, U2},
@@ -76,14 +77,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, Error> {
         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())
     }
 }
 
@@ -165,7 +167,7 @@ mod test {
             .unwrap();
 
             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);
         }