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