Group revamp (#261)

* Revamp `KeGroup` trait

* Update dependencies

* Fix `hash_to_scalar` using `OprfGroup` instead of `KeGroup`

* Relax constraints on associated types of `KeGroup`

* Improve `KeGroup` implementation on `Curve`

* Improve `KeyExchange` trait

* Fix new Clippy 1.59 warnings

src/key_exchange/group/elliptic_curve.rs

@@ -0,0 +1,85 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+//
+// This source code is licensed under both the MIT license found in the
+// LICENSE-MIT file in the root directory of this source tree and the Apache
+// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
+// of this source tree.
+
+use digest::core_api::BlockSizeUser;
+use digest::Digest;
+use elliptic_curve::group::cofactor::CofactorGroup;
+use elliptic_curve::hash2curve::{ExpandMsgXmd, FromOkm, GroupDigest};
+use elliptic_curve::sec1::{FromEncodedPoint, ModulusSize, ToEncodedPoint};
+use elliptic_curve::{
+    AffinePoint, Curve, FieldSize, NonZeroScalar, ProjectiveArithmetic, ProjectivePoint, PublicKey,
+    Scalar, SecretKey,
+};
+use generic_array::typenum::{IsLess, IsLessOrEqual, U256};
+use generic_array::GenericArray;
+use rand::{CryptoRng, RngCore};
+
+use super::KeGroup;
+use crate::errors::InternalError;
+
+impl<G: Curve + GroupDigest + ProjectiveArithmetic> KeGroup for G
+where
+    FieldSize<Self>: ModulusSize,
+    AffinePoint<Self>: FromEncodedPoint<Self> + ToEncodedPoint<Self>,
+    ProjectivePoint<Self>: CofactorGroup + ToEncodedPoint<Self>,
+    Scalar<Self>: FromOkm,
+{
+    type Pk = PublicKey<Self>;
+
+    type PkLen = <FieldSize<Self> as ModulusSize>::CompressedPointSize;
+
+    type Sk = SecretKey<Self>;
+
+    type SkLen = FieldSize<Self>;
+
+    fn serialize_pk(pk: &Self::Pk) -> GenericArray<u8, Self::PkLen> {
+        GenericArray::clone_from_slice(pk.to_encoded_point(true).as_bytes())
+    }
+
+    fn deserialize_pk(bytes: &GenericArray<u8, Self::PkLen>) -> Result<Self::Pk, InternalError> {
+        PublicKey::from_sec1_bytes(bytes).map_err(|_| InternalError::PointError)
+    }
+
+    fn random_sk<R: RngCore + CryptoRng>(rng: &mut R) -> Self::Sk {
+        SecretKey::random(rng)
+    }
+
+    // Implements the `HashToScalar()` function
+    fn hash_to_scalar<H>(input: &[&[u8]], dst: &[u8]) -> Result<Self::Sk, InternalError>
+    where
+        H: Digest + BlockSizeUser,
+        H::OutputSize: IsLess<U256> + IsLessOrEqual<H::BlockSize>,
+    {
+        Self::hash_to_scalar::<ExpandMsgXmd<H>>(input, dst)
+            .ok()
+            .and_then(|scalar| Option::<NonZeroScalar<Self>>::from(NonZeroScalar::new(scalar)))
+            .map(SecretKey::from)
+            .ok_or(InternalError::HashToScalar)
+    }
+
+    fn public_key(sk: &Self::Sk) -> Self::Pk {
+        sk.public_key()
+    }
+
+    fn diffie_hellman(pk: &Self::Pk, sk: &Self::Sk) -> GenericArray<u8, Self::PkLen> {
+        GenericArray::clone_from_slice(
+            (pk.to_projective() * sk.to_nonzero_scalar().as_ref())
+                .to_encoded_point(true)
+                .as_bytes(),
+        )
+    }
+
+    fn zeroize_sk_on_drop(_sk: &mut Self::Sk) {}
+
+    fn serialize_sk(sk: &Self::Sk) -> GenericArray<u8, Self::SkLen> {
+        sk.to_be_bytes()
+    }
+
+    fn deserialize_sk(bytes: &GenericArray<u8, Self::SkLen>) -> Result<Self::Sk, InternalError> {
+        SecretKey::from_be_bytes(bytes).map_err(|_| InternalError::PointError)
+    }
+}

src/key_exchange/group/mod.rs

@@ -7,37 +7,62 @@
 
 //! Includes the KeGroup trait and definitions for the key exchange groups
 
+mod elliptic_curve;
+#[cfg(feature = "ristretto255")]
+pub mod ristretto255;
+#[cfg(feature = "x25519")]
+pub mod x25519;
+
+use digest::core_api::BlockSizeUser;
+use digest::Digest;
+use generic_array::typenum::{IsLess, IsLessOrEqual, U256};
 use generic_array::{ArrayLength, GenericArray};
 use rand::{CryptoRng, RngCore};
 
 use crate::errors::InternalError;
 
 /// A group representation for use in the key exchange
-pub trait KeGroup: Sized + Clone {
+pub trait KeGroup {
+    /// Public key
+    type Pk: Clone;
     /// Length of the public key
-    type PkLen: ArrayLength<u8> + 'static;
+    type PkLen: ArrayLength<u8>;
+    /// Secret key
+    type Sk: Clone;
     /// Length of the secret key
-    type SkLen: ArrayLength<u8> + 'static;
-
-    /// Return a public key from its fixed-length bytes representation
-    fn from_pk_slice(element_bits: &GenericArray<u8, Self::PkLen>) -> Result<Self, InternalError>;
-
-    /// Generate a random secret key
-    fn random_sk<R: RngCore + CryptoRng>(rng: &mut R) -> GenericArray<u8, Self::SkLen>;
-
-    /// Return a public key from its secret key
-    fn public_key(sk: &GenericArray<u8, Self::SkLen>) -> Self;
+    type SkLen: ArrayLength<u8>;
 
     /// Serializes `self`
-    fn to_arr(&self) -> GenericArray<u8, Self::PkLen>;
+    fn serialize_pk(pk: &Self::Pk) -> GenericArray<u8, Self::PkLen>;
+
+    /// Return a public key from its fixed-length bytes representation
+    fn deserialize_pk(bytes: &GenericArray<u8, Self::PkLen>) -> Result<Self::Pk, InternalError>;
+
+    /// Generate a random secret key
+    fn random_sk<R: RngCore + CryptoRng>(rng: &mut R) -> Self::Sk;
+
+    /// Hashes a slice of pseudo-random bytes to a scalar
+    ///
+    /// # Errors
+    /// [`InternalError::HashToScalar`] if the `input` is empty or longer then
+    /// [`u16::MAX`].
+    fn hash_to_scalar<H>(input: &[&[u8]], dst: &[u8]) -> Result<Self::Sk, InternalError>
+    where
+        H: Digest + BlockSizeUser,
+        H::OutputSize: IsLess<U256> + IsLessOrEqual<H::BlockSize>;
+
+    /// Return a public key from its secret key
+    fn public_key(sk: &Self::Sk) -> Self::Pk;
 
     /// Diffie-Hellman key exchange
-    fn diffie_hellman(&self, sk: &GenericArray<u8, Self::SkLen>) -> GenericArray<u8, Self::PkLen>;
-}
+    fn diffie_hellman(pk: &Self::Pk, sk: &Self::Sk) -> GenericArray<u8, Self::PkLen>;
 
-#[cfg(feature = "p256")]
-pub mod p256;
-#[cfg(feature = "ristretto255")]
-pub mod ristretto255;
-#[cfg(feature = "x25519")]
-pub mod x25519;
+    /// Zeroize secret key on drop.
+    fn zeroize_sk_on_drop(sk: &mut Self::Sk);
+
+    /// Serializes `self`
+    fn serialize_sk(sk: &Self::Sk) -> GenericArray<u8, Self::SkLen>;
+
+    /// Return a public key from its fixed-length bytes representation
+    fn deserialize_sk(bytes: &GenericArray<u8, Self::SkLen>) -> Result<Self::Sk, InternalError>;
+}

src/key_exchange/group/p256.rs

@@ -1,52 +0,0 @@
-// Copyright (c) Facebook, Inc. and its affiliates.
-//
-// This source code is licensed under both the MIT license found in the
-// LICENSE-MIT file in the root directory of this source tree and the Apache
-// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
-// of this source tree.
-
-//! Key Exchange group implementation for p256
-
-use generic_array::typenum::{U32, U33};
-use generic_array::GenericArray;
-use p256_::elliptic_curve::group::GroupEncoding;
-use p256_::elliptic_curve::sec1::ToEncodedPoint;
-use p256_::elliptic_curve::{PublicKey, SecretKey};
-use p256_::NistP256;
-use rand::{CryptoRng, RngCore};
-
-use super::KeGroup;
-use crate::errors::InternalError;
-
-impl KeGroup for PublicKey<NistP256> {
-    type PkLen = U33;
-    type SkLen = U32;
-
-    fn from_pk_slice(element_bits: &GenericArray<u8, Self::PkLen>) -> Result<Self, InternalError> {
-        Self::from_sec1_bytes(element_bits).map_err(|_| InternalError::PointError)
-    }
-
-    fn random_sk<R: RngCore + CryptoRng>(rng: &mut R) -> GenericArray<u8, Self::SkLen> {
-        SecretKey::<NistP256>::random(rng).to_be_bytes()
-    }
-
-    fn public_key(sk: &GenericArray<u8, Self::SkLen>) -> Self {
-        SecretKey::<NistP256>::from_be_bytes(sk)
-            .unwrap()
-            .public_key()
-    }
-
-    fn to_arr(&self) -> GenericArray<u8, Self::PkLen> {
-        GenericArray::clone_from_slice(self.to_encoded_point(true).as_bytes())
-    }
-
-    fn diffie_hellman(&self, sk: &GenericArray<u8, Self::SkLen>) -> GenericArray<u8, Self::PkLen> {
-        (self.to_projective()
-            * SecretKey::<NistP256>::from_be_bytes(sk)
-                .unwrap()
-                .to_nonzero_scalar()
-                .as_ref())
-        .to_affine()
-        .to_bytes()
-    }
-}

src/key_exchange/group/ristretto255.rs

@@ -10,24 +10,40 @@
 use curve25519_dalek::constants::RISTRETTO_BASEPOINT_POINT;
 use curve25519_dalek::ristretto::{CompressedRistretto, RistrettoPoint};
 use curve25519_dalek::scalar::Scalar;
-use generic_array::typenum::U32;
+use digest::core_api::BlockSizeUser;
+use digest::{Digest, OutputSizeUser};
+use elliptic_curve::hash2curve::{ExpandMsg, ExpandMsgXmd, Expander};
+use generic_array::typenum::{IsLess, IsLessOrEqual, U256, U32, U64};
 use generic_array::GenericArray;
 use rand::{CryptoRng, RngCore};
+use voprf::Group;
+use zeroize::Zeroize;
 
 use super::KeGroup;
 use crate::errors::InternalError;
 
-impl KeGroup for RistrettoPoint {
+/// Implementation for Ristretto255.
+// This is necessary because Rust lacks specialization, otherwise we could
+// implement `KeGroup` for `voprf::Ristretto255`.
+pub struct Ristretto255;
+
+impl KeGroup for Ristretto255 {
+    type Pk = RistrettoPoint;
     type PkLen = U32;
+    type Sk = Scalar;
     type SkLen = U32;
 
-    fn from_pk_slice(element_bits: &GenericArray<u8, Self::PkLen>) -> Result<Self, InternalError> {
-        CompressedRistretto::from_slice(element_bits)
+    fn serialize_pk(pk: &Self::Pk) -> GenericArray<u8, Self::PkLen> {
+        pk.compress().to_bytes().into()
+    }
+
+    fn deserialize_pk(bytes: &GenericArray<u8, Self::PkLen>) -> Result<Self::Pk, InternalError> {
+        CompressedRistretto::from_slice(bytes)
             .decompress()
             .ok_or(InternalError::PointError)
     }
 
-    fn random_sk<R: RngCore + CryptoRng>(rng: &mut R) -> GenericArray<u8, Self::SkLen> {
+    fn random_sk<R: RngCore + CryptoRng>(rng: &mut R) -> Self::Sk {
         loop {
             let scalar = {
                 #[cfg(not(test))]
@@ -47,21 +63,121 @@ impl KeGroup for RistrettoPoint {
                 }
             };
 
-            if scalar != Scalar::zero() {
-                break scalar.to_bytes().into();
+            if scalar != Scalar::zero() && scalar.is_canonical() {
+                break scalar;
             }
         }
     }
 
-    fn public_key(sk: &GenericArray<u8, Self::SkLen>) -> Self {
-        RISTRETTO_BASEPOINT_POINT * Scalar::from_bits(*sk.as_ref())
+    // Implements the `HashToScalar()` function from
+    // https://www.ietf.org/archive/id/draft-irtf-cfrg-voprf-08.html#section-4.1
+    fn hash_to_scalar<'a, H>(input: &[&[u8]], dst: &[u8]) -> Result<Self::Sk, InternalError>
+    where
+        H: Digest + BlockSizeUser,
+        H::OutputSize: IsLess<U256> + IsLessOrEqual<H::BlockSize>,
+    {
+        let mut uniform_bytes = GenericArray::<_, U64>::default();
+        ExpandMsgXmd::<H>::expand_message(input, dst, 64)
+            .map_err(|_| InternalError::HashToScalar)?
+            .fill_bytes(&mut uniform_bytes);
+
+        Ok(Scalar::from_bytes_mod_order_wide(&uniform_bytes.into()))
     }
 
-    fn to_arr(&self) -> GenericArray<u8, Self::PkLen> {
-        self.compress().to_bytes().into()
+    fn public_key(sk: &Self::Sk) -> Self::Pk {
+        RISTRETTO_BASEPOINT_POINT * sk
     }
 
-    fn diffie_hellman(&self, sk: &GenericArray<u8, Self::SkLen>) -> GenericArray<u8, Self::SkLen> {
-        (self * Scalar::from_bits(*sk.as_ref())).to_arr()
+    fn diffie_hellman(pk: &Self::Pk, sk: &Self::Sk) -> GenericArray<u8, Self::PkLen> {
+        Self::serialize_pk(&(pk * sk))
+    }
+
+    fn zeroize_sk_on_drop(sk: &mut Self::Sk) {
+        sk.zeroize()
+    }
+
+    fn serialize_sk(sk: &Self::Sk) -> GenericArray<u8, Self::SkLen> {
+        sk.to_bytes().into()
+    }
+
+    fn deserialize_sk(bytes: &GenericArray<u8, Self::PkLen>) -> Result<Self::Sk, InternalError> {
+        Scalar::from_canonical_bytes((*bytes).into()).ok_or(InternalError::PointError)
+    }
+}
+
+#[cfg(feature = "ristretto255_voprf")]
+impl voprf::CipherSuite for Ristretto255 {
+    const ID: u16 = voprf::Ristretto255::ID;
+
+    type Group = <voprf::Ristretto255 as voprf::CipherSuite>::Group;
+
+    type Hash = <voprf::Ristretto255 as voprf::CipherSuite>::Hash;
+}
+
+impl Group for Ristretto255 {
+    type Elem = <voprf::Ristretto255 as Group>::Elem;
+
+    type ElemLen = <voprf::Ristretto255 as Group>::ElemLen;
+
+    type Scalar = <voprf::Ristretto255 as Group>::Scalar;
+
+    type ScalarLen = <voprf::Ristretto255 as Group>::ScalarLen;
+
+    fn hash_to_curve<CS: voprf::CipherSuite>(
+        input: &[&[u8]],
+        dst: &[u8],
+    ) -> voprf::Result<Self::Elem, voprf::InternalError>
+    where
+        <CS::Hash as OutputSizeUser>::OutputSize:
+            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
+    {
+        <voprf::Ristretto255 as Group>::hash_to_curve::<CS>(input, dst)
+    }
+
+    fn hash_to_scalar<CS: voprf::CipherSuite>(
+        input: &[&[u8]],
+        dst: &[u8],
+    ) -> voprf::Result<Self::Scalar, voprf::InternalError>
+    where
+        <CS::Hash as OutputSizeUser>::OutputSize:
+            IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
+    {
+        <voprf::Ristretto255 as Group>::hash_to_scalar::<CS>(input, dst)
+    }
+
+    fn base_elem() -> Self::Elem {
+        <voprf::Ristretto255 as Group>::base_elem()
+    }
+
+    fn identity_elem() -> Self::Elem {
+        <voprf::Ristretto255 as Group>::identity_elem()
+    }
+
+    fn serialize_elem(elem: Self::Elem) -> GenericArray<u8, Self::ElemLen> {
+        <voprf::Ristretto255 as Group>::serialize_elem(elem)
+    }
+
+    fn deserialize_elem(element_bits: &[u8]) -> voprf::Result<Self::Elem> {
+        <voprf::Ristretto255 as Group>::deserialize_elem(element_bits)
+    }
+
+    fn random_scalar<R: RngCore + CryptoRng>(rng: &mut R) -> Self::Scalar {
+        <voprf::Ristretto255 as Group>::random_scalar(rng)
+    }
+
+    fn invert_scalar(scalar: Self::Scalar) -> Self::Scalar {
+        <voprf::Ristretto255 as Group>::invert_scalar(scalar)
+    }
+
+    fn is_zero_scalar(scalar: Self::Scalar) -> subtle::Choice {
+        <voprf::Ristretto255 as Group>::is_zero_scalar(scalar)
+    }
+
+    fn serialize_scalar(scalar: Self::Scalar) -> GenericArray<u8, Self::ScalarLen> {
+        <voprf::Ristretto255 as Group>::serialize_scalar(scalar)
+    }
+
+    fn deserialize_scalar(scalar_bits: &[u8]) -> voprf::Result<Self::Scalar> {
+        <voprf::Ristretto255 as Group>::deserialize_scalar(scalar_bits)
     }
 }

src/key_exchange/group/x25519.rs

@@ -7,47 +7,97 @@
 
 //! Key Exchange group implementation for X25519
 
-use generic_array::typenum::U32;
+use curve25519_dalek_3::scalar::Scalar;
+use digest::core_api::BlockSizeUser;
+use digest::Digest;
+use elliptic_curve::hash2curve::{ExpandMsg, ExpandMsgXmd, Expander};
+use generic_array::typenum::{IsLess, IsLessOrEqual, U256, U32, U64};
 use generic_array::GenericArray;
 use rand::{CryptoRng, RngCore};
 use x25519_dalek::{PublicKey, StaticSecret};
+use zeroize::Zeroize;
 
 use super::KeGroup;
 use crate::errors::InternalError;
 
+/// Implementation for X25519.
+pub struct X25519;
+
 /// The implementation of such a subgroup for Ristretto
-impl KeGroup for PublicKey {
+impl KeGroup for X25519 {
+    type Pk = PublicKey;
     type PkLen = U32;
+    type Sk = StaticSecret;
     type SkLen = U32;
 
-    fn from_pk_slice(element_bits: &GenericArray<u8, Self::PkLen>) -> Result<Self, InternalError> {
-        Ok(Self::from(<[u8; 32]>::from(*element_bits)))
+    fn serialize_pk(pk: &Self::Pk) -> GenericArray<u8, Self::PkLen> {
+        pk.to_bytes().into()
     }
 
-    fn random_sk<R: RngCore + CryptoRng>(rng: &mut R) -> GenericArray<u8, Self::SkLen> {
+    fn deserialize_pk(bytes: &GenericArray<u8, Self::PkLen>) -> Result<Self::Pk, InternalError> {
+        if **bytes == [0; 32] {
+            Err(InternalError::PointError)
+        } else {
+            Ok(PublicKey::from(<[_; 32]>::from(*bytes)))
+        }
+    }
+
+    fn random_sk<R: RngCore + CryptoRng>(rng: &mut R) -> Self::Sk {
         let mut scalar_bytes = [0u8; 32];
 
         loop {
             rng.fill_bytes(&mut scalar_bytes);
 
             if scalar_bytes != [0u8; 32] {
-                break StaticSecret::from(scalar_bytes).to_bytes().into();
+                break StaticSecret::from(scalar_bytes);
             }
         }
     }
 
-    fn public_key(sk: &GenericArray<u8, Self::SkLen>) -> Self {
-        Self::from(&StaticSecret::from(<[u8; 32]>::from(*sk)))
+    // Implements the `HashToScalar()` function from
+    // https://www.ietf.org/archive/id/draft-irtf-cfrg-voprf-08.html#section-4.1
+    fn hash_to_scalar<'a, H>(input: &[&[u8]], dst: &[u8]) -> Result<Self::Sk, InternalError>
+    where
+        H: Digest + BlockSizeUser,
+        H::OutputSize: IsLess<U256> + IsLessOrEqual<H::BlockSize>,
+    {
+        let mut uniform_bytes = GenericArray::<_, U64>::default();
+        ExpandMsgXmd::<H>::expand_message(input, dst, 64)
+            .map_err(|_| InternalError::HashToScalar)?
+            .fill_bytes(&mut uniform_bytes);
+
+        Ok(StaticSecret::from(
+            Scalar::from_bytes_mod_order_wide(&uniform_bytes.into()).to_bytes(),
+        ))
     }
 
-    fn to_arr(&self) -> GenericArray<u8, Self::PkLen> {
-        self.to_bytes().into()
+    fn public_key(sk: &Self::Sk) -> Self::Pk {
+        PublicKey::from(sk)
     }
 
-    fn diffie_hellman(&self, sk: &GenericArray<u8, Self::SkLen>) -> GenericArray<u8, Self::SkLen> {
-        StaticSecret::from(<[u8; 32]>::from(*sk))
-            .diffie_hellman(self)
-            .to_bytes()
-            .into()
+    fn diffie_hellman(pk: &Self::Pk, sk: &Self::Sk) -> GenericArray<u8, Self::PkLen> {
+        sk.diffie_hellman(pk).to_bytes().into()
+    }
+
+    fn zeroize_sk_on_drop(sk: &mut Self::Sk) {
+        sk.zeroize()
+    }
+
+    fn serialize_sk(sk: &Self::Sk) -> GenericArray<u8, Self::SkLen> {
+        sk.to_bytes().into()
+    }
+
+    fn deserialize_sk(bytes: &GenericArray<u8, Self::PkLen>) -> Result<Self::Sk, InternalError> {
+        if **bytes == [0; 32] {
+            Err(InternalError::PointError)
+        } else {
+            let sk = StaticSecret::from(<[u8; 32]>::from(*bytes));
+
+            if sk.to_bytes() == **bytes {
+                Ok(sk)
+            } else {
+                Err(InternalError::PointError)
+            }
+        }
     }
 }

src/key_exchange/traits.rs

@@ -10,9 +10,9 @@ use digest::Output;
 use generic_array::typenum::{IsLess, Le, NonZero, U256};
 use generic_array::{ArrayLength, GenericArray};
 use rand::{CryptoRng, RngCore};
-use zeroize::Zeroize;
+use zeroize::ZeroizeOnDrop;
 
-use crate::ciphersuite::CipherSuite;
+use crate::ciphersuite::{CipherSuite, OprfHash};
 use crate::errors::ProtocolError;
 use crate::hash::{Hash, ProxyHash};
 use crate::key_exchange::group::KeGroup;
@@ -46,17 +46,17 @@ where
     <D::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
     Le<<D::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
 {
-    type KE1State: FromBytes + ToBytes + Zeroize + Clone;
-    type KE2State: FromBytes + ToBytes + Zeroize + Clone;
-    type KE1Message: FromBytes + ToBytes + Zeroize + Clone;
-    type KE2Message: FromBytes + ToBytes + Clone;
-    type KE3Message: FromBytes + ToBytes + Clone;
+    type KE1State: FromBytes + ToBytes + ZeroizeOnDrop + Clone;
+    type KE2State: FromBytes + ToBytes + ZeroizeOnDrop + Clone;
+    type KE1Message: FromBytes + ToBytes + ZeroizeOnDrop + Clone;
+    type KE2Message: FromBytes + ToBytes + ZeroizeOnDrop + Clone;
+    type KE3Message: FromBytes + ToBytes + ZeroizeOnDrop + Clone;
 
     fn generate_ke1<R: RngCore + CryptoRng>(
         rng: &mut R,
     ) -> Result<(Self::KE1State, Self::KE1Message), ProtocolError>;
 
-    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
+    #[allow(clippy::too_many_arguments)]
     fn generate_ke2<'a, 'b, 'c, 'd, R: RngCore + CryptoRng, S: SecretKey<G>>(
         rng: &mut R,
         l1_bytes: impl Iterator<Item = &'a [u8]>,
@@ -69,7 +69,7 @@ where
         context: &[u8],
     ) -> Result<GenerateKe2Result<Self, D, G>, ProtocolError<S::Error>>;
 
-    #[allow(clippy::too_many_arguments, clippy::type_complexity)]
+    #[allow(clippy::too_many_arguments)]
     fn generate_ke3<'a, 'b, 'c, 'd>(
         l2_component: impl Iterator<Item = &'a [u8]>,
         ke2_message: Self::KE2Message,
@@ -82,13 +82,10 @@ where
         context: &[u8],
     ) -> Result<GenerateKe3Result<Self, D, G>, ProtocolError>;
 
-    #[allow(clippy::type_complexity)]
     fn finish_ke(
         ke3_message: Self::KE3Message,
         ke2_state: &Self::KE2State,
     ) -> Result<Output<D>, ProtocolError>;
-
-    fn ke2_message_size() -> usize;
 }
 
 pub trait FromBytes: Sized {
@@ -101,14 +98,13 @@ pub trait ToBytes {
     fn to_bytes(&self) -> GenericArray<u8, Self::Len>;
 }
 
-#[allow(dead_code)]
 pub type Ke1StateLen<CS: CipherSuite> =
-    <<CS::KeyExchange as KeyExchange<CS::Hash, CS::KeGroup>>::KE1State as ToBytes>::Len;
+    <<CS::KeyExchange as KeyExchange<OprfHash<CS>, CS::KeGroup>>::KE1State as ToBytes>::Len;
 pub type Ke1MessageLen<CS: CipherSuite> =
-    <<CS::KeyExchange as KeyExchange<CS::Hash, CS::KeGroup>>::KE1Message as ToBytes>::Len;
+    <<CS::KeyExchange as KeyExchange<OprfHash<CS>, CS::KeGroup>>::KE1Message as ToBytes>::Len;
 pub type Ke2StateLen<CS: CipherSuite> =
-    <<CS::KeyExchange as KeyExchange<CS::Hash, CS::KeGroup>>::KE2State as ToBytes>::Len;
+    <<CS::KeyExchange as KeyExchange<OprfHash<CS>, CS::KeGroup>>::KE2State as ToBytes>::Len;
 pub type Ke2MessageLen<CS: CipherSuite> =
-    <<CS::KeyExchange as KeyExchange<CS::Hash, CS::KeGroup>>::KE2Message as ToBytes>::Len;
+    <<CS::KeyExchange as KeyExchange<OprfHash<CS>, CS::KeGroup>>::KE2Message as ToBytes>::Len;
 pub type Ke3MessageLen<CS: CipherSuite> =
-    <<CS::KeyExchange as KeyExchange<CS::Hash, CS::KeGroup>>::KE3Message as ToBytes>::Len;
+    <<CS::KeyExchange as KeyExchange<OprfHash<CS>, CS::KeGroup>>::KE3Message as ToBytes>::Len;

src/key_exchange/tripledh.rs

@@ -9,7 +9,7 @@
 use core::convert::TryFrom;
 use core::ops::Add;
 
-use derive_where::DeriveWhere;
+use derive_where::derive_where;
 use digest::core_api::BlockSizeUser;
 use digest::{Digest, Output};
 use generic_array::sequence::Concat;
@@ -18,6 +18,7 @@ use generic_array::{ArrayLength, GenericArray};
 use hkdf::{Hkdf, HkdfExtract};
 use hmac::{Hmac, Mac};
 use rand::{CryptoRng, RngCore};
+use zeroize::{Zeroize, ZeroizeOnDrop};
 
 use crate::errors::utils::{check_slice_size, check_slice_size_atleast};
 use crate::errors::{InternalError, ProtocolError};
@@ -55,36 +56,63 @@ pub struct TripleDH;
 #[cfg_attr(
     feature = "serde",
     derive(serde_::Deserialize, serde_::Serialize),
-    serde(bound = "", crate = "serde_")
+    serde(
+        bound(
+            deserialize = "KG::Sk: serde_::Deserialize<'de>",
+            serialize = "KG::Sk: serde_::Serialize",
+        ),
+        crate = "serde_"
+    )
 )]
-#[derive(DeriveWhere)]
-#[derive_where(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Zeroize(drop))]
+#[derive_where(Clone)]
+#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; KG::Sk)]
 pub struct Ke1State<KG: KeGroup> {
     client_e_sk: PrivateKey<KG>,
     client_nonce: GenericArray<u8, NonceLen>,
 }
 
+impl<KG: KeGroup> Drop for Ke1State<KG> {
+    fn drop(&mut self) {
+        self.client_nonce.zeroize();
+    }
+}
+
+impl<KG: KeGroup> ZeroizeOnDrop for Ke1State<KG> {}
+
 /// The first key exchange message
 #[cfg_attr(
     feature = "serde",
     derive(serde_::Deserialize, serde_::Serialize),
-    serde(bound = "", crate = "serde_")
+    serde(
+        bound(
+            deserialize = "KG::Pk: serde_::Deserialize<'de>",
+            serialize = "KG::Pk: serde_::Serialize",
+        ),
+        crate = "serde_"
+    )
 )]
-#[derive(DeriveWhere)]
-#[derive_where(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Zeroize)]
+#[derive_where(Clone)]
+#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; KG::Pk)]
 pub struct Ke1Message<KG: KeGroup> {
     pub(crate) client_nonce: GenericArray<u8, NonceLen>,
     pub(crate) client_e_pk: PublicKey<KG>,
 }
 
+impl<KG: KeGroup> Drop for Ke1Message<KG> {
+    fn drop(&mut self) {
+        self.client_nonce.zeroize();
+    }
+}
+
+impl<KG: KeGroup> ZeroizeOnDrop for Ke1Message<KG> {}
+
 /// The server state produced after the second key exchange message
 #[cfg_attr(
     feature = "serde",
     derive(serde_::Deserialize, serde_::Serialize),
     serde(bound = "", crate = "serde_")
 )]
-#[derive(DeriveWhere)]
-#[derive_where(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Zeroize(drop))]
+#[derive_where(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
 pub struct Ke2State<D: Hash>
 where
     D::Core: ProxyHash,
@@ -96,14 +124,41 @@ where
     session_key: Output<D>,
 }
 
+impl<D: Hash> Drop for Ke2State<D>
+where
+    D::Core: ProxyHash,
+    <D::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<D::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+{
+    fn drop(&mut self) {
+        self.km3.zeroize();
+        self.hashed_transcript.zeroize();
+        self.session_key.zeroize();
+    }
+}
+
+impl<D: Hash> ZeroizeOnDrop for Ke2State<D>
+where
+    D::Core: ProxyHash,
+    <D::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<D::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+{
+}
+
 /// The second key exchange message
 #[cfg_attr(
     feature = "serde",
     derive(serde_::Deserialize, serde_::Serialize),
-    serde(bound = "", crate = "serde_")
+    serde(
+        bound(
+            deserialize = "KG::Pk: serde_::Deserialize<'de>",
+            serialize = "KG::Pk: serde_::Serialize",
+        ),
+        crate = "serde_"
+    )
 )]
-#[derive(DeriveWhere)]
-#[derive_where(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
+#[derive_where(Clone)]
+#[derive_where(Debug, Eq, Hash, Ord, PartialEq, PartialOrd; KG::Pk)]
 pub struct Ke2Message<D: Hash, KG: KeGroup>
 where
     D::Core: ProxyHash,
@@ -115,13 +170,32 @@ where
     mac: Output<D>,
 }
 
+impl<D: Hash, KG: KeGroup> Drop for Ke2Message<D, KG>
+where
+    D::Core: ProxyHash,
+    <D::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<D::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+{
+    fn drop(&mut self) {
+        self.server_nonce.zeroize();
+        self.mac.zeroize();
+    }
+}
+
+impl<D: Hash, KG: KeGroup> ZeroizeOnDrop for Ke2Message<D, KG>
+where
+    D::Core: ProxyHash,
+    <D::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<D::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+{
+}
+
 /// The third key exchange message
 #[cfg_attr(
     feature = "serde",
     derive(serde_::Deserialize, serde_::Serialize),
     serde(bound = "", crate = "serde_")
 )]
-#[derive(DeriveWhere)]
 #[derive_where(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
 pub struct Ke3Message<D: Hash>
 where
@@ -132,6 +206,25 @@ where
     mac: Output<D>,
 }
 
+impl<D: Hash> Drop for Ke3Message<D>
+where
+    D::Core: ProxyHash,
+    <D::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<D::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+{
+    fn drop(&mut self) {
+        self.mac.zeroize();
+    }
+}
+
+impl<D: Hash> ZeroizeOnDrop for Ke3Message<D>
+where
+    D::Core: ProxyHash,
+    <D::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<D::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+{
+}
+
 ////////////////////////////////
 // High-level Implementations //
 // ========================== //
@@ -210,13 +303,13 @@ where
             .chain_iter(id_s.into_iter())
             .chain_iter(l2_bytes)
             .chain(server_nonce)
-            .chain(&server_e_kp.public().to_arr());
+            .chain(&server_e_kp.public().to_bytes());
 
         let result = derive_3dh_keys::<D, KG, S>(
             TripleDHComponents {
                 pk1: ke1_message.client_e_pk.clone(),
                 sk1: server_e_kp.private().clone(),
-                pk2: ke1_message.client_e_pk,
+                pk2: ke1_message.client_e_pk.clone(),
                 sk2: server_s_sk,
                 pk3: client_s_pk,
                 sk3: server_e_kp.private().clone(),
@@ -268,7 +361,7 @@ where
             .chain_iter(serialized_credential_request)
             .chain_iter(id_s)
             .chain_iter(l2_component)
-            .chain_iter(ke2_message.to_bytes_without_info_or_mac());
+            .chain(ke2_message.to_bytes_without_mac());
 
         let result = derive_3dh_keys::<D, KG, PrivateKey<KG>>(
             TripleDHComponents {
@@ -322,10 +415,6 @@ where
 
         Ok(ke2_state.session_key.clone())
     }
-
-    fn ke2_message_size() -> usize {
-        NonceLen::USIZE + <KG as KeGroup>::PkLen::USIZE + OutputSize::<D>::USIZE
-    }
 }
 
 /////////////////////////
@@ -489,7 +578,7 @@ impl<KG: KeGroup> FromBytes for Ke1State<KG> {
         let checked_bytes = check_slice_size_atleast(bytes, key_len + nonce_len, "ke1_state")?;
 
         Ok(Self {
-            client_e_sk: PrivateKey::from_bytes(&checked_bytes[..key_len])?,
+            client_e_sk: PrivateKey::deserialize(&checked_bytes[..key_len])?,
             client_nonce: GenericArray::clone_from_slice(
                 &checked_bytes[key_len..key_len + nonce_len],
             ),
@@ -506,7 +595,7 @@ where
     type Len = Sum<KG::SkLen, NonceLen>;
 
     fn to_bytes(&self) -> GenericArray<u8, Self::Len> {
-        self.client_e_sk.to_arr().concat(self.client_nonce)
+        self.client_e_sk.serialize().concat(self.client_nonce)
     }
 }
 
@@ -521,7 +610,7 @@ impl<KG: KeGroup> FromBytes for Ke1Message<KG> {
 
         Ok(Self {
             client_nonce: GenericArray::clone_from_slice(&checked_nonce[..nonce_len]),
-            client_e_pk: PublicKey::from_bytes(&checked_nonce[nonce_len..])?,
+            client_e_pk: PublicKey::deserialize(&checked_nonce[nonce_len..])?,
         })
     }
 }
@@ -535,7 +624,7 @@ where
     type Len = Sum<NonceLen, KG::PkLen>;
 
     fn to_bytes(&self) -> GenericArray<u8, Self::Len> {
-        self.client_nonce.concat(self.client_e_pk.to_arr())
+        self.client_nonce.concat(self.client_e_pk.to_bytes())
     }
 }
 
@@ -602,13 +691,11 @@ where
         )?;
 
         // Check the public key bytes
-        let server_e_pk = KeyPair::<KG>::check_public_key(PublicKey::from_bytes(
-            &unchecked_server_e_pk[..key_len],
-        )?)?;
+        let server_e_pk = PublicKey::deserialize(&unchecked_server_e_pk[..key_len])?;
 
         Ok(Self {
             server_nonce: GenericArray::clone_from_slice(&checked_nonce[..nonce_len]),
-            server_e_pk: PublicKey::from_bytes(&server_e_pk)?,
+            server_e_pk,
             mac: GenericArray::clone_from_slice(checked_mac),
         })
     }
@@ -628,7 +715,7 @@ where
 
     fn to_bytes(&self) -> GenericArray<u8, Self::Len> {
         self.server_nonce
-            .concat(self.server_e_pk.to_arr())
+            .concat(self.server_e_pk.to_bytes())
             .concat(self.mac.clone())
     }
 }
@@ -638,9 +725,11 @@ where
     D::Core: ProxyHash,
     <D::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
     Le<<D::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+    NonceLen: Add<KG::PkLen>,
+    Sum<NonceLen, KG::PkLen>: ArrayLength<u8>,
 {
-    fn to_bytes_without_info_or_mac(&self) -> impl Iterator<Item = &[u8]> {
-        [self.server_nonce.as_slice(), self.server_e_pk.as_slice()].into_iter()
+    fn to_bytes_without_mac(&self) -> GenericArray<u8, Sum<NonceLen, KG::PkLen>> {
+        self.server_nonce.concat(self.server_e_pk.to_bytes())
     }
 }