Repository: facebook/opaque-ke (language: rust)

Issue:
`Group` implementation support
I have done some work towards implementing P256 support in opaque-ke. In this process I discovered some potential shortcomings in some of the interfaces:
- `PublicKey` and `PrivateKey` are currently fixed to be 32-bytes, which doesn't apply to other prime-order groups, like P256
- [p256's `Scalar` implementation](https://docs.rs/p256/0.9.0/p256/struct.Scalar.html) for example, doesn't support returning a reference to it's underlying memory, but currently `Group::scalar_as_bytes` requires this, I figured that `Group::Scalar` might just implement `Copy` to solve this issue, which might be reasonable?
- `expand_message_xmd` is currently not exposed, which makes sense, but I understand if opaque-ke expects implementors to bring their own implementation
- testing, for example the test vectors given in the RFC, is hard to impossible from outside the crate

I'm not actually suggesting opaque-ke to provide a P256 implementation, though I don't mind if it does and I'm willing to provide the PR. But it seems that the current provided interface has to be improved either way to make implementations work.


Discussion / hints:

Fix patch:
diff --git a/src/envelope.rs b/src/envelope.rs
index e1f71cb7..172a17c3 100644
--- a/src/envelope.rs
+++ b/src/envelope.rs
@@ -32,14 +32,14 @@ const NONCE_LEN: usize = 32;
 fn build_inner_envelope_internal<CS: CipherSuite>(
     random_pwd: &[u8],
     nonce: &[u8],
-) -> Result<PublicKey, InternalPakeError> {
+) -> Result<PublicKey<CS::Group>, InternalPakeError> {
     let h = Hkdf::<CS::Hash>::new(None, random_pwd);
-    let mut keypair_seed = vec![0u8; <PrivateKey as SizedBytes>::Len::to_usize()];
+    let mut keypair_seed = vec![0u8; <PrivateKey<CS::Group> as SizedBytes>::Len::to_usize()];
     h.expand(&[nonce, STR_PRIVATE_KEY].concat(), &mut keypair_seed)
         .map_err(|_| InternalPakeError::HkdfError)?;
     let client_static_keypair =
-        KeyPair::<CS::Group>::from_private_key_slice(CS::Group::scalar_as_bytes(
-            &CS::Group::hash_to_scalar::<CS::Hash>(&keypair_seed[..], STR_OPAQUE_HASH_TO_SCALAR)?,
+        KeyPair::<CS::Group>::from_private_key_slice(&CS::Group::scalar_as_bytes(
+            CS::Group::hash_to_scalar::<CS::Hash>(&keypair_seed[..], STR_OPAQUE_HASH_TO_SCALAR)?,
         ))?;
 
     Ok(client_static_keypair.public().clone())
@@ -50,12 +50,12 @@ fn recover_keys_internal<CS: CipherSuite>(
     nonce: &[u8],
 ) -> Result<KeyPair<CS::Group>, InternalPakeError> {
     let h = Hkdf::<CS::Hash>::new(None, random_pwd);
-    let mut keypair_seed = vec![0u8; <PrivateKey as SizedBytes>::Len::to_usize()];
+    let mut keypair_seed = vec![0u8; <PrivateKey<CS::Group> as SizedBytes>::Len::to_usize()];
     h.expand(&[nonce, STR_PRIVATE_KEY].concat(), &mut keypair_seed)
         .map_err(|_| InternalPakeError::HkdfError)?;
     let client_static_keypair =
-        KeyPair::<CS::Group>::from_private_key_slice(CS::Group::scalar_as_bytes(
-            &CS::Group::hash_to_scalar::<CS::Hash>(&keypair_seed[..], STR_OPAQUE_HASH_TO_SCALAR)?,
+        KeyPair::<CS::Group>::from_private_key_slice(&CS::Group::scalar_as_bytes(
+            CS::Group::hash_to_scalar::<CS::Hash>(&keypair_seed[..], STR_OPAQUE_HASH_TO_SCALAR)?,
         ))?;
 
     Ok(client_static_keypair)
@@ -185,7 +185,7 @@ impl<CS: CipherSuite> Envelope<CS> {
     ) -> Result<
         (
             Self,
-            PublicKey,
+            PublicKey<CS::Group>,
             GenericArray<u8, <CS::Hash as Digest>::OutputSize>,
         ),
         InternalPakeError,
diff --git a/src/group.rs b/src/group.rs
index 2d02a112..6e59b89c 100644
--- a/src/group.rs
+++ b/src/group.rs
@@ -28,9 +28,9 @@ use zeroize::Zeroize;
 /// subgroup is noted additively — as in the draft RFC — in this trait.
 pub trait Group: Copy + Sized + for<'a> Mul<&'a <Self as Group>::Scalar, Output = Self> {
     /// The type of base field scalars
-    type Scalar: Zeroize + Clone;
+    type Scalar: Zeroize + Copy;
     /// The byte length necessary to represent scalars
-    type ScalarLen: ArrayLength<u8>;
+    type ScalarLen: ArrayLength<u8> + 'static;
     /// Return a scalar from its fixed-length bytes representation
     fn from_scalar_slice(
         scalar_bits: &GenericArray<u8, Self::ScalarLen>,
@@ -38,12 +38,12 @@ pub trait Group: Copy + Sized + for<'a> Mul<&'a <Self as Group>::Scalar, Output
     /// picks a scalar at random
     fn random_nonzero_scalar<R: RngCore + CryptoRng>(rng: &mut R) -> Self::Scalar;
     /// Serializes a scalar to bytes
-    fn scalar_as_bytes(scalar: &Self::Scalar) -> &GenericArray<u8, Self::ScalarLen>;
+    fn scalar_as_bytes(scalar: Self::Scalar) -> GenericArray<u8, Self::ScalarLen>;
     /// The multiplicative inverse of this scalar
     fn scalar_invert(scalar: &Self::Scalar) -> Self::Scalar;
 
     /// The byte length necessary to represent group elements
-    type ElemLen: ArrayLength<u8>;
+    type ElemLen: ArrayLength<u8> + 'static;
     /// Return an element from its fixed-length bytes representation
     fn from_element_slice(
         element_bits: &GenericArray<u8, Self::ElemLen>,
@@ -58,7 +58,9 @@ pub trait Group: Copy + Sized + for<'a> Mul<&'a <Self as Group>::Scalar, Output
     type UniformBytesLen: ArrayLength<u8>;
 
     /// Hashes a slice of pseudo-random bytes of the correct length to a curve point
-    fn hash_to_curve(uniform_bytes: &GenericArray<u8, Self::UniformBytesLen>) -> Self;
+    fn hash_to_curve(
+        uniform_bytes: &GenericArray<u8, Self::UniformBytesLen>,
+    ) -> Result<Self, InternalPakeError>;
 
     /// Get the base point for the group
     fn base_point() -> Self;
@@ -105,8 +107,8 @@ impl Group for RistrettoPoint {
             }
         }
     }
-    fn scalar_as_bytes(scalar: &Self::Scalar) -> &GenericArray<u8, Self::ScalarLen> {
-        GenericArray::from_slice(scalar.as_bytes())
+    fn scalar_as_bytes(scalar: Self::Scalar) -> GenericArray<u8, Self::ScalarLen> {
+        scalar.to_bytes().into()
     }
     fn scalar_invert(scalar: &Self::Scalar) -> Self::Scalar {
         scalar.invert()
@@ -128,14 +130,16 @@ impl Group for RistrettoPoint {
     }
 
     type UniformBytesLen = U64;
-    fn hash_to_curve(uniform_bytes: &GenericArray<u8, Self::UniformBytesLen>) -> Self {
+    fn hash_to_curve(
+        uniform_bytes: &GenericArray<u8, Self::UniformBytesLen>,
+    ) -> Result<Self, InternalPakeError> {
         // https://caniuse.rs/features/array_gt_32_impls
         let bits: [u8; 64] = {
             let mut bytes = [0u8; 64];
             bytes.copy_from_slice(uniform_bytes);
             bytes
         };
-        RistrettoPoint::from_uniform_bytes(&bits)
+        Ok(RistrettoPoint::from_uniform_bytes(&bits))
     }
 
     fn base_point() -> Self {
diff --git a/src/key_exchange/traits.rs b/src/key_exchange/traits.rs
index e8a62c13..a02eb54f 100644
--- a/src/key_exchange/traits.rs
+++ b/src/key_exchange/traits.rs
@@ -30,8 +30,8 @@ pub trait KeyExchange<D: Hash, G: Group> {
         l1_bytes: Vec<u8>,
         l2_bytes: Vec<u8>,
         ke1_message: Self::KE1Message,
-        client_s_pk: PublicKey,
-        server_s_sk: PrivateKey,
+        client_s_pk: PublicKey<G>,
+        server_s_sk: PrivateKey<G>,
         id_u: Vec<u8>,
         id_s: Vec<u8>,
         context: Vec<u8>,
@@ -43,8 +43,8 @@ pub trait KeyExchange<D: Hash, G: Group> {
         ke2_message: Self::KE2Message,
         ke1_state: &Self::KE1State,
         serialized_credential_request: &[u8],
-        server_s_pk: PublicKey,
-        client_s_sk: PrivateKey,
+        server_s_pk: PublicKey<G>,
+        client_s_sk: PrivateKey<G>,
         id_u: Vec<u8>,
         id_s: Vec<u8>,
         context: Vec<u8>,
diff --git a/src/key_exchange/tripledh.rs b/src/key_exchange/tripledh.rs
index fa7debf1..5d0da1f0 100644
--- a/src/key_exchange/tripledh.rs
+++ b/src/key_exchange/tripledh.rs
@@ -43,10 +43,10 @@ static STR_OPAQUE: &[u8] = b"OPAQUE-";
 pub struct TripleDH;
 
 impl<D: Hash, G: Group> KeyExchange<D, G> for TripleDH {
-    type KE1State = Ke1State;
+    type KE1State = Ke1State<G>;
     type KE2State = Ke2State<<D as FixedOutput>::OutputSize>;
-    type KE1Message = Ke1Message;
-    type KE2Message = Ke2Message<<D as FixedOutput>::OutputSize>;
+    type KE1Message = Ke1Message<G>;
+    type KE2Message = Ke2Message<G, <D as FixedOutput>::OutputSize>;
     type KE3Message = Ke3Message<<D as FixedOutput>::OutputSize>;
 
     fn generate_ke1<R: RngCore + CryptoRng>(
@@ -75,8 +75,8 @@ impl<D: Hash, G: Group> KeyExchange<D, G> for TripleDH {
         serialized_credential_request: Vec<u8>,
         l2_bytes: Vec<u8>,
         ke1_message: Self::KE1Message,
-        client_s_pk: PublicKey,
-        server_s_sk: PrivateKey,
+        client_s_pk: PublicKey<G>,
+        server_s_sk: PrivateKey<G>,
         id_u: Vec<u8>,
         id_s: Vec<u8>,
         context: Vec<u8>,
@@ -133,8 +133,8 @@ impl<D: Hash, G: Group> KeyExchange<D, G> for TripleDH {
         ke2_message: Self::KE2Message,
         ke1_state: &Self::KE1State,
         serialized_credential_request: &[u8],
-        server_s_pk: PublicKey,
-        client_s_sk: PrivateKey,
+        server_s_pk: PublicKey<G>,
+        client_s_sk: PrivateKey<G>,
         id_u: Vec<u8>,
         id_s: Vec<u8>,
         context: Vec<u8>,
@@ -208,23 +208,44 @@ impl<D: Hash, G: Group> KeyExchange<D, G> for TripleDH {
 }
 
 /// The client state produced after the first key exchange message
-#[derive(PartialEq, Eq, Debug, Hash, Zeroize, Clone)]
 #[cfg_attr(feature = "serialize", derive(serde::Deserialize, serde::Serialize))]
-#[zeroize(drop)]
-pub struct Ke1State {
-    client_e_sk: PrivateKey,
+pub struct Ke1State<G: Group> {
+    client_e_sk: PrivateKey<G>,
     client_nonce: GenericArray<u8, NonceLen>,
 }
 
+impl_clone_for!(
+    struct Ke1State<G: Group>,
+    [client_e_sk, client_nonce],
+);
+impl_debug_eq_hash_for!(
+    struct Ke1State<G: Group>,
+    [client_e_sk, client_nonce],
+);
+
+// This can't be derived because of the use of a generic parameter
+impl<G: Group> Zeroize for Ke1State<G> {
+    fn zeroize(&mut self) {
+        self.client_e_sk.zeroize();
+        self.client_nonce.zeroize();
+    }
+}
+
+impl<G: Group> Drop for Ke1State<G> {
+    fn drop(&mut self) {
+        self.zeroize();
+    }
+}
+
 /// The first key exchange message
 #[derive(PartialEq, Eq, Debug, Hash, Clone)]
 #[cfg_attr(feature = "serialize", derive(serde::Deserialize, serde::Serialize))]
-pub struct Ke1Message {
+pub struct Ke1Message<G: Group> {
     pub(crate) client_nonce: GenericArray<u8, NonceLen>,
-    pub(crate) client_e_pk: PublicKey,
+    pub(crate) client_e_pk: PublicKey<G>,
 }
 
-impl FromBytes for Ke1State {
+impl<G: Group> FromBytes for Ke1State<G> {
     fn from_bytes<CS: CipherSuite>(bytes: &[u8]) -> Result<Self, PakeError> {
         let nonce_len = NonceLen::to_usize();
         let checked_bytes = check_slice_size_atleast(bytes, KEY_LEN + nonce_len, "ke1_state")?;
@@ -238,7 +259,7 @@ impl FromBytes for Ke1State {
     }
 }
 
-impl ToBytesWithPointers for Ke1State {
+impl<G: Group> ToBytesWithPointers for Ke1State<G> {
     fn to_bytes(&self) -> Vec<u8> {
         let output: Vec<u8> = [&self.client_e_sk.to_arr(), &self.client_nonce[..]].concat();
         output
@@ -249,20 +270,20 @@ impl ToBytesWithPointers for Ke1State {
         vec![
             (
                 self.client_e_sk.as_ptr(),
-                <PrivateKey as SizedBytes>::Len::to_usize(),
+                <PrivateKey<G> as SizedBytes>::Len::to_usize(),
             ),
             (self.client_nonce.as_ptr(), NonceLen::to_usize()),
         ]
     }
 }
 
-impl ToBytes for Ke1Message {
+impl<G: Group> ToBytes for Ke1Message<G> {
     fn to_bytes(&self) -> Vec<u8> {
         [&self.client_nonce[..], &self.client_e_pk.to_arr()].concat()
     }
 }
 
-impl FromBytes for Ke1Message {
+impl<G: Group> FromBytes for Ke1Message<G> {
     fn from_bytes<CS: CipherSuite>(ke1_message_bytes: &[u8]) -> Result<Self, PakeError> {
         let nonce_len = NonceLen::to_usize();
         let checked_nonce =
@@ -323,9 +344,9 @@ impl<HashLen: ArrayLength<u8>> ToBytesWithPointers for Ke2State<HashLen> {
 #[derive(Clone, Debug, Eq, Hash, PartialEq)]
 #[cfg_attr(feature = "serialize", derive(serde::Deserialize, serde::Serialize))]
 #[cfg_attr(feature = "serialize", serde(bound = ""))]
-pub struct Ke2Message<HashLen: ArrayLength<u8>> {
+pub struct Ke2Message<G: Group, HashLen: ArrayLength<u8>> {
     server_nonce: GenericArray<u8, NonceLen>,
-    server_e_pk: PublicKey,
+    server_e_pk: PublicKey<G>,
     mac: GenericArray<u8, HashLen>,
 }
 
@@ -344,19 +365,19 @@ impl<HashLen: ArrayLength<u8>> FromBytes for Ke2State<HashLen> {
     }
 }
 
-impl<HashLen: ArrayLength<u8>> ToBytes for Ke2Message<HashLen> {
+impl<G: Group, HashLen: ArrayLength<u8>> ToBytes for Ke2Message<G, HashLen> {
     fn to_bytes(&self) -> Vec<u8> {
         [&self.to_bytes_without_info_or_mac(), &self.mac[..]].concat()
     }
 }
 
-impl<HashLen: ArrayLength<u8>> Ke2Message<HashLen> {
+impl<G: Group, HashLen: ArrayLength<u8>> Ke2Message<G, HashLen> {
     fn to_bytes_without_info_or_mac(&self) -> Vec<u8> {
         [&self.server_nonce[..], &self.server_e_pk.to_arr()].concat()
     }
 }
 
-impl<HashLen: ArrayLength<u8>> FromBytes for Ke2Message<HashLen> {
+impl<G: Group, HashLen: ArrayLength<u8>> FromBytes for Ke2Message<G, HashLen> {
     fn from_bytes<CS: CipherSuite>(input: &[u8]) -> Result<Self, PakeError> {
         let nonce_len = NonceLen::to_usize();
         let checked_nonce = check_slice_size_atleast(input, nonce_len, "ke2_message nonce")?;
@@ -387,13 +408,13 @@ impl<HashLen: ArrayLength<u8>> FromBytes for Ke2Message<HashLen> {
 
 #[allow(clippy::upper_case_acronyms)]
 // The triple of public and private components used in the 3DH computation
-struct TripleDHComponents {
-    pk1: PublicKey,
-    sk1: PrivateKey,
-    pk2: PublicKey,
-    sk2: PrivateKey,
-    pk3: PublicKey,
-    sk3: PrivateKey,
+struct TripleDHComponents<G: Group> {
+    pk1: PublicKey<G>,
+    sk1: PrivateKey<G>,
+    pk2: PublicKey<G>,
+    sk2: PrivateKey<G>,
+    pk3: PublicKey<G>,
+    sk3: PrivateKey<G>,
 }
 
 #[allow(clippy::upper_case_acronyms)]
@@ -433,7 +454,7 @@ impl<HashLen: ArrayLength<u8>> FromBytes for Ke3Message<HashLen> {
 // 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, G: Group>(
-    dh: TripleDHComponents,
+    dh: TripleDHComponents<G>,
     hashed_derivation_transcript: &[u8],
 ) -> Result<TripleDHDerivationResult<D>, ProtocolError> {
     let ikm: Vec<u8> = [
diff --git a/src/keypair.rs b/src/keypair.rs
index 798487dc..b3ecdac7 100644
--- a/src/keypair.rs
+++ b/src/keypair.rs
@@ -11,7 +11,7 @@ use crate::errors::InternalPakeError;
 use crate::group::Group;
 #[cfg(test)]
 use generic_array::typenum::Unsigned;
-use generic_array::{typenum::U32, GenericArray};
+use generic_array::{ArrayLength, GenericArray};
 use generic_bytes::{SizedBytes, TryFromSizedBytesError};
 #[cfg(test)]
 use proptest::prelude::*;
@@ -19,7 +19,6 @@ use proptest::prelude::*;
 use rand::{rngs::StdRng, SeedableRng};
 use rand::{CryptoRng, RngCore};
 use std::fmt::Debug;
-use std::marker::PhantomData;
 use std::ops::Deref;
 use zeroize::Zeroize;
 
@@ -35,31 +34,34 @@ pub trait SizedBytesExt: SizedBytes {
 impl<T> SizedBytesExt for T where T: SizedBytes {}
 
 /// A Keypair trait with public-private verification
-#[cfg_attr(feature = "serialize", derive(serde::Deserialize, serde::Serialize))]
-pub struct KeyPair<G> {
-    pk: PublicKey,
-    sk: PrivateKey,
-    _g: PhantomData<G>,
+#[cfg_attr(
+    feature = "serialize",
+    derive(serde::Deserialize, serde::Serialize),
+    serde(bound = "")
+)]
+pub struct KeyPair<G: Group> {
+    pk: PublicKey<G>,
+    sk: PrivateKey<G>,
 }
 
 impl_clone_for!(
-    struct KeyPair<G>,
-    [pk, sk, _g],
+    struct KeyPair<G: Group>,
+    [pk, sk],
 );
 impl_debug_eq_hash_for!(
-    struct KeyPair<G>,
-    [pk, sk, _g],
+    struct KeyPair<G: Group>,
+    [pk, sk],
 );
 
-// This can't be derived because of the use of a phantom parameter
-impl<G> Zeroize for KeyPair<G> {
+// This can't be derived because of the use of a generic parameter
+impl<G: Group> Zeroize for KeyPair<G> {
     fn zeroize(&mut self) {
         self.pk.zeroize();
         self.sk.zeroize();
     }
 }
 
-impl<G> Drop for KeyPair<G> {
+impl<G: Group> Drop for KeyPair<G> {
     fn drop(&mut self) {
         self.zeroize();
     }
@@ -67,49 +69,45 @@ impl<G> Drop for KeyPair<G> {
 
 impl<G: Group> KeyPair<G> {
     /// The public key component
-    pub fn public(&self) -> &PublicKey {
+    pub fn public(&self) -> &PublicKey<G> {
         &self.pk
     }
 
     /// The private key component
-    pub fn private(&self) -> &PrivateKey {
+    pub fn private(&self) -> &PrivateKey<G> {
         &self.sk
     }
 
     /// Generating a random key pair given a cryptographic rng
     pub(crate) fn generate_random<R: RngCore + CryptoRng>(rng: &mut R) -> Self {
         let sk = G::random_nonzero_scalar(rng);
-        let sk_bytes = G::scalar_as_bytes(&sk);
-        let pk = G::base_point().mult_by_slice(sk_bytes);
+        let sk_bytes = G::scalar_as_bytes(sk);
+        let pk = G::base_point().mult_by_slice(&sk_bytes);
         Self {
-            pk: PublicKey(Key(pk.to_arr().to_vec())),
-            sk: PrivateKey(Key(sk_bytes.to_vec())),
-            _g: PhantomData,
+            pk: PublicKey(Key(pk.to_arr())),
+            sk: PrivateKey(Key(sk_bytes)),
         }
     }
 
     /// Obtaining a public key from secret bytes. At all times, we should have
     /// &public_from_private(self.private()) == self.public()
-    pub(crate) fn public_from_private(bytes: &PrivateKey) -> PublicKey {
+    pub(crate) fn public_from_private(bytes: &PrivateKey<G>) -> PublicKey<G> {
         let bytes_data = GenericArray::<u8, G::ScalarLen>::from_slice(&bytes.0[..]);
-        PublicKey(Key(G::base_point()
-            .mult_by_slice(bytes_data)
-            .to_arr()
-            .to_vec()))
+        PublicKey(Key(G::base_point().mult_by_slice(bytes_data).to_arr()))
     }
 
     /// Check whether a public key is valid. This is meant to be applied on
     /// material provided through the network which fits the key
     /// representation (i.e. can be mapped to a curve point), but presents
     /// some risk - e.g. small subgroup check
-    pub(crate) fn check_public_key(key: PublicKey) -> Result<PublicKey, InternalPakeError> {
+    pub(crate) fn check_public_key(key: PublicKey<G>) -> Result<PublicKey<G>, InternalPakeError> {
         G::from_element_slice(GenericArray::from_slice(&key.0)).map(|_| key)
     }
 
     /// Computes the diffie hellman function on a public key and private key
     pub(crate) fn diffie_hellman(
-        pk: PublicKey,
-        sk: PrivateKey,
+        pk: PublicKey<G>,
+        sk: PrivateKey<G>,
     ) -> Result<Vec<u8>, InternalPakeError> {
         let pk_data = GenericArray::<u8, G::ElemLen>::from_slice(&pk.0[..]);
         let point = G::from_element_slice(pk_data)?;
@@ -119,20 +117,16 @@ impl<G: Group> KeyPair<G> {
 
     /// Obtains a KeyPair from a slice representing the private key
     pub fn from_private_key_slice(input: &[u8]) -> Result<Self, InternalPakeError> {
-        let sk = PrivateKey(Key::from_arr(GenericArray::from_slice(input))?);
+        let sk = PrivateKey(Key(GenericArray::clone_from_slice(input)));
         let pk = Self::public_from_private(&sk);
-        Ok(Self {
-            pk,
-            sk,
-            _g: PhantomData,
-        })
+        Ok(Self { pk, sk })
     }
 
     #[cfg(test)]
     pub fn as_byte_ptrs(&self) -> Vec<(*const u8, usize)> {
         vec![
-            (self.pk.as_ptr(), KeyLen::to_usize()),
-            (self.sk.as_ptr(), KeyLen::to_usize()),
+            (self.pk.as_ptr(), G::ElemLen::to_usize()),
+            (self.sk.as_ptr(), G::ScalarLen::to_usize()),
         ]
     }
 }
@@ -154,18 +148,56 @@ impl<G: Group + Debug> KeyPair<G> {
     }
 }
 
-type KeyLen = U32;
-
 /// A minimalist key type built around a \[u8; 32\]
-#[derive(Debug, PartialEq, Eq, Clone, Hash, Zeroize)]
-#[cfg_attr(feature = "serialize", derive(serde::Deserialize, serde::Serialize))]
-// Ensure Key material is zeroed after use.
-#[zeroize(drop)]
+#[cfg_attr(
+    feature = "serialize",
+    derive(serde::Deserialize, serde::Serialize),
+    serde(bound = "")
+)]
 #[repr(transparent)]
-pub struct Key(Vec<u8>);
+pub struct Key<L: ArrayLength<u8>>(GenericArray<u8, L>);
+
+impl<L: ArrayLength<u8>> Clone for Key<L> {
+    fn clone(&self) -> Self {
+        Self(self.0.clone())
+    }
+}
+
+impl<L: ArrayLength<u8>> Debug for Key<L> {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("Key").field(&self.0).finish()
+    }
+}
 
-impl Deref for Key {
-    type Target = Vec<u8>;
+impl<L: ArrayLength<u8>> Eq for Key<L> {}
+
+impl<L: ArrayLength<u8>> PartialEq for Key<L> {
+    fn eq(&self, other: &Self) -> bool {
+        self.0.eq(&other.0)
+    }
+}
+
+impl<L: ArrayLength<u8>> std::hash::Hash for Key<L> {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.0.hash(state);
+    }
+}
+
+// This can't be derived because of the use of a generic parameter
+impl<L: ArrayLength<u8>> Zeroize for Key<L> {
+    fn zeroize(&mut self) {
+        self.0.zeroize();
+    }
+}
+
+impl<L: ArrayLength<u8>> Drop for Key<L> {
+    fn drop(&mut self) {
+        self.zeroize();
+    }
+}
+
+impl<L: ArrayLength<u8>> Deref for Key<L> {
+    type Target = GenericArray<u8, L>;
 
     fn deref(&self) -> &Self::Target {
         &self.0
@@ -173,35 +205,54 @@ impl Deref for Key {
 }
 
 // Don't make it implement SizedBytes so that it's not constructible outside of this module.
-impl Key {
-    fn to_arr(&self) -> GenericArray<u8, KeyLen> {
+impl<L: ArrayLength<u8>> Key<L> {
+    fn to_arr(&self) -> GenericArray<u8, L> {
         GenericArray::clone_from_slice(&self.0[..])
     }
 
     #[allow(clippy::unnecessary_wraps)]
-    fn from_arr(key_bytes: &GenericArray<u8, KeyLen>) -> Result<Self, TryFromSizedBytesError> {
-        Ok(Key(key_bytes.to_vec()))
+    fn from_arr(key_bytes: &GenericArray<u8, L>) -> Result<Self, TryFromSizedBytesError> {
+        Ok(Key(key_bytes.to_owned()))
     }
 }
 
 /// Wrapper around a Key to enforce that it's a private one.
-#[derive(Debug, PartialEq, Eq, Clone, Hash, Zeroize)]
 #[cfg_attr(feature = "serialize", derive(serde::Deserialize, serde::Serialize))]
-// Ensure Key material is zeroed after use.
-#[zeroize(drop)]
 #[repr(transparent)]
-pub struct PrivateKey(Key);
+pub struct PrivateKey<G: Group>(Key<G::ScalarLen>);
 
-impl Deref for PrivateKey {
-    type Target = Key;
+impl_clone_for!(
+    tuple PrivateKey<G: Group>,
+    [0],
+);
+impl_debug_eq_hash_for!(
+    tuple PrivateKey<G: Group>,
+    [0],
+);
+
+// This can't be derived because of the use of a generic parameter
+impl<G: Group> Zeroize for PrivateKey<G> {
+    fn zeroize(&mut self) {
+        self.0.zeroize();
+    }
+}
+
+impl<G: Group> Drop for PrivateKey<G> {
+    fn drop(&mut self) {
+        self.zeroize();
+    }
+}
+
+impl<G: Group> Deref for PrivateKey<G> {
+    type Target = Key<G::ScalarLen>;
 
     fn deref(&self) -> &Self::Target {
         &self.0
     }
 }
 
-impl SizedBytes for PrivateKey {
-    type Len = KeyLen;
+impl<G: Group> SizedBytes for PrivateKey<G> {
+    type Len = G::ScalarLen;
 
     fn to_arr(&self) -> GenericArray<u8, Self::Len> {
         self.0.to_arr()
@@ -213,23 +264,42 @@ impl SizedBytes for PrivateKey {
 }
 
 /// Wrapper around a Key to enforce that it's a public one.
-#[derive(Debug, PartialEq, Eq, Clone, Hash, Zeroize)]
 #[cfg_attr(feature = "serialize", derive(serde::Deserialize, serde::Serialize))]
-// Ensure Key material is zeroed after use.
-#[zeroize(drop)]
 #[repr(transparent)]
-pub struct PublicKey(Key);
+pub struct PublicKey<G: Group>(Key<G::ElemLen>);
+
+impl_clone_for!(
+    tuple PublicKey<G: Group>,
+    [0],
+);
+impl_debug_eq_hash_for!(
+    tuple PublicKey<G: Group>,
+    [0],
+);
+
+// This can't be derived because of the use of a generic parameter
+impl<G: Group> Zeroize for PublicKey<G> {
+    fn zeroize(&mut self) {
+        self.0.zeroize();
+    }
+}
+
+impl<G: Group> Drop for PublicKey<G> {
+    fn drop(&mut self) {
+        self.zeroize();
+    }
+}
 
-impl Deref for PublicKey {
-    type Target = Key;
+impl<G: Group> Deref for PublicKey<G> {
+    type Target = Key<G::ElemLen>;
 
     fn deref(&self) -> &Self::Target {
         &self.0
     }
 }
 
-impl SizedBytes for PublicKey {
-    type Len = KeyLen;
+impl<G: Group> SizedBytes for PublicKey<G> {
+    type Len = G::ElemLen;
 
     fn to_arr(&self) -> GenericArray<u8, Self::Len> {
         self.0.to_arr()
@@ -251,8 +321,12 @@ mod tests {
 
     #[test]
     fn test_zeroize_key() -> Result<(), ProtocolError> {
-        let key_len = KeyLen::to_usize();
-        let mut key = Key(vec![1u8; key_len]);
+        let key_len = <RistrettoPoint as Group>::ElemLen::to_usize();
+        let mut key =
+            Key::<<RistrettoPoint as Group>::ElemLen>(GenericArray::clone_from_slice(&vec![
+                1u8;
+                key_len
+            ]));
         let ptr = key.as_ptr();
 
         key.zeroize();
diff --git a/src/map_to_curve.rs b/src/map_to_curve.rs
index 24b05253..7bbb40a5 100644
--- a/src/map_to_curve.rs
+++ b/src/map_to_curve.rs
@@ -43,9 +43,7 @@ impl GroupWithMapToCurve for RistrettoPoint {
     fn map_to_curve<H: Hash>(msg: &[u8], dst: &[u8]) -> Result<Self, InternalPakeError> {
         let uniform_bytes =
             expand_message_xmd::<H>(msg, dst, <H as Digest>::OutputSize::to_usize())?;
-        Ok(<Self as Group>::hash_to_curve(
-            &GenericArray::clone_from_slice(&uniform_bytes[..]),
-        ))
+        <Self as Group>::hash_to_curve(&GenericArray::clone_from_slice(&uniform_bytes[..]))
     }
 
     fn hash_to_scalar<H: Hash>(
@@ -75,9 +73,9 @@ fn xor(x: &[u8], y: &[u8]) -> Result<Vec<u8>, InternalPakeError> {
     Ok(x.iter().zip(y).map(|(&x1, &x2)| x1 ^ x2).collect())
 }
 
-// Corresponds to the expand_message_xmd() function defined in
-// https://www.ietf.org/archive/id/draft-irtf-cfrg-hash-to-curve-10.txt
-pub(crate) fn expand_message_xmd<H: Hash>(
+/// Corresponds to the expand_message_xmd() function defined in
+/// https://www.ietf.org/archive/id/draft-irtf-cfrg-hash-to-curve-10.txt
+pub fn expand_message_xmd<H: Hash>(
     msg: &[u8],
     dst: &[u8],
     len_in_bytes: usize,
diff --git a/src/messages.rs b/src/messages.rs
index 04bf4fcc..f14dfa0d 100644
--- a/src/messages.rs
+++ b/src/messages.rs
@@ -79,7 +79,7 @@ pub struct RegistrationResponse<CS: CipherSuite> {
     /// The server's oprf output
     pub(crate) beta: CS::Group,
     /// Server's static public key
-    pub(crate) server_s_pk: PublicKey,
+    pub(crate) server_s_pk: PublicKey<CS::Group>,
 }
 
 // Cannot be derived because it would require for CS to be Clone.
@@ -107,7 +107,7 @@ impl<CS: CipherSuite> RegistrationResponse<CS> {
     /// Deserialization from bytes
     pub fn deserialize(input: &[u8]) -> Result<Self, ProtocolError> {
         let elem_len = <CS::Group as Group>::ElemLen::to_usize();
-        let key_len = <PublicKey as SizedBytes>::Len::to_usize();
+        let key_len = <PublicKey<CS::Group> as SizedBytes>::Len::to_usize();
         let checked_slice =
             check_slice_size(input, elem_len + key_len, "registration_response_bytes")?;
 
@@ -141,7 +141,7 @@ pub struct RegistrationUpload<CS: CipherSuite> {
     /// The masking key used to mask the envelope
     pub(crate) masking_key: GenericArray<u8, <CS::Hash as Digest>::OutputSize>,
     /// The user's public key
-    pub(crate) client_s_pk: PublicKey,
+    pub(crate) client_s_pk: PublicKey<CS::Group>,
 }
 
 impl_clone_for!(
@@ -166,7 +166,7 @@ impl<CS: CipherSuite> RegistrationUpload<CS> {
 
     /// Deserialization from bytes
     pub fn deserialize(input: &[u8]) -> Result<Self, ProtocolError> {
-        let key_len = <PublicKey as SizedBytes>::Len::to_usize();
+        let key_len = <PublicKey<CS::Group> as SizedBytes>::Len::to_usize();
         let hash_len = <CS::Hash as Digest>::OutputSize::to_usize();
         let checked_slice =
             check_slice_size_atleast(input, key_len + hash_len, "registration_upload_bytes")?;
@@ -311,7 +311,7 @@ impl<CS: CipherSuite> CredentialResponse<CS> {
     /// Deserialization from bytes
     pub fn deserialize(input: &[u8]) -> Result<Self, ProtocolError> {
         let elem_len = <CS::Group as Group>::ElemLen::to_usize();
-        let key_len = <PublicKey as SizedBytes>::Len::to_usize();
+        let key_len = <PublicKey<CS::Group> as SizedBytes>::Len::to_usize();
         let nonce_len: usize = 32;
         let envelope_len = Envelope::<CS>::len();
         let masked_response_len = key_len + envelope_len;
diff --git a/src/opaque.rs b/src/opaque.rs
index 1209461a..48c9f40d 100644
--- a/src/opaque.rs
+++ b/src/opaque.rs
@@ -69,7 +69,7 @@ impl<CS: CipherSuite> ServerSetup<CS> {
     /// Deserialization from bytes
     pub fn deserialize(input: &[u8]) -> Result<Self, ProtocolError> {
         let seed_len = <CS::Hash as Digest>::OutputSize::to_usize();
-        let key_len = <PrivateKey as SizedBytes>::Len::to_usize();
+        let key_len = <PrivateKey<CS::Group> as SizedBytes>::Len::to_usize();
         let checked_slice = check_slice_size(input, seed_len + key_len + key_len, "server_setup")?;
 
         Ok(Self {
@@ -115,7 +115,7 @@ impl<CS: CipherSuite> ClientRegistration<CS> {
     /// Serialization into bytes
     pub fn serialize(&self) -> Vec<u8> {
         [
-            &CS::Group::scalar_as_bytes(&self.token.blind)[..],
+            &CS::Group::scalar_as_bytes(self.token.blind)[..],
             &self.token.data,
         ]
         .concat()
@@ -242,7 +242,7 @@ pub struct ClientRegistrationFinishResult<CS: CipherSuite> {
     /// The export key output by client registration
     pub export_key: GenericArray<u8, <CS::Hash as Digest>::OutputSize>,
     /// The server's static public key
-    pub server_s_pk: PublicKey,
+    pub server_s_pk: PublicKey<CS::Group>,
     /// Instance of the ClientRegistration, only used in tests for checking zeroize
     #[cfg(test)]
     pub state: ClientRegistration<CS>,
@@ -415,7 +415,7 @@ impl<CS: CipherSuite> ClientLogin<CS> {
     /// Serialization into bytes
     pub fn serialize(&self) -> Vec<u8> {
         let output: Vec<u8> = [
-            &CS::Group::scalar_as_bytes(&self.token.blind)[..],
+            &CS::Group::scalar_as_bytes(self.token.blind)[..],
             &serialize(&self.serialized_credential_request, 2),
             &serialize(&self.ke1_state.to_bytes(), 2),
             &self.token.data,
@@ -517,7 +517,7 @@ pub struct ClientLoginFinishResult<CS: CipherSuite> {
     /// The client-side export key
     pub export_key: GenericArray<u8, <CS::Hash as Digest>::OutputSize>,
     /// The server's static public key
-    pub server_s_pk: PublicKey,
+    pub server_s_pk: PublicKey<CS::Group>,
     /// Instance of the ClientLogin, only used in tests for checking zeroize
     #[cfg(test)]
     pub state: ClientLogin<CS>,
@@ -915,7 +915,7 @@ fn oprf_key_from_seed<G: GroupWithMapToCurve, D: Hash>(
     oprf_seed: &GenericArray<u8, D::OutputSize>,
     credential_identifier: &[u8],
 ) -> Result<G::Scalar, InternalPakeError> {
-    let mut oprf_key_bytes = vec![0u8; <PrivateKey as SizedBytes>::Len::to_usize()];
+    let mut oprf_key_bytes = vec![0u8; <PrivateKey<G> as SizedBytes>::Len::to_usize()];
     Hkdf::<D>::from_prk(oprf_seed)
         .map_err(|_| InternalPakeError::HkdfError)?
         .expand(
@@ -929,10 +929,11 @@ fn oprf_key_from_seed<G: GroupWithMapToCurve, D: Hash>(
 fn mask_response<CS: CipherSuite>(
     masking_key: &[u8],
     masking_nonce: &[u8],
-    server_s_pk: &PublicKey,
+    server_s_pk: &PublicKey<CS::Group>,
     envelope: &Envelope<CS>,
 ) -> Result<Vec<u8>, ProtocolError> {
-    let mut xor_pad = vec![0u8; <PublicKey as SizedBytes>::Len::to_usize() + Envelope::<CS>::len()];
+    let mut xor_pad =
+        vec![0u8; <PublicKey<CS::Group> as SizedBytes>::Len::to_usize() + Envelope::<CS>::len()];
     Hkdf::<CS::Hash>::from_prk(masking_key)
         .map_err(|_| InternalPakeError::HkdfError)?
         .expand(
@@ -954,8 +955,9 @@ fn unmask_response<CS: CipherSuite>(
     masking_key: &[u8],
     masking_nonce: &[u8],
     masked_response: &[u8],
-) -> Result<(PublicKey, Envelope<CS>), ProtocolError> {
-    let mut xor_pad = vec![0u8; <PublicKey as SizedBytes>::Len::to_usize() + Envelope::<CS>::len()];
+) -> Result<(PublicKey<CS::Group>, Envelope<CS>), ProtocolError> {
+    let mut xor_pad =
+        vec![0u8; <PublicKey<CS::Group> as SizedBytes>::Len::to_usize() + Envelope::<CS>::len()];
     Hkdf::<CS::Hash>::from_prk(masking_key)
         .map_err(|_| InternalPakeError::HkdfError)?
         .expand(
@@ -968,7 +970,7 @@ fn unmask_response<CS: CipherSuite>(
         .zip(masked_response.iter())
         .map(|(&x1, &x2)| x1 ^ x2)
         .collect();
-    let key_len = <PublicKey as SizedBytes>::Len::to_usize();
+    let key_len = <PublicKey<CS::Group> as SizedBytes>::Len::to_usize();
     let unchecked_server_s_pk =
         PublicKey::from_arr(&GenericArray::clone_from_slice(&plaintext[..key_len]))?;
     let envelope = Envelope::deserialize(&plaintext[key_len..])?;


Test patch:
diff --git a/src/serialization/tests.rs b/src/serialization/tests.rs
index aea155dc..202dcf7e 100644
--- a/src/serialization/tests.rs
+++ b/src/serialization/tests.rs
@@ -211,8 +211,11 @@ fn credential_response_roundtrip() {
     let mut masking_nonce = vec![0u8; 32];
     rng.fill_bytes(&mut masking_nonce);
 
-    let mut masked_response =
-        vec![0u8; <PublicKey as SizedBytes>::Len::to_usize() + Envelope::<Default>::len()];
+    let mut masked_response = vec![
+        0u8;
+        <PublicKey<RistrettoPoint> as SizedBytes>::Len::to_usize()
+            + Envelope::<Default>::len()
+    ];
     rng.fill_bytes(&mut masked_response);
 
     let server_e_kp = KeyPair::<<Default as CipherSuite>::Group>::generate_random(&mut rng);
diff --git a/src/tests/full_test.rs b/src/tests/full_test.rs
index 4f86c2b1..5ec44c75 100644
--- a/src/tests/full_test.rs
+++ b/src/tests/full_test.rs
@@ -307,13 +307,13 @@ fn generate_parameters<CS: CipherSuite>() -> TestVectorParameters {
     .unwrap();
 
     let blinding_factor = CS::Group::random_nonzero_scalar(&mut rng);
-    let blinding_factor_bytes = CS::Group::scalar_as_bytes(&blinding_factor).clone();
+    let blinding_factor_bytes = CS::Group::scalar_as_bytes(blinding_factor);
 
     let mut blinding_factor_registration_rng = CycleRng::new(blinding_factor_bytes.to_vec());
     let client_registration_start_result =
         ClientRegistration::<CS>::start(&mut blinding_factor_registration_rng, password).unwrap();
     let blinding_factor_bytes_returned =
-        CS::Group::scalar_as_bytes(&client_registration_start_result.state.token.blind).clone();
+        CS::Group::scalar_as_bytes(client_registration_start_result.state.token.blind);
     assert_eq!(
         hex::encode(&blinding_factor_bytes),
         hex::encode(&blinding_factor_bytes_returned)
diff --git a/src/tests/opaque_test_vectors.rs b/src/tests/opaque_test_vectors.rs
index e6c6d3b3..a3d78bd2 100644
--- a/src/tests/opaque_test_vectors.rs
+++ b/src/tests/opaque_test_vectors.rs
@@ -445,7 +445,7 @@ fn populate_test_vectors(values: &Value) -> TestVectorParameters {
         dummy_private_key: parse_default!(
             values,
             "client_private_key",
-            vec![0u8; <PrivateKey as SizedBytes>::Len::to_usize()]
+            vec![0u8; <PrivateKey<RistrettoPoint> as SizedBytes>::Len::to_usize()]
         ),
         dummy_masking_key: parse_default!(values, "masking_key", vec![0u8; 64]),
         context: parse!(values, "Context"),
diff --git a/src/tests/voprf_test_vectors.rs b/src/tests/voprf_test_vectors.rs
index dbfe361c..1d8a4493 100644
--- a/src/tests/voprf_test_vectors.rs
+++ b/src/tests/voprf_test_vectors.rs
@@ -73,7 +73,7 @@ fn test_blind() -> Result<(), PakeError> {
 
         assert_eq!(
             &parameters.blind,
-            &RistrettoPoint::scalar_as_bytes(&token.blind).to_vec()
+            &RistrettoPoint::scalar_as_bytes(token.blind).to_vec()
         );
         assert_eq!(
             &parameters.blinded_element,