kem: Defined an encapsulate_in_place method

kem/src/lib.rs

@@ -21,6 +21,14 @@ pub trait Encapsulate<EK, SS> {
 
     /// Encapsulates a fresh shared secret
     fn encapsulate(&self, rng: &mut impl CryptoRngCore) -> Result<(EK, SS), Self::Error>;
+
+    /// 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>;
 }
 
 /// A value that can be used to decapsulate an encapsulated key. Often, this will just be a secret

kem/tests/hpke.rs

@@ -21,7 +21,17 @@ 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))
+    }
+
+    fn encapsulate_in_place(
+        &self,
+        csprng: &mut impl CryptoRngCore,
+        encapsulated_key: &mut EncappedKey,
+    ) -> Result<SharedSecret, HpkeError> {
+        let (ek, ss) = self.encapsulate(csprng)?;
+        *encapsulated_key = ek;
+        Ok(ss)
     }
 }
 

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,16 +14,25 @@ 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))
     }
+
+    fn encapsulate_in_place(
+        &self,
+        csprng: &mut impl CryptoRngCore,
+        encapsulated_key: &mut SaberEncappedKey,
+    ) -> Result<SaberSharedSecret, Infallible> {
+        let (ek, ss) = self.encapsulate(csprng)?;
+        *encapsulated_key = ek;
+        Ok(ss)
+    }
 }
 
 impl Decapsulate<SaberEncappedKey, SaberSharedSecret> for SaberPrivateKey {

kem/tests/x3dh.rs

@@ -68,6 +68,16 @@ impl Encapsulate<EphemeralKey, SharedSecret> for EncapContext {
 
         Ok((ek, shared_secret))
     }
+
+    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