Add SimpleKEM and FullKEM traits to kem

kem/Cargo.toml

@@ -18,7 +18,7 @@ zeroize = { version = "1.7", default-features = false }
 
 [dev-dependencies]
 hpke = "0.10"
-p256 = { version = "0.9", features = ["ecdsa"] }
+p256 = { version = "0.9", features = ["ecdh", "ecdsa"] }
 pqcrypto = { version = "0.15", default-features = false, features = [
     "pqcrypto-saber",
 ] }

kem/src/lib.rs

@@ -33,3 +33,129 @@ pub trait Decapsulate<EK, SS> {
     /// Decapsulates the given encapsulated key
     fn decapsulate(&self, encapsulated_key: &EK) -> Result<SS, Self::Error>;
 }
+
+// Helper type alias for SimpleKEM encapsulate errors
+type SimpleEncapError<X> = <<X as SimpleKEM>::EncapsulatingKey as Encapsulate<
+    <X as SimpleKEM>::EncapsulatedKey,
+    <X as SimpleKEM>::SharedSecret,
+>>::Error;
+
+// Helper type alias for SimpleKEM decapsulate errors
+type SimpleDecapError<X> = <<X as SimpleKEM>::DecapsulatingKey as Decapsulate<
+    <X as SimpleKEM>::EncapsulatedKey,
+    <X as SimpleKEM>::SharedSecret,
+>>::Error;
+
+/// This trait represents a simplified KEM model, where the encapsulating key and public key are the
+/// same type.
+pub trait SimpleKEM {
+    /// The type that will implement [`Decapsulate`]
+    type DecapsulatingKey: Decapsulate<Self::EncapsulatedKey, Self::SharedSecret>;
+
+    /// The type that will implement [`Encapsulate`]
+    type EncapsulatingKey: Encapsulate<Self::EncapsulatedKey, Self::SharedSecret>;
+
+    /// The type of the encapsulated key
+    type EncapsulatedKey;
+
+    /// The type of the shared secret
+    type SharedSecret;
+
+    /// Generates a new (decapsulating key, encapsulating key) keypair for the KEM model
+    fn random_keypair(
+        rng: &mut impl CryptoRngCore,
+    ) -> (Self::DecapsulatingKey, Self::EncapsulatingKey);
+
+    /// Forwards a call to [`encapsulate`](Encapsulate::encapsulate)
+    fn encapsulate(
+        ek: &Self::EncapsulatingKey,
+        rng: &mut impl CryptoRngCore,
+    ) -> Result<(Self::EncapsulatedKey, Self::SharedSecret), SimpleEncapError<Self>> {
+        ek.encapsulate(rng)
+    }
+
+    /// Forwards a call to [`decapsulate`](Decapsulate::decapsulate)
+    fn decapsulate(
+        dk: &Self::DecapsulatingKey,
+        ek: &Self::EncapsulatedKey,
+    ) -> Result<Self::SharedSecret, SimpleDecapError<Self>> {
+        dk.decapsulate(ek)
+    }
+}
+
+// Helper type alias for FullKEM encapsulate errors
+type FullEncapError<X> = <<X as FullKEM>::EncapsulatingKey as Encapsulate<
+    <X as FullKEM>::EncapsulatedKey,
+    <X as FullKEM>::SharedSecret,
+>>::Error;
+
+// Helper type alias for FullKEM decapsulate errors
+type FullDecapError<X> = <<X as FullKEM>::DecapsulatingKey as Decapsulate<
+    <X as FullKEM>::EncapsulatedKey,
+    <X as FullKEM>::SharedSecret,
+>>::Error;
+
+/// This is a trait that all KEM models should implement. It represents all the stages and types
+/// necessary for a KEM.
+///
+/// In particular,
+///
+/// 1. `KeyGen() -> (PrivateKey, PublicKey)`
+/// 2. `Encaps(EncapsulatingKey) -> (EncappedKey, SharedSecret)`
+/// 3. `Decaps(DecapsulatingKey, EncappedKey) -> SharedSecret`
+///
+/// Promotion from [`PrivateKey`](FullKEM::PrivateKey) to
+/// [`DecapsulatingKey`](FullKEM::DecapsulatingKey) and [`PublicKey`](FullKEM::PublicKey) to
+/// [`EncapsulatingKey`](FullKEM::EncapsulatingKey) is context dependent.
+pub trait FullKEM {
+    /// The private key produced by [`random_keypair`](FullKEM::random_keypair)
+    type PrivateKey;
+
+    /// The public key produced by [`random_keypair`](FullKEM::random_keypair)
+    type PublicKey;
+
+    /// The type that will implement [`Decapsulate`]
+    type DecapsulatingKey: Decapsulate<Self::EncapsulatedKey, Self::SharedSecret>;
+
+    /// The type that will implement [`Encapsulate`]
+    type EncapsulatingKey: Encapsulate<Self::EncapsulatedKey, Self::SharedSecret>;
+
+    /// The type of the encapsulated key
+    type EncapsulatedKey;
+
+    /// The type of the shared secret
+    type SharedSecret;
+
+    /// Generates a new ([`PrivateKey`](FullKEM::PrivateKey), [`PublicKey`](FullKEM::PublicKey))
+    /// keypair for the KEM model
+    fn random_keypair(rng: &mut impl CryptoRngCore) -> (Self::PrivateKey, Self::PublicKey);
+
+    /// Forwards a call to [`encapsulate`](Encapsulate::encapsulate)
+    fn encapsulate(
+        ek: &Self::EncapsulatingKey,
+        rng: &mut impl CryptoRngCore,
+    ) -> Result<(Self::EncapsulatedKey, Self::SharedSecret), FullEncapError<Self>> {
+        ek.encapsulate(rng)
+    }
+
+    /// Forwards a call to [`decapsulate`](Decapsulate::decapsulate)
+    fn decapsulate(
+        dk: &Self::DecapsulatingKey,
+        ek: &Self::EncapsulatedKey,
+    ) -> Result<Self::SharedSecret, FullDecapError<Self>> {
+        dk.decapsulate(ek)
+    }
+}
+
+impl<K: SimpleKEM> FullKEM for K {
+    type PrivateKey = K::DecapsulatingKey;
+    type PublicKey = K::EncapsulatingKey;
+    type DecapsulatingKey = K::DecapsulatingKey;
+    type EncapsulatingKey = K::EncapsulatingKey;
+    type EncapsulatedKey = K::EncapsulatedKey;
+    type SharedSecret = K::SharedSecret;
+
+    fn random_keypair(rng: &mut impl CryptoRngCore) -> (Self::PrivateKey, Self::PublicKey) {
+        Self::random_keypair(rng)
+    }
+}

kem/tests/kemtrait_p256.rs

@@ -0,0 +1,78 @@
+use kem::{Decapsulate, Encapsulate, SimpleKEM};
+use p256::{
+    ecdh::{EphemeralSecret, SharedSecret},
+    PublicKey,
+};
+use rand_core::CryptoRngCore;
+
+struct KemNistP256;
+struct DecapsulatorP256(EphemeralSecret);
+struct EncapsulatorP256(PublicKey);
+struct Secret(SharedSecret);
+
+impl Decapsulate<PublicKey, Secret> for DecapsulatorP256 {
+    type Error = ();
+
+    fn decapsulate(&self, encapsulated_key: &PublicKey) -> Result<Secret, Self::Error> {
+        Ok(Secret(self.0.diffie_hellman(encapsulated_key)))
+    }
+}
+
+impl Encapsulate<PublicKey, Secret> for EncapsulatorP256 {
+    type Error = ();
+
+    fn encapsulate(
+        &self,
+        rng: &mut impl CryptoRngCore,
+    ) -> Result<(PublicKey, Secret), Self::Error> {
+        let sk = EphemeralSecret::random(rng);
+        let pk = sk.public_key();
+
+        Ok((pk, Secret(sk.diffie_hellman(&self.0))))
+    }
+}
+
+impl SimpleKEM for KemNistP256 {
+    type DecapsulatingKey = DecapsulatorP256;
+    type EncapsulatingKey = EncapsulatorP256;
+    type EncapsulatedKey = PublicKey;
+    type SharedSecret = Secret;
+
+    fn random_keypair(
+        rng: &mut impl CryptoRngCore,
+    ) -> (Self::DecapsulatingKey, Self::EncapsulatingKey) {
+        let sk = EphemeralSecret::random(rng);
+        let pk = sk.public_key();
+
+        (DecapsulatorP256(sk), EncapsulatorP256(pk))
+    }
+}
+
+// Helper trait so that shared secrets can be more easily tested for equality during testing
+pub trait SecretBytes {
+    fn as_slice(&self) -> &[u8];
+}
+
+impl SecretBytes for Secret {
+    fn as_slice(&self) -> &[u8] {
+        self.0.as_bytes().as_slice()
+    }
+}
+
+// use a generic SimpleKEM function to ensure correctness
+fn test_kemtrait_basic<K: SimpleKEM>()
+where
+    <K as SimpleKEM>::SharedSecret: SecretBytes,
+{
+    let mut rng = rand::thread_rng();
+    let (sk, pk) = K::random_keypair(&mut rng);
+    let (ek, ss1) = K::encapsulate(&pk, &mut rng).expect("never fails");
+    let ss2 = K::decapsulate(&sk, &ek).expect("never fails");
+
+    assert_eq!(ss1.as_slice(), ss2.as_slice());
+}
+
+#[test]
+fn test_kemtrait_p256() {
+    test_kemtrait_basic::<KemNistP256>();
+}

kem/tests/x3dh.rs

@@ -1,9 +1,12 @@
-use kem::{Decapsulate, Encapsulate};
+use kem::{Decapsulate, Encapsulate, FullKEM};
 
 use p256::ecdsa::Signature;
 use rand_core::CryptoRngCore;
 use x3dh_ke::{x3dh_a, x3dh_b, EphemeralKey, IdentityKey, Key, OneTimePreKey, SignedPreKey};
 
+/// A struct representing the x3dh KEM
+struct X3Dh;
+
 /// The shared secret type defined by x3dh_ke
 type SharedSecret = [u8; 32];
 
@@ -92,6 +95,22 @@ impl Decapsulate<EphemeralKey, SharedSecret> for DecapContext {
     }
 }
 
+impl FullKEM for X3Dh {
+    type PrivateKey = X3DhPrivkeyBundle;
+    type PublicKey = X3DhPubkeyBundle;
+    type DecapsulatingKey = DecapContext;
+    type EncapsulatingKey = EncapContext;
+    type EncapsulatedKey = EphemeralKey;
+    type SharedSecret = SharedSecret;
+
+    fn random_keypair(_: &mut impl CryptoRngCore) -> (Self::PrivateKey, Self::PublicKey) {
+        let sk = Self::PrivateKey::gen();
+        let pk = sk.as_pubkeys();
+
+        (sk, pk)
+    }
+}
+
 #[test]
 fn test_x3dh() {
     let mut rng = rand::thread_rng();
@@ -111,3 +130,23 @@ fn test_x3dh() {
     let ss2 = decap_context.decapsulate(&encapped_key).unwrap();
     assert_eq!(ss1, ss2);
 }
+
+#[test]
+fn test_kemtrait_x3dh() {
+    let mut rng = rand::thread_rng();
+
+    // We use _a and _b suffixes to denote whether a key belongs to Alice or Bob. Alice is the
+    // sender in this case.
+    let sk_ident_a = IdentityKey::default();
+    let pk_ident_a = sk_ident_a.strip();
+    let (sk_bundle_b, pk_bundle_b) = X3Dh::random_keypair(&mut rng);
+
+    let encap_context = EncapContext(pk_bundle_b, sk_ident_a);
+    let decap_context = DecapContext(sk_bundle_b, pk_ident_a);
+
+    // Now do an authenticated encap
+    let (encapped_key, ss1) = X3Dh::encapsulate(&encap_context, &mut rng).unwrap();
+    let ss2 = X3Dh::decapsulate(&decap_context, &encapped_key).unwrap();
+
+    assert_eq!(ss1, ss2);
+}