Skip to main content

core_crypto_keystore/
mls.rs

1use openmls::prelude::Ciphersuite;
2use openmls_basic_credential::SignatureKeyPair;
3use openmls_traits::key_store::{MlsEntity, MlsEntityId};
4
5use crate::{
6    CryptoKeystoreError, CryptoKeystoreResult, Sha256Hash,
7    entities::{
8        PersistedMlsGroup, PersistedMlsPendingGroup, StoredCredential, StoredE2eiEnrollment, StoredEncryptionKeyPair,
9        StoredEpochEncryptionKeypair, StoredHpkePrivateKey, StoredKeypackage, StoredPskBundle,
10    },
11    traits::{Entity as _, FetchFromDatabase},
12};
13
14impl crate::Database {
15    /// Fetches Keypackages
16    ///
17    /// # Arguments
18    /// * `count` - amount of entries to be returned
19    ///
20    /// # Errors
21    /// Any common error that can happen during a database connection. IoError being a common error
22    /// for example.
23    pub async fn mls_fetch_key_packages<V: MlsEntity>(&self, count: u32) -> CryptoKeystoreResult<Vec<V>> {
24        let keypackages = StoredKeypackage::load_all(&*self.conn().await)?;
25        Ok(keypackages
26            .into_iter()
27            .filter_map(|kpb| postcard::from_bytes(&kpb.keypackage).ok())
28            .take(count as _)
29            .collect())
30    }
31
32    /// Checks if the given MLS group id exists in the keystore
33    /// Note: in case of any error, this will return false
34    ///
35    /// # Arguments
36    /// * `group_id` - group/conversation id
37    pub async fn mls_group_exists(&self, group_id: impl AsRef<[u8]> + Send) -> bool {
38        matches!(
39            self.get_borrowed::<PersistedMlsGroup>(group_id.as_ref()).await,
40            Ok(Some(_))
41        )
42    }
43
44    /// Persists a `MlsGroup`
45    ///
46    /// # Arguments
47    /// * `group_id` - group/conversation id
48    /// * `state` - the group state
49    ///
50    /// # Errors
51    /// Any common error that can happen during a database connection. IoError being a common error
52    /// for example.
53    pub async fn mls_group_persist(
54        &self,
55        group_id: impl AsRef<[u8]> + Send,
56        state: &[u8],
57        parent_group_id: Option<&[u8]>,
58    ) -> CryptoKeystoreResult<()> {
59        self.with_transaction(async |tx| {
60            tx.save(PersistedMlsGroup {
61                id: group_id.as_ref().to_owned(),
62                state: state.into(),
63                parent_id: parent_group_id.map(Into::into),
64            })
65            .await
66        })
67        .await
68    }
69
70    /// Loads `MlsGroups` from the database. It will be returned as a `HashMap` where the key is
71    /// the group/conversation id and the value the group state
72    ///
73    /// # Errors
74    /// Any common error that can happen during a database connection. IoError being a common error
75    /// for example.
76    pub async fn mls_groups_restore(
77        &self,
78    ) -> CryptoKeystoreResult<std::collections::HashMap<Vec<u8>, (Option<Vec<u8>>, Vec<u8>)>> {
79        let groups = PersistedMlsGroup::load_all(&*self.conn().await)?;
80        Ok(groups
81            .into_iter()
82            .map(|mut group| {
83                let id = std::mem::take(&mut group.id);
84                let parent_id = std::mem::take(&mut group.parent_id);
85                let state = std::mem::take(&mut group.state);
86                (id, (parent_id, state))
87            })
88            .collect())
89    }
90
91    /// Deletes `MlsGroups` from the database.
92    /// # Errors
93    /// Any common error that can happen during a database connection. IoError being a common error
94    /// for example.
95    pub async fn mls_group_delete(&self, group_id: impl AsRef<[u8]> + Send) -> CryptoKeystoreResult<()> {
96        self.with_transaction(async |tx| tx.remove_borrowed::<PersistedMlsGroup>(group_id.as_ref()).await)
97            .await?;
98        Ok(())
99    }
100
101    /// Saves a `MlsGroup` in a temporary table (typically used in scenarios where the group cannot
102    /// be committed until the backend acknowledges it, like external commits)
103    ///
104    /// # Arguments
105    /// * `group_id` - group/conversation id
106    /// * `mls_group` - the group/conversation state
107    /// * `custom_configuration` - local group configuration
108    ///
109    /// # Errors
110    /// Any common error that can happen during a database connection. IoError being a common error
111    /// for example.
112    pub async fn mls_pending_groups_save(
113        &self,
114        group_id: impl AsRef<[u8]> + Send,
115        mls_group: &[u8],
116        custom_configuration: &[u8],
117        parent_group_id: Option<&[u8]>,
118    ) -> CryptoKeystoreResult<()> {
119        self.with_transaction(async |tx| {
120            tx.save(PersistedMlsPendingGroup {
121                id: group_id.as_ref().to_owned(),
122                state: mls_group.into(),
123                custom_configuration: custom_configuration.into(),
124                parent_id: parent_group_id.map(Into::into),
125            })
126            .await
127        })
128        .await?;
129        Ok(())
130    }
131
132    /// Loads a temporary `MlsGroup` and its configuration from the database
133    ///
134    /// # Arguments
135    /// * `id` - group/conversation id
136    ///
137    /// # Errors
138    /// Any common error that can happen during a database connection. IoError being a common error
139    /// for example.
140    pub async fn mls_pending_groups_load(
141        &self,
142        group_id: impl AsRef<[u8]> + Send,
143    ) -> CryptoKeystoreResult<Option<(Vec<u8>, Vec<u8>)>> {
144        let optional = self.get_borrowed::<PersistedMlsPendingGroup>(group_id.as_ref()).await?;
145        Ok(optional.map(|pending_group| (pending_group.state.clone(), pending_group.custom_configuration.clone())))
146    }
147
148    /// Deletes a temporary `MlsGroup` from the database
149    ///
150    /// # Arguments
151    /// * `id` - group/conversation id
152    ///
153    /// # Errors
154    /// Any common error that can happen during a database connection. IoError being a common error
155    /// for example.
156    pub async fn mls_pending_groups_delete(&self, group_id: impl AsRef<[u8]> + Send) -> CryptoKeystoreResult<()> {
157        self.with_transaction(async |tx| tx.remove_borrowed::<PersistedMlsPendingGroup>(group_id.as_ref()).await)
158            .await
159            .map(|_| ())
160    }
161
162    /// Persists an enrollment instance
163    ///
164    /// # Arguments
165    /// * `id` - hash of the enrollment and unique identifier
166    /// * `content` - serialized enrollment
167    pub async fn save_e2ei_enrollment(&self, id: &[u8], content: &[u8]) -> CryptoKeystoreResult<()> {
168        let id = id.into();
169        let content = content.into();
170        self.with_transaction(async |tx| tx.save(StoredE2eiEnrollment { id, content }).await)
171            .await
172    }
173
174    /// Fetches and delete the enrollment instance
175    ///
176    /// # Arguments
177    /// * `id` - hash of the enrollment and unique identifier
178    pub async fn pop_e2ei_enrollment(&self, id: &[u8]) -> CryptoKeystoreResult<Option<Vec<u8>>> {
179        // someone who has time could try to optimize this but honestly it's really on the cold path
180        let Some(mut enrollment) = self.get_borrowed::<StoredE2eiEnrollment>(id).await? else {
181            return Ok(None);
182        };
183        self.with_transaction(async |tx| tx.remove_borrowed::<StoredE2eiEnrollment>(id).await)
184            .await?;
185        Ok(Some(std::mem::take(&mut enrollment.content)))
186    }
187}
188
189#[inline(always)]
190pub fn deser<T: MlsEntity>(bytes: &[u8]) -> Result<T, CryptoKeystoreError> {
191    Ok(postcard::from_bytes(bytes)?)
192}
193
194#[inline(always)]
195pub fn ser<T: MlsEntity>(value: &T) -> Result<Vec<u8>, CryptoKeystoreError> {
196    Ok(postcard::to_stdvec(value)?)
197}
198
199#[cfg_attr(target_os = "unknown", async_trait::async_trait(?Send))]
200#[cfg_attr(not(target_os = "unknown"), async_trait::async_trait)]
201impl openmls_traits::key_store::OpenMlsKeyStore for crate::Database {
202    type Error = CryptoKeystoreError;
203
204    async fn store<V: MlsEntity + Sync>(&self, id: &[u8], value: &V) -> Result<(), Self::Error>
205    where
206        Self: Sized,
207    {
208        if id.is_empty() {
209            return Err(CryptoKeystoreError::MlsKeyStoreError(
210                "The provided key is empty".into(),
211            ));
212        }
213
214        let data = ser(value)?;
215
216        match V::ID {
217            MlsEntityId::GroupState => {
218                return Err(CryptoKeystoreError::IncorrectApiUsage(
219                    "Groups must not be saved using OpenMLS's APIs. You should use the keystore's provided methods",
220                ));
221            }
222            MlsEntityId::SignatureKeyPair => {
223                return Err(CryptoKeystoreError::IncorrectApiUsage(
224                    "Signature keys must not be saved using OpenMLS's APIs. Save a credential via the keystore API
225                    instead.",
226                ));
227            }
228            MlsEntityId::KeyPackage => {
229                let kp = StoredKeypackage {
230                    keypackage_ref: id.into(),
231                    keypackage: data,
232                };
233                self.save(kp).await?;
234            }
235            MlsEntityId::HpkePrivateKey => {
236                let kp = StoredHpkePrivateKey {
237                    pk: id.into(),
238                    sk: data,
239                };
240                self.save(kp).await?;
241            }
242            MlsEntityId::PskBundle => {
243                let kp = StoredPskBundle {
244                    psk_id: id.into(),
245                    psk: data,
246                };
247                self.save(kp).await?;
248            }
249            MlsEntityId::EncryptionKeyPair => {
250                let kp = StoredEncryptionKeyPair {
251                    pk: id.into(),
252                    sk: data,
253                };
254                self.save(kp).await?;
255            }
256            MlsEntityId::EpochEncryptionKeyPair => {
257                let kp = StoredEpochEncryptionKeypair {
258                    id: id.into(),
259                    keypairs: data,
260                };
261                self.save(kp).await?;
262            }
263        }
264
265        Ok(())
266    }
267
268    async fn read<V: MlsEntity>(&self, id: &[u8]) -> Option<V>
269    where
270        Self: Sized,
271    {
272        if id.is_empty() {
273            return None;
274        }
275
276        match V::ID {
277            MlsEntityId::GroupState => {
278                let v = self.get_borrowed::<PersistedMlsGroup>(id).await.ok().flatten()?;
279                deser(&v.state).ok()
280            }
281            MlsEntityId::SignatureKeyPair => {
282                let conn = &*self.conn().await;
283                let hash = Sha256Hash::from_existing_hash(id).ok()?;
284                let stored_credential = StoredCredential::get(conn, &hash).ok().flatten()?;
285                let ciphersuite = Ciphersuite::try_from(stored_credential.ciphersuite).ok()?;
286                let signature_scheme = ciphersuite.signature_algorithm();
287
288                let mls_keypair = SignatureKeyPair::from_raw(
289                    signature_scheme,
290                    stored_credential.private_key.to_vec(),
291                    stored_credential.public_key.to_vec(),
292                );
293
294                // In a well designed interface, something like this should not be necessary. However, we don't have
295                // a well-designed interface.
296                let data = ser(&mls_keypair).ok()?;
297                deser(&data).ok()
298            }
299            MlsEntityId::KeyPackage => {
300                let v = self.get_borrowed::<StoredKeypackage>(id).await.ok().flatten()?;
301                deser(&v.keypackage).ok()
302            }
303            MlsEntityId::HpkePrivateKey => {
304                let v = self.get_borrowed::<StoredHpkePrivateKey>(id).await.ok().flatten()?;
305                deser(&v.sk).ok()
306            }
307            MlsEntityId::PskBundle => {
308                let v = self.get_borrowed::<StoredPskBundle>(id).await.ok().flatten()?;
309                deser(&v.psk).ok()
310            }
311            MlsEntityId::EncryptionKeyPair => {
312                let v = self.get_borrowed::<StoredEncryptionKeyPair>(id).await.ok().flatten()?;
313                deser(&v.sk).ok()
314            }
315            MlsEntityId::EpochEncryptionKeyPair => {
316                let v = self
317                    .get_borrowed::<StoredEpochEncryptionKeypair>(id)
318                    .await
319                    .ok()
320                    .flatten()?;
321                deser(&v.keypairs).ok()
322            }
323        }
324    }
325
326    async fn delete<V: MlsEntity>(&self, id: &[u8]) -> Result<(), Self::Error> {
327        match V::ID {
328            MlsEntityId::GroupState => self.remove_borrowed::<PersistedMlsGroup>(id).await?,
329            MlsEntityId::SignatureKeyPair => unimplemented!(
330                "Deleting a signature key pair should not be done through this API, any keypair should be deleted via
331                deleting a credential."
332            ),
333            MlsEntityId::HpkePrivateKey => self.remove_borrowed::<StoredHpkePrivateKey>(id).await?,
334            MlsEntityId::KeyPackage => self.remove_borrowed::<StoredKeypackage>(id).await?,
335            MlsEntityId::PskBundle => self.remove_borrowed::<StoredPskBundle>(id).await?,
336            MlsEntityId::EncryptionKeyPair => self.remove_borrowed::<StoredEncryptionKeyPair>(id).await?,
337            MlsEntityId::EpochEncryptionKeyPair => self.remove_borrowed::<StoredEpochEncryptionKeypair>(id).await?,
338        }
339
340        Ok(())
341    }
342}