Define a Encapsulate::encapsulate_in_place method

kem/src/lib.rs

@@ -12,9 +12,9 @@
 use core::fmt::Debug;
 use rand_core::CryptoRngCore;
 
-/// A value that can be encapsulated to. Often, this will just be a public key. However, it can
-/// also be a bundle of public keys, or it can include a sender's private key for authenticated
-/// encapsulation.
+/// This trait implements encapsulation. Often, the implementer will just be a public key. However,
+/// it can also be a bundle of public keys, or it can include a sender's private key for
+/// authenticated encapsulation.
 pub trait Encapsulate<EK, SS> {
     /// Encapsulation error
     type Error: Debug;
@@ -23,9 +23,25 @@ pub trait Encapsulate<EK, SS> {
     fn encapsulate(&self, rng: &mut impl CryptoRngCore) -> Result<(EK, SS), Self::Error>;
 }
 
-/// A value that can be used to decapsulate an encapsulated key. Often, this will just be a secret
-/// key. But, as with [`Encapsulate`], it can be a bundle of secret keys, or it can include a
-/// sender's private key for authenticated encapsulation.
+/// This trait implements in-place encapsulation. Often, the implementer will just be a public key.
+/// However, it can also be a bundle of public keys, or it can include a sender's private key for
+/// authenticated encapsulation.
+pub trait EncapsulateInPlace<EK, SS> {
+    /// Encapsulation error
+    type Error: Debug;
+
+    /// Encapsulates a fresh shared secret, placing the encapsulated key into the given mut ref.
+    /// If this errors, the final value of `encapsulated_key` MUST equal its original value.
+    fn encapsulate_in_place(
+        &self,
+        rng: &mut impl CryptoRngCore,
+        encapsulated_key: &mut EK,
+    ) -> Result<SS, Self::Error>;
+}
+
+/// This trait implements decapsulation. Often, the implementer will just be a secret key. But, as
+/// with [`Encapsulate`], it can be a bundle of secret keys, or it can include a sender's private
+/// key for authenticated encapsulation.
 pub trait Decapsulate<EK, SS> {
     /// Decapsulation error
     type Error: Debug;

kem/tests/hpke.rs

@@ -21,7 +21,7 @@ impl Encapsulate<EncappedKey, SharedSecret> for PublicKey {
         &self,
         mut csprng: &mut impl CryptoRngCore,
     ) -> Result<(EncappedKey, SharedSecret), HpkeError> {
-        <X25519HkdfSha256 as KemTrait>::encap(&self.0, None, &mut csprng).map(|(ek, ss)| (ss, ek))
+        <X25519HkdfSha256 as KemTrait>::encap(&self.0, None, &mut csprng).map(|(ss, ek)| (ek, ss))
     }
 }
 
@@ -50,4 +50,6 @@ fn test_hpke() {
     let (ek, ss1) = pk_recip.encapsulate(&mut rng).unwrap();
     let ss2 = sk_recip.decapsulate(&ek).unwrap();
     assert_eq!(ss1.0, ss2.0);
+
+    // Can't use encapsulate_in_place for this crate, because EncappedKey has no constructor
 }

kem/tests/saber.rs

@@ -1,5 +1,7 @@
 use kem::{Decapsulate, Encapsulate};
 
+use core::convert::Infallible;
+
 use pqcrypto::kem::firesaber::{
     decapsulate, encapsulate, keypair, Ciphertext as SaberEncappedKey, PublicKey, SecretKey,
     SharedSecret as SaberSharedSecret,
@@ -12,23 +14,21 @@ struct SaberPublicKey(PublicKey);
 struct SaberPrivateKey(SecretKey);
 
 impl Encapsulate<SaberEncappedKey, SaberSharedSecret> for SaberPublicKey {
-    // TODO: Encapsulation is infallible. Make this the never type once it's available
-    type Error = ();
+    type Error = Infallible;
 
     fn encapsulate(
         &self,
         _: &mut impl CryptoRngCore,
-    ) -> Result<(SaberEncappedKey, SaberSharedSecret), ()> {
+    ) -> Result<(SaberEncappedKey, SaberSharedSecret), Infallible> {
         let (ss, ek) = encapsulate(&self.0);
         Ok((ek, ss))
     }
 }
 
 impl Decapsulate<SaberEncappedKey, SaberSharedSecret> for SaberPrivateKey {
-    // TODO: Decapsulation is infallible. Make this the never type once it's available
-    type Error = ();
+    type Error = Infallible;
 
-    fn decapsulate(&self, ek: &SaberEncappedKey) -> Result<SaberSharedSecret, ()> {
+    fn decapsulate(&self, ek: &SaberEncappedKey) -> Result<SaberSharedSecret, Infallible> {
         Ok(decapsulate(ek, &self.0))
     }
 }
@@ -50,4 +50,6 @@ fn test_saber() {
     let (ek, ss1) = pk_recip.encapsulate(&mut rng).unwrap();
     let ss2 = sk_recip.decapsulate(&ek).unwrap();
     assert_eq!(ss1.as_bytes(), ss2.as_bytes());
+
+    // Can't use encapsulate_in_place for this crate, because Ciphertext has no constructor
 }

kem/tests/x3dh.rs

@@ -1,4 +1,4 @@
-use kem::{Decapsulate, Encapsulate};
+use kem::{Decapsulate, Encapsulate, EncapsulateInPlace};
 
 use p256::ecdsa::Signature;
 use rand_core::CryptoRngCore;
@@ -70,6 +70,22 @@ impl Encapsulate<EphemeralKey, SharedSecret> for EncapContext {
     }
 }
 
+// Same thing but in-place
+impl EncapsulateInPlace<EphemeralKey, SharedSecret> for EncapContext {
+    type Error = &'static str;
+
+    // Do the dumb thing and just call encapsulate() and copy the result into the mut ref
+    fn encapsulate_in_place(
+        &self,
+        rng: &mut impl CryptoRngCore,
+        encapsulated_key: &mut EphemeralKey,
+    ) -> Result<SharedSecret, Self::Error> {
+        let (ek, ss) = self.encapsulate(rng)?;
+        *encapsulated_key = ek;
+        Ok(ss)
+    }
+}
+
 // Define an decapsulator. Since authenticated and unauthenticated encapped keys are represented by
 // the same type (which, outside of testing, should not be the case), this can do both auth'd and
 // unauth'd decapsulation.
@@ -110,4 +126,12 @@ fn test_x3dh() {
     let (encapped_key, ss1) = encap_context.encapsulate(&mut rng).unwrap();
     let ss2 = decap_context.decapsulate(&encapped_key).unwrap();
     assert_eq!(ss1, ss2);
+
+    // Now do the same but with encapsulate_in_place
+    let mut encapped_key = EphemeralKey::default();
+    let ss1 = encap_context
+        .encapsulate_in_place(&mut rng, &mut encapped_key)
+        .unwrap();
+    let ss2 = decap_context.decapsulate(&encapped_key).unwrap();
+    assert_eq!(ss1, ss2);
 }