ml-kem: replace EncodedSizeUser with ExpandedKeyEncoding

First, this commit completely migrates `EncapsulationKey` to using only
the `kem`/`crypto-common` traits: `KeySizeUser`, `TryKeyInit`,
`KeyExport`, and changes any remaining uses to use only the new traits.

Since `DecapsulationKey` uses those same traits for handling `Seed`s,
the only remaining use of the old `EncodedSizeUser` trait is handling
the expanded form of `DecapsulationKey`. So this commit repurposes it
into an `ExpandedKeyEncoding` trait.

Like `DecapsulationKey::from_expanded`, the trait has been marked
deprecated with a rationale given in the documentation for
`ExpandedKeyEncoding`, namely that the expanded form has only
disadvantages when compared to seeds which are significantly smaller,
uniformly sized, and avoid the need to do expanded key validation.
It also notes several ML-KEM libraries have dropped support entirely.

In the `ml-kem` crate, for now, we still need this functionality if only
for tests which have been written generically, including but not limited
to the ones that run the NIST ACVP vectors.

Author: tarcieri

ml-kem/benches/mlkem.rs

@@ -1,4 +1,4 @@
-use ::kem::{Decapsulate, Encapsulate, Kem, KeyExport, KeyInit, TryKeyInit};
+use ::kem::{Decapsulate, Encapsulate, Kem, KeyExport, KeyInit};
 use core::hint::black_box;
 use criterion::{Criterion, criterion_group, criterion_main};
 use getrandom::SysRng;

ml-kem/src/decapsulation_key.rs

@@ -1,11 +1,14 @@
 use crate::{
-    B32, EncapsulationKey, Encoded, EncodedSizeUser, ExpandedDecapsulationKey, Seed, SharedKey,
+    B32, EncapsulationKey, Seed, SharedKey,
     crypto::{G, J},
-    kem::{Generate, InvalidKey, Kem, KeyInit, KeySizeUser},
-    param::{DecapsulationKeySize, KemParams},
+    kem::{Generate, InvalidKey, Kem, KeyExport, KeyInit, KeySizeUser},
+    param::{DecapsulationKeySize, ExpandedDecapsulationKey, KemParams},
     pke::{DecryptionKey, EncryptionKey},
 };
-use array::sizes::{U32, U64};
+use array::{
+    Array, ArraySize,
+    sizes::{U32, U64},
+};
 use kem::{Ciphertext, Decapsulate};
 use rand_core::{TryCryptoRng, TryRng};
 use subtle::{ConditionallySelectable, ConstantTimeEq};
@@ -41,7 +44,7 @@ where
     /// Initialize a [`DecapsulationKey`] from the serialized expanded key form.
     ///
     /// Note that this form is deprecated in practice; prefer to use
-    /// [`DecapsulationKey::from_seed`].
+    /// [`DecapsulationKey::from_seed`]. See [`ExpandedKeyEncoding`] for more information.
     ///
     /// # Errors
     /// - Returns [`InvalidKey`] in the event the expanded key failed validation
@@ -164,24 +167,6 @@ where
     }
 }
 
-impl<P> EncodedSizeUser for DecapsulationKey<P>
-where
-    P: KemParams,
-{
-    type EncodedSize = DecapsulationKeySize<P>;
-
-    fn from_encoded_bytes(expanded: &Encoded<Self>) -> Result<Self, InvalidKey> {
-        #[allow(deprecated)]
-        Self::from_expanded(expanded)
-    }
-
-    fn to_encoded_bytes(&self) -> Encoded<Self> {
-        let dk_pke = self.dk_pke.to_bytes();
-        let ek = self.ek.to_encoded_bytes();
-        P::concat_dk(dk_pke, ek, self.ek.h(), self.z.clone())
-    }
-}
-
 impl<P> Generate for DecapsulationKey<P>
 where
     P: KemParams,
@@ -210,3 +195,68 @@ where
         Self::from_seed(*seed)
     }
 }
+
+/// DEPRECATED: support for encoding and decoding [`DecapsulationKey`]s in the legacy expanded form,
+/// as opposed to the more widely adopted [`Seed`] form.
+///
+/// The expanded encoding format is problematic for several reasons, notably they need to validated
+/// whereas generation from seeds is always correct, meaning there is no performance advantage to
+/// using them, only additional complexity.
+///
+/// They are significantly larger than seeds (which are 64-bytes) and their sizes vary depending on
+/// security level whereas the size of a seed is constant:
+/// - ML-KEM-512: 1632 bytes
+/// - ML-KEM-768: 2400 bytes
+/// - ML-KEM-1024: 3168 bytes
+///
+/// Many ML-KEM libraries have dropped support for this format entirely.
+#[deprecated(since = "0.3.0", note = "use `DecapsulationKey::from_seed` instead")]
+pub trait ExpandedKeyEncoding: Sized {
+    /// The size of an expanded decapsulation key.
+    type EncodedSize: ArraySize;
+
+    /// Parse a [`DecapsulationKey`] from its legacy expanded form.
+    ///
+    /// # Errors
+    /// - If the key fails to validate successfully.
+    fn from_expanded_bytes(enc: &Array<u8, Self::EncodedSize>) -> Result<Self, InvalidKey>;
+
+    /// Serialize a [`DecapsulationKey`] to its legacy expanded form.
+    fn to_expanded_bytes(&self) -> Array<u8, Self::EncodedSize>;
+}
+
+#[allow(deprecated)]
+impl<P> ExpandedKeyEncoding for DecapsulationKey<P>
+where
+    P: KemParams,
+{
+    type EncodedSize = DecapsulationKeySize<P>;
+
+    fn from_expanded_bytes(expanded: &ExpandedDecapsulationKey<P>) -> Result<Self, InvalidKey> {
+        Self::from_expanded(expanded)
+    }
+
+    fn to_expanded_bytes(&self) -> ExpandedDecapsulationKey<P> {
+        let dk_pke = self.dk_pke.to_bytes();
+        let ek = self.ek.to_bytes();
+        P::concat_dk(dk_pke, ek, self.ek.h(), self.z.clone())
+    }
+}
+
+/// Initialize a KEM from a seed.
+pub trait FromSeed: Kem {
+    /// Using the provided [`Seed`] value, create a KEM keypair.
+    fn from_seed(seed: &Seed) -> (Self::DecapsulationKey, Self::EncapsulationKey);
+}
+
+impl<K> FromSeed for K
+where
+    K: Kem,
+    K::DecapsulationKey: KeyInit + KeySizeUser<KeySize = U64>,
+{
+    fn from_seed(seed: &Seed) -> (K::DecapsulationKey, K::EncapsulationKey) {
+        let dk = K::DecapsulationKey::new(seed);
+        let ek = dk.as_ref().clone();
+        (dk, ek)
+    }
+}

ml-kem/src/encapsulation_key.rs

@@ -1,5 +1,5 @@
 use crate::{
-    B32, Encoded, EncodedSizeUser, SharedKey,
+    B32, SharedKey,
     crypto::{G, H},
     kem::{InvalidKey, Kem, Key, KeyExport, KeySizeUser, TryKeyInit},
     param::{EncapsulationKeySize, KemParams},
@@ -24,6 +24,16 @@ impl<P> EncapsulationKey<P>
 where
     P: Kem<SharedKeySize = U32> + KemParams,
 {
+    /// Create a new [`EncapsulationKey`] from its serialized form.
+    ///
+    /// # Errors
+    /// If the key failed validation during decoding.
+    pub fn new(encapsulation_key: &Key<Self>) -> Result<Self, InvalidKey> {
+        EncryptionKey::from_bytes(encapsulation_key)
+            .map(Self::from_encryption_key)
+            .map_err(|_| InvalidKey)
+    }
+
     /// Encapsulates with the given randomness. This is useful for testing against known vectors.
     ///
     /// # Warning
@@ -66,21 +76,6 @@ where
     }
 }
 
-impl<P> EncodedSizeUser for EncapsulationKey<P>
-where
-    P: KemParams,
-{
-    type EncodedSize = EncapsulationKeySize<P>;
-
-    fn from_encoded_bytes(enc: &Encoded<Self>) -> Result<Self, InvalidKey> {
-        Ok(Self::from_encryption_key(EncryptionKey::from_bytes(enc)?))
-    }
-
-    fn to_encoded_bytes(&self) -> Encoded<Self> {
-        self.ek_pke.to_bytes()
-    }
-}
-
 impl<P> KeyExport for EncapsulationKey<P>
 where
     P: KemParams,
@@ -102,9 +97,7 @@ where
     P: KemParams,
 {
     fn new(encapsulation_key: &Key<Self>) -> Result<Self, InvalidKey> {
-        EncryptionKey::from_bytes(encapsulation_key)
-            .map(Self::from_encryption_key)
-            .map_err(|_| InvalidKey)
+        Self::new(encapsulation_key)
     }
 }
 

ml-kem/src/lib.rs

@@ -68,13 +68,13 @@ mod pke;
 /// Section 7. Parameter Sets
 mod param;
 
+/// PKCS#8 key encoding support
 pub mod pkcs8;
 
-/// Trait definitions
-mod traits;
-
 pub use array;
-pub use decapsulation_key::DecapsulationKey;
+#[allow(deprecated)]
+pub use decapsulation_key::ExpandedKeyEncoding;
+pub use decapsulation_key::{DecapsulationKey, FromSeed};
 pub use encapsulation_key::EncapsulationKey;
 pub use kem::{
     self, Ciphertext, Decapsulate, Encapsulate, Generate, InvalidKey, Kem, Key, KeyExport, KeyInit,
@@ -85,7 +85,6 @@ pub use ml_kem_768::MlKem768;
 pub use ml_kem_1024::MlKem1024;
 pub use module_lattice::encoding::ArraySize;
 pub use param::{ExpandedDecapsulationKey, ParameterSet};
-pub use traits::*;
 
 use array::{
     Array,
@@ -297,21 +296,41 @@ mod test {
         round_trip_test::<MlKem1024>();
     }
 
+    fn seed_test<P>()
+    where
+        P: KemParams,
+    {
+        let mut rng = UnwrapErr(SysRng);
+        let mut seed = Seed::default();
+        rng.try_fill_bytes(&mut seed).unwrap();
+
+        let dk = DecapsulationKey::<P>::from_seed(seed.clone());
+        let seed_encoded = dk.to_seed().unwrap();
+        assert_eq!(seed, seed_encoded);
+
+        let ek_original = dk.encapsulation_key();
+        let ek_encoded = ek_original.to_bytes();
+        let ek_decoded = EncapsulationKey::new(&ek_encoded).unwrap();
+        assert_eq!(ek_original, &ek_decoded);
+    }
+
+    #[test]
+    fn seed() {
+        seed_test::<MlKem512>();
+        seed_test::<MlKem768>();
+        seed_test::<MlKem1024>();
+    }
+
+    #[allow(deprecated)]
     fn expanded_key_test<P>()
     where
         P: KemParams,
     {
         let mut rng = UnwrapErr(SysRng);
         let dk_original = DecapsulationKey::<P>::generate_from_rng(&mut rng);
-        let ek_original = dk_original.encapsulation_key().clone();
-
-        let dk_encoded = dk_original.to_encoded_bytes();
-        let dk_decoded = DecapsulationKey::from_encoded_bytes(&dk_encoded).unwrap();
+        let dk_encoded = dk_original.to_expanded_bytes();
+        let dk_decoded = DecapsulationKey::from_expanded_bytes(&dk_encoded).unwrap();
         assert_eq!(dk_original, dk_decoded);
-
-        let ek_encoded = ek_original.to_encoded_bytes();
-        let ek_decoded = EncapsulationKey::from_encoded_bytes(&ek_encoded).unwrap();
-        assert_eq!(ek_original, ek_decoded);
     }
 
     #[test]
@@ -328,13 +347,17 @@ mod test {
         let mut rng = UnwrapErr(SysRng);
         let dk_original = DecapsulationKey::<P>::generate_from_rng(&mut rng);
 
-        let mut dk_encoded = dk_original.to_encoded_bytes();
+        #[allow(deprecated)]
+        let mut dk_encoded = dk_original.to_expanded_bytes();
+
         // Corrupt the hash value
         let hash_offset = P::NttVectorSize::USIZE + P::EncryptionKeySize::USIZE;
         dk_encoded[hash_offset] ^= 0xFF;
 
+        #[allow(deprecated)]
         let dk_decoded: Result<DecapsulationKey<P>, InvalidKey> =
-            DecapsulationKey::from_encoded_bytes(&dk_encoded);
+            DecapsulationKey::from_expanded_bytes(&dk_encoded);
+
         assert!(dk_decoded.is_err());
     }
 
@@ -345,26 +368,6 @@ mod test {
         invalid_hash_expanded_key_test::<MlKem1024>();
     }
 
-    fn seed_test<P>()
-    where
-        P: KemParams,
-    {
-        let mut rng = UnwrapErr(SysRng);
-        let mut seed = Seed::default();
-        rng.try_fill_bytes(&mut seed).unwrap();
-
-        let dk = DecapsulationKey::<P>::from_seed(seed.clone());
-        let seed_encoded = dk.to_seed().unwrap();
-        assert_eq!(seed, seed_encoded);
-    }
-
-    #[test]
-    fn seed() {
-        seed_test::<MlKem512>();
-        seed_test::<MlKem768>();
-        seed_test::<MlKem1024>();
-    }
-
     fn key_inequality_test<P>()
     where
         P: KemParams,

ml-kem/src/pkcs8.rs

@@ -8,6 +8,8 @@
 //! [`EncapsulationKey`].
 
 #![cfg(feature = "pkcs8")]
+// TODO: debug why rustdoc can't find the `pkcs8::{Decode*, Encode*}` traits. They're right there!
+#![allow(rustdoc::broken_intra_doc_links)]
 
 pub use ::pkcs8::{DecodePrivateKey, DecodePublicKey, spki::AssociatedAlgorithmIdentifier};
 pub use const_oid::AssociatedOid;
@@ -31,7 +33,7 @@ use array::Array;
 
 #[cfg(feature = "alloc")]
 use {
-    crate::EncodedSizeUser,
+    ::kem::KeyExport,
     ::pkcs8::der::{Encode, TagMode, asn1::BitStringRef},
 };
 
@@ -100,7 +102,7 @@ where
     /// Serialize the given `EncapsulationKey` into DER format.
     /// Returns a `Document` which wraps the DER document in case of success.
     fn to_public_key_der(&self) -> spki::Result<pkcs8::Document> {
-        let public_key = self.to_encoded_bytes();
+        let public_key = self.to_bytes();
         let subject_public_key = BitStringRef::new(0, &public_key)?;
 
         ::pkcs8::SubjectPublicKeyInfo {