Experiment: enforce constraints of ExpandMsg on the type level

elliptic-curve/src/hash2curve/group_digest.rs

@@ -1,5 +1,11 @@
+use core::ops::Mul;
+
 use super::MapToCurve;
 use crate::hash2field::{hash_to_field, ExpandMsg, FromOkm};
+use generic_array::{
+    typenum::{Prod, U1, U2},
+    ArrayLength,
+};
 use group::cofactor::CofactorGroup;
 
 /// Adds hashing arbitrary byte sequences to a valid group element
@@ -34,9 +40,13 @@ pub trait GroupDigest {
     /// let pt = ProjectivePoint::hash_from_bytes::<hash2field::ExpandMsgXof<sha3::Shake256>>(b"test data", b"CURVE_XOF:SHAKE-256_SSWU_RO_");
     /// ```
     ///
-    fn hash_from_bytes<X: ExpandMsg>(msg: &[u8], dst: &'static [u8]) -> Self::Output {
-        let mut u = [Self::FieldElement::default(), Self::FieldElement::default()];
-        hash_to_field::<X, _>(msg, dst, &mut u);
+    fn hash_from_bytes<X>(msg: &[u8], dst: &'static [u8]) -> Self::Output
+    where
+        X: ExpandMsg<Prod<<Self::FieldElement as FromOkm>::Length, U2>>,
+        <Self::FieldElement as FromOkm>::Length: Mul<U2>,
+        Prod<<Self::FieldElement as FromOkm>::Length, U2>: ArrayLength<u8>,
+    {
+        let u = hash_to_field::<X, _, U2>(msg, dst);
         let q0 = Self::Output::map_to_curve(u[0]);
         let q1 = Self::Output::map_to_curve(u[1]);
         // Ideally we could add and then clear cofactor once
@@ -60,9 +70,13 @@ pub trait GroupDigest {
     /// uniformly random in G: the set of possible outputs of
     /// encode_to_curve is only a fraction of the points in G, and some
     /// points in this set are more likely to be output than others.
-    fn encode_from_bytes<X: ExpandMsg>(msg: &[u8], dst: &'static [u8]) -> Self::Output {
-        let mut u = [Self::FieldElement::default()];
-        hash_to_field::<X, _>(msg, dst, &mut u);
+    fn encode_from_bytes<X>(msg: &[u8], dst: &'static [u8]) -> Self::Output
+    where
+        X: ExpandMsg<Prod<<Self::FieldElement as FromOkm>::Length, U1>>,
+        <Self::FieldElement as FromOkm>::Length: Mul<U1>,
+        Prod<<Self::FieldElement as FromOkm>::Length, U1>: ArrayLength<u8>,
+    {
+        let u = hash_to_field::<X, _, U1>(msg, dst);
         let q0 = Self::Output::map_to_curve(u[0]);
         q0.clear_cofactor()
     }

elliptic-curve/src/hash2field.rs

@@ -2,10 +2,13 @@ mod expand_msg;
 mod expand_msg_xmd;
 mod expand_msg_xof;
 
+use core::ops::Mul;
+
 pub use expand_msg::*;
 pub use expand_msg_xmd::*;
 pub use expand_msg_xof::*;
-use generic_array::{typenum::Unsigned, ArrayLength, GenericArray};
+use generic_array::typenum::Prod;
+use generic_array::{ArrayLength, GenericArray};
 
 /// The trait for helping to convert to a scalar
 pub trait FromOkm {
@@ -18,16 +21,22 @@ pub trait FromOkm {
 
 /// Convert an arbitrary byte sequence according to
 /// <https://tools.ietf.org/html/draft-irtf-cfrg-hash-to-curve-11#section-5.3>
-pub fn hash_to_field<E, T>(data: &[u8], domain: &'static [u8], out: &mut [T])
+pub fn hash_to_field<E, T, O: ArrayLength<T>>(
+    data: &[u8],
+    domain: &'static [u8],
+) -> GenericArray<T, O>
 where
-    E: ExpandMsg,
+    E: ExpandMsg<Prod<T::Length, O>>,
     T: FromOkm + Default,
+    T::Length: Mul<O>,
+    Prod<T::Length, O>: ArrayLength<u8>,
 {
-    let len_in_bytes = T::Length::to_usize() * out.len();
     let mut tmp = GenericArray::<u8, <T as FromOkm>::Length>::default();
-    let mut expander = E::expand_message(data, domain, len_in_bytes);
+    let mut expander = E::expand_message(data, domain);
+    let mut out = GenericArray::default();
     for o in out.iter_mut() {
         expander.fill_bytes(&mut tmp);
         *o = T::from_okm(&tmp);
     }
+    out
 }

elliptic-curve/src/hash2field/expand_msg.rs

@@ -1,4 +1,5 @@
 use digest::{Digest, ExtendableOutputDirty, Update, XofReader};
+use generic_array::typenum::{IsLess, U256};
 use generic_array::{ArrayLength, GenericArray};
 
 /// Salt when the DST is too long
@@ -7,11 +8,11 @@ const OVERSIZE_DST_SALT: &[u8] = b"H2C-OVERSIZE-DST-";
 const MAX_DST_LEN: usize = 255;
 
 /// Trait for types implementing expand_message interface for hash_to_field
-pub trait ExpandMsg {
+pub trait ExpandMsg<L: ArrayLength<u8>> {
     /// Expands `msg` to the required number of bytes
     /// Returns an expander that can be used to call `read` until enough
     /// bytes have been consumed
-    fn expand_message(msg: &[u8], dst: &'static [u8], len_in_bytes: usize) -> Self;
+    fn expand_message(msg: &[u8], dst: &'static [u8]) -> Self;
 
     /// Fill the array with the expanded bytes
     fn fill_bytes(&mut self, okm: &mut [u8]);
@@ -22,14 +23,20 @@ pub trait ExpandMsg {
 /// 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,
@@ -65,10 +72,11 @@ 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(),
+            Self::Hashed(_) => L::to_u8(),
+            // Can't overflow because enforced on a type level.
+            Self::Array(d) => d.len() as u8,
         }
     }
 }

elliptic-curve/src/hash2field/expand_msg_xmd.rs

@@ -1,11 +1,9 @@
+use core::ops::Mul;
+
 use super::{Domain, ExpandMsg};
-use digest::{
-    generic_array::{
-        typenum::{IsLess, IsLessOrEqual, Unsigned, U256},
-        GenericArray,
-    },
-    BlockInput, Digest,
-};
+use digest::{BlockInput, Digest};
+use generic_array::typenum::{IsLess, IsLessOrEqual, NonZero, Prod, Unsigned, U255, U256, U65536};
+use generic_array::{ArrayLength, GenericArray};
 
 /// Placeholder type for implementing expand_message_xmd based on a hash function
 pub struct ExpandMsgXmd<HashT>
@@ -17,9 +15,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>
@@ -41,9 +39,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 {
@@ -53,32 +51,45 @@ where
 }
 
 /// ExpandMsgXmd implements expand_message_xmd for the ExpandMsg trait
-impl<HashT> ExpandMsg for ExpandMsgXmd<HashT>
+impl<HashT, L> ExpandMsg<L> for ExpandMsgXmd<HashT>
 where
     HashT: Digest + BlockInput,
+    L: ArrayLength<u8>,
+    U255: Mul<HashT::OutputSize>,
+    // 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>,
+    // Constraint set by `expand_message_xmd`:
+    // https://www.ietf.org/archive/id/draft-irtf-cfrg-hash-to-curve-13.html#section-5.4.1-6
+    L: NonZero + IsLess<Prod<U255, HashT::OutputSize>> + IsLess<U65536>,
 {
-    fn expand_message(msg: &[u8], dst: &'static [u8], len_in_bytes: usize) -> Self {
-        let b_in_bytes = HashT::OutputSize::to_usize();
-        let ell = (len_in_bytes + b_in_bytes - 1) / b_in_bytes;
+    fn expand_message(msg: &[u8], dst: &'static [u8]) -> Self {
+        let b_in_bytes = HashT::OutputSize::to_u16();
+        // Can't overflow because enforced on a type level.
+        let ell = ((L::to_u16() + b_in_bytes - 1) / b_in_bytes) as u8;
+        // Enforced on the type level
         // if ell > 255 {
         //     panic!("ell was too big in expand_message_xmd");
         // }
         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(L::to_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();
 
         Self {

elliptic-curve/src/hash2field/expand_msg_xof.rs

@@ -1,28 +1,33 @@
 use super::ExpandMsg;
 use crate::hash2field::Domain;
-use digest::{ExtendableOutput, ExtendableOutputDirty, Update, XofReader};
-use generic_array::typenum::U32;
+use digest::{ExtendableOutput, ExtendableOutputDirty, FixedOutput, Update, XofReader};
+use generic_array::typenum::{IsLessOrEqual, NonZero, U32, U65536};
+use generic_array::ArrayLength;
 
 /// Placeholder type for implementing expand_message_xof based on an extendable output function
 pub struct ExpandMsgXof<HashT>
 where
-    HashT: Default + ExtendableOutput + ExtendableOutputDirty + Update,
+    HashT: Default + ExtendableOutput + ExtendableOutputDirty + FixedOutput + Update,
 {
     reader: <HashT as ExtendableOutput>::Reader,
 }
 
 /// ExpandMsgXof implements expand_message_xof for the ExpandMsg trait
-impl<HashT> ExpandMsg for ExpandMsgXof<HashT>
+impl<HashT, L> ExpandMsg<L> for ExpandMsgXof<HashT>
 where
-    HashT: Default + ExtendableOutput + ExtendableOutputDirty + Update,
+    HashT: Default + ExtendableOutput + ExtendableOutputDirty + FixedOutput + Update,
+    L: ArrayLength<u8>,
+    // Constraint set by `expand_message_xof`:
+    // https://www.ietf.org/archive/id/draft-irtf-cfrg-hash-to-curve-13.html#section-5.4.2-5
+    L: NonZero + IsLessOrEqual<U65536>,
 {
-    fn expand_message(msg: &[u8], dst: &'static [u8], len_in_bytes: usize) -> Self {
+    fn expand_message(msg: &[u8], dst: &'static [u8]) -> Self {
         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(L::to_u16().to_be_bytes())
             .chain(domain.data())
-            .chain([domain.len() as u8])
+            .chain([domain.len()])
             .finalize_xof();
         Self { reader }
     }