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
14impl 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 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 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 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 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 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 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 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 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 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}