Apply constraints and some improvements

Author: daxpedda

elliptic-curve/src/hash2field/expand_msg.rs

@@ -1,5 +1,6 @@
 use crate::Result;
-use digest::{Digest, ExtendableOutputDirty, Update, XofReader};
+use digest::{Digest, ExtendableOutput, Update, XofReader};
+use generic_array::typenum::{IsLess, U256};
 use generic_array::{ArrayLength, GenericArray};
 
 /// Salt when the DST is too long
@@ -23,24 +24,30 @@ pub trait ExpandMsg: Sized {
 /// Implements [section 5.4.3 of `draft-irtf-cfrg-hash-to-curve-13`][dst].
 ///
 /// [dst]: https://datatracker.ietf.org/doc/html/draft-irtf-cfrg-hash-to-curve-13#section-5.4.3
-pub(crate) enum Domain<L: ArrayLength<u8>> {
+pub(crate) enum Domain<L>
+where
+    L: ArrayLength<u8> + IsLess<U256>,
+{
     /// > 255
     Hashed(GenericArray<u8, L>),
     /// <= 255
     Array(&'static [u8]),
 }
 
-impl<L: ArrayLength<u8>> Domain<L> {
+impl<L> Domain<L>
+where
+    L: ArrayLength<u8> + IsLess<U256>,
+{
     pub fn xof<X>(dst: &'static [u8]) -> Self
     where
-        X: Default + ExtendableOutputDirty + Update,
+        X: Default + ExtendableOutput + Update,
     {
         if dst.len() > MAX_DST_LEN {
             let mut data = GenericArray::<u8, L>::default();
             X::default()
                 .chain(OVERSIZE_DST_SALT)
                 .chain(dst)
-                .finalize_xof_dirty()
+                .finalize_xof()
                 .read(&mut data);
             Self::Hashed(data)
         } else {
@@ -66,10 +73,12 @@ impl<L: ArrayLength<u8>> Domain<L> {
         }
     }
 
-    pub fn len(&self) -> usize {
+    pub fn len(&self) -> u8 {
         match self {
-            Self::Hashed(_) => L::to_usize(),
-            Self::Array(d) => d.len(),
+            // Can't overflow because it's enforced on a type level.
+            Self::Hashed(_) => L::to_u8(),
+            // Can't overflow because it's checked on creation.
+            Self::Array(d) => d.len() as u8,
         }
     }
 }

elliptic-curve/src/hash2field/expand_msg_xmd.rs

@@ -18,9 +18,9 @@ where
     b_0: GenericArray<u8, HashT::OutputSize>,
     b_vals: GenericArray<u8, HashT::OutputSize>,
     domain: Domain<HashT::OutputSize>,
-    index: usize,
+    index: u8,
     offset: usize,
-    ell: usize,
+    ell: u8,
 }
 
 impl<HashT> ExpandMsgXmd<HashT>
@@ -42,9 +42,9 @@ where
                 .for_each(|(j, (b0val, bi1val))| tmp[j] = b0val ^ bi1val);
             self.b_vals = HashT::new()
                 .chain(tmp)
-                .chain([self.index as u8])
+                .chain([self.index])
                 .chain(self.domain.data())
-                .chain([self.domain.len() as u8])
+                .chain([self.domain.len()])
                 .finalize();
             true
         } else {
@@ -57,34 +57,39 @@ where
 impl<HashT> ExpandMsg for ExpandMsgXmd<HashT>
 where
     HashT: Digest + BlockInput,
-    HashT::OutputSize: IsLess<U256> + IsLessOrEqual<HashT::BlockSize>,
+    // 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:
+    // 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>,
 {
     fn expand_message(msg: &[u8], dst: &'static [u8], len_in_bytes: usize) -> Result<Self> {
-        if len_in_bytes > 0xFFFF {
+        if len_in_bytes == 0 {
             return Err(Error);
         }
 
-        let b_in_bytes = HashT::OutputSize::to_usize();
-        let ell = (len_in_bytes + b_in_bytes - 1) / b_in_bytes;
+        let len_in_bytes_u16 = u16::try_from(len_in_bytes).map_err(|_| Error)?;
 
-        if ell > 255 {
-            return Err(Error);
-        }
+        let b_in_bytes = HashT::OutputSize::to_usize();
+        let ell = u8::try_from((len_in_bytes + b_in_bytes - 1) / b_in_bytes).map_err(|_| Error)?;
 
         let domain = Domain::xmd::<HashT>(dst);
         let b_0 = HashT::new()
             .chain(GenericArray::<u8, HashT::BlockSize>::default())
             .chain(msg)
-            .chain([(len_in_bytes >> 8) as u8, len_in_bytes as u8, 0u8])
+            .chain(len_in_bytes_u16.to_be_bytes())
+            .chain([0])
             .chain(domain.data())
-            .chain([domain.len() as u8])
+            .chain([domain.len()])
             .finalize();
 
         let b_vals = HashT::new()
             .chain(&b_0[..])
             .chain([1u8])
             .chain(domain.data())
-            .chain([domain.len() as u8])
+            .chain([domain.len()])
             .finalize();
 
         Ok(Self {

elliptic-curve/src/hash2field/expand_msg_xof.rs

@@ -1,5 +1,5 @@
 use super::ExpandMsg;
-use crate::{hash2field::Domain, Result};
+use crate::{hash2field::Domain, Error, Result};
 use digest::{ExtendableOutput, ExtendableOutputDirty, Update, XofReader};
 use generic_array::typenum::U32;
 
@@ -17,12 +17,18 @@ where
     HashT: Default + ExtendableOutput + ExtendableOutputDirty + Update,
 {
     fn expand_message(msg: &[u8], dst: &'static [u8], len_in_bytes: usize) -> Result<Self> {
+        if len_in_bytes == 0 {
+            return Err(Error);
+        }
+
+        let len_in_bytes = u16::try_from(len_in_bytes).map_err(|_| Error)?;
+
         let domain = Domain::<U32>::xof::<HashT>(dst);
         let reader = HashT::default()
             .chain(msg)
-            .chain([(len_in_bytes >> 8) as u8, len_in_bytes as u8])
+            .chain(len_in_bytes.to_be_bytes())
             .chain(domain.data())
-            .chain([domain.len() as u8])
+            .chain([domain.len()])
             .finalize_xof();
         Ok(Self { reader })
     }