Allow items in KeyStore to expire

xmtp · Feb 22, 2024 · 3992d09 · 3992d09
1 parent 0da7dcb
commit 3992d09
Show file tree

Hide file tree

Showing 16 changed files with 41 additions and 28 deletions.
diff --git a/basic_credential/src/lib.rs b/basic_credential/src/lib.rs
@@ -121,7 +121,7 @@ impl SignatureKeyPair {
     where
         T: OpenMlsKeyStore,
     {
-        key_store.store(&self.id(), self)
+        key_store.store(&self.id(), self, None)
     }
 
     /// Read a signature key pair from the key store.

diff --git a/cli/src/persistent_key_store.rs b/cli/src/persistent_key_store.rs
@@ -28,7 +28,7 @@ impl OpenMlsKeyStore for PersistentKeyStore {
     /// serialization for ID `k`.
     ///
     /// Returns an error if storing fails.
-    fn store<V: MlsEntity>(&self, k: &[u8], v: &V) -> Result<(), Self::Error> {
+    fn store<V: MlsEntity>(&self, k: &[u8], v: &V, _: Option<u64>) -> Result<(), Self::Error> {
         let value =
             serde_json::to_vec(v).map_err(|_| PersistentKeyStoreError::SerializationError)?;
         // We unwrap here, because this is the only function claiming a write

diff --git a/interop_client/src/main.rs b/interop_client/src/main.rs
@@ -401,7 +401,7 @@ impl MlsClient for MlsClientImpl {
         // Store keys so OpenMLS can find them.
         crypto_provider
             .key_store()
-            .store(my_key_package.hpke_init_key().as_slice(), &private_key)
+            .store(my_key_package.hpke_init_key().as_slice(), &private_key, None)
             .map_err(|_| Status::aborted("failed to interact with the key store"))?;
 
         // Store the key package in the key store with the hash reference as id
@@ -414,6 +414,7 @@ impl MlsClient for MlsClientImpl {
                     .map_err(into_status)?
                     .as_slice(),
                 &my_key_package,
+                None,
             )
             .map_err(into_status)?;
 
@@ -424,7 +425,7 @@ impl MlsClient for MlsClientImpl {
         // The key is the public key.
         crypto_provider
             .key_store()
-            .store::<HpkePrivateKey>(my_key_package.hpke_init_key().as_slice(), &private_key)
+            .store::<HpkePrivateKey>(my_key_package.hpke_init_key().as_slice(), &private_key, None)
             .map_err(into_status)?;
 
         let welcome_msg = MlsMessageIn::tls_deserialize(&mut request.welcome.as_slice())

diff --git a/memory_keystore/src/lib.rs b/memory_keystore/src/lib.rs
@@ -14,7 +14,7 @@ impl OpenMlsKeyStore for MemoryKeyStore {
     /// serialization for ID `k`.
     ///
     /// Returns an error if storing fails.
-    fn store<V: MlsEntity>(&self, k: &[u8], v: &V) -> Result<(), Self::Error> {
+    fn store<V: MlsEntity>(&self, k: &[u8], v: &V, _: Option<u64>) -> Result<(), Self::Error> {
         let value = serde_json::to_vec(v).map_err(|_| MemoryKeyStoreError::SerializationError)?;
         // We unwrap here, because this is the only function claiming a write
         // lock on `credential_bundles`. It only holds the lock very briefly and

diff --git a/openmls/src/group/core_group/kat_passive_client.rs b/openmls/src/group/core_group/kat_passive_client.rs
@@ -263,6 +263,7 @@ impl PassiveClient {
                     .unwrap()
                     .as_slice(),
                 &key_package,
+                None,
             )
             .unwrap();
 
@@ -272,6 +273,7 @@ impl PassiveClient {
             .store::<HpkePrivateKey>(
                 key_package.hpke_init_key().as_slice(),
                 key_package_bundle.private_key(),
+                None,
             )
             .unwrap();
 
@@ -282,7 +284,7 @@ impl PassiveClient {
         ));
 
         key_pair
-            .write_to_key_store(self.provider.key_store())
+            .write_to_key_store(self.provider.key_store(), None)
             .unwrap();
     }
 

diff --git a/openmls/src/group/core_group/kat_welcome.rs b/openmls/src/group/core_group/kat_welcome.rs
@@ -171,6 +171,7 @@ pub fn run_test_vector(test_vector: WelcomeTestVector) -> Result<(), &'static st
         .store(
             key_package.hash_ref(provider.crypto()).unwrap().as_slice(),
             &key_package,
+            None,
         )
         .unwrap();
 
@@ -179,6 +180,7 @@ pub fn run_test_vector(test_vector: WelcomeTestVector) -> Result<(), &'static st
         .store::<HpkePrivateKey>(
             key_package.hpke_init_key().as_slice(),
             key_package_bundle.private_key(),
+            None,
         )
         .unwrap();
 

diff --git a/openmls/src/group/core_group/mod.rs b/openmls/src/group/core_group/mod.rs
@@ -779,7 +779,7 @@ impl CoreGroup {
             self.context().epoch().as_u64(),
             self.own_leaf_index(),
         );
-        store.store(&k.0, &keypair_references.to_vec())
+        store.store(&k.0, &keypair_references.to_vec(), None)
     }
 
     /// Read the [`EncryptionKeyPair`]s of this group and its current

diff --git a/openmls/src/group/mls_group/mod.rs b/openmls/src/group/mls_group/mod.rs
@@ -305,7 +305,7 @@ impl MlsGroup {
         &mut self,
         store: &KeyStore,
     ) -> Result<(), KeyStore::Error> {
-        store.store(self.group_id().as_slice(), &*self)?;
+        store.store(self.group_id().as_slice(), &*self, None)?;
 
         self.state_changed = InnerState::Persisted;
         Ok(())

diff --git a/openmls/src/group/mls_group/updates.rs b/openmls/src/group/mls_group/updates.rs
@@ -101,7 +101,7 @@ impl MlsGroup {
             )?;
             // TODO #1207: Move to the top of the function.
             keypair
-                .write_to_key_store(provider.key_store())
+                .write_to_key_store(provider.key_store(), None)
                 .map_err(ProposeSelfUpdateError::KeyStoreError)?;
         };
 

diff --git a/openmls/src/group/tests/test_proposal_validation.rs b/openmls/src/group/tests/test_proposal_validation.rs
@@ -1714,7 +1714,7 @@ fn test_valsem110(ciphersuite: Ciphersuite, provider: &impl OpenMlsProvider) {
         .find(|keypair| keypair.public_key() == &alice_encryption_key)
         .unwrap();
     leaf_keypair
-        .write_to_key_store(provider.key_store())
+        .write_to_key_store(provider.key_store(), None)
         .unwrap();
 
     // Have bob process the resulting plaintext

diff --git a/openmls/src/key_packages/mod.rs b/openmls/src/key_packages/mod.rs
@@ -422,6 +422,7 @@ impl KeyPackage {
             leaf_node_extensions,
             init_key,
         )?;
+        let expiration = Some(key_package.life_time().not_after());
 
         // Store the key package in the key store with the hash reference as id
         // for retrieval when parsing welcome messages.
@@ -430,12 +431,13 @@ impl KeyPackage {
             .store(
                 key_package.hash_ref(provider.crypto())?.as_slice(),
                 &key_package,
+                expiration,
             )
             .map_err(KeyPackageNewError::KeyStoreError)?;
 
         // Store the encryption key pair in the key store.
         encryption_key_pair
-            .write_to_key_store(provider.key_store())
+            .write_to_key_store(provider.key_store(), expiration)
             .map_err(KeyPackageNewError::KeyStoreError)?;
 
         Ok(key_package)
@@ -461,13 +463,6 @@ impl KeyPackage {
             .crypto()
             .derive_hpke_keypair(config.ciphersuite.hpke_config(), ikm.as_slice());
 
-        // Store the private part of the init_key into the key store.
-        // The key is the public key.
-        provider
-            .key_store()
-            .store::<HpkePrivateKey>(&init_key.public, &init_key.private)
-            .map_err(KeyPackageNewError::KeyStoreError)?;
-
         // We don't need the private key here. It's stored in the key store for
         // use later when creating a group with this key package.
         let leaf_node = LeafNode::create_new_with_key(
@@ -484,12 +479,20 @@ impl KeyPackage {
         let key_package = KeyPackageTbs {
             protocol_version: config.version,
             ciphersuite: config.ciphersuite,
-            init_key: init_key.public.into(),
+            init_key: init_key.public.clone().into(),
             leaf_node,
             extensions,
         };
 
         let key_package = key_package.sign(signer)?;
+        let expiration = Some(key_package.life_time().not_after());
+
+        // Store the private part of the init_key into the key store.
+        // The key is the public key.
+        provider
+            .key_store()
+            .store::<HpkePrivateKey>(&init_key.public, &init_key.private, expiration)
+            .map_err(KeyPackageNewError::KeyStoreError)?;
 
         // Store the key package in the key store with the hash reference as id
         // for retrieval when parsing welcome messages.
@@ -498,6 +501,7 @@ impl KeyPackage {
             .store(
                 key_package.hash_ref(provider.crypto())?.as_slice(),
                 &key_package,
+                expiration,
             )
             .map_err(KeyPackageNewError::KeyStoreError)?;
 
@@ -643,6 +647,7 @@ impl KeyPackageBuilder {
             self.leaf_node_capabilities.unwrap_or_default(),
             self.leaf_node_extensions.unwrap_or_default(),
         )?;
+        let expiration = Some(key_package.life_time().not_after());
 
         // Store the key package in the key store with the hash reference as id
         // for retrieval when parsing welcome messages.
@@ -651,19 +656,20 @@ impl KeyPackageBuilder {
             .store(
                 key_package.hash_ref(provider.crypto())?.as_slice(),
                 &key_package,
+                expiration,
             )
             .map_err(KeyPackageNewError::KeyStoreError)?;
 
         // Store the encryption key pair in the key store.
         encryption_keypair
-            .write_to_key_store(provider.key_store())
+            .write_to_key_store(provider.key_store(), expiration)
             .map_err(KeyPackageNewError::KeyStoreError)?;
 
         // Store the private part of the init_key into the key store.
         // The key is the public key.
         provider
             .key_store()
-            .store::<HpkePrivateKey>(key_package.hpke_init_key().as_slice(), &init_private_key)
+            .store::<HpkePrivateKey>(key_package.hpke_init_key().as_slice(), &init_private_key, expiration)
             .map_err(KeyPackageNewError::KeyStoreError)?;
 
         Ok(key_package)

diff --git a/openmls/src/messages/tests/test_welcome.rs b/openmls/src/messages/tests/test_welcome.rs
@@ -183,15 +183,16 @@ fn test_welcome_context_mismatch(ciphersuite: Ciphersuite, provider: &impl OpenM
         .store(
             bob_kp.hash_ref(provider.crypto()).unwrap().as_slice(),
             bob_kp,
+            None,
         )
         .unwrap();
     provider
         .key_store()
-        .store::<HpkePrivateKey>(bob_kp.hpke_init_key().as_slice(), bob_private_key)
+        .store::<HpkePrivateKey>(bob_kp.hpke_init_key().as_slice(), bob_private_key, None)
         .unwrap();
 
     encryption_keypair
-        .write_to_key_store(provider.key_store())
+        .write_to_key_store(provider.key_store(), None)
         .unwrap();
 
     let _group = MlsGroup::new_from_welcome(

diff --git a/openmls/src/schedule/psk.rs b/openmls/src/schedule/psk.rs
@@ -300,7 +300,7 @@ impl PreSharedKeyId {
 
         provider
             .key_store()
-            .store(&keystore_id, &psk_bundle)
+            .store(&keystore_id, &psk_bundle, None)
             .map_err(|_| PskError::KeyStore)
     }
 

diff --git a/openmls/src/treesync/node/encryption_keys.rs b/openmls/src/treesync/node/encryption_keys.rs
@@ -169,8 +169,9 @@ impl EncryptionKeyPair {
     pub(crate) fn write_to_key_store<KeyStore: OpenMlsKeyStore>(
         &self,
         store: &KeyStore,
+        expiration: Option<u64>,
     ) -> Result<(), KeyStore::Error> {
-        store.store(&self.public_key().to_bytes_with_prefix(), self)
+        store.store(&self.public_key().to_bytes_with_prefix(), self, expiration)
     }
 
     /// Read the [`EncryptionKeyPair`] from the key store of the `provider`. This
@@ -239,7 +240,7 @@ pub mod test_utils {
     pub fn write_keys_from_key_store(provider: &impl OpenMlsProvider, encryption_key: HpkeKeyPair) {
         let keypair = EncryptionKeyPair::from(encryption_key);
 
-        keypair.write_to_key_store(provider.key_store()).unwrap();
+        keypair.write_to_key_store(provider.key_store(), None).unwrap();
     }
 }
 

diff --git a/openmls/src/treesync/node/leaf_node.rs b/openmls/src/treesync/node/leaf_node.rs
@@ -229,7 +229,7 @@ impl LeafNode {
 
         // Store the encryption key pair in the key store.
         encryption_key_pair
-            .write_to_key_store(provider.key_store())
+            .write_to_key_store(provider.key_store(), None)
             .map_err(LeafNodeGenerationError::KeyStoreError)?;
 
         Ok(leaf_node)

diff --git a/traits/src/key_store.rs b/traits/src/key_store.rs
@@ -34,7 +34,7 @@ pub trait OpenMlsKeyStore {
     /// serialization for ID `k`.
     ///
     /// Returns an error if storing fails.
-    fn store<V: MlsEntity>(&self, k: &[u8], v: &V) -> Result<(), Self::Error>
+    fn store<V: MlsEntity>(&self, k: &[u8], v: &V, expiration: Option<u64>) -> Result<(), Self::Error>
     where
         Self: Sized;