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 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 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 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 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 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 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 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 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 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 pub async fn pop_e2ei_enrollment(&self, id: &[u8]) -> CryptoKeystoreResult<Option<Vec<u8>>> {
179 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 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}