Publicly Expose trait KeyExchange (#379)

* Clean up types used in `trait KeyExchange`

* Expose `trait KeyExchange` publicly

src/ciphersuite.rs

@@ -18,7 +18,7 @@ use generic_array::ArrayLength;
 use crate::envelope::NonceLen;
 use crate::hash::{Hash, OutputSize, ProxyHash};
 use crate::key_exchange::group::Group;
-use crate::key_exchange::traits::KeyExchange;
+use crate::key_exchange::KeyExchange;
 use crate::ksf::Ksf;
 use crate::opaque::MaskedResponseLen;
 

src/envelope.rs

@@ -22,7 +22,7 @@ use crate::ciphersuite::{CipherSuite, KeGroup, OprfHash};
 use crate::errors::{InternalError, ProtocolError};
 use crate::hash::OutputSize;
 use crate::key_exchange::group::Group;
-use crate::key_exchange::traits::SerializedIdentifiers;
+use crate::key_exchange::SerializedIdentifiers;
 use crate::keypair::{KeyPair, PrivateKey, PublicKey};
 use crate::opaque::Identifiers;
 use crate::serialization::{GenericArrayExt, SliceExt, UpdateExt};

src/key_exchange/mod.rs

@@ -12,7 +12,434 @@
 pub mod group;
 pub(crate) mod shared;
 pub mod sigma_i;
-pub(crate) mod traits;
 pub mod tripledh;
 
-pub use crate::key_exchange::traits::KeyExchange;
+use core::iter;
+use core::ops::Add;
+
+use derive_where::derive_where;
+use digest::core_api::{BlockSizeUser, CoreProxy};
+use digest::Output;
+use generic_array::sequence::Concat;
+use generic_array::typenum::{IsLess, Le, NonZero, Sum, U2, U256};
+use generic_array::{ArrayLength, GenericArray};
+use rand::{CryptoRng, RngCore};
+use voprf::{BlindedElement, EvaluationElement};
+use zeroize::{Zeroize, ZeroizeOnDrop};
+
+#[cfg(test)]
+use crate::ciphersuite::KeHash;
+use crate::ciphersuite::{CipherSuite, OprfGroup};
+use crate::errors::ProtocolError;
+use crate::hash::{Hash, ProxyHash};
+use crate::key_exchange::group::Group;
+use crate::key_exchange::shared::{NonceLen, STR_CONTEXT};
+use crate::keypair::{PrivateKey, PublicKey};
+use crate::opaque::{Identifiers, MaskedResponse, MaskedResponseLen};
+use crate::serialization::{i2osp, SliceExt};
+
+/// The key exchange trait.
+pub trait KeyExchange
+where
+    <Self::Hash as CoreProxy>::Core: ProxyHash,
+    <<Self::Hash as CoreProxy>::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<<Self::Hash as CoreProxy>::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+{
+    /// The group used for the key exchange.
+    type Group: Group;
+    /// The hash used for the key exchange.
+    type Hash: Hash;
+
+    /// Client state.
+    type KE1State: ZeroizeOnDrop + Clone;
+    /// Server state.
+    type KE2State<CS: CipherSuite>: ZeroizeOnDrop + Clone;
+    /// First message sent by the client.
+    type KE1Message: ZeroizeOnDrop + Clone;
+    /// Server state builder.
+    type KE2Builder<'a, CS: CipherSuite<KeyExchange = Self>>: ZeroizeOnDrop + Clone;
+    /// Server data for the remote key interaction.
+    type KE2BuilderData<'a, CS: 'static + CipherSuite>;
+    /// Server remote key input.
+    type KE2BuilderInput<CS: CipherSuite>;
+    /// Message sent by the server.
+    type KE2Message: ZeroizeOnDrop + Clone;
+    /// Second message sent by the client.
+    type KE3Message: ZeroizeOnDrop + Clone;
+
+    /// Client generates [`KE1Message`](Self::KE1Message) and
+    /// [`KE1State`](Self::KE1State).
+    fn generate_ke1<R: RngCore + CryptoRng>(
+        rng: &mut R,
+    ) -> Result<GenerateKe1Result<Self>, ProtocolError>;
+
+    /// Server generates [`KE2Builder`](Self::KE2Builder).
+    fn ke2_builder<'a, CS: CipherSuite<KeyExchange = Self>, R: RngCore + CryptoRng>(
+        rng: &mut R,
+        credential_request: SerializedCredentialRequest<CS>,
+        ke1_message: Self::KE1Message,
+        credential_response: SerializedCredentialResponse<CS>,
+        client_s_pk: PublicKey<Self::Group>,
+        identifiers: SerializedIdentifiers<'a, Self::Group>,
+        context: SerializedContext<'a>,
+    ) -> Result<Self::KE2Builder<'a, CS>, ProtocolError>;
+
+    /// Server returns the data for the remote key interaction.
+    fn ke2_builder_data<'a, CS: CipherSuite<KeyExchange = Self>>(
+        builder: &'a Self::KE2Builder<'_, CS>,
+    ) -> Self::KE2BuilderData<'a, CS>;
+
+    /// Server generates the input without a remote key.
+    fn generate_ke2_input<CS: CipherSuite<KeyExchange = Self>, R: CryptoRng + RngCore>(
+        builder: &Self::KE2Builder<'_, CS>,
+        rng: &mut R,
+        server_s_sk: &PrivateKey<Self::Group>,
+    ) -> Self::KE2BuilderInput<CS>;
+
+    /// Server generates [`KE2Message`](Self::KE2Message) and
+    /// [`KE2State`](Self::KE2State).
+    fn build_ke2<CS: CipherSuite<KeyExchange = Self>>(
+        builder: Self::KE2Builder<'_, CS>,
+        input: Self::KE2BuilderInput<CS>,
+    ) -> Result<GenerateKe2Result<CS>, ProtocolError>;
+
+    /// Client generates [`KE3Message`](Self::KE3Message) and the session key.
+    #[allow(clippy::too_many_arguments)]
+    fn generate_ke3<CS: CipherSuite<KeyExchange = Self>, R: CryptoRng + RngCore>(
+        rng: &mut R,
+        credential_request: SerializedCredentialRequest<CS>,
+        ke1_message: Self::KE1Message,
+        credential_response: SerializedCredentialResponse<CS>,
+        ke1_state: &Self::KE1State,
+        ke2_message: Self::KE2Message,
+        server_s_pk: PublicKey<Self::Group>,
+        client_s_sk: PrivateKey<Self::Group>,
+        identifiers: SerializedIdentifiers<'_, Self::Group>,
+        context: SerializedContext<'_>,
+    ) -> Result<GenerateKe3Result<Self>, ProtocolError>;
+
+    /// Server generates the session key.
+    fn finish_ke<CS: CipherSuite<KeyExchange = Self>>(
+        ke2_state: &Self::KE2State<CS>,
+        ke3_message: Self::KE3Message,
+        identifiers: Identifiers<'_>,
+        context: SerializedContext<'_>,
+    ) -> Result<Output<Self::Hash>, ProtocolError>;
+}
+
+/// Serialized form of [`CredentialRequest`](crate::CredentialRequest).
+#[cfg_attr(
+    feature = "serde",
+    derive(serde::Deserialize, serde::Serialize),
+    serde(bound = "")
+)]
+#[derive_where(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Zeroize)]
+pub struct SerializedCredentialRequest<CS: CipherSuite>(
+    GenericArray<u8, <OprfGroup<CS> as voprf::Group>::ElemLen>,
+);
+
+impl<CS: CipherSuite> SerializedCredentialRequest<CS> {
+    pub(crate) fn new(blinded_element: &BlindedElement<CS::OprfCs>) -> Self {
+        Self(blinded_element.serialize())
+    }
+
+    /// Returns the serialized form of
+    /// [`CredentialRequest`](crate::CredentialRequest) in multiple byte slices.
+    pub fn iter(&self) -> impl Clone + Iterator<Item = &[u8]> {
+        iter::once(self.0.as_slice())
+    }
+
+    /// Returns a [`SerializedCredentialRequest`] deserialized from the given
+    /// `bytes`.
+    pub fn deserialize_take(bytes: &mut &[u8]) -> Result<Self, ProtocolError> {
+        Ok(Self(bytes.take_array("blinded element")?))
+    }
+}
+
+type SerializedCredentialRequestLen<CS: CipherSuite> = <OprfGroup<CS> as voprf::Group>::ElemLen;
+
+impl<CS: CipherSuite> Serialize for SerializedCredentialRequest<CS> {
+    type Len = SerializedCredentialRequestLen<CS>;
+
+    fn serialize(&self) -> GenericArray<u8, Self::Len> {
+        self.0.clone()
+    }
+}
+
+/// Serialized form of [`CredentialResponse`](crate::CredentialResponse).
+#[cfg_attr(
+    feature = "serde",
+    derive(serde::Deserialize, serde::Serialize),
+    serde(bound = "")
+)]
+#[derive_where(Clone, Debug, Eq, Hash, PartialEq, Zeroize)]
+pub struct SerializedCredentialResponse<CS: CipherSuite> {
+    evaluation_element: GenericArray<u8, <OprfGroup<CS> as voprf::Group>::ElemLen>,
+    masking_nonce: GenericArray<u8, NonceLen>,
+    masked_response: MaskedResponse<CS>,
+}
+
+impl<CS: CipherSuite> SerializedCredentialResponse<CS> {
+    pub(crate) fn new(
+        evaluation_element: &EvaluationElement<CS::OprfCs>,
+        masking_nonce: GenericArray<u8, NonceLen>,
+        masked_response: MaskedResponse<CS>,
+    ) -> Self {
+        Self {
+            evaluation_element: evaluation_element.serialize(),
+            masking_nonce,
+            masked_response,
+        }
+    }
+
+    /// Returns the serialized form of
+    /// [`CredentialResponse`](crate::CredentialResponse) in multiple byte
+    /// slices.
+    pub fn iter(&self) -> impl Clone + Iterator<Item = &[u8]> {
+        [self.evaluation_element.as_slice(), &self.masking_nonce]
+            .into_iter()
+            .chain(self.masked_response.iter())
+    }
+
+    /// Returns a [`SerializedCredentialRequest`] deserialized from the given
+    /// `bytes`.
+    pub fn deserialize_take(input: &mut &[u8]) -> Result<Self, ProtocolError> {
+        Ok(Self {
+            evaluation_element: input.take_array("evaluation element")?,
+            masking_nonce: input.take_array("masking nonce")?,
+            masked_response: MaskedResponse::deserialize_take(input)?,
+        })
+    }
+}
+
+type SerializedCredentialResponseLen<CS: CipherSuite> =
+    Sum<Sum<<OprfGroup<CS> as voprf::Group>::ElemLen, NonceLen>, MaskedResponseLen<CS>>;
+
+impl<CS: CipherSuite> Serialize for SerializedCredentialResponse<CS>
+where
+    <OprfGroup<CS> as voprf::Group>::ElemLen: Add<NonceLen>,
+    Sum<<OprfGroup<CS> as voprf::Group>::ElemLen, NonceLen>:
+        ArrayLength<u8> + Add<MaskedResponseLen<CS>>,
+    SerializedCredentialResponseLen<CS>: ArrayLength<u8>,
+{
+    type Len = SerializedCredentialResponseLen<CS>;
+
+    fn serialize(&self) -> GenericArray<u8, Self::Len> {
+        self.evaluation_element
+            .clone()
+            .concat(self.masking_nonce)
+            .concat(self.masked_response.serialize())
+    }
+}
+
+/// Serialized form of a `context` given in
+/// [`ClientLoginFinishParameters`](crate::ClientLoginFinishParameters) or
+/// [`ServerLoginParameters`](crate::ServerLoginParameters).
+#[cfg_attr(
+    feature = "serde",
+    derive(serde::Deserialize, serde::Serialize),
+    serde(bound = "")
+)]
+#[derive(Clone, Debug, Eq, Hash, PartialEq, Zeroize)]
+pub struct SerializedContext<'a> {
+    length: GenericArray<u8, U2>,
+    #[zeroize(skip)]
+    context: &'a [u8],
+}
+
+impl<'a> SerializedContext<'a> {
+    pub(crate) fn from(context: Option<&'a [u8]>) -> Result<Self, ProtocolError> {
+        let context = context.unwrap_or(&[]);
+
+        Ok(Self {
+            length: i2osp::<U2>(context.len())?,
+            context,
+        })
+    }
+
+    /// Returns the serialized form of `context` in multiple byte slices.
+    pub fn iter(&self) -> impl Clone + Iterator<Item = &[u8]> {
+        iter::once(STR_CONTEXT).chain([self.length.as_slice(), self.context])
+    }
+}
+
+/// Serialized form of [`Identifiers`](crate::Identifiers).
+#[cfg_attr(
+    feature = "serde",
+    derive(serde::Deserialize, serde::Serialize),
+    serde(bound(deserialize = "'de: 'a", serialize = ""))
+)]
+#[derive_where(Clone, Debug, Eq, Hash, PartialEq, Zeroize)]
+pub struct SerializedIdentifiers<'a, G: Group> {
+    /// Client identifiers.
+    pub client: SerializedIdentifier<'a, G>,
+    /// Server identifiers.
+    pub server: SerializedIdentifier<'a, G>,
+}
+
+/// Serialized form of a single identifier from
+/// [`Identifiers`](crate::Identifiers).
+#[cfg_attr(
+    feature = "serde",
+    derive(serde::Deserialize, serde::Serialize),
+    serde(bound(deserialize = "'de: 'a", serialize = ""))
+)]
+#[derive_where(Clone, Debug, Eq, Hash, PartialEq, Zeroize)]
+pub struct SerializedIdentifier<'a, G: Group> {
+    length: GenericArray<u8, U2>,
+    identifier: Identifier<'a, G>,
+}
+
+#[cfg_attr(
+    feature = "serde",
+    derive(serde::Deserialize, serde::Serialize),
+    serde(bound = "")
+)]
+#[derive_where(Clone, Debug, Eq, Hash, PartialEq, Zeroize)]
+enum Identifier<'a, G: Group> {
+    Owned(GenericArray<u8, G::PkLen>),
+    #[derive_where(skip_inner(Zeroize))]
+    Borrowed(&'a [u8]),
+}
+
+impl<'a, G: Group> SerializedIdentifiers<'a, G> {
+    pub(crate) fn from_identifiers(
+        ids: Identifiers<'a>,
+        client_s_pk: GenericArray<u8, G::PkLen>,
+        server_s_pk: GenericArray<u8, G::PkLen>,
+    ) -> Result<Self, ProtocolError> {
+        let client = SerializedIdentifier::from_identifier(ids.client, client_s_pk)?;
+        let server = SerializedIdentifier::from_identifier(ids.server, server_s_pk)?;
+
+        Ok(Self { client, server })
+    }
+}
+
+impl<'a, G: Group> SerializedIdentifier<'a, G> {
+    /// Creates a [`SerializedIdentifier`] an identifier or the corresponding
+    /// static public key.
+    pub fn from_identifier(
+        id: Option<&'a [u8]>,
+        s_pk: GenericArray<u8, G::PkLen>,
+    ) -> Result<Self, ProtocolError> {
+        if let Some(id) = id {
+            Ok(SerializedIdentifier {
+                length: i2osp::<U2>(id.len())?,
+                identifier: Identifier::Borrowed(id),
+            })
+        } else {
+            Ok(SerializedIdentifier {
+                length: i2osp::<U2>(s_pk.len())?,
+                identifier: Identifier::Owned(s_pk),
+            })
+        }
+    }
+
+    /// Returns the serialized form of an identifier in multiple byte slices.
+    pub fn iter(&self) -> impl Clone + Iterator<Item = &[u8]> {
+        [self.length.as_slice()]
+            .into_iter()
+            .chain(match &self.identifier {
+                Identifier::Owned(bytes) => [bytes.as_slice()],
+                Identifier::Borrowed(bytes) => [*bytes],
+            })
+    }
+}
+
+/// Deserialization trait for key exchange types.
+pub trait Deserialize: Sized {
+    /// Deserialize [`Self`] from the given `bytes`.
+    ///
+    /// The deserialized bytes must be taken from `bytes`.
+    fn deserialize_take(bytes: &mut &[u8]) -> Result<Self, ProtocolError>;
+}
+
+/// Serialization trait for key exchange types.
+pub trait Serialize {
+    /// The length of the serialized types.
+    type Len: ArrayLength<u8>;
+
+    /// Serialize [`Self`] to a fixed-length byte array.
+    fn serialize(&self) -> GenericArray<u8, Self::Len>;
+}
+
+/// Result type of [`KeyExchange::generate_ke1()`].
+pub struct GenerateKe1Result<KE: KeyExchange + ?Sized> {
+    /// The client state.
+    pub state: KE::KE1State,
+    /// The first client message.
+    pub message: KE::KE1Message,
+}
+
+/// Result type of [`KeyExchange::build_ke2()`].
+pub struct GenerateKe2Result<CS: CipherSuite> {
+    /// The server state.
+    pub state: <CS::KeyExchange as KeyExchange>::KE2State<CS>,
+    /// The server message.
+    pub message: <CS::KeyExchange as KeyExchange>::KE2Message,
+    #[cfg(test)]
+    pub(crate) handshake_secret: Output<KeHash<CS>>,
+    #[cfg(test)]
+    pub(crate) km2: Output<KeHash<CS>>,
+}
+
+/// Result type of [`KeyExchange::generate_ke3()`].
+pub struct GenerateKe3Result<KE: KeyExchange + ?Sized> {
+    /// The session key.
+    pub session_key: Output<KE::Hash>,
+    /// The second client message.
+    pub message: KE::KE3Message,
+    #[cfg(test)]
+    pub(crate) handshake_secret: Output<KE::Hash>,
+    #[cfg(test)]
+    pub(crate) km3: Output<KE::Hash>,
+}
+
+pub(crate) type Ke1StateLen<CS: CipherSuite> =
+    <<CS::KeyExchange as KeyExchange>::KE1State as Serialize>::Len;
+pub(crate) type Ke1MessageLen<CS: CipherSuite> =
+    <<CS::KeyExchange as KeyExchange>::KE1Message as Serialize>::Len;
+pub(crate) type Ke2StateLen<CS: CipherSuite> =
+    <<CS::KeyExchange as KeyExchange>::KE2State<CS> as Serialize>::Len;
+pub(crate) type Ke2MessageLen<CS: CipherSuite> =
+    <<CS::KeyExchange as KeyExchange>::KE2Message as Serialize>::Len;
+pub(crate) type Ke3MessageLen<CS: CipherSuite> =
+    <<CS::KeyExchange as KeyExchange>::KE3Message as Serialize>::Len;
+
+//////////////////////////
+// Test Implementations //
+//===================== //
+//////////////////////////
+
+#[cfg(test)]
+use crate::serialization::AssertZeroized;
+
+#[cfg(test)]
+impl<CS: CipherSuite> AssertZeroized for SerializedCredentialRequest<CS> {
+    fn assert_zeroized(&self) {
+        let Self(blinded_element) = self;
+
+        for byte in blinded_element.iter() {
+            assert_eq!(byte, &0);
+        }
+    }
+}
+
+#[cfg(test)]
+impl<CS: CipherSuite> AssertZeroized for SerializedCredentialResponse<CS> {
+    fn assert_zeroized(&self) {
+        let Self {
+            evaluation_element,
+            masking_nonce,
+            masked_response,
+        } = self;
+
+        for byte in evaluation_element
+            .iter()
+            .chain(masking_nonce)
+            .chain(masked_response.iter().flatten())
+        {
+            assert_eq!(byte, &0);
+        }
+    }
+}

src/key_exchange/shared.rs

@@ -17,14 +17,14 @@ use generic_array::{ArrayLength, GenericArray};
 use hkdf::{Hkdf, HkdfExtract};
 use rand::{CryptoRng, RngCore};
 
+use super::{
+    Deserialize, GenerateKe1Result, KeyExchange, Serialize, SerializedContext,
+    SerializedCredentialRequest, SerializedCredentialResponse, SerializedIdentifiers,
+};
 use crate::ciphersuite::{CipherSuite, KeGroup, KeHash};
 use crate::errors::{InternalError, ProtocolError};
 use crate::hash::{Hash, OutputSize, ProxyHash};
 use crate::key_exchange::group::Group;
-use crate::key_exchange::traits::{
-    CredentialRequestParts, CredentialResponseParts, Deserialize, Serialize, SerializedContext,
-    SerializedIdentifiers,
-};
 use crate::keypair::{KeyPair, PrivateKey, PublicKey};
 use crate::serialization::{i2osp, SliceExt, UpdateExt};
 
@@ -100,9 +100,13 @@ pub(super) struct DerivedKeys<H: OutputSizeUser> {
 
 // Helper functions
 
-pub(super) fn generate_ke1<R: RngCore + CryptoRng, G: Group>(
+pub(super) fn generate_ke1<
+    R: RngCore + CryptoRng,
+    KE: KeyExchange<KE1State = Ke1State<G>, KE1Message = Ke1Message<G>>,
+    G: Group,
+>(
     rng: &mut R,
-) -> Result<(Ke1State<G>, Ke1Message<G>), ProtocolError> {
+) -> Result<GenerateKe1Result<KE>, ProtocolError> {
     let client_e_kp = KeyPair::<G>::derive_random(rng);
     let client_nonce = generate_nonce::<R>(rng);
 
@@ -111,13 +115,13 @@ pub(super) fn generate_ke1<R: RngCore + CryptoRng, G: Group>(
         client_e_pk: client_e_kp.public().clone(),
     };
 
-    Ok((
-        Ke1State {
+    Ok(GenerateKe1Result {
+        state: Ke1State {
             client_e_sk: client_e_kp.private().clone(),
             client_nonce,
         },
-        ke1_message,
-    ))
+        message: ke1_message,
+    })
 }
 
 // Generate a random nonce up to NonceLen::USIZE bytes.
@@ -130,9 +134,9 @@ pub(super) fn generate_nonce<R: RngCore + CryptoRng>(rng: &mut R) -> GenericArra
 pub(super) fn transcript<CS: CipherSuite, KE: Group>(
     context: &SerializedContext<'_>,
     identifiers: &SerializedIdentifiers<'_, KeGroup<CS>>,
-    credential_request: &CredentialRequestParts<CS>,
+    credential_request: &SerializedCredentialRequest<CS>,
     ke1_message: &Ke1MessageIter<KE>,
-    credential_response: &CredentialResponseParts<CS>,
+    credential_response: &SerializedCredentialResponse<CS>,
     server_nonce: GenericArray<u8, NonceLen>,
     server_e_pk: &GenericArray<u8, KE::PkLen>,
 ) -> KeHash<CS> {

src/key_exchange/sigma_i/message.rs

@@ -20,10 +20,10 @@ use crate::errors::ProtocolError;
 use crate::hash::OutputSize;
 use crate::key_exchange::group::Group;
 use crate::key_exchange::shared::{Ke1MessageIter, Ke1MessageIterLen, NonceLen};
-use crate::key_exchange::traits::{
-    CredentialRequestParts, CredentialRequestPartsLen, CredentialResponseParts,
-    CredentialResponsePartsLen, Deserialize, Serialize, SerializedContext, SerializedIdentifier,
-    SerializedIdentifiers,
+use crate::key_exchange::{
+    Deserialize, Serialize, SerializedContext, SerializedCredentialRequest,
+    SerializedCredentialRequestLen, SerializedCredentialResponse, SerializedCredentialResponseLen,
+    SerializedIdentifier, SerializedIdentifiers,
 };
 use crate::opaque::MaskedResponseLen;
 use crate::serialization::{SliceExt, UpdateExt};
@@ -99,9 +99,9 @@ pub struct MessageBuilder<'a, CS: CipherSuite> {
 )]
 #[derive_where(Clone, Debug, Eq, Hash, PartialEq, Zeroize, ZeroizeOnDrop)]
 pub struct CachedMessage<CS: CipherSuite, KE: Group> {
-    pub(super) credential_request: CredentialRequestParts<CS>,
+    pub(super) credential_request: SerializedCredentialRequest<CS>,
     pub(super) ke1_message: Ke1MessageIter<KE>,
-    pub(super) credential_response: CredentialResponseParts<CS>,
+    pub(super) credential_response: SerializedCredentialResponse<CS>,
     pub(super) server_nonce: GenericArray<u8, NonceLen>,
     pub(super) server_e_pk: GenericArray<u8, KE::PkLen>,
     pub(super) server_mac: Output<KeHash<CS>>,
@@ -231,9 +231,9 @@ impl<'a, CS: CipherSuite> MessageBuilder<'a, CS> {
 impl<CS: CipherSuite, KE: Group> Deserialize for CachedMessage<CS, KE> {
     fn deserialize_take(input: &mut &[u8]) -> Result<Self, ProtocolError> {
         Ok(Self {
-            credential_request: CredentialRequestParts::deserialize_take(input)?,
+            credential_request: SerializedCredentialRequest::deserialize_take(input)?,
             ke1_message: Ke1MessageIter::deserialize_take(input)?,
-            credential_response: CredentialResponseParts::deserialize_take(input)?,
+            credential_response: SerializedCredentialResponse::deserialize_take(input)?,
             server_nonce: input.take_array("server nonce")?,
             server_e_pk: input.take_array("serialized server ephemeral key")?,
             server_mac: input.take_array("server mac")?,
@@ -246,8 +246,8 @@ type CachedMessageLen<CS: CipherSuite, KE: Group> = Sum<
     Sum<
         Sum<
             Sum<
-                Sum<CredentialRequestPartsLen<CS>, Ke1MessageIterLen<KE>>,
-                CredentialResponsePartsLen<CS>,
+                Sum<SerializedCredentialRequestLen<CS>, Ke1MessageIterLen<KE>>,
+                SerializedCredentialResponseLen<CS>,
             >,
             NonceLen,
         >,
@@ -258,23 +258,25 @@ type CachedMessageLen<CS: CipherSuite, KE: Group> = Sum<
 
 impl<CS: CipherSuite, KE: Group> Serialize for CachedMessage<CS, KE>
 where
-    CredentialRequestPartsLen<CS>: ArrayLength<u8> + Add<Ke1MessageIterLen<KE>>,
-    Sum<CredentialRequestPartsLen<CS>, Ke1MessageIterLen<KE>>:
-        ArrayLength<u8> + Add<CredentialResponsePartsLen<CS>>,
-    Sum<Sum<CredentialRequestPartsLen<CS>, Ke1MessageIterLen<KE>>, CredentialResponsePartsLen<CS>>:
-        ArrayLength<u8> + Add<NonceLen>,
+    SerializedCredentialRequestLen<CS>: ArrayLength<u8> + Add<Ke1MessageIterLen<KE>>,
+    Sum<SerializedCredentialRequestLen<CS>, Ke1MessageIterLen<KE>>:
+        ArrayLength<u8> + Add<SerializedCredentialResponseLen<CS>>,
+    Sum<
+        Sum<SerializedCredentialRequestLen<CS>, Ke1MessageIterLen<KE>>,
+        SerializedCredentialResponseLen<CS>,
+    >: ArrayLength<u8> + Add<NonceLen>,
     Sum<
         Sum<
-            Sum<CredentialRequestPartsLen<CS>, Ke1MessageIterLen<KE>>,
-            CredentialResponsePartsLen<CS>,
+            Sum<SerializedCredentialRequestLen<CS>, Ke1MessageIterLen<KE>>,
+            SerializedCredentialResponseLen<CS>,
         >,
         NonceLen,
     >: ArrayLength<u8> + Add<KE::PkLen>,
     Sum<
         Sum<
             Sum<
-                Sum<CredentialRequestPartsLen<CS>, Ke1MessageIterLen<KE>>,
-                CredentialResponsePartsLen<CS>,
+                Sum<SerializedCredentialRequestLen<CS>, Ke1MessageIterLen<KE>>,
+                SerializedCredentialResponseLen<CS>,
             >,
             NonceLen,
         >,
@@ -288,7 +290,7 @@ where
     <OprfGroup<CS> as voprf::Group>::ElemLen: Add<NonceLen>,
     Sum<<OprfGroup<CS> as voprf::Group>::ElemLen, NonceLen>:
         ArrayLength<u8> + Add<MaskedResponseLen<CS>>,
-    CredentialResponsePartsLen<CS>: ArrayLength<u8>,
+    SerializedCredentialResponseLen<CS>: ArrayLength<u8>,
 {
     type Len = CachedMessageLen<CS, KE>;
 

src/key_exchange/sigma_i/mod.rs

@@ -35,6 +35,11 @@ use zeroize::Zeroize;
 
 use self::message::Role;
 pub use self::message::{CachedMessage, HashOutput, Message, MessageBuilder, VerifyMessage};
+use super::{
+    Deserialize, GenerateKe1Result, GenerateKe2Result, GenerateKe3Result, KeyExchange, Serialize,
+    SerializedContext, SerializedCredentialRequest, SerializedCredentialResponse,
+    SerializedIdentifier, SerializedIdentifiers,
+};
 use crate::ciphersuite::{CipherSuite, KeGroup, KeHash};
 use crate::envelope::NonceLen;
 use crate::errors::{InternalError, ProtocolError};
@@ -42,11 +47,6 @@ use crate::hash::{Hash, OutputSize, ProxyHash};
 use crate::key_exchange::group::Group;
 use crate::key_exchange::shared::{derive_keys, generate_ke1, generate_nonce, transcript};
 pub use crate::key_exchange::shared::{DiffieHellman, Ke1Message, Ke1State};
-use crate::key_exchange::traits::{
-    CredentialRequestParts, CredentialResponseParts, Deserialize, GenerateKe2Result,
-    GenerateKe3Result, KeyExchange, Sealed, Serialize, SerializedContext, SerializedIdentifier,
-    SerializedIdentifiers,
-};
 use crate::keypair::{KeyPair, PrivateKey, PublicKey};
 use crate::opaque::Identifiers;
 use crate::serialization::{SliceExt, UpdateExt};
@@ -156,9 +156,9 @@ pub struct Ke2Builder<'a, CS: CipherSuite, KE: Group> {
     expected_mac: Output<KeHash<CS>>,
     session_key: Output<KeHash<CS>>,
     #[cfg(test)]
-    km3: Output<KeHash<CS>>,
-    #[cfg(test)]
     handshake_secret: Output<KeHash<CS>>,
+    #[cfg(test)]
+    km2: Output<KeHash<CS>>,
 }
 
 /// The server state produced after the second key exchange message
@@ -239,15 +239,15 @@ where
 
     fn generate_ke1<R: RngCore + CryptoRng>(
         rng: &mut R,
-    ) -> Result<(Self::KE1State, Self::KE1Message), ProtocolError> {
+    ) -> Result<GenerateKe1Result<Self>, ProtocolError> {
         generate_ke1(rng)
     }
 
     fn ke2_builder<'a, CS: CipherSuite<KeyExchange = Self>, R: RngCore + CryptoRng>(
         rng: &mut R,
-        credential_request: CredentialRequestParts<CS>,
+        credential_request: SerializedCredentialRequest<CS>,
         ke1_message: Self::KE1Message,
-        credential_response: CredentialResponseParts<CS>,
+        credential_response: SerializedCredentialResponse<CS>,
         client_s_pk: PublicKey<Self::Group>,
         identifiers: SerializedIdentifiers<'a, KeGroup<CS>>,
         context: SerializedContext<'a>,
@@ -309,9 +309,9 @@ where
             expected_mac: client_mac,
             session_key: derived_keys.session_key,
             #[cfg(test)]
-            km3: derived_keys.km3,
-            #[cfg(test)]
             handshake_secret: derived_keys.handshake_secret,
+            #[cfg(test)]
+            km2: derived_keys.km2,
         })
     }
 
@@ -333,33 +333,33 @@ where
         builder: Self::KE2Builder<'_, CS>,
         input: Self::KE2BuilderInput<CS>,
     ) -> Result<GenerateKe2Result<CS>, ProtocolError> {
-        Ok((
-            Ke2State {
+        Ok(GenerateKe2Result {
+            state: Ke2State {
                 client_s_pk: builder.client_s_pk.clone(),
                 session_key: builder.session_key.clone(),
                 verify_state: input.1,
                 expected_mac: builder.expected_mac.clone(),
             },
-            Ke2Message {
+            message: Ke2Message {
                 server_nonce: builder.server_nonce,
                 server_e_pk: builder.server_e_pk.clone(),
                 signature: input.0,
                 mac: builder.transcript.cache.server_mac.clone(),
             },
             #[cfg(test)]
-            builder.handshake_secret.clone(),
+            handshake_secret: builder.handshake_secret.clone(),
             #[cfg(test)]
-            builder.km3.clone(),
-        ))
+            km2: builder.km2.clone(),
+        })
     }
 
     fn generate_ke3<CS: CipherSuite<KeyExchange = Self>, R: CryptoRng + RngCore>(
         rng: &mut R,
-        credential_request: CredentialRequestParts<CS>,
+        credential_request: SerializedCredentialRequest<CS>,
         ke1_message: Self::KE1Message,
-        credential_response: CredentialResponseParts<CS>,
-        ke2_message: Self::KE2Message,
+        credential_response: SerializedCredentialResponse<CS>,
         ke1_state: &Self::KE1State,
+        ke2_message: Self::KE2Message,
         server_s_pk: PublicKey<Self::Group>,
         client_s_sk: PrivateKey<Self::Group>,
         identifiers: SerializedIdentifiers<'_, KeGroup<CS>>,
@@ -427,22 +427,22 @@ where
             &ke2_message.signature,
         )?;
 
-        Ok((
-            derived_keys.session_key,
-            Ke3Message {
+        Ok(GenerateKe3Result {
+            session_key: derived_keys.session_key,
+            message: Ke3Message {
                 signature,
                 mac: client_mac,
             },
             #[cfg(test)]
-            derived_keys.handshake_secret,
+            handshake_secret: derived_keys.handshake_secret,
             #[cfg(test)]
-            derived_keys.km3,
-        ))
+            km3: derived_keys.km3,
+        })
     }
 
     fn finish_ke<CS: CipherSuite<KeyExchange = Self>>(
-        ke3_message: Self::KE3Message,
         ke2_state: &Self::KE2State<CS>,
+        ke3_message: Self::KE3Message,
         identifiers: Identifiers<'_>,
         context: SerializedContext<'_>,
     ) -> Result<Output<KEH>, ProtocolError> {
@@ -468,14 +468,6 @@ where
     }
 }
 
-impl<SIG: SignatureProtocol, KE: 'static + Group, KEH: Hash> Sealed for SigmaI<SIG, KE, KEH>
-where
-    KEH::Core: ProxyHash,
-    <KEH::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
-    Le<<KEH::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
-{
-}
-
 impl<CS: CipherSuite, SIG: SignatureProtocol, KE: Group> Deserialize for Ke2State<CS, SIG, KE>
 where
     SIG::VerifyState<CS, KE>: Deserialize,

src/key_exchange/sigma_i/shared.rs

@@ -11,7 +11,7 @@ use digest::{Output, OutputSizeUser};
 use generic_array::{ArrayLength, GenericArray};
 
 use crate::errors::ProtocolError;
-use crate::key_exchange::traits::{Deserialize, Serialize};
+use crate::key_exchange::{Deserialize, Serialize};
 use crate::serialization::SliceExt;
 
 /// Pre-hash of the message to be verified.

src/key_exchange/traits.rs

@@ -1,402 +0,0 @@
-// Copyright (c) Meta Platforms, Inc. and affiliates.
-//
-// This source code is dual-licensed under either the MIT license found in the
-// LICENSE-MIT file in the root directory of this source tree or the Apache
-// License, Version 2.0 found in the LICENSE-APACHE file in the root directory
-// of this source tree. You may select, at your option, one of the above-listed
-// licenses.
-
-use core::iter;
-use core::ops::Add;
-
-use derive_where::derive_where;
-use digest::core_api::{BlockSizeUser, CoreProxy};
-use digest::Output;
-use generic_array::sequence::Concat;
-use generic_array::typenum::{IsLess, Le, NonZero, Sum, U2, U256};
-use generic_array::{ArrayLength, GenericArray};
-use rand::{CryptoRng, RngCore};
-use voprf::{BlindedElement, EvaluationElement};
-use zeroize::{Zeroize, ZeroizeOnDrop};
-
-#[cfg(test)]
-use crate::ciphersuite::KeHash;
-use crate::ciphersuite::{CipherSuite, OprfGroup};
-use crate::errors::ProtocolError;
-use crate::hash::{Hash, ProxyHash};
-use crate::key_exchange::group::Group;
-use crate::key_exchange::shared::{NonceLen, STR_CONTEXT};
-use crate::keypair::{PrivateKey, PublicKey};
-use crate::opaque::{Identifiers, MaskedResponse, MaskedResponseLen};
-use crate::serialization::{i2osp, SliceExt};
-
-/// The key exchange trait. This is only exposed so users can use it in generics
-/// and qualified bounds.
-#[allow(private_bounds)]
-pub trait KeyExchange: Sealed
-where
-    <Self::Hash as CoreProxy>::Core: ProxyHash,
-    <<Self::Hash as CoreProxy>::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
-    Le<<<Self::Hash as CoreProxy>::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
-{
-    /// The group used for the key exchange.
-    type Group: Group;
-    /// The has used for the key exchange.
-    type Hash: Hash;
-
-    #[doc(hidden)]
-    type KE1State: ZeroizeOnDrop + Clone;
-    #[doc(hidden)]
-    type KE2State<CS: CipherSuite>: ZeroizeOnDrop + Clone;
-    #[doc(hidden)]
-    type KE1Message: ZeroizeOnDrop + Clone;
-    #[doc(hidden)]
-    type KE2Builder<'a, CS: CipherSuite<KeyExchange = Self>>: ZeroizeOnDrop + Clone;
-    #[doc(hidden)]
-    type KE2BuilderData<'a, CS: 'static + CipherSuite>;
-    #[doc(hidden)]
-    type KE2BuilderInput<CS: CipherSuite>;
-    #[doc(hidden)]
-    type KE2Message: ZeroizeOnDrop + Clone;
-    #[doc(hidden)]
-    type KE3Message: ZeroizeOnDrop + Clone;
-
-    #[doc(hidden)]
-    fn generate_ke1<R: RngCore + CryptoRng>(
-        rng: &mut R,
-    ) -> Result<(Self::KE1State, Self::KE1Message), ProtocolError>;
-
-    #[doc(hidden)]
-    fn ke2_builder<'a, CS: CipherSuite<KeyExchange = Self>, R: RngCore + CryptoRng>(
-        rng: &mut R,
-        credential_request: CredentialRequestParts<CS>,
-        ke1_message: Self::KE1Message,
-        credential_response: CredentialResponseParts<CS>,
-        client_s_pk: PublicKey<Self::Group>,
-        identifiers: SerializedIdentifiers<'a, Self::Group>,
-        context: SerializedContext<'a>,
-    ) -> Result<Self::KE2Builder<'a, CS>, ProtocolError>;
-
-    #[doc(hidden)]
-    fn ke2_builder_data<'a, CS: CipherSuite<KeyExchange = Self>>(
-        builder: &'a Self::KE2Builder<'_, CS>,
-    ) -> Self::KE2BuilderData<'a, CS>;
-
-    #[doc(hidden)]
-    fn generate_ke2_input<CS: CipherSuite<KeyExchange = Self>, R: CryptoRng + RngCore>(
-        builder: &Self::KE2Builder<'_, CS>,
-        rng: &mut R,
-        server_s_sk: &PrivateKey<Self::Group>,
-    ) -> Self::KE2BuilderInput<CS>;
-
-    #[doc(hidden)]
-    fn build_ke2<CS: CipherSuite<KeyExchange = Self>>(
-        builder: Self::KE2Builder<'_, CS>,
-        input: Self::KE2BuilderInput<CS>,
-    ) -> Result<GenerateKe2Result<CS>, ProtocolError>;
-
-    #[doc(hidden)]
-    #[allow(clippy::too_many_arguments)]
-    fn generate_ke3<CS: CipherSuite<KeyExchange = Self>, R: CryptoRng + RngCore>(
-        rng: &mut R,
-        credential_request: CredentialRequestParts<CS>,
-        ke1_message: Self::KE1Message,
-        credential_response: CredentialResponseParts<CS>,
-        ke2_message: Self::KE2Message,
-        ke1_state: &Self::KE1State,
-        server_s_pk: PublicKey<Self::Group>,
-        client_s_sk: PrivateKey<Self::Group>,
-        identifiers: SerializedIdentifiers<'_, Self::Group>,
-        context: SerializedContext<'_>,
-    ) -> Result<GenerateKe3Result<Self>, ProtocolError>;
-
-    #[doc(hidden)]
-    fn finish_ke<CS: CipherSuite<KeyExchange = Self>>(
-        ke3_message: Self::KE3Message,
-        ke2_state: &Self::KE2State<CS>,
-        identifiers: Identifiers<'_>,
-        context: SerializedContext<'_>,
-    ) -> Result<Output<Self::Hash>, ProtocolError>;
-}
-
-pub(super) trait Sealed {}
-
-#[cfg_attr(
-    feature = "serde",
-    derive(serde::Deserialize, serde::Serialize),
-    serde(bound = "")
-)]
-#[derive_where(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd, Zeroize)]
-pub struct CredentialRequestParts<CS: CipherSuite>(
-    GenericArray<u8, <OprfGroup<CS> as voprf::Group>::ElemLen>,
-);
-
-impl<CS: CipherSuite> CredentialRequestParts<CS> {
-    pub(crate) fn new(blinded_element: &BlindedElement<CS::OprfCs>) -> Self {
-        Self(blinded_element.serialize())
-    }
-
-    pub fn iter(&self) -> impl Clone + Iterator<Item = &[u8]> {
-        iter::once(self.0.as_slice())
-    }
-
-    pub fn deserialize_take(input: &mut &[u8]) -> Result<Self, ProtocolError> {
-        Ok(Self(input.take_array("blinded element")?))
-    }
-}
-
-pub type CredentialRequestPartsLen<CS: CipherSuite> = <OprfGroup<CS> as voprf::Group>::ElemLen;
-
-impl<CS: CipherSuite> Serialize for CredentialRequestParts<CS> {
-    type Len = CredentialRequestPartsLen<CS>;
-
-    fn serialize(&self) -> GenericArray<u8, Self::Len> {
-        self.0.clone()
-    }
-}
-
-#[cfg_attr(
-    feature = "serde",
-    derive(serde::Deserialize, serde::Serialize),
-    serde(bound = "")
-)]
-#[derive_where(Clone, Debug, Eq, Hash, PartialEq, Zeroize)]
-pub struct CredentialResponseParts<CS: CipherSuite> {
-    evaluation_element: GenericArray<u8, <OprfGroup<CS> as voprf::Group>::ElemLen>,
-    masking_nonce: GenericArray<u8, NonceLen>,
-    masked_response: MaskedResponse<CS>,
-}
-
-impl<CS: CipherSuite> CredentialResponseParts<CS> {
-    pub(crate) fn new(
-        evaluation_element: &EvaluationElement<CS::OprfCs>,
-        masking_nonce: GenericArray<u8, NonceLen>,
-        masked_response: MaskedResponse<CS>,
-    ) -> Self {
-        Self {
-            evaluation_element: evaluation_element.serialize(),
-            masking_nonce,
-            masked_response,
-        }
-    }
-
-    pub fn iter(&self) -> impl Clone + Iterator<Item = &[u8]> {
-        [self.evaluation_element.as_slice(), &self.masking_nonce]
-            .into_iter()
-            .chain(self.masked_response.iter())
-    }
-
-    pub fn deserialize_take(input: &mut &[u8]) -> Result<Self, ProtocolError> {
-        Ok(Self {
-            evaluation_element: input.take_array("evaluation element")?,
-            masking_nonce: input.take_array("masking nonce")?,
-            masked_response: MaskedResponse::deserialize_take(input)?,
-        })
-    }
-}
-
-pub type CredentialResponsePartsLen<CS: CipherSuite> =
-    Sum<Sum<<OprfGroup<CS> as voprf::Group>::ElemLen, NonceLen>, MaskedResponseLen<CS>>;
-
-impl<CS: CipherSuite> Serialize for CredentialResponseParts<CS>
-where
-    <OprfGroup<CS> as voprf::Group>::ElemLen: Add<NonceLen>,
-    Sum<<OprfGroup<CS> as voprf::Group>::ElemLen, NonceLen>:
-        ArrayLength<u8> + Add<MaskedResponseLen<CS>>,
-    CredentialResponsePartsLen<CS>: ArrayLength<u8>,
-{
-    type Len = CredentialResponsePartsLen<CS>;
-
-    fn serialize(&self) -> GenericArray<u8, Self::Len> {
-        self.evaluation_element
-            .clone()
-            .concat(self.masking_nonce)
-            .concat(self.masked_response.serialize())
-    }
-}
-
-#[cfg_attr(
-    feature = "serde",
-    derive(serde::Deserialize, serde::Serialize),
-    serde(bound = "")
-)]
-#[derive(Clone, Debug, Eq, Hash, PartialEq, Zeroize)]
-pub struct SerializedContext<'a> {
-    length: GenericArray<u8, U2>,
-    #[zeroize(skip)]
-    context: &'a [u8],
-}
-
-impl<'a> SerializedContext<'a> {
-    pub(crate) fn from(context: Option<&'a [u8]>) -> Result<Self, ProtocolError> {
-        let context = context.unwrap_or(&[]);
-
-        Ok(Self {
-            length: i2osp::<U2>(context.len())?,
-            context,
-        })
-    }
-
-    pub fn iter(&self) -> impl Clone + Iterator<Item = &[u8]> {
-        iter::once(STR_CONTEXT).chain([self.length.as_slice(), self.context])
-    }
-}
-
-#[cfg_attr(
-    feature = "serde",
-    derive(serde::Deserialize, serde::Serialize),
-    serde(bound(deserialize = "'de: 'a", serialize = ""))
-)]
-#[derive_where(Clone, Debug, Eq, Hash, PartialEq, Zeroize)]
-pub struct SerializedIdentifiers<'a, G: Group> {
-    pub client: SerializedIdentifier<'a, G>,
-    pub server: SerializedIdentifier<'a, G>,
-}
-
-/// Computes `I2OSP(len(input), max_bytes) || input` and helps hold output
-/// without allocation.
-#[cfg_attr(
-    feature = "serde",
-    derive(serde::Deserialize, serde::Serialize),
-    serde(bound(deserialize = "'de: 'a", serialize = ""))
-)]
-#[derive_where(Clone, Debug, Eq, Hash, PartialEq, Zeroize)]
-pub struct SerializedIdentifier<'a, G: Group> {
-    length: GenericArray<u8, U2>,
-    identifier: Identifier<'a, G>,
-}
-
-#[cfg_attr(
-    feature = "serde",
-    derive(serde::Deserialize, serde::Serialize),
-    serde(bound = "")
-)]
-#[derive_where(Clone, Debug, Eq, Hash, PartialEq, Zeroize)]
-enum Identifier<'a, G: Group> {
-    Owned(GenericArray<u8, G::PkLen>),
-    #[derive_where(skip_inner(Zeroize))]
-    Borrowed(&'a [u8]),
-}
-
-impl<'a, G: Group> SerializedIdentifiers<'a, G> {
-    pub(crate) fn from_identifiers(
-        ids: Identifiers<'a>,
-        client_s_pk: GenericArray<u8, G::PkLen>,
-        server_s_pk: GenericArray<u8, G::PkLen>,
-    ) -> Result<Self, ProtocolError> {
-        let client = SerializedIdentifier::from_identifier(ids.client, client_s_pk)?;
-        let server = SerializedIdentifier::from_identifier(ids.server, server_s_pk)?;
-
-        Ok(Self { client, server })
-    }
-}
-
-impl<'a, G: Group> SerializedIdentifier<'a, G> {
-    pub fn from_identifier(
-        id: Option<&'a [u8]>,
-        s_pk: GenericArray<u8, G::PkLen>,
-    ) -> Result<Self, ProtocolError> {
-        if let Some(id) = id {
-            Ok(SerializedIdentifier {
-                length: i2osp::<U2>(id.len())?,
-                identifier: Identifier::Borrowed(id),
-            })
-        } else {
-            Ok(SerializedIdentifier {
-                length: i2osp::<U2>(s_pk.len())?,
-                identifier: Identifier::Owned(s_pk),
-            })
-        }
-    }
-
-    pub(crate) fn iter(&self) -> impl Clone + Iterator<Item = &[u8]> {
-        // Some magic to make it output the same type in all branches.
-        [self.length.as_slice()]
-            .into_iter()
-            .chain(match &self.identifier {
-                Identifier::Owned(bytes) => [bytes.as_slice()],
-                Identifier::Borrowed(bytes) => [*bytes],
-            })
-    }
-}
-
-pub trait Deserialize: Sized {
-    fn deserialize_take(input: &mut &[u8]) -> Result<Self, ProtocolError>;
-}
-
-pub trait Serialize {
-    type Len: ArrayLength<u8>;
-
-    fn serialize(&self) -> GenericArray<u8, Self::Len>;
-}
-
-#[cfg(not(test))]
-pub type GenerateKe2Result<CS: CipherSuite> = (
-    <CS::KeyExchange as KeyExchange>::KE2State<CS>,
-    <CS::KeyExchange as KeyExchange>::KE2Message,
-);
-#[cfg(test)]
-pub type GenerateKe2Result<CS: CipherSuite> = (
-    <CS::KeyExchange as KeyExchange>::KE2State<CS>,
-    <CS::KeyExchange as KeyExchange>::KE2Message,
-    Output<KeHash<CS>>,
-    Output<KeHash<CS>>,
-);
-#[cfg(not(test))]
-pub type GenerateKe3Result<K: KeyExchange> = (Output<K::Hash>, K::KE3Message);
-#[cfg(test)]
-pub type GenerateKe3Result<K: KeyExchange> = (
-    Output<K::Hash>,
-    K::KE3Message,
-    Output<K::Hash>,
-    Output<K::Hash>,
-);
-
-pub type Ke1StateLen<CS: CipherSuite> =
-    <<CS::KeyExchange as KeyExchange>::KE1State as Serialize>::Len;
-pub type Ke1MessageLen<CS: CipherSuite> =
-    <<CS::KeyExchange as KeyExchange>::KE1Message as Serialize>::Len;
-pub type Ke2StateLen<CS: CipherSuite> =
-    <<CS::KeyExchange as KeyExchange>::KE2State<CS> as Serialize>::Len;
-pub type Ke2MessageLen<CS: CipherSuite> =
-    <<CS::KeyExchange as KeyExchange>::KE2Message as Serialize>::Len;
-pub type Ke3MessageLen<CS: CipherSuite> =
-    <<CS::KeyExchange as KeyExchange>::KE3Message as Serialize>::Len;
-
-//////////////////////////
-// Test Implementations //
-//===================== //
-//////////////////////////
-
-#[cfg(test)]
-use crate::serialization::AssertZeroized;
-
-#[cfg(test)]
-impl<CS: CipherSuite> AssertZeroized for CredentialRequestParts<CS> {
-    fn assert_zeroized(&self) {
-        let Self(blinded_element) = self;
-
-        for byte in blinded_element.iter() {
-            assert_eq!(byte, &0);
-        }
-    }
-}
-
-#[cfg(test)]
-impl<CS: CipherSuite> AssertZeroized for CredentialResponseParts<CS> {
-    fn assert_zeroized(&self) {
-        let Self {
-            evaluation_element,
-            masking_nonce,
-            masked_response,
-        } = self;
-
-        for byte in evaluation_element
-            .iter()
-            .chain(masking_nonce)
-            .chain(masked_response.iter().flatten())
-        {
-            assert_eq!(byte, &0);
-        }
-    }
-}

src/key_exchange/tripledh.rs

@@ -22,16 +22,17 @@ use rand::{CryptoRng, RngCore};
 use subtle::{ConstantTimeEq, CtOption};
 use zeroize::{Zeroize, ZeroizeOnDrop};
 
+use super::{
+    Deserialize, GenerateKe1Result, GenerateKe2Result, GenerateKe3Result, KeyExchange, Serialize,
+    SerializedContext, SerializedCredentialRequest, SerializedCredentialResponse,
+    SerializedIdentifiers,
+};
 use crate::ciphersuite::{CipherSuite, KeGroup};
 use crate::errors::{InternalError, ProtocolError};
 use crate::hash::{Hash, OutputSize, ProxyHash};
 use crate::key_exchange::group::Group;
 use crate::key_exchange::shared::{self, NonceLen};
 pub use crate::key_exchange::shared::{DiffieHellman, Ke1Message, Ke1State};
-use crate::key_exchange::traits::{
-    CredentialRequestParts, CredentialResponseParts, Deserialize, GenerateKe2Result,
-    GenerateKe3Result, KeyExchange, Sealed, Serialize, SerializedContext, SerializedIdentifiers,
-};
 use crate::keypair::{KeyPair, PrivateKey, PublicKey};
 use crate::opaque::Identifiers;
 use crate::serialization::SliceExt;
@@ -151,15 +152,15 @@ where
 
     fn generate_ke1<R: RngCore + CryptoRng>(
         rng: &mut R,
-    ) -> Result<(Self::KE1State, Self::KE1Message), ProtocolError> {
+    ) -> Result<GenerateKe1Result<Self>, ProtocolError> {
         shared::generate_ke1(rng)
     }
 
     fn ke2_builder<'a, CS: CipherSuite<KeyExchange = Self>, R: RngCore + CryptoRng>(
         rng: &mut R,
-        credential_request: CredentialRequestParts<CS>,
+        credential_request: SerializedCredentialRequest<CS>,
         ke1_message: Self::KE1Message,
-        credential_response: CredentialResponseParts<CS>,
+        credential_response: SerializedCredentialResponse<CS>,
         client_s_pk: PublicKey<G>,
         identifiers: SerializedIdentifiers<'_, KeGroup<CS>>,
         context: SerializedContext<'a>,
@@ -240,30 +241,30 @@ where
         );
         let expected_mac = mac_hasher.finalize().into_bytes();
 
-        Ok((
-            Ke2State {
+        Ok(GenerateKe2Result {
+            state: Ke2State {
                 session_key: derived_keys.session_key,
                 expected_mac,
             },
-            Ke2Message {
+            message: Ke2Message {
                 server_nonce: builder.server_nonce,
                 server_e_pk: builder.server_e_pk.clone(),
                 mac,
             },
             #[cfg(test)]
-            derived_keys.handshake_secret,
+            handshake_secret: derived_keys.handshake_secret,
             #[cfg(test)]
-            derived_keys.km2,
-        ))
+            km2: derived_keys.km2,
+        })
     }
 
     fn generate_ke3<CS: CipherSuite<KeyExchange = Self>, R: CryptoRng + RngCore>(
         _: &mut R,
-        credential_request: CredentialRequestParts<CS>,
+        credential_request: SerializedCredentialRequest<CS>,
         ke1_message: Self::KE1Message,
-        credential_response: CredentialResponseParts<CS>,
-        ke2_message: Self::KE2Message,
+        credential_response: SerializedCredentialResponse<CS>,
         ke1_state: &Self::KE1State,
+        ke2_message: Self::KE2Message,
         server_s_pk: PublicKey<G>,
         client_s_sk: PrivateKey<G>,
         identifiers: SerializedIdentifiers<'_, KeGroup<CS>>,
@@ -309,21 +310,21 @@ where
             Hmac::<H>::new_from_slice(&derived_keys.km3).map_err(|_| InternalError::HmacError)?;
         Mac::update(&mut client_mac, &transcript_hasher.finalize());
 
-        Ok((
-            derived_keys.session_key,
-            Ke3Message {
+        Ok(GenerateKe3Result {
+            session_key: derived_keys.session_key,
+            message: Ke3Message {
                 mac: client_mac.finalize().into_bytes(),
             },
             #[cfg(test)]
-            derived_keys.handshake_secret,
+            handshake_secret: derived_keys.handshake_secret,
             #[cfg(test)]
-            derived_keys.km3,
-        ))
+            km3: derived_keys.km3,
+        })
     }
 
     fn finish_ke<CS: CipherSuite>(
-        ke3_message: Self::KE3Message,
         ke2_state: &Self::KE2State<CS>,
+        ke3_message: Self::KE3Message,
         _: Identifiers<'_>,
         _: SerializedContext<'_>,
     ) -> Result<Output<H>, ProtocolError> {
@@ -336,14 +337,6 @@ where
     }
 }
 
-impl<G: Group + 'static, H: Hash> Sealed for TripleDh<G, H>
-where
-    H::Core: ProxyHash,
-    <H::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
-    Le<<H::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
-{
-}
-
 ////////////////////////////////////////////////
 // Trait Implementations //
 // ========================================== //

src/keypair.rs

@@ -97,7 +97,8 @@ impl<G: Group> PrivateKey<G> {
         PublicKey(G::public_key(self.0))
     }
 
-    pub(crate) fn serialize(&self) -> GenericArray<u8, G::SkLen> {
+    /// Serializes this private key to a fixed-length byte array.
+    pub fn serialize(&self) -> GenericArray<u8, G::SkLen> {
         G::serialize_sk(self.0)
     }
 
@@ -149,8 +150,10 @@ pub trait PrivateKeySerialization<G: Group>: Clone {
     fn serialize_key_pair(key_pair: &KeyPair<G, Self>) -> GenericArray<u8, Self::Len>;
 
     /// Deserialization from bytes
+    ///
+    /// The deserialized bytes must be taken from `bytes`.
     fn deserialize_take_key_pair(
-        input: &mut &[u8],
+        bytes: &mut &[u8],
     ) -> Result<KeyPair<G, Self>, ProtocolError<Self::Error>>;
 }
 
@@ -281,7 +284,9 @@ pub trait OprfSeedSerialization<H, E>: Sized {
     fn serialize(&self) -> GenericArray<u8, Self::Len>;
 
     /// Deserialization from bytes
-    fn deserialize_take(input: &mut &[u8]) -> Result<Self, ProtocolError<E>>;
+    ///
+    /// The deserialized bytes must be taken from `bytes`.
+    fn deserialize_take(bytes: &mut &[u8]) -> Result<Self, ProtocolError<E>>;
 }
 
 impl<H: OutputSizeUser, E> OprfSeedSerialization<H, E> for OprfSeed<H> {

src/lib.rs

@@ -1349,9 +1349,9 @@ mod tests;
 
 // Exports
 
-pub use ciphersuite::CipherSuite;
 pub use rand;
 
+pub use crate::ciphersuite::CipherSuite;
 #[cfg(feature = "curve25519")]
 pub use crate::key_exchange::group::curve25519::Curve25519;
 #[cfg(feature = "ed25519")]

src/messages.rs

@@ -25,9 +25,9 @@ use crate::errors::ProtocolError;
 use crate::hash::OutputSize;
 use crate::key_exchange::group::Group;
 use crate::key_exchange::shared::NonceLen;
-use crate::key_exchange::traits::{
-    CredentialRequestParts, CredentialResponseParts, Deserialize, Ke1MessageLen, Ke2MessageLen,
-    Ke3MessageLen, KeyExchange, Serialize,
+use crate::key_exchange::{
+    Deserialize, Ke1MessageLen, Ke2MessageLen, Ke3MessageLen, KeyExchange, Serialize,
+    SerializedCredentialRequest, SerializedCredentialResponse,
 };
 use crate::keypair::PublicKey;
 use crate::opaque::{
@@ -359,8 +359,8 @@ impl<CS: CipherSuite> CredentialRequest<CS> {
         })
     }
 
-    pub(crate) fn to_parts(&self) -> CredentialRequestParts<CS> {
-        CredentialRequestParts::new(&self.blinded_element)
+    pub(crate) fn to_parts(&self) -> SerializedCredentialRequest<CS> {
+        SerializedCredentialRequest::new(&self.blinded_element)
     }
 
     /// Only used for testing purposes
@@ -415,8 +415,8 @@ impl<CS: CipherSuite> CredentialResponse<CS> {
         })
     }
 
-    pub(crate) fn to_parts(&self) -> CredentialResponseParts<CS> {
-        CredentialResponseParts::new(
+    pub(crate) fn to_parts(&self) -> SerializedCredentialResponse<CS> {
+        SerializedCredentialResponse::new(
             &self.evaluation_element,
             self.masking_nonce,
             self.masked_response.clone(),

src/opaque.rs

@@ -27,9 +27,9 @@ use crate::errors::{InternalError, ProtocolError};
 use crate::hash::OutputSize;
 use crate::key_exchange::group::Group;
 use crate::key_exchange::shared::NonceLen;
-use crate::key_exchange::traits::{
-    CredentialResponseParts, Deserialize, Ke1MessageLen, Ke1StateLen, Ke2StateLen, KeyExchange,
-    Serialize, SerializedContext, SerializedIdentifiers,
+use crate::key_exchange::{
+    Deserialize, Ke1MessageLen, Ke1StateLen, Ke2StateLen, KeyExchange, Serialize,
+    SerializedContext, SerializedCredentialResponse, SerializedIdentifiers,
 };
 use crate::keypair::{
     KeyPair, OprfSeed, OprfSeedSerialization, PrivateKey, PrivateKeySerialization, PublicKey,
@@ -503,18 +503,18 @@ impl<CS: CipherSuite> ClientLogin<CS> {
         password: &[u8],
     ) -> Result<ClientLoginStartResult<CS>, ProtocolError> {
         let blind_result = blind::<CS, _>(rng, password)?;
-        let (ke1_state, ke1_message) = CS::KeyExchange::generate_ke1(rng)?;
+        let ke1_result = CS::KeyExchange::generate_ke1(rng)?;
 
         let credential_request = CredentialRequest {
             blinded_element: blind_result.message,
-            ke1_message,
+            ke1_message: ke1_result.message,
         };
 
         Ok(ClientLoginStartResult {
             message: credential_request.clone(),
             state: Self {
                 oprf_client: blind_result.state,
-                ke1_state,
+                ke1_state: ke1_result.state,
                 credential_request,
             },
         })
@@ -582,8 +582,8 @@ impl<CS: CipherSuite> ClientLogin<CS> {
             self.credential_request.to_parts(),
             self.credential_request.ke1_message.clone(),
             credential_response.to_parts(),
-            credential_response.ke2_message,
             &self.ke1_state,
+            credential_response.ke2_message,
             server_s_pk.clone(),
             opened_envelope.client_static_keypair.private().clone(),
             opened_envelope.identifiers,
@@ -592,17 +592,17 @@ impl<CS: CipherSuite> ClientLogin<CS> {
 
         Ok(ClientLoginFinishResult {
             message: CredentialFinalization {
-                ke3_message: result.1,
+                ke3_message: result.message,
             },
-            session_key: result.0,
+            session_key: result.session_key,
             export_key: opened_envelope.export_key,
             server_s_pk,
             #[cfg(test)]
             state: self,
             #[cfg(test)]
-            handshake_secret: result.2,
+            handshake_secret: result.handshake_secret,
             #[cfg(test)]
-            client_mac_key: result.3,
+            client_mac_key: result.km3,
         })
     }
 }
@@ -679,7 +679,7 @@ impl<CS: CipherSuite> ServerLogin<CS> {
         let server = voprf::OprfServer::new_with_key(&oprf_key).map_err(ProtocolError::from)?;
         let evaluation_element = server.blind_evaluate(&credential_request.blinded_element);
 
-        let credential_response = CredentialResponseParts::new(
+        let credential_response = SerializedCredentialResponse::new(
             &evaluation_element,
             masking_nonce,
             masked_response.clone(),
@@ -743,18 +743,18 @@ impl<CS: CipherSuite> ServerLogin<CS> {
             evaluation_element: builder.evaluation_element.clone(),
             masking_nonce: *builder.masking_nonce.deref(),
             masked_response: builder.masked_response.clone(),
-            ke2_message: result.1,
+            ke2_message: result.message,
         };
 
         Ok(ServerLoginStartResult {
             message: credential_response,
             state: Self {
-                ke2_state: result.0,
+                ke2_state: result.state,
             },
             #[cfg(test)]
-            handshake_secret: result.2,
+            handshake_secret: result.handshake_secret,
             #[cfg(test)]
-            server_mac_key: result.3,
+            server_mac_key: result.km2,
             #[cfg(test)]
             oprf_key: builder.oprf_key.deref().clone(),
         })
@@ -797,8 +797,8 @@ impl<CS: CipherSuite> ServerLogin<CS> {
         let context = SerializedContext::from(parameters.context)?;
 
         let session_key = <CS::KeyExchange as KeyExchange>::finish_ke(
-            message.ke3_message,
             &self.ke2_state,
+            message.ke3_message,
             parameters.identifiers,
             context,
         )?;

src/serialization/tests.rs

@@ -25,7 +25,7 @@ use crate::errors::*;
 use crate::hash::OutputSize;
 use crate::key_exchange::group::Group;
 use crate::key_exchange::shared::NonceLen;
-use crate::key_exchange::traits::{
+use crate::key_exchange::{
     Deserialize, Ke1MessageLen, Ke1StateLen, Ke2MessageLen, KeyExchange, Serialize,
 };
 use crate::keypair::KeyPair;

src/tests/full_test.rs

@@ -27,7 +27,7 @@ use crate::errors::*;
 use crate::hash::OutputSize;
 use crate::key_exchange::group::Group;
 use crate::key_exchange::shared::NonceLen;
-use crate::key_exchange::traits::{
+use crate::key_exchange::{
     Deserialize, Ke1MessageLen, Ke1StateLen, Ke2MessageLen, KeyExchange, Serialize,
 };
 use crate::ksf::Identity;

src/tests/test_opaque_vectors.rs

@@ -23,9 +23,7 @@ use crate::errors::*;
 use crate::hash::OutputSize;
 use crate::key_exchange::group::Group;
 use crate::key_exchange::shared::NonceLen;
-use crate::key_exchange::traits::{
-    Deserialize, Ke1MessageLen, Ke2MessageLen, KeyExchange, Serialize,
-};
+use crate::key_exchange::{Deserialize, Ke1MessageLen, Ke2MessageLen, KeyExchange, Serialize};
 use crate::ksf::Identity;
 use crate::messages::{
     CredentialRequestLen, CredentialResponseLen, CredentialResponseWithoutKeLen,