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