Add kem::{TryKeyInit + KeyExport} impls for encapsulation keys

Cargo.toml

@@ -12,3 +12,6 @@ debug = true
 
 [patch.crates-io]
 ml-kem = { path = "./ml-kem" }
+
+kem = { git = "https://github.com/RustCrypto/traits" }
+sec1 = { git = "https://github.com/RustCrypto/formats" }

dhkem/Cargo.toml

@@ -17,6 +17,9 @@ readme = "README.md"
 kem = "0.4.0-rc.4"
 rand_core = "0.10.0-rc-5"
 
+# TODO(tarcieri): remove this and get these from `kem`
+common = { package = "crypto-common", version = "0.2.0-rc.12" }
+
 # optional dependencies
 elliptic-curve = { version = "0.14.0-rc.23", optional = true, default-features = false }
 k256 = { version = "0.14.0-rc.5", optional = true, default-features = false, features = ["arithmetic"] }

dhkem/src/ecdh_kem.rs

@@ -3,10 +3,18 @@
 use crate::{DhDecapsulator, DhEncapsulator, DhKem};
 use core::{convert::Infallible, marker::PhantomData};
 use elliptic_curve::{
-    CurveArithmetic, Generate, PublicKey,
+    AffinePoint,
+    CurveArithmetic,
+    FieldBytesSize,
+    Generate,
+    PublicKey,
+    common::InvalidKey, // TODO(tarcieri): get from `kem` crate
     ecdh::{EphemeralSecret, SharedSecret},
+    sec1::{
+        FromEncodedPoint, ModulusSize, ToEncodedPoint, UncompressedPoint, UncompressedPointSize,
+    },
 };
-use kem::{Decapsulate, Encapsulate};
+use kem::{Decapsulate, Encapsulate, KeyExport, KeySizeUser, TryKeyInit};
 use rand_core::{CryptoRng, TryCryptoRng};
 
 /// Generic Elliptic Curve Diffie-Hellman KEM adapter compatible with curves implemented using
@@ -15,9 +23,68 @@ use rand_core::{CryptoRng, TryCryptoRng};
 /// Implements a KEM interface that internally uses ECDH.
 pub struct EcdhKem<C: CurveArithmetic>(PhantomData<C>);
 
+/// From [RFC9810 §7.1.1]: `SerializePublicKey` and `DeserializePublicKey`:
+///
+/// > For P-256, P-384, and P-521, the SerializePublicKey() function of the
+/// > KEM performs the uncompressed Elliptic-Curve-Point-to-Octet-String
+/// > conversion according to [SECG].
+///
+/// [RFC9810 §7.1.1]: https://datatracker.ietf.org/doc/html/rfc9180#name-serializepublickey-and-dese
+/// [SECG]: https://www.secg.org/sec1-v2.pdf
+impl<C> KeySizeUser for DhEncapsulator<PublicKey<C>>
+where
+    C: CurveArithmetic,
+    FieldBytesSize<C>: ModulusSize,
+{
+    type KeySize = UncompressedPointSize<C>;
+}
+
+/// From [RFC9810 §7.1.1]: `SerializePublicKey` and `DeserializePublicKey`:
+///
+/// > DeserializePublicKey() performs the uncompressed
+/// > Octet-String-to-Elliptic-Curve-Point conversion.
+///
+/// [RFC9810 §7.1.1]: https://datatracker.ietf.org/doc/html/rfc9180#name-serializepublickey-and-dese
+impl<C> TryKeyInit for DhEncapsulator<PublicKey<C>>
+where
+    C: CurveArithmetic,
+    FieldBytesSize<C>: ModulusSize,
+    AffinePoint<C>: FromEncodedPoint<C> + ToEncodedPoint<C>,
+{
+    fn new(encapsulation_key: &UncompressedPoint<C>) -> Result<Self, InvalidKey> {
+        PublicKey::<C>::from_sec1_bytes(encapsulation_key)
+            .map(Into::into)
+            .map_err(|_| InvalidKey)
+    }
+}
+
+/// From [RFC9810 §7.1.1]: `SerializePublicKey` and `DeserializePublicKey`:
+///
+/// > For P-256, P-384, and P-521, the SerializePublicKey() function of the
+/// > KEM performs the uncompressed Elliptic-Curve-Point-to-Octet-String
+/// > conversion according to [SECG].
+///
+/// [RFC9810 §7.1.1]: https://datatracker.ietf.org/doc/html/rfc9180#name-serializepublickey-and-dese
+/// [SECG]: https://www.secg.org/sec1-v2.pdf
+impl<C> KeyExport for DhEncapsulator<PublicKey<C>>
+where
+    C: CurveArithmetic,
+    FieldBytesSize<C>: ModulusSize,
+    AffinePoint<C>: FromEncodedPoint<C> + ToEncodedPoint<C>,
+{
+    fn to_bytes(&self) -> UncompressedPoint<C> {
+        // TODO(tarcieri): use `ToEncodedPoint::to_uncompressed_point` (RustCrypto/traits#2221)
+        let mut ret = UncompressedPoint::<C>::default();
+        ret.copy_from_slice(self.0.to_encoded_point(false).as_bytes());
+        ret
+    }
+}
+
 impl<C> Encapsulate<PublicKey<C>, SharedSecret<C>> for DhEncapsulator<PublicKey<C>>
 where
     C: CurveArithmetic,
+    FieldBytesSize<C>: ModulusSize,
+    AffinePoint<C>: FromEncodedPoint<C> + ToEncodedPoint<C>,
 {
     type Error = Infallible;
 
@@ -38,6 +105,8 @@ where
 impl<C> Decapsulate<PublicKey<C>, SharedSecret<C>> for DhDecapsulator<EphemeralSecret<C>>
 where
     C: CurveArithmetic,
+    FieldBytesSize<C>: ModulusSize,
+    AffinePoint<C>: FromEncodedPoint<C> + ToEncodedPoint<C>,
 {
     type Encapsulator = DhEncapsulator<PublicKey<C>>;
     type Error = Infallible;
@@ -56,6 +125,8 @@ where
 impl<C> DhKem for EcdhKem<C>
 where
     C: CurveArithmetic,
+    FieldBytesSize<C>: ModulusSize,
+    AffinePoint<C>: FromEncodedPoint<C> + ToEncodedPoint<C>,
 {
     type DecapsulatingKey = DhDecapsulator<EphemeralSecret<C>>;
     type EncapsulatingKey = DhEncapsulator<PublicKey<C>>;

dhkem/src/x25519_kem.rs

@@ -1,14 +1,54 @@
 use crate::{DhDecapsulator, DhEncapsulator, DhKem};
 use core::convert::Infallible;
-use kem::{Decapsulate, Encapsulate};
+use kem::{Decapsulate, Encapsulate, KeyExport, KeySizeUser, TryKeyInit, consts::U32};
 use rand_core::{CryptoRng, TryCryptoRng, UnwrapErr};
 use x25519::{PublicKey, ReusableSecret, SharedSecret};
 
+// TODO(tarcieri): get these from `kem`
+use common::{InvalidKey, Key};
+
 /// X22519 Diffie-Hellman KEM adapter.
 ///
 /// Implements a KEM interface that internally uses X25519 ECDH.
 pub struct X25519Kem;
 
+/// From [RFC9810 §7.1.1]: `SerializePublicKey` and `DeserializePublicKey`:
+///
+/// > For X25519 and X448, the SerializePublicKey() and
+/// > DeserializePublicKey() functions are the identity function, since
+/// > these curves already use fixed-length byte strings for public keys.
+///
+/// [RFC9810 §7.1.1]: https://datatracker.ietf.org/doc/html/rfc9180#name-serializepublickey-and-dese
+impl KeySizeUser for DhEncapsulator<PublicKey> {
+    type KeySize = U32;
+}
+
+/// From [RFC9810 §7.1.1]: `SerializePublicKey` and `DeserializePublicKey`:
+///
+/// > For X25519 and X448, the SerializePublicKey() and
+/// > DeserializePublicKey() functions are the identity function, since
+/// > these curves already use fixed-length byte strings for public keys.
+///
+/// [RFC9810 §7.1.1]: https://datatracker.ietf.org/doc/html/rfc9180#name-serializepublickey-and-dese
+impl TryKeyInit for DhEncapsulator<PublicKey> {
+    fn new(encapsulation_key: &Key<Self>) -> Result<Self, InvalidKey> {
+        Ok(Self(PublicKey::from(encapsulation_key.0)))
+    }
+}
+
+/// From [RFC9810 §7.1.1]: `SerializePublicKey` and `DeserializePublicKey`:
+///
+/// > For X25519 and X448, the SerializePublicKey() and
+/// > DeserializePublicKey() functions are the identity function, since
+/// > these curves already use fixed-length byte strings for public keys.
+///
+/// [RFC9810 §7.1.1]: https://datatracker.ietf.org/doc/html/rfc9180#name-serializepublickey-and-dese
+impl KeyExport for DhEncapsulator<PublicKey> {
+    fn to_bytes(&self) -> Key<Self> {
+        self.0.to_bytes().into()
+    }
+}
+
 impl Encapsulate<PublicKey, SharedSecret> for DhEncapsulator<PublicKey> {
     type Error = Infallible;
 
@@ -30,9 +70,7 @@ impl Decapsulate<PublicKey, SharedSecret> for DhDecapsulator<ReusableSecret> {
     type Error = Infallible;
 
     fn decapsulate(&self, encapsulated_key: &PublicKey) -> Result<SharedSecret, Self::Error> {
-        let ss = self.0.diffie_hellman(encapsulated_key);
-
-        Ok(ss)
+        Ok(self.0.diffie_hellman(encapsulated_key))
     }
 
     fn encapsulator(&self) -> DhEncapsulator<PublicKey> {

ml-kem/Cargo.toml

@@ -31,6 +31,9 @@ rand_core = "0.10.0-rc-5"
 sha3 = { version = "0.11.0-rc.3", default-features = false }
 subtle = { version = "2", default-features = false }
 
+# TODO(tarcieri): remove this and get these from `kem`
+common = { package = "crypto-common", version = "0.2.0-rc.12" }
+
 # optional dependencies
 const-oid = { version = "0.10.1", optional = true, default-features = false, features = ["db"] }
 pkcs8 = { version = "0.11.0-rc.8", optional = true, default-features = false }

ml-kem/benches/mlkem.rs

@@ -11,15 +11,15 @@ fn criterion_benchmark(c: &mut Criterion) {
     c.bench_function("keygen", |b| {
         b.iter(|| {
             let dk = ml_kem_768::DecapsulationKey::generate_from_rng(&mut rng);
-            let _dk_bytes = dk.to_bytes();
-            let _ek_bytes = dk.encapsulator().to_bytes();
+            let _dk_bytes = dk.to_encoded_bytes();
+            let _ek_bytes = dk.encapsulator().to_encoded_bytes();
         })
     });
 
     let dk = ml_kem_768::DecapsulationKey::generate_from_rng(&mut rng);
-    let dk_bytes = dk.to_bytes();
-    let ek_bytes = dk.encapsulator().to_bytes();
-    let ek = ml_kem_768::EncapsulationKey::from_bytes(&ek_bytes).unwrap();
+    let dk_bytes = dk.to_encoded_bytes();
+    let ek_bytes = dk.encapsulator().to_encoded_bytes();
+    let ek = ml_kem_768::EncapsulationKey::from_encoded_bytes(&ek_bytes).unwrap();
 
     // Encapsulation
     c.bench_function("encapsulate", |b| {
@@ -28,7 +28,7 @@ fn criterion_benchmark(c: &mut Criterion) {
     let (ct, _ss) = ek.encapsulate_with_rng(&mut rng).unwrap();
 
     // Decapsulation
-    let dk = <MlKem768 as KemCore>::DecapsulationKey::from_bytes(&dk_bytes).unwrap();
+    let dk = <MlKem768 as KemCore>::DecapsulationKey::from_encoded_bytes(&dk_bytes).unwrap();
 
     c.bench_function("decapsulate", |b| {
         b.iter(|| {