Rework SecretKey API to facilitate async (#371)

* Rework `SecretKey` API to facilitate async

* Remove left-over constraints

src/key_exchange/group/curve25519.rs

@@ -9,7 +9,7 @@
 //! Key Exchange group implementation for Curve25519
 
 use curve25519_dalek::montgomery::MontgomeryPoint;
-use curve25519_dalek::scalar::{self, Scalar};
+use curve25519_dalek::scalar;
 use curve25519_dalek::traits::Identity;
 use digest::core_api::BlockSizeUser;
 use digest::{FixedOutput, HashMarker, OutputSizeUser};
@@ -17,9 +17,11 @@ use generic_array::typenum::{IsLess, IsLessOrEqual, U256, U32};
 use generic_array::GenericArray;
 use rand::{CryptoRng, RngCore};
 use subtle::ConstantTimeEq;
+use zeroize::Zeroize;
 
 use super::KeGroup;
-use crate::errors::InternalError;
+use crate::errors::{InternalError, ProtocolError};
+use crate::key_exchange::tripledh::DiffieHellman;
 
 /// Implementation for Curve25519.
 pub struct Curve25519;
@@ -28,20 +30,20 @@ pub struct Curve25519;
 impl KeGroup for Curve25519 {
     type Pk = MontgomeryPoint;
     type PkLen = U32;
-    type Sk = [u8; 32];
+    type Sk = Scalar;
     type SkLen = U32;
 
     fn serialize_pk(pk: Self::Pk) -> GenericArray<u8, Self::PkLen> {
         pk.to_bytes().into()
     }
 
-    fn deserialize_pk(bytes: &[u8]) -> Result<Self::Pk, InternalError> {
+    fn deserialize_pk(bytes: &[u8]) -> Result<Self::Pk, ProtocolError> {
         bytes
             .try_into()
             .ok()
             .map(MontgomeryPoint)
             .filter(|pk| pk != &MontgomeryPoint::identity())
-            .ok_or(InternalError::PointError)
+            .ok_or(ProtocolError::SerializationError)
     }
 
     fn random_sk<R: RngCore + CryptoRng>(rng: &mut R) -> Self::Sk {
@@ -51,8 +53,8 @@ impl KeGroup for Curve25519 {
             rng.fill_bytes(&mut scalar_bytes);
             let scalar = scalar::clamp_integer(scalar_bytes);
 
-            if scalar != Scalar::ZERO.to_bytes() {
-                break scalar;
+            if scalar != curve25519_dalek::Scalar::ZERO.to_bytes() {
+                break Scalar(scalar);
             }
         }
     }
@@ -72,26 +74,22 @@ impl KeGroup for Curve25519 {
         <CS::Hash as OutputSizeUser>::OutputSize:
             IsLess<U256> + IsLessOrEqual<<CS::Hash as BlockSizeUser>::BlockSize>,
     {
-        Ok(scalar::clamp_integer(seed.into()))
+        Ok(Scalar(scalar::clamp_integer(seed.into())))
     }
 
     fn is_zero_scalar(scalar: Self::Sk) -> subtle::Choice {
-        scalar.ct_eq(&Scalar::ZERO.to_bytes())
+        scalar.0.ct_eq(&curve25519_dalek::Scalar::ZERO.to_bytes())
     }
 
     fn public_key(sk: Self::Sk) -> Self::Pk {
-        MontgomeryPoint::mul_base_clamped(sk)
-    }
-
-    fn diffie_hellman(pk: Self::Pk, sk: Self::Sk) -> GenericArray<u8, Self::PkLen> {
-        Self::serialize_pk(pk.mul_clamped(sk))
+        MontgomeryPoint::mul_base_clamped(sk.0)
     }
 
     fn serialize_sk(sk: Self::Sk) -> GenericArray<u8, Self::SkLen> {
-        sk.into()
+        sk.0.into()
     }
 
-    fn deserialize_sk(bytes: &[u8]) -> Result<Self::Sk, InternalError> {
+    fn deserialize_sk(bytes: &[u8]) -> Result<Self::Sk, ProtocolError> {
         bytes
             .try_into()
             .ok()
@@ -99,7 +97,18 @@ impl KeGroup for Curve25519 {
                 let scalar = scalar::clamp_integer(bytes);
                 (scalar == bytes).then_some(scalar)
             })
-            .filter(|scalar| scalar != &Scalar::ZERO.to_bytes())
-            .ok_or(InternalError::PointError)
+            .filter(|scalar| scalar != &curve25519_dalek::Scalar::ZERO.to_bytes())
+            .map(Scalar)
+            .ok_or(ProtocolError::SerializationError)
+    }
+}
+
+/// Curve25519 scalar.
+#[derive(Clone, Copy, Zeroize)]
+pub struct Scalar([u8; 32]);
+
+impl DiffieHellman<Curve25519> for Scalar {
+    fn diffie_hellman(self, pk: MontgomeryPoint) -> GenericArray<u8, U32> {
+        Curve25519::serialize_pk(pk.mul_clamped(self.0))
     }
 }

src/key_exchange/group/elliptic_curve.rs

@@ -19,7 +19,8 @@ use generic_array::GenericArray;
 use rand::{CryptoRng, RngCore};
 
 use super::KeGroup;
-use crate::errors::InternalError;
+use crate::errors::{InternalError, ProtocolError};
+use crate::key_exchange::tripledh::DiffieHellman;
 
 impl<G> KeGroup for G
 where
@@ -41,10 +42,10 @@ where
         GenericArray::clone_from_slice(pk.to_encoded_point(true).as_bytes())
     }
 
-    fn deserialize_pk(bytes: &[u8]) -> Result<Self::Pk, InternalError> {
+    fn deserialize_pk(bytes: &[u8]) -> Result<Self::Pk, ProtocolError> {
         PublicKey::<Self>::from_sec1_bytes(bytes)
             .map(|public_key| public_key.to_projective())
-            .map_err(|_| InternalError::PointError)
+            .map_err(|_| ProtocolError::SerializationError)
     }
 
     fn random_sk<R: RngCore + CryptoRng>(rng: &mut R) -> Self::Sk {
@@ -77,17 +78,29 @@ where
         scalar.is_zero()
     }
 
-    fn diffie_hellman(pk: Self::Pk, sk: Self::Sk) -> GenericArray<u8, Self::PkLen> {
-        Self::serialize_pk(pk * sk)
-    }
-
     fn serialize_sk(sk: Self::Sk) -> GenericArray<u8, Self::SkLen> {
         sk.into()
     }
 
-    fn deserialize_sk(bytes: &[u8]) -> Result<Self::Sk, InternalError> {
+    fn deserialize_sk(bytes: &[u8]) -> Result<Self::Sk, ProtocolError> {
         SecretKey::<Self>::from_slice(bytes)
             .map(|secret_key| *secret_key.to_nonzero_scalar())
-            .map_err(|_| InternalError::PointError)
+            .map_err(|_| ProtocolError::SerializationError)
+    }
+}
+
+impl<G> DiffieHellman<G> for Scalar<G>
+where
+    G: GroupDigest,
+    FieldBytesSize<G>: ModulusSize,
+    AffinePoint<G>: FromEncodedPoint<G> + ToEncodedPoint<G>,
+    ProjectivePoint<G>: CofactorGroup + ToEncodedPoint<G>,
+    Scalar<G>: FromOkm,
+{
+    fn diffie_hellman(
+        self,
+        pk: ProjectivePoint<G>,
+    ) -> GenericArray<u8, <FieldBytesSize<G> as ModulusSize>::CompressedPointSize> {
+        GenericArray::clone_from_slice((pk * self).to_encoded_point(true).as_bytes())
     }
 }

src/key_exchange/group/mod.rs

@@ -22,7 +22,7 @@ use generic_array::{ArrayLength, GenericArray};
 use rand::{CryptoRng, RngCore};
 use zeroize::Zeroize;
 
-use crate::errors::InternalError;
+use crate::errors::{InternalError, ProtocolError};
 
 const STR_OPAQUE_DERIVE_AUTH_KEY_PAIR: [u8; 33] = *b"OPAQUE-DeriveDiffieHellmanKeyPair";
 
@@ -41,7 +41,7 @@ pub trait KeGroup {
     fn serialize_pk(pk: Self::Pk) -> GenericArray<u8, Self::PkLen>;
 
     /// Return a public key from its fixed-length bytes representation
-    fn deserialize_pk(bytes: &[u8]) -> Result<Self::Pk, InternalError>;
+    fn deserialize_pk(bytes: &[u8]) -> Result<Self::Pk, ProtocolError>;
 
     /// Generate a random secret key
     fn random_sk<R: RngCore + CryptoRng>(rng: &mut R) -> Self::Sk;
@@ -104,14 +104,11 @@ pub trait KeGroup {
     /// Return a public key from its secret key
     fn public_key(sk: Self::Sk) -> Self::Pk;
 
-    /// Diffie-Hellman key exchange
-    fn diffie_hellman(pk: Self::Pk, sk: Self::Sk) -> GenericArray<u8, Self::PkLen>;
-
     /// 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: &[u8]) -> Result<Self::Sk, InternalError>;
+    fn deserialize_sk(bytes: &[u8]) -> Result<Self::Sk, ProtocolError>;
 }
 
 // Helper functions used to compute DeriveAuthKeyPair() (taken from the voprf

src/key_exchange/group/ristretto255.rs

@@ -21,7 +21,8 @@ use subtle::ConstantTimeEq;
 use voprf::Group;
 
 use super::KeGroup;
-use crate::errors::InternalError;
+use crate::errors::{InternalError, ProtocolError};
+use crate::key_exchange::tripledh::DiffieHellman;
 
 /// Implementation for Ristretto255.
 // This is necessary because Rust lacks specialization, otherwise we could
@@ -38,12 +39,12 @@ impl KeGroup for Ristretto255 {
         pk.compress().to_bytes().into()
     }
 
-    fn deserialize_pk(bytes: &[u8]) -> Result<Self::Pk, InternalError> {
+    fn deserialize_pk(bytes: &[u8]) -> Result<Self::Pk, ProtocolError> {
         CompressedRistretto::from_slice(bytes)
-            .map_err(|_| InternalError::PointError)?
+            .map_err(|_| ProtocolError::SerializationError)?
             .decompress()
             .filter(|point| point != &RistrettoPoint::identity())
-            .ok_or(InternalError::PointError)
+            .ok_or(ProtocolError::SerializationError)
     }
 
     fn random_sk<R: RngCore + CryptoRng>(rng: &mut R) -> Self::Sk {
@@ -89,21 +90,17 @@ impl KeGroup for Ristretto255 {
         RISTRETTO_BASEPOINT_POINT * sk
     }
 
-    fn diffie_hellman(pk: Self::Pk, sk: Self::Sk) -> GenericArray<u8, Self::PkLen> {
-        Self::serialize_pk(pk * sk)
-    }
-
     fn serialize_sk(sk: Self::Sk) -> GenericArray<u8, Self::SkLen> {
         sk.to_bytes().into()
     }
 
-    fn deserialize_sk(bytes: &[u8]) -> Result<Self::Sk, InternalError> {
+    fn deserialize_sk(bytes: &[u8]) -> Result<Self::Sk, ProtocolError> {
         bytes
             .try_into()
             .ok()
             .and_then(|bytes| Scalar::from_canonical_bytes(bytes).into())
             .filter(|scalar| scalar != &Scalar::ZERO)
-            .ok_or(InternalError::PointError)
+            .ok_or(ProtocolError::SerializationError)
     }
 }
 
@@ -183,3 +180,9 @@ impl Group for Ristretto255 {
         <voprf::Ristretto255 as Group>::deserialize_scalar(scalar_bits)
     }
 }
+
+impl DiffieHellman<Ristretto255> for Scalar {
+    fn diffie_hellman(self, pk: RistrettoPoint) -> GenericArray<u8, U32> {
+        Ristretto255::serialize_pk(pk * self)
+    }
+}

src/key_exchange/traits.rs

@@ -17,7 +17,7 @@ use crate::ciphersuite::{CipherSuite, OprfHash};
 use crate::errors::ProtocolError;
 use crate::hash::{Hash, ProxyHash};
 use crate::key_exchange::group::KeGroup;
-use crate::keypair::{PrivateKey, PublicKey, SecretKey};
+use crate::keypair::{PrivateKey, PublicKey};
 
 pub trait KeyExchange<D: Hash, G: KeGroup>
 where
@@ -28,6 +28,9 @@ where
     type KE1State: Deserialize + Serialize + ZeroizeOnDrop + Clone;
     type KE2State: Deserialize + Serialize + ZeroizeOnDrop + Clone;
     type KE1Message: Deserialize + Serialize + ZeroizeOnDrop + Clone;
+    type KE2Builder: ZeroizeOnDrop + Clone;
+    type KE2BuilderData<'a>;
+    type KE2BuilderInput;
     type KE2Message: Deserialize + Serialize + ZeroizeOnDrop + Clone;
     type KE3Message: Deserialize + Serialize + ZeroizeOnDrop + Clone;
 
@@ -36,25 +39,28 @@ where
     ) -> Result<(Self::KE1State, Self::KE1Message), ProtocolError>;
 
     #[allow(clippy::too_many_arguments)]
-    fn generate_ke2<
-        'a,
-        'b,
-        'c,
-        'd,
-        OprfCs: voprf::CipherSuite,
-        R: RngCore + CryptoRng,
-        S: SecretKey<G>,
-    >(
+    fn ke2_builder<'a, 'b, 'c, 'd, OprfCs: voprf::CipherSuite, R: RngCore + CryptoRng>(
         rng: &mut R,
-        l1_bytes: impl Iterator<Item = &'a [u8]>,
-        l2_bytes: impl Iterator<Item = &'b [u8]>,
+        serialized_credential_request: impl Iterator<Item = &'a [u8]>,
+        serialized_credential_response: impl Iterator<Item = &'b [u8]>,
         ke1_message: Self::KE1Message,
         client_s_pk: PublicKey<G>,
-        server_s_sk: S,
         id_u: impl Iterator<Item = &'c [u8]>,
         id_s: impl Iterator<Item = &'d [u8]>,
         context: &[u8],
-    ) -> Result<GenerateKe2Result<Self, D, G>, ProtocolError<S::Error>>;
+    ) -> Result<Self::KE2Builder, ProtocolError>;
+
+    fn ke2_builder_data(builder: &Self::KE2Builder) -> Self::KE2BuilderData<'_>;
+
+    fn generate_ke2_input(
+        builder: &Self::KE2Builder,
+        server_s_sk: &PrivateKey<G>,
+    ) -> Self::KE2BuilderInput;
+
+    fn build_ke2(
+        builder: Self::KE2Builder,
+        input: Self::KE2BuilderInput,
+    ) -> Result<GenerateKe2Result<Self, D, G>, ProtocolError>;
 
     #[allow(clippy::too_many_arguments)]
     fn generate_ke3<'a, 'b, 'c, 'd>(

src/key_exchange/tripledh.rs

@@ -19,6 +19,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};
@@ -27,7 +28,7 @@ use crate::key_exchange::group::KeGroup;
 use crate::key_exchange::traits::{
     Deserialize, GenerateKe2Result, GenerateKe3Result, KeyExchange, Serialize,
 };
-use crate::keypair::{KeyPair, PrivateKey, PublicKey, SecretKey};
+use crate::keypair::{KeyPair, PrivateKey, PublicKey};
 use crate::serialization::{Input, UpdateExt};
 
 ///////////////
@@ -49,6 +50,14 @@ static STR_OPAQUE: &[u8] = b"OPAQUE-";
 ////////////////////////////
 
 /// The Triple Diffie-Hellman key exchange implementation
+///
+/// # Remote Key
+///
+/// [`ServerLoginBuilder::data()`](crate::ServerLoginBuilder::data()) will
+/// return the client's ephemeral public key.
+/// [`ServerLoginBuilder::build()`](crate::ServerLoginBuilder::build()) expects
+/// a shared secret computed through Diffie-Hellman from the server's private
+/// key and the given public key.
 pub struct TripleDh;
 
 /// The client state produced after the first key exchange message
@@ -95,6 +104,31 @@ where
     session_key: Output<D>,
 }
 
+/// Builder for the second key exchange message
+#[cfg_attr(
+    feature = "serde",
+    derive(serde::Deserialize, serde::Serialize),
+    serde(bound(
+        deserialize = "D: serde::Deserialize<'de>,  PublicKey<KG>: serde::Deserialize<'de>",
+        serialize = "D: serde::Serialize,  PublicKey<KG>: serde::Serialize",
+    ))
+)]
+#[derive_where(Clone)]
+#[derive_where(Debug, Eq, Hash, PartialEq; D, PublicKey<KG>)]
+pub struct Ke2Builder<D: Hash, KG: KeGroup>
+where
+    D::Core: ProxyHash,
+    <D::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<D::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+{
+    server_nonce: GenericArray<u8, NonceLen>,
+    transcript_hasher: D,
+    client_e_pk: PublicKey<KG>,
+    server_e_pk: PublicKey<KG>,
+    shared_secret_1: GenericArray<u8, KG::PkLen>,
+    shared_secret_3: GenericArray<u8, KG::PkLen>,
+}
+
 /// The second key exchange message
 #[cfg_attr(
     feature = "serde",
@@ -130,13 +164,20 @@ where
     mac: Output<D>,
 }
 
+/// Trait required by [`KeGroup::Sk`] to be compatible with [`TripleDh`].
+pub trait DiffieHellman<KG: KeGroup> {
+    /// Diffie-Hellman key exchange.
+    fn diffie_hellman(self, pk: KG::Pk) -> GenericArray<u8, KG::PkLen>;
+}
+
 ////////////////////////////////
 // High-level Implementations //
 // ========================== //
 ////////////////////////////////
 
-impl<D: Hash, KG: KeGroup> KeyExchange<D, KG> for TripleDh
+impl<D: Hash, KG: KeGroup + 'static> KeyExchange<D, KG> for TripleDh
 where
+    KG::Sk: DiffieHellman<KG>,
     D::Core: ProxyHash,
     <D::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
     Le<<D::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
@@ -158,6 +199,9 @@ where
     type KE1State = Ke1State<KG>;
     type KE2State = Ke2State<D>;
     type KE1Message = Ke1Message<KG>;
+    type KE2Builder = Ke2Builder<D, KG>;
+    type KE2BuilderData<'a> = &'a PublicKey<KG>;
+    type KE2BuilderInput = GenericArray<u8, KG::PkLen>;
     type KE2Message = Ke2Message<D, KG>;
     type KE3Message = Ke3Message<D>;
 
@@ -181,71 +225,82 @@ where
         ))
     }
 
-    #[allow(clippy::type_complexity)]
-    fn generate_ke2<
-        'a,
-        'b,
-        'c,
-        'd,
-        OprfCs: voprf::CipherSuite,
-        R: RngCore + CryptoRng,
-        S: SecretKey<KG>,
-    >(
+    fn ke2_builder<'a, 'b, 'c, 'd, OprfCs: voprf::CipherSuite, R: RngCore + CryptoRng>(
         rng: &mut R,
         serialized_credential_request: impl Iterator<Item = &'a [u8]>,
-        l2_bytes: impl Iterator<Item = &'b [u8]>,
+        serialized_credential_response: impl Iterator<Item = &'b [u8]>,
         ke1_message: Self::KE1Message,
         client_s_pk: PublicKey<KG>,
-        server_s_sk: S,
         id_u: impl Iterator<Item = &'c [u8]>,
         id_s: impl Iterator<Item = &'d [u8]>,
         context: &[u8],
-    ) -> Result<GenerateKe2Result<Self, D, KG>, ProtocolError<S::Error>> {
-        let server_e_kp = KeyPair::<KG>::generate_random::<OprfCs, _>(rng);
+    ) -> Result<Self::KE2Builder, ProtocolError> {
+        let server_e = KeyPair::<KG>::generate_random::<OprfCs, _>(rng);
         let server_nonce = generate_nonce::<R>(rng);
 
-        let mut transcript_hasher = D::new()
+        let transcript_hasher = D::new()
             .chain(STR_CONTEXT)
-            .chain_iter(
-                Input::<U2>::from(context)
-                    .map_err(ProtocolError::into_custom)?
-                    .iter(),
-            )
+            .chain_iter(Input::<U2>::from(context)?.iter())
             .chain_iter(id_u.into_iter())
             .chain_iter(serialized_credential_request)
             .chain_iter(id_s.into_iter())
-            .chain_iter(l2_bytes)
+            .chain_iter(serialized_credential_response)
             .chain(server_nonce)
-            .chain(server_e_kp.public().serialize());
+            .chain(server_e.public().serialize());
 
-        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.clone(),
-                sk2: server_s_sk,
-                pk3: client_s_pk,
-                sk3: server_e_kp.private().clone(),
-            },
-            &transcript_hasher.clone().finalize(),
+        let shared_secret_1 = server_e
+            .private()
+            .ke_diffie_hellman(&ke1_message.client_e_pk);
+        let shared_secret_3 = server_e.private().ke_diffie_hellman(&client_s_pk);
+
+        Ok(Ke2Builder {
+            server_nonce,
+            transcript_hasher,
+            client_e_pk: ke1_message.client_e_pk.clone(),
+            server_e_pk: server_e.public().clone(),
+            shared_secret_1,
+            shared_secret_3,
+        })
+    }
+
+    fn ke2_builder_data(builder: &Self::KE2Builder) -> Self::KE2BuilderData<'_> {
+        &builder.client_e_pk
+    }
+
+    fn generate_ke2_input(
+        builder: &Self::KE2Builder,
+        server_s_sk: &PrivateKey<KG>,
+    ) -> Self::KE2BuilderInput {
+        server_s_sk.ke_diffie_hellman(&builder.client_e_pk)
+    }
+
+    fn build_ke2(
+        mut builder: Self::KE2Builder,
+        shared_secret_2: Self::KE2BuilderInput,
+    ) -> Result<GenerateKe2Result<Self, D, KG>, ProtocolError> {
+        let result = derive_3dh_keys::<D, KG>(
+            builder.shared_secret_1.clone(),
+            shared_secret_2,
+            builder.shared_secret_3.clone(),
+            &builder.transcript_hasher.clone().finalize(),
         )?;
 
         let mut mac_hasher =
             Hmac::<D>::new_from_slice(&result.1).map_err(|_| InternalError::HmacError)?;
-        mac_hasher.update(&transcript_hasher.clone().finalize());
+        mac_hasher.update(&builder.transcript_hasher.clone().finalize());
         let mac = mac_hasher.finalize().into_bytes();
 
-        Digest::update(&mut transcript_hasher, &mac);
+        Digest::update(&mut builder.transcript_hasher, &mac);
 
         Ok((
             Ke2State {
                 km3: result.2,
-                hashed_transcript: transcript_hasher.finalize(),
+                hashed_transcript: builder.transcript_hasher.clone().finalize(),
                 session_key: result.0,
             },
             Ke2Message {
-                server_nonce,
-                server_e_pk: server_e_kp.public().clone(),
+                server_nonce: builder.server_nonce,
+                server_e_pk: builder.server_e_pk.clone(),
                 mac,
             },
             #[cfg(test)]
@@ -276,15 +331,12 @@ where
             .chain_iter(l2_component)
             .chain(ke2_message.to_bytes_without_mac());
 
-        let result = derive_3dh_keys::<D, KG, PrivateKey<KG>>(
-            TripleDhComponents {
-                pk1: ke2_message.server_e_pk.clone(),
-                sk1: ke1_state.client_e_sk.clone(),
-                pk2: server_s_pk,
-                sk2: ke1_state.client_e_sk.clone(),
-                pk3: ke2_message.server_e_pk.clone(),
-                sk3: client_s_sk,
-            },
+        let result = derive_3dh_keys::<D, KG>(
+            ke1_state
+                .client_e_sk
+                .ke_diffie_hellman(&ke2_message.server_e_pk),
+            ke1_state.client_e_sk.ke_diffie_hellman(&server_s_pk),
+            client_s_sk.ke_diffie_hellman(&ke2_message.server_e_pk),
             &transcript_hasher.clone().finalize(),
         )?;
 
@@ -335,16 +387,6 @@ where
 //==================== //
 /////////////////////////
 
-// The triple of public and private components used in the 3DH computation
-struct TripleDhComponents<KG: KeGroup, S: SecretKey<KG>> {
-    pk1: PublicKey<KG>,
-    sk1: PrivateKey<KG>,
-    pk2: PublicKey<KG>,
-    sk2: S,
-    pk3: PublicKey<KG>,
-    sk3: PrivateKey<KG>,
-}
-
 // Consists of a session key, followed by two mac keys: (session_key, km2, km3)
 #[cfg(not(test))]
 type TripleDhDerivationResult<D> = (Output<D>, Output<D>, Output<D>);
@@ -361,10 +403,12 @@ type TripleDhDerivationResult<D> = (Output<D>, Output<D>, Output<D>, Output<D>);
 // Internal function which takes the public and private components of the client
 // and server keypairs, along with some auxiliary metadata, to produce the
 // session key and two MAC keys
-fn derive_3dh_keys<D: Hash, KG: KeGroup, S: SecretKey<KG>>(
-    dh: TripleDhComponents<KG, S>,
+fn derive_3dh_keys<D: Hash, KG: KeGroup>(
+    shared_secret_1: GenericArray<u8, KG::PkLen>,
+    shared_secret_2: GenericArray<u8, KG::PkLen>,
+    shared_secret_3: GenericArray<u8, KG::PkLen>,
     hashed_derivation_transcript: &[u8],
-) -> Result<TripleDhDerivationResult<D>, ProtocolError<S::Error>>
+) -> Result<TripleDhDerivationResult<D>, ProtocolError>
 where
     D::Core: ProxyHash,
     <D::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
@@ -372,36 +416,24 @@ where
 {
     let mut hkdf = HkdfExtract::<D>::new(None);
 
-    hkdf.input_ikm(
-        &dh.sk1
-            .diffie_hellman(dh.pk1)
-            .map_err(InternalError::into_custom)?,
-    );
-    hkdf.input_ikm(&dh.sk2.diffie_hellman(dh.pk2)?);
-    hkdf.input_ikm(
-        &dh.sk3
-            .diffie_hellman(dh.pk3)
-            .map_err(InternalError::into_custom)?,
-    );
+    hkdf.input_ikm(&shared_secret_1);
+    hkdf.input_ikm(&shared_secret_2);
+    hkdf.input_ikm(&shared_secret_3);
 
     let (_, extracted_ikm) = hkdf.finalize();
     let handshake_secret = derive_secrets::<D>(
         &extracted_ikm,
         STR_HANDSHAKE_SECRET,
         hashed_derivation_transcript,
-    )
-    .map_err(ProtocolError::into_custom)?;
+    )?;
     let session_key = derive_secrets::<D>(
         &extracted_ikm,
         STR_SESSION_KEY,
         hashed_derivation_transcript,
-    )
-    .map_err(ProtocolError::into_custom)?;
+    )?;
 
-    let km2 = hkdf_expand_label::<D>(&handshake_secret, STR_SERVER_MAC, b"")
-        .map_err(ProtocolError::into_custom)?;
-    let km3 = hkdf_expand_label::<D>(&handshake_secret, STR_CLIENT_MAC, b"")
-        .map_err(ProtocolError::into_custom)?;
+    let km2 = hkdf_expand_label::<D>(&handshake_secret, STR_SERVER_MAC, b"")?;
+    let km3 = hkdf_expand_label::<D>(&handshake_secret, STR_CLIENT_MAC, b"")?;
 
     Ok((
         session_key,
@@ -579,6 +611,32 @@ where
     }
 }
 
+impl<KG: KeGroup, D: Hash> Drop for Ke2Builder<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) {
+        struct AssertZeroizeOnDrop<'a, T: ZeroizeOnDrop>(#[allow(unused)] &'a T);
+
+        self.server_nonce.zeroize();
+        self.transcript_hasher.reset();
+        let _ = AssertZeroizeOnDrop(&self.client_e_pk);
+        let _ = AssertZeroizeOnDrop(&self.server_e_pk);
+        self.shared_secret_1.zeroize();
+        self.shared_secret_3.zeroize();
+    }
+}
+
+impl<KG: KeGroup, D: Hash> ZeroizeOnDrop for Ke2Builder<D, KG>
+where
+    D::Core: ProxyHash,
+    <D::Core as BlockSizeUser>::BlockSize: IsLess<U256>,
+    Le<<D::Core as BlockSizeUser>::BlockSize, U256>: NonZero,
+{
+}
+
 impl<KG: KeGroup, D: Hash> Deserialize for Ke2Message<D, KG>
 where
     D::Core: ProxyHash,