Various fixes and improvements to hash2curve

daxpedda · daxpedda · commit 5f3658cefdbb · 2025-04-11T21:49:00.000+02:00
diff --git a/elliptic-curve/src/hash2curve/hash2field.rs b/elliptic-curve/src/hash2curve/hash2field.rs
@@ -4,15 +4,20 @@
 
 mod expand_msg;
 
+use core::num::NonZeroUsize;
+
 pub use expand_msg::{xmd::*, xof::*, *};
 
 use crate::{Error, Result};
-use hybrid_array::{Array, ArraySize, typenum::Unsigned};
+use hybrid_array::{
+    Array, ArraySize,
+    typenum::{NonZero, Unsigned},
+};
 
 /// The trait for helping to convert to a field element.
 pub trait FromOkm {
     /// The number of bytes needed to convert to a field element.
-    type Length: ArraySize;
+    type Length: ArraySize + NonZero;
 
     /// Convert a byte sequence into a field element.
     fn from_okm(data: &Array<u8, Self::Length>) -> Self;
@@ -37,7 +42,10 @@ where
     E: ExpandMsg<'a>,
     T: FromOkm + Default,
 {
-    let len_in_bytes = T::Length::to_usize().checked_mul(out.len()).ok_or(Error)?;
+    let len_in_bytes = T::Length::to_usize()
+        .checked_mul(out.len())
+        .and_then(NonZeroUsize::new)
+        .ok_or(Error)?;
     let mut tmp = Array::<u8, <T as FromOkm>::Length>::default();
     let mut expander = E::expand_message(data, domain, len_in_bytes)?;
     for o in out.iter_mut() {
diff --git a/elliptic-curve/src/hash2curve/hash2field/expand_msg.rs b/elliptic-curve/src/hash2curve/hash2field/expand_msg.rs
@@ -3,6 +3,8 @@
 pub(super) mod xmd;
 pub(super) mod xof;
 
+use core::num::NonZero;
+
 use crate::{Error, Result};
 use digest::{Digest, ExtendableOutput, Update, XofReader};
 use hybrid_array::typenum::{IsLess, U256};
@@ -28,7 +30,7 @@ pub trait ExpandMsg<'a> {
     fn expand_message(
         msgs: &[&[u8]],
         dsts: &'a [&'a [u8]],
-        len_in_bytes: usize,
+        len_in_bytes: NonZero<usize>,
     ) -> Result<Self::Expander>;
 }
 
diff --git a/elliptic-curve/src/hash2curve/hash2field/expand_msg/xmd.rs b/elliptic-curve/src/hash2curve/hash2field/expand_msg/xmd.rs
@@ -1,59 +1,71 @@
 //! `expand_message_xmd` based on a hash function.
 
-use core::marker::PhantomData;
+use core::{marker::PhantomData, num::NonZero, ops::Mul};
 
 use super::{Domain, ExpandMsg, Expander};
 use crate::{Error, Result};
 use digest::{
     FixedOutput, HashMarker,
     array::{
         Array,
-        typenum::{IsLess, IsLessOrEqual, U256, Unsigned},
+        typenum::{IsGreaterOrEqual, IsLess, IsLessOrEqual, U2, U8, U256, Unsigned},
     },
     core_api::BlockSizeUser,
 };
 
-/// Placeholder type for implementing `expand_message_xmd` based on a hash function
+/// Implements `expand_message_xof` via the [`ExpandMsg`] trait:
+/// https://www.rfc-editor.org/rfc/rfc9380.html#name-expand_message_xmd
+///
+/// `K` is the target security level in bits:
+/// https://www.rfc-editor.org/rfc/rfc9380.html#section-8.9-2.2
+/// https://www.rfc-editor.org/rfc/rfc9380.html#name-target-security-levels
 ///
 /// # Errors
 /// - `dst.is_empty()`
-/// - `len_in_bytes == 0`
 /// - `len_in_bytes > u16::MAX`
 /// - `len_in_bytes > 255 * HashT::OutputSize`
 #[derive(Debug)]
-pub struct ExpandMsgXmd<HashT>(PhantomData<HashT>)
+pub struct ExpandMsgXmd<HashT, K>(PhantomData<(HashT, K)>)
 where
     HashT: BlockSizeUser + Default + FixedOutput + HashMarker,
     HashT::OutputSize: IsLess<U256>,
-    HashT::OutputSize: IsLessOrEqual<HashT::BlockSize>;
+    HashT::OutputSize: IsLessOrEqual<HashT::BlockSize>,
+    HashT::OutputSize: Mul<U8>,
+    U2: Mul<K>,
+    <HashT::OutputSize as Mul<U8>>::Output: IsGreaterOrEqual<<U2 as Mul<K>>::Output>;
 
-/// ExpandMsgXmd implements expand_message_xmd for the ExpandMsg trait
-impl<'a, HashT> ExpandMsg<'a> for ExpandMsgXmd<HashT>
+impl<'a, HashT, K> ExpandMsg<'a> for ExpandMsgXmd<HashT, K>
 where
     HashT: BlockSizeUser + Default + FixedOutput + HashMarker,
-    // If `len_in_bytes` is bigger then 256, length of the `DST` will depend on
-    // the output size of the hash, which is still not allowed to be bigger then 256:
+    // If DST is larger than 255 bytes, the length of the computed DST will depend on the output
+    // size of the hash, which is still not allowed to be larger than 256:
     // https://www.ietf.org/archive/id/draft-irtf-cfrg-hash-to-curve-13.html#section-5.4.1-6
     HashT::OutputSize: IsLess<U256>,
     // Constraint set by `expand_message_xmd`:
     // https://www.ietf.org/archive/id/draft-irtf-cfrg-hash-to-curve-13.html#section-5.4.1-4
     HashT::OutputSize: IsLessOrEqual<HashT::BlockSize>,
+    // The number of bits output by `HashT` MUST be larger or equal to `2 * K`:
+    // https://www.rfc-editor.org/rfc/rfc9380.html#section-5.3.1-2.1
+    HashT::OutputSize: Mul<U8>,
+    U2: Mul<K>,
+    <HashT::OutputSize as Mul<U8>>::Output: IsGreaterOrEqual<<U2 as Mul<K>>::Output>,
 {
     type Expander = ExpanderXmd<'a, HashT>;
 
     fn expand_message(
         msgs: &[&[u8]],
         dsts: &'a [&'a [u8]],
-        len_in_bytes: usize,
+        len_in_bytes: NonZero<usize>,
     ) -> Result<Self::Expander> {
-        if len_in_bytes == 0 {
+        let len_in_bytes_u16 = u16::try_from(len_in_bytes.get()).map_err(|_| Error)?;
+
+        // `255 * <b_in_bytes>` can not exceed `u16::MAX`
+        if len_in_bytes_u16 > 255 * HashT::OutputSize::to_u16() {
             return Err(Error);
         }
 
-        let len_in_bytes_u16 = u16::try_from(len_in_bytes).map_err(|_| Error)?;
-
         let b_in_bytes = HashT::OutputSize::to_usize();
-        let ell = u8::try_from(len_in_bytes.div_ceil(b_in_bytes)).map_err(|_| Error)?;
+        let ell = u8::try_from(len_in_bytes.get().div_ceil(b_in_bytes)).map_err(|_| Error)?;
 
         let domain = Domain::xmd::<HashT>(dsts)?;
         let mut b_0 = HashT::default();
@@ -157,7 +169,7 @@ mod test {
     use hex_literal::hex;
     use hybrid_array::{
         ArraySize,
-        typenum::{U32, U128},
+        typenum::{U8, U32, U128},
     };
     use sha2::Sha256;
 
@@ -209,13 +221,18 @@ mod test {
         ) -> Result<()>
         where
             HashT: BlockSizeUser + Default + FixedOutput + HashMarker,
-            HashT::OutputSize: IsLess<U256> + IsLessOrEqual<HashT::BlockSize>,
+            HashT::OutputSize: IsLess<U256> + IsLessOrEqual<HashT::BlockSize> + Mul<U8>,
+            U2: Mul<U32>,
+            <HashT::OutputSize as Mul<U8>>::Output: IsGreaterOrEqual<<U2 as Mul<U32>>::Output>,
         {
             assert_message::<HashT>(self.msg, domain, L::to_u16(), self.msg_prime);
 
             let dst = [dst];
-            let mut expander =
-                ExpandMsgXmd::<HashT>::expand_message(&[self.msg], &dst, L::to_usize())?;
+            let mut expander = ExpandMsgXmd::<HashT, U32>::expand_message(
+                &[self.msg],
+                &dst,
+                NonZero::new(L::to_usize()).ok_or(Error)?,
+            )?;
 
             let mut uniform_bytes = Array::<u8, L>::default();
             expander.fill_bytes(&mut uniform_bytes);
diff --git a/elliptic-curve/src/hash2curve/hash2field/expand_msg/xof.rs b/elliptic-curve/src/hash2curve/hash2field/expand_msg/xof.rs
@@ -2,26 +2,45 @@
 
 use super::{Domain, ExpandMsg, Expander};
 use crate::{Error, Result};
-use core::fmt;
-use digest::{ExtendableOutput, Update, XofReader};
-use hybrid_array::typenum::U32;
-
-/// Placeholder type for implementing `expand_message_xof` based on an extendable output function
+use core::{
+    fmt,
+    marker::PhantomData,
+    num::NonZero,
+    ops::{Div, Mul},
+};
+use digest::{ExtendableOutput, HashMarker, Update, XofReader};
+use hybrid_array::{
+    ArraySize,
+    typenum::{IsLess, U2, U8, U256},
+};
+
+/// Implements `expand_message_xof` via the [`ExpandMsg`] trait:
+/// https://www.rfc-editor.org/rfc/rfc9380.html#name-expand_message_xof
+///
+/// `K` is the target security level in bits:
+/// https://www.rfc-editor.org/rfc/rfc9380.html#section-8.9-2.2
+/// https://www.rfc-editor.org/rfc/rfc9380.html#name-target-security-levels
 ///
 /// # Errors
 /// - `dst.is_empty()`
-/// - `len_in_bytes == 0`
 /// - `len_in_bytes > u16::MAX`
-pub struct ExpandMsgXof<HashT>
+pub struct ExpandMsgXof<HashT, K>
 where
-    HashT: Default + ExtendableOutput + Update,
+    HashT: Default + ExtendableOutput + Update + HashMarker,
+    U2: Mul<K>,
+    <U2 as Mul<K>>::Output: Div<U8>,
+    HashSize<K>: ArraySize + IsLess<U256>,
 {
     reader: <HashT as ExtendableOutput>::Reader,
+    _k: PhantomData<K>,
 }
 
-impl<HashT> fmt::Debug for ExpandMsgXof<HashT>
+impl<HashT, K> fmt::Debug for ExpandMsgXof<HashT, K>
 where
-    HashT: Default + ExtendableOutput + Update,
+    HashT: Default + ExtendableOutput + Update + HashMarker,
+    U2: Mul<K>,
+    <U2 as Mul<K>>::Output: Div<U8>,
+    HashSize<K>: ArraySize + IsLess<U256>,
     <HashT as ExtendableOutput>::Reader: fmt::Debug,
 {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
@@ -31,25 +50,28 @@ where
     }
 }
 
-/// ExpandMsgXof implements `expand_message_xof` for the [`ExpandMsg`] trait
-impl<'a, HashT> ExpandMsg<'a> for ExpandMsgXof<HashT>
+type HashSize<K> = <<U2 as Mul<K>>::Output as Div<U8>>::Output;
+
+impl<'a, HashT, K> ExpandMsg<'a> for ExpandMsgXof<HashT, K>
 where
-    HashT: Default + ExtendableOutput + Update,
+    HashT: Default + ExtendableOutput + Update + HashMarker,
+    // If DST is larger than 255 bytes, the length of the computed DST is calculated by
+    // `2 * k / 8`.
+    // https://www.rfc-editor.org/rfc/rfc9380.html#section-5.3.1-2.1
+    U2: Mul<K>,
+    <U2 as Mul<K>>::Output: Div<U8>,
+    HashSize<K>: ArraySize + IsLess<U256>,
 {
     type Expander = Self;
 
     fn expand_message(
         msgs: &[&[u8]],
         dsts: &'a [&'a [u8]],
-        len_in_bytes: usize,
+        len_in_bytes: NonZero<usize>,
     ) -> Result<Self::Expander> {
-        if len_in_bytes == 0 {
-            return Err(Error);
-        }
-
-        let len_in_bytes = u16::try_from(len_in_bytes).map_err(|_| Error)?;
+        let len_in_bytes = u16::try_from(len_in_bytes.get()).map_err(|_| Error)?;
 
-        let domain = Domain::<U32>::xof::<HashT>(dsts)?;
+        let domain = Domain::<HashSize<K>>::xof::<HashT>(dsts)?;
         let mut reader = HashT::default();
 
         for msg in msgs {
@@ -60,13 +82,19 @@ where
         domain.update_hash(&mut reader);
         reader.update(&[domain.len()]);
         let reader = reader.finalize_xof();
-        Ok(Self { reader })
+        Ok(Self {
+            reader,
+            _k: PhantomData,
+        })
     }
 }
 
-impl<HashT> Expander for ExpandMsgXof<HashT>
+impl<HashT, K> Expander for ExpandMsgXof<HashT, K>
 where
-    HashT: Default + ExtendableOutput + Update,
+    HashT: Default + ExtendableOutput + Update + HashMarker,
+    U2: Mul<K>,
+    <U2 as Mul<K>>::Output: Div<U8>,
+    HashSize<K>: ArraySize + IsLess<U256>,
 {
     fn fill_bytes(&mut self, okm: &mut [u8]) {
         self.reader.read(okm);
@@ -78,7 +106,10 @@ mod test {
     use super::*;
     use core::mem::size_of;
     use hex_literal::hex;
-    use hybrid_array::{Array, ArraySize, typenum::U128};
+    use hybrid_array::{
+        Array, ArraySize,
+        typenum::{U32, U128},
+    };
     use sha3::Shake128;
 
     fn assert_message(msg: &[u8], domain: &Domain<'_, U32>, len_in_bytes: u16, bytes: &[u8]) {
@@ -110,13 +141,16 @@ mod test {
         #[allow(clippy::panic_in_result_fn)]
         fn assert<HashT, L>(&self, dst: &'static [u8], domain: &Domain<'_, U32>) -> Result<()>
         where
-            HashT: Default + ExtendableOutput + Update,
+            HashT: Default + ExtendableOutput + Update + HashMarker,
             L: ArraySize,
         {
             assert_message(self.msg, domain, L::to_u16(), self.msg_prime);
 
-            let mut expander =
-                ExpandMsgXof::<HashT>::expand_message(&[self.msg], &[dst], L::to_usize())?;
+            let mut expander = ExpandMsgXof::<HashT, U128>::expand_message(
+                &[self.msg],
+                &[dst],
+                NonZero::new(L::to_usize()).ok_or(Error)?,
+            )?;
 
             let mut uniform_bytes = Array::<u8, L>::default();
             expander.fill_bytes(&mut uniform_bytes);
