core_crypto/
proteus.rs

1use crate::{
2    CoreCrypto, Error, KeystoreError, LeafError, ProteusError, Result,
3    group_store::{GroupStore, GroupStoreEntity, GroupStoreValue},
4};
5use core_crypto_keystore::{
6    Connection as CryptoKeystore,
7    connection::FetchFromDatabase,
8    entities::{ProteusIdentity, ProteusSession},
9};
10use proteus_wasm::{
11    keys::{IdentityKeyPair, PreKeyBundle},
12    message::Envelope,
13    session::Session,
14};
15use std::{collections::HashMap, sync::Arc};
16
17/// Proteus session IDs, it seems it's basically a string
18pub type SessionIdentifier = String;
19
20/// Proteus Session wrapper, that contains the identifier and the associated proteus Session
21#[derive(Debug)]
22pub struct ProteusConversationSession {
23    pub(crate) identifier: SessionIdentifier,
24    pub(crate) session: Session<Arc<IdentityKeyPair>>,
25}
26
27impl ProteusConversationSession {
28    /// Encrypts a message for this Proteus session
29    pub fn encrypt(&mut self, plaintext: &[u8]) -> Result<Vec<u8>> {
30        self.session
31            .encrypt(plaintext)
32            .and_then(|e| e.serialise())
33            .map_err(ProteusError::wrap("encrypting message for proteus session"))
34            .map_err(Into::into)
35    }
36
37    /// Decrypts a message for this Proteus session
38    pub async fn decrypt(
39        &mut self,
40        store: &mut core_crypto_keystore::Connection,
41        ciphertext: &[u8],
42    ) -> Result<Vec<u8>> {
43        let envelope = Envelope::deserialise(ciphertext).map_err(ProteusError::wrap("deserializing envelope"))?;
44        self.session
45            .decrypt(store, &envelope)
46            .await
47            .map_err(ProteusError::wrap("decrypting message for proteus session"))
48            .map_err(Into::into)
49    }
50
51    /// Returns the session identifier
52    pub fn identifier(&self) -> &str {
53        &self.identifier
54    }
55
56    /// Returns the public key fingerprint of the local identity (= self identity)
57    pub fn fingerprint_local(&self) -> String {
58        self.session.local_identity().fingerprint()
59    }
60
61    /// Returns the public key fingerprint of the remote identity (= client you're communicating with)
62    pub fn fingerprint_remote(&self) -> String {
63        self.session.remote_identity().fingerprint()
64    }
65}
66
67#[cfg_attr(target_family = "wasm", async_trait::async_trait(?Send))]
68#[cfg_attr(not(target_family = "wasm"), async_trait::async_trait)]
69impl GroupStoreEntity for ProteusConversationSession {
70    type RawStoreValue = core_crypto_keystore::entities::ProteusSession;
71    type IdentityType = Arc<proteus_wasm::keys::IdentityKeyPair>;
72
73    async fn fetch_from_id(
74        id: &[u8],
75        identity: Option<Self::IdentityType>,
76        keystore: &impl FetchFromDatabase,
77    ) -> crate::Result<Option<Self>> {
78        let result = keystore
79            .find::<Self::RawStoreValue>(id)
80            .await
81            .map_err(KeystoreError::wrap("finding raw group store entity by id"))?;
82        let Some(store_value) = result else {
83            return Ok(None);
84        };
85
86        let Some(identity) = identity else {
87            return Err(crate::Error::ProteusNotInitialized);
88        };
89
90        let session = proteus_wasm::session::Session::deserialise(identity, &store_value.session)
91            .map_err(ProteusError::wrap("deserializing session"))?;
92
93        Ok(Some(Self {
94            identifier: store_value.id.clone(),
95            session,
96        }))
97    }
98}
99
100impl CoreCrypto {
101    /// Proteus session accessor
102    ///
103    /// Warning: The Proteus client **MUST** be initialized with
104    /// [crate::transaction_context::TransactionContext::proteus_init] first or an error will be
105    /// returned
106    pub async fn proteus_session(
107        &self,
108        session_id: &str,
109    ) -> Result<Option<GroupStoreValue<ProteusConversationSession>>> {
110        let mut mutex = self.proteus.lock().await;
111        let proteus = mutex.as_mut().ok_or(Error::ProteusNotInitialized)?;
112        let keystore = self.mls.crypto_provider.keystore();
113        proteus.session(session_id, &keystore).await
114    }
115
116    /// Proteus session exists
117    ///
118    /// Warning: The Proteus client **MUST** be initialized with
119    /// [crate::transaction_context::TransactionContext::proteus_init] first or an error will be
120    /// returned
121    pub async fn proteus_session_exists(&self, session_id: &str) -> Result<bool> {
122        let mut mutex = self.proteus.lock().await;
123        let proteus = mutex.as_mut().ok_or(Error::ProteusNotInitialized)?;
124        let keystore = self.mls.crypto_provider.keystore();
125        Ok(proteus.session_exists(session_id, &keystore).await)
126    }
127
128    /// Returns the proteus last resort prekey id (u16::MAX = 65535)
129    pub fn proteus_last_resort_prekey_id() -> u16 {
130        ProteusCentral::last_resort_prekey_id()
131    }
132
133    /// Returns the proteus identity's public key fingerprint
134    ///
135    /// Warning: The Proteus client **MUST** be initialized with
136    /// [crate::transaction_context::TransactionContext::proteus_init] first or an error will be
137    /// returned
138    pub async fn proteus_fingerprint(&self) -> Result<String> {
139        let mutex = self.proteus.lock().await;
140        let proteus = mutex.as_ref().ok_or(Error::ProteusNotInitialized)?;
141        Ok(proteus.fingerprint())
142    }
143
144    /// Returns the proteus identity's public key fingerprint
145    ///
146    /// Warning: The Proteus client **MUST** be initialized with
147    /// [crate::transaction_context::TransactionContext::proteus_init] first or an error will be
148    /// returned
149    pub async fn proteus_fingerprint_local(&self, session_id: &str) -> Result<String> {
150        let mut mutex = self.proteus.lock().await;
151        let proteus = mutex.as_mut().ok_or(Error::ProteusNotInitialized)?;
152        let keystore = self.mls.crypto_provider.keystore();
153        proteus.fingerprint_local(session_id, &keystore).await
154    }
155
156    /// Returns the proteus identity's public key fingerprint
157    ///
158    /// Warning: The Proteus client **MUST** be initialized with
159    /// [crate::transaction_context::TransactionContext::proteus_init] first or an error will be
160    /// returned
161    pub async fn proteus_fingerprint_remote(&self, session_id: &str) -> Result<String> {
162        let mut mutex = self.proteus.lock().await;
163        let proteus = mutex.as_mut().ok_or(Error::ProteusNotInitialized)?;
164        let keystore = self.mls.crypto_provider.keystore();
165        proteus.fingerprint_remote(session_id, &keystore).await
166    }
167}
168
169/// Proteus counterpart of [crate::mls::session::Session]
170///
171/// The big difference is that [ProteusCentral] doesn't *own* its own keystore but must borrow it from the outside.
172/// Whether it's exclusively for this struct's purposes or it's shared with our main struct, [crate::mls::session::Session]
173#[derive(Debug)]
174pub struct ProteusCentral {
175    proteus_identity: Arc<IdentityKeyPair>,
176    proteus_sessions: GroupStore<ProteusConversationSession>,
177}
178
179impl ProteusCentral {
180    /// Initializes the [ProteusCentral]
181    pub async fn try_new(keystore: &CryptoKeystore) -> Result<Self> {
182        let proteus_identity: Arc<IdentityKeyPair> = Arc::new(Self::load_or_create_identity(keystore).await?);
183        let proteus_sessions = Self::restore_sessions(keystore, &proteus_identity).await?;
184
185        Ok(Self {
186            proteus_identity,
187            proteus_sessions,
188        })
189    }
190
191    /// Restore proteus sessions from disk
192    pub(crate) async fn reload_sessions(&mut self, keystore: &CryptoKeystore) -> Result<()> {
193        self.proteus_sessions = Self::restore_sessions(keystore, &self.proteus_identity).await?;
194        Ok(())
195    }
196
197    /// This function will try to load a proteus Identity from our keystore; If it cannot, it will create a new one
198    /// This means this function doesn't fail except in cases of deeper errors (such as in the Keystore and other crypto errors)
199    async fn load_or_create_identity(keystore: &CryptoKeystore) -> Result<IdentityKeyPair> {
200        let Some(identity) = keystore
201            .find::<ProteusIdentity>(ProteusIdentity::ID)
202            .await
203            .map_err(KeystoreError::wrap("finding proteus identity"))?
204        else {
205            return Self::create_identity(keystore).await;
206        };
207
208        let sk = identity.sk_raw();
209        let pk = identity.pk_raw();
210
211        // SAFETY: Byte lengths are ensured at the keystore level so this function is safe to call, despite being cursed
212        IdentityKeyPair::from_raw_key_pair(*sk, *pk)
213            .map_err(ProteusError::wrap("constructing identity keypair"))
214            .map_err(Into::into)
215    }
216
217    /// Internal function to create and save a new Proteus Identity
218    async fn create_identity(keystore: &CryptoKeystore) -> Result<IdentityKeyPair> {
219        let kp = IdentityKeyPair::new();
220        let pk = kp.public_key.public_key.as_slice().to_vec();
221
222        let ks_identity = ProteusIdentity {
223            sk: kp.secret_key.to_keypair_bytes().into(),
224            pk,
225        };
226        keystore
227            .save(ks_identity)
228            .await
229            .map_err(KeystoreError::wrap("saving new proteus identity"))?;
230
231        Ok(kp)
232    }
233
234    /// Restores the saved sessions in memory. This is performed automatically on init
235    async fn restore_sessions(
236        keystore: &core_crypto_keystore::Connection,
237        identity: &Arc<IdentityKeyPair>,
238    ) -> Result<GroupStore<ProteusConversationSession>> {
239        let mut proteus_sessions = GroupStore::new_with_limit(crate::group_store::ITEM_LIMIT * 2);
240        for session in keystore
241            .find_all::<ProteusSession>(Default::default())
242            .await
243            .map_err(KeystoreError::wrap("finding all proteus sessions"))?
244            .into_iter()
245        {
246            let proteus_session = Session::deserialise(identity.clone(), &session.session)
247                .map_err(ProteusError::wrap("deserializing session"))?;
248
249            let identifier = session.id.clone();
250
251            let proteus_conversation = ProteusConversationSession {
252                identifier: identifier.clone(),
253                session: proteus_session,
254            };
255
256            if proteus_sessions
257                .try_insert(identifier.into_bytes(), proteus_conversation)
258                .is_err()
259            {
260                break;
261            }
262        }
263
264        Ok(proteus_sessions)
265    }
266
267    /// Creates a new session from a prekey
268    pub async fn session_from_prekey(
269        &mut self,
270        session_id: &str,
271        key: &[u8],
272    ) -> Result<GroupStoreValue<ProteusConversationSession>> {
273        let prekey = PreKeyBundle::deserialise(key).map_err(ProteusError::wrap("deserializing prekey bundle"))?;
274        // Note on the `::<>` turbofish below:
275        //
276        // `init_from_prekey` returns an error type which is parametric over some wrapped `E`,
277        // because one variant (not relevant to this particular operation) wraps an error type based
278        // on a parameter of a different function entirely.
279        //
280        // Rust complains here, because it can't figure out what type that `E` should be. After all, it's
281        // not inferrable from this function call! It is also entirely irrelevant in this case.
282        //
283        // We can derive two general rules about error-handling in Rust from this example:
284        //
285        // 1. It's better to make smaller error types where possible, encapsulating fallible operations
286        //    with their own error variants, and then wrapping those errors where required, as opposed to
287        //    creating giant catch-all errors. Doing so also has knock-on benefits with regard to tracing
288        //    the precise origin of the error.
289        // 2. One should never make an error wrapper parametric. If you need to wrap an unknown error,
290        //    it's always better to wrap a `Box<dyn std::error::Error>` than to make your error type parametric.
291        //    The allocation cost of creating the `Box` is utterly trivial in an error-handling path, and
292        //    it avoids parametric virality. (`init_from_prekey` is itself only generic because it returns
293        //    this error type with a type-parametric variant, which the function never returns.)
294        //
295        // In this case, we have the out of band knowledge that `ProteusErrorKind` has a `#[from]` implementation
296        // for `proteus_wasm::session::Error<core_crypto_keystore::CryptoKeystoreError>` and for no other kinds
297        // of session error. So we can safely say that the type of error we are meant to catch here, and
298        // therefore pass in that otherwise-irrelevant type, to ensure that error handling works properly.
299        //
300        // Some people say that if it's stupid but it works, it's not stupid. I disagree. If it's stupid but
301        // it works, that's our cue to seek out even better, non-stupid ways to get things done. I reiterate:
302        // the actual type referred to in this turbofish is nothing but a magic incantation to make error
303        // handling work; it has no bearing on the error retured from this function. How much better would it
304        // have been if `session::Error` were not parametric and we could have avoided the turbofish entirely?
305        let proteus_session = Session::init_from_prekey::<core_crypto_keystore::CryptoKeystoreError>(
306            self.proteus_identity.clone(),
307            prekey,
308        )
309        .map_err(ProteusError::wrap("initializing session from prekey"))?;
310
311        let proteus_conversation = ProteusConversationSession {
312            identifier: session_id.into(),
313            session: proteus_session,
314        };
315
316        self.proteus_sessions.insert(session_id.into(), proteus_conversation);
317
318        Ok(self.proteus_sessions.get(session_id.as_bytes()).unwrap().clone())
319    }
320
321    /// Creates a new proteus Session from a received message
322    pub(crate) async fn session_from_message(
323        &mut self,
324        keystore: &mut CryptoKeystore,
325        session_id: &str,
326        envelope: &[u8],
327    ) -> Result<(GroupStoreValue<ProteusConversationSession>, Vec<u8>)> {
328        let message = Envelope::deserialise(envelope).map_err(ProteusError::wrap("deserialising envelope"))?;
329        let (session, payload) = Session::init_from_message(self.proteus_identity.clone(), keystore, &message)
330            .await
331            .map_err(ProteusError::wrap("initializing session from message"))?;
332
333        let proteus_conversation = ProteusConversationSession {
334            identifier: session_id.into(),
335            session,
336        };
337
338        self.proteus_sessions.insert(session_id.into(), proteus_conversation);
339
340        Ok((
341            self.proteus_sessions.get(session_id.as_bytes()).unwrap().clone(),
342            payload,
343        ))
344    }
345
346    /// Persists a session in store
347    ///
348    /// **Note**: This isn't usually needed as persisting sessions happens automatically when decrypting/encrypting messages and initializing Sessions
349    pub(crate) async fn session_save(&mut self, keystore: &CryptoKeystore, session_id: &str) -> Result<()> {
350        if let Some(session) = self
351            .proteus_sessions
352            .get_fetch(session_id.as_bytes(), keystore, Some(self.proteus_identity.clone()))
353            .await?
354        {
355            Self::session_save_by_ref(keystore, session).await?;
356        }
357
358        Ok(())
359    }
360
361    pub(crate) async fn session_save_by_ref(
362        keystore: &CryptoKeystore,
363        session: GroupStoreValue<ProteusConversationSession>,
364    ) -> Result<()> {
365        let session = session.read().await;
366        let db_session = ProteusSession {
367            id: session.identifier().to_string(),
368            session: session
369                .session
370                .serialise()
371                .map_err(ProteusError::wrap("serializing session"))?,
372        };
373        keystore
374            .save(db_session)
375            .await
376            .map_err(KeystoreError::wrap("saving proteus session"))?;
377        Ok(())
378    }
379
380    /// Deletes a session in the store
381    pub(crate) async fn session_delete(&mut self, keystore: &CryptoKeystore, session_id: &str) -> Result<()> {
382        if keystore.remove::<ProteusSession, _>(session_id).await.is_ok() {
383            let _ = self.proteus_sessions.remove(session_id.as_bytes());
384        }
385        Ok(())
386    }
387
388    /// Session accessor
389    pub(crate) async fn session(
390        &mut self,
391        session_id: &str,
392        keystore: &CryptoKeystore,
393    ) -> Result<Option<GroupStoreValue<ProteusConversationSession>>> {
394        self.proteus_sessions
395            .get_fetch(session_id.as_bytes(), keystore, Some(self.proteus_identity.clone()))
396            .await
397    }
398
399    /// Session exists
400    pub(crate) async fn session_exists(&mut self, session_id: &str, keystore: &CryptoKeystore) -> bool {
401        self.session(session_id, keystore).await.ok().flatten().is_some()
402    }
403
404    /// Decrypt a proteus message for an already existing session
405    /// Note: This cannot be used for handshake messages, see [ProteusCentral::session_from_message]
406    pub(crate) async fn decrypt(
407        &mut self,
408        keystore: &mut CryptoKeystore,
409        session_id: &str,
410        ciphertext: &[u8],
411    ) -> Result<Vec<u8>> {
412        let session = self
413            .proteus_sessions
414            .get_fetch(session_id.as_bytes(), keystore, Some(self.proteus_identity.clone()))
415            .await?
416            .ok_or(LeafError::ConversationNotFound(session_id.as_bytes().into()))
417            .map_err(ProteusError::wrap("getting session"))?;
418
419        let plaintext = session.write().await.decrypt(keystore, ciphertext).await?;
420        ProteusCentral::session_save_by_ref(keystore, session).await?;
421
422        Ok(plaintext)
423    }
424
425    /// Encrypt a message for a session
426    pub(crate) async fn encrypt(
427        &mut self,
428        keystore: &mut CryptoKeystore,
429        session_id: &str,
430        plaintext: &[u8],
431    ) -> Result<Vec<u8>> {
432        let session = self
433            .session(session_id, keystore)
434            .await?
435            .ok_or(LeafError::ConversationNotFound(session_id.as_bytes().into()))
436            .map_err(ProteusError::wrap("getting session"))?;
437
438        let ciphertext = session.write().await.encrypt(plaintext)?;
439        ProteusCentral::session_save_by_ref(keystore, session).await?;
440
441        Ok(ciphertext)
442    }
443
444    /// Encrypts a message for a list of sessions
445    /// This is mainly used for conversations with multiple clients, this allows to minimize FFI roundtrips
446    pub(crate) async fn encrypt_batched(
447        &mut self,
448        keystore: &mut CryptoKeystore,
449        sessions: &[impl AsRef<str>],
450        plaintext: &[u8],
451    ) -> Result<HashMap<String, Vec<u8>>> {
452        let mut acc = HashMap::new();
453        for session_id in sessions {
454            if let Some(session) = self.session(session_id.as_ref(), keystore).await? {
455                let mut session_w = session.write().await;
456                acc.insert(session_w.identifier.clone(), session_w.encrypt(plaintext)?);
457                drop(session_w);
458
459                ProteusCentral::session_save_by_ref(keystore, session).await?;
460            }
461        }
462        Ok(acc)
463    }
464
465    /// Generates a new Proteus PreKey, stores it in the keystore and returns a serialized PreKeyBundle to be consumed externally
466    pub(crate) async fn new_prekey(&self, id: u16, keystore: &CryptoKeystore) -> Result<Vec<u8>> {
467        use proteus_wasm::keys::{PreKey, PreKeyId};
468
469        let prekey_id = PreKeyId::new(id);
470        let prekey = PreKey::new(prekey_id);
471        let keystore_prekey = core_crypto_keystore::entities::ProteusPrekey::from_raw(
472            id,
473            prekey.serialise().map_err(ProteusError::wrap("serialising prekey"))?,
474        );
475        let bundle = PreKeyBundle::new(self.proteus_identity.as_ref().public_key.clone(), &prekey);
476        let bundle = bundle
477            .serialise()
478            .map_err(ProteusError::wrap("serialising prekey bundle"))?;
479        keystore
480            .save(keystore_prekey)
481            .await
482            .map_err(KeystoreError::wrap("saving keystore prekey"))?;
483        Ok(bundle)
484    }
485
486    /// Generates a new Proteus Prekey, with an automatically auto-incremented ID.
487    ///
488    /// See [ProteusCentral::new_prekey]
489    pub(crate) async fn new_prekey_auto(&self, keystore: &CryptoKeystore) -> Result<(u16, Vec<u8>)> {
490        let id = core_crypto_keystore::entities::ProteusPrekey::get_free_id(keystore)
491            .await
492            .map_err(KeystoreError::wrap("getting proteus prekey by id"))?;
493        Ok((id, self.new_prekey(id, keystore).await?))
494    }
495
496    /// Returns the Proteus last resort prekey ID (u16::MAX = 65535 = 0xFFFF)
497    pub fn last_resort_prekey_id() -> u16 {
498        proteus_wasm::keys::MAX_PREKEY_ID.value()
499    }
500
501    /// Returns the Proteus last resort prekey
502    /// If it cannot be found, one will be created.
503    pub(crate) async fn last_resort_prekey(&self, keystore: &CryptoKeystore) -> Result<Vec<u8>> {
504        let last_resort = if let Some(last_resort) = keystore
505            .find::<core_crypto_keystore::entities::ProteusPrekey>(
506                Self::last_resort_prekey_id().to_le_bytes().as_slice(),
507            )
508            .await
509            .map_err(KeystoreError::wrap("finding proteus prekey"))?
510        {
511            proteus_wasm::keys::PreKey::deserialise(&last_resort.prekey)
512                .map_err(ProteusError::wrap("deserialising proteus prekey"))?
513        } else {
514            let last_resort = proteus_wasm::keys::PreKey::last_resort();
515
516            use core_crypto_keystore::CryptoKeystoreProteus as _;
517            keystore
518                .proteus_store_prekey(
519                    Self::last_resort_prekey_id(),
520                    &last_resort
521                        .serialise()
522                        .map_err(ProteusError::wrap("serialising last resort prekey"))?,
523                )
524                .await
525                .map_err(KeystoreError::wrap("storing proteus prekey"))?;
526
527            last_resort
528        };
529
530        let bundle = PreKeyBundle::new(self.proteus_identity.as_ref().public_key.clone(), &last_resort);
531        let bundle = bundle
532            .serialise()
533            .map_err(ProteusError::wrap("serialising prekey bundle"))?;
534
535        Ok(bundle)
536    }
537
538    /// Proteus identity keypair
539    pub fn identity(&self) -> &IdentityKeyPair {
540        self.proteus_identity.as_ref()
541    }
542
543    /// Proteus Public key hex-encoded fingerprint
544    pub fn fingerprint(&self) -> String {
545        self.proteus_identity.as_ref().public_key.fingerprint()
546    }
547
548    /// Proteus Session local hex-encoded fingerprint
549    ///
550    /// # Errors
551    /// When the session is not found
552    pub(crate) async fn fingerprint_local(&mut self, session_id: &str, keystore: &CryptoKeystore) -> Result<String> {
553        let session = self
554            .session(session_id, keystore)
555            .await?
556            .ok_or(LeafError::ConversationNotFound(session_id.as_bytes().into()))
557            .map_err(ProteusError::wrap("getting session"))?;
558        let fingerprint = session.read().await.fingerprint_local();
559        Ok(fingerprint)
560    }
561
562    /// Proteus Session remote hex-encoded fingerprint
563    ///
564    /// # Errors
565    /// When the session is not found
566    pub(crate) async fn fingerprint_remote(&mut self, session_id: &str, keystore: &CryptoKeystore) -> Result<String> {
567        let session = self
568            .session(session_id, keystore)
569            .await?
570            .ok_or(LeafError::ConversationNotFound(session_id.as_bytes().into()))
571            .map_err(ProteusError::wrap("getting session"))?;
572        let fingerprint = session.read().await.fingerprint_remote();
573        Ok(fingerprint)
574    }
575
576    /// Hex-encoded fingerprint of the given prekey
577    ///
578    /// # Errors
579    /// If the prekey cannot be deserialized
580    pub fn fingerprint_prekeybundle(prekey: &[u8]) -> Result<String> {
581        let prekey = PreKeyBundle::deserialise(prekey).map_err(ProteusError::wrap("deserialising prekey bundle"))?;
582        Ok(prekey.identity_key.fingerprint())
583    }
584}
585
586#[cfg(test)]
587mod tests {
588    use crate::{
589        prelude::{CertificateBundle, ClientIdentifier, MlsCredentialType, Session, SessionConfig},
590        test_utils::{proteus_utils::*, x509::X509TestChain, *},
591    };
592
593    use crate::prelude::INITIAL_KEYING_MATERIAL_COUNT;
594
595    use super::*;
596
597    use core_crypto_keystore::{ConnectionType, DatabaseKey};
598
599    #[apply(all_cred_cipher)]
600    async fn cc_can_init(case: TestContext) {
601        #[cfg(not(target_family = "wasm"))]
602        let (path, db_file) = tmp_db_file();
603        #[cfg(target_family = "wasm")]
604        let (path, _) = tmp_db_file();
605        let client_id = "alice".into();
606        let cfg = SessionConfig::builder()
607            .persistent(&path)
608            .database_key(DatabaseKey::generate())
609            .client_id(client_id)
610            .ciphersuites([case.ciphersuite()])
611            .build()
612            .validate()
613            .unwrap();
614
615        let cc: CoreCrypto = Session::try_new(cfg).await.unwrap().into();
616        let context = cc.new_transaction().await.unwrap();
617        assert!(context.proteus_init().await.is_ok());
618        assert!(context.proteus_new_prekey(1).await.is_ok());
619        context.finish().await.unwrap();
620        #[cfg(not(target_family = "wasm"))]
621        drop(db_file);
622    }
623
624    #[apply(all_cred_cipher)]
625    async fn cc_can_2_phase_init(case: TestContext) {
626        #[cfg(not(target_family = "wasm"))]
627        let (path, db_file) = tmp_db_file();
628        #[cfg(target_family = "wasm")]
629        let (path, _) = tmp_db_file();
630        // we are deferring MLS initialization here, not passing a MLS 'client_id' yet
631        let cfg = SessionConfig::builder()
632            .persistent(&path)
633            .database_key(DatabaseKey::generate())
634            .ciphersuites([case.ciphersuite()])
635            .build()
636            .validate()
637            .unwrap();
638
639        let cc: CoreCrypto = Session::try_new(cfg).await.unwrap().into();
640        let transaction = cc.new_transaction().await.unwrap();
641        let x509_test_chain = X509TestChain::init_empty(case.signature_scheme());
642        x509_test_chain.register_with_central(&transaction).await;
643        assert!(transaction.proteus_init().await.is_ok());
644        // proteus is initialized, prekeys can be generated
645        assert!(transaction.proteus_new_prekey(1).await.is_ok());
646        // 👇 and so a unique 'client_id' can be fetched from wire-server
647        let client_id = "alice";
648        let identifier = match case.credential_type {
649            MlsCredentialType::Basic => ClientIdentifier::Basic(client_id.into()),
650            MlsCredentialType::X509 => {
651                CertificateBundle::rand_identifier(client_id, &[x509_test_chain.find_local_intermediate_ca()])
652            }
653        };
654        transaction
655            .mls_init(
656                identifier,
657                vec![case.ciphersuite()],
658                Some(INITIAL_KEYING_MATERIAL_COUNT),
659            )
660            .await
661            .unwrap();
662        // expect MLS to work
663        assert_eq!(
664            transaction
665                .get_or_create_client_keypackages(case.ciphersuite(), case.credential_type, 2)
666                .await
667                .unwrap()
668                .len(),
669            2
670        );
671        #[cfg(not(target_family = "wasm"))]
672        drop(db_file);
673    }
674
675    #[macro_rules_attribute::apply(smol_macros::test)]
676    async fn can_init() {
677        #[cfg(not(target_family = "wasm"))]
678        let (path, db_file) = tmp_db_file();
679        #[cfg(target_family = "wasm")]
680        let (path, _) = tmp_db_file();
681        let key = DatabaseKey::generate();
682        let keystore = core_crypto_keystore::Connection::open(ConnectionType::Persistent(&path), &key)
683            .await
684            .unwrap();
685        keystore.new_transaction().await.unwrap();
686        let central = ProteusCentral::try_new(&keystore).await.unwrap();
687        let identity = (*central.proteus_identity).clone();
688        keystore.commit_transaction().await.unwrap();
689
690        let keystore = core_crypto_keystore::Connection::open(ConnectionType::Persistent(&path), &key)
691            .await
692            .unwrap();
693        keystore.new_transaction().await.unwrap();
694        let central = ProteusCentral::try_new(&keystore).await.unwrap();
695        keystore.commit_transaction().await.unwrap();
696        assert_eq!(identity, *central.proteus_identity);
697
698        keystore.wipe().await.unwrap();
699        #[cfg(not(target_family = "wasm"))]
700        drop(db_file);
701    }
702
703    #[macro_rules_attribute::apply(smol_macros::test)]
704    async fn can_talk_with_proteus() {
705        #[cfg(not(target_family = "wasm"))]
706        let (path, db_file) = tmp_db_file();
707        #[cfg(target_family = "wasm")]
708        let (path, _) = tmp_db_file();
709
710        let session_id = uuid::Uuid::new_v4().hyphenated().to_string();
711
712        let key = DatabaseKey::generate();
713        let mut keystore = core_crypto_keystore::Connection::open(ConnectionType::Persistent(&path), &key)
714            .await
715            .unwrap();
716        keystore.new_transaction().await.unwrap();
717
718        let mut alice = ProteusCentral::try_new(&keystore).await.unwrap();
719
720        let mut bob = CryptoboxLike::init();
721        let bob_pk_bundle = bob.new_prekey();
722
723        alice
724            .session_from_prekey(&session_id, &bob_pk_bundle.serialise().unwrap())
725            .await
726            .unwrap();
727
728        let message = b"Hello world";
729
730        let encrypted = alice.encrypt(&mut keystore, &session_id, message).await.unwrap();
731        let decrypted = bob.decrypt(&session_id, &encrypted).await;
732        assert_eq!(decrypted, message);
733
734        let encrypted = bob.encrypt(&session_id, message);
735        let decrypted = alice.decrypt(&mut keystore, &session_id, &encrypted).await.unwrap();
736        assert_eq!(decrypted, message);
737
738        keystore.commit_transaction().await.unwrap();
739        keystore.wipe().await.unwrap();
740        #[cfg(not(target_family = "wasm"))]
741        drop(db_file);
742    }
743
744    #[macro_rules_attribute::apply(smol_macros::test)]
745    async fn can_produce_proteus_consumed_prekeys() {
746        #[cfg(not(target_family = "wasm"))]
747        let (path, db_file) = tmp_db_file();
748        #[cfg(target_family = "wasm")]
749        let (path, _) = tmp_db_file();
750
751        let session_id = uuid::Uuid::new_v4().hyphenated().to_string();
752
753        let key = DatabaseKey::generate();
754        let mut keystore = core_crypto_keystore::Connection::open(ConnectionType::Persistent(&path), &key)
755            .await
756            .unwrap();
757        keystore.new_transaction().await.unwrap();
758        let mut alice = ProteusCentral::try_new(&keystore).await.unwrap();
759
760        let mut bob = CryptoboxLike::init();
761
762        let alice_prekey_bundle_ser = alice.new_prekey(1, &keystore).await.unwrap();
763
764        bob.init_session_from_prekey_bundle(&session_id, &alice_prekey_bundle_ser);
765        let message = b"Hello world!";
766        let encrypted = bob.encrypt(&session_id, message);
767
768        let (_, decrypted) = alice
769            .session_from_message(&mut keystore, &session_id, &encrypted)
770            .await
771            .unwrap();
772
773        assert_eq!(message, decrypted.as_slice());
774
775        let encrypted = alice.encrypt(&mut keystore, &session_id, message).await.unwrap();
776        let decrypted = bob.decrypt(&session_id, &encrypted).await;
777
778        assert_eq!(message, decrypted.as_slice());
779        keystore.commit_transaction().await.unwrap();
780        keystore.wipe().await.unwrap();
781        #[cfg(not(target_family = "wasm"))]
782        drop(db_file);
783    }
784
785    #[macro_rules_attribute::apply(smol_macros::test)]
786    async fn auto_prekeys_are_sequential() {
787        use core_crypto_keystore::entities::ProteusPrekey;
788        const GAP_AMOUNT: u16 = 5;
789        const ID_TEST_RANGE: std::ops::RangeInclusive<u16> = 1..=30;
790
791        #[cfg(not(target_family = "wasm"))]
792        let (path, db_file) = tmp_db_file();
793        #[cfg(target_family = "wasm")]
794        let (path, _) = tmp_db_file();
795
796        let key = DatabaseKey::generate();
797        let keystore = core_crypto_keystore::Connection::open(ConnectionType::Persistent(&path), &key)
798            .await
799            .unwrap();
800        keystore.new_transaction().await.unwrap();
801        let alice = ProteusCentral::try_new(&keystore).await.unwrap();
802
803        for i in ID_TEST_RANGE {
804            let (pk_id, pkb) = alice.new_prekey_auto(&keystore).await.unwrap();
805            assert_eq!(i, pk_id);
806            let prekey = proteus_wasm::keys::PreKeyBundle::deserialise(&pkb).unwrap();
807            assert_eq!(prekey.prekey_id.value(), pk_id);
808        }
809
810        use rand::Rng as _;
811        let mut rng = rand::thread_rng();
812        let mut gap_ids: Vec<u16> = (0..GAP_AMOUNT).map(|_| rng.gen_range(ID_TEST_RANGE)).collect();
813        gap_ids.sort();
814        gap_ids.dedup();
815        while gap_ids.len() < GAP_AMOUNT as usize {
816            gap_ids.push(rng.gen_range(ID_TEST_RANGE));
817            gap_ids.sort();
818            gap_ids.dedup();
819        }
820        for gap_id in gap_ids.iter() {
821            keystore.remove::<ProteusPrekey, _>(gap_id.to_le_bytes()).await.unwrap();
822        }
823
824        gap_ids.sort();
825
826        for gap_id in gap_ids.iter() {
827            let (pk_id, pkb) = alice.new_prekey_auto(&keystore).await.unwrap();
828            assert_eq!(pk_id, *gap_id);
829            let prekey = proteus_wasm::keys::PreKeyBundle::deserialise(&pkb).unwrap();
830            assert_eq!(prekey.prekey_id.value(), *gap_id);
831        }
832
833        let mut gap_ids: Vec<u16> = (0..GAP_AMOUNT).map(|_| rng.gen_range(ID_TEST_RANGE)).collect();
834        gap_ids.sort();
835        gap_ids.dedup();
836        while gap_ids.len() < GAP_AMOUNT as usize {
837            gap_ids.push(rng.gen_range(ID_TEST_RANGE));
838            gap_ids.sort();
839            gap_ids.dedup();
840        }
841        for gap_id in gap_ids.iter() {
842            keystore.remove::<ProteusPrekey, _>(gap_id.to_le_bytes()).await.unwrap();
843        }
844
845        let potential_range = *ID_TEST_RANGE.end()..=(*ID_TEST_RANGE.end() * 2);
846        let potential_range_check = potential_range.clone();
847        for _ in potential_range {
848            let (pk_id, pkb) = alice.new_prekey_auto(&keystore).await.unwrap();
849            assert!(gap_ids.contains(&pk_id) || potential_range_check.contains(&pk_id));
850            let prekey = proteus_wasm::keys::PreKeyBundle::deserialise(&pkb).unwrap();
851            assert_eq!(prekey.prekey_id.value(), pk_id);
852        }
853        keystore.commit_transaction().await.unwrap();
854        keystore.wipe().await.unwrap();
855        #[cfg(not(target_family = "wasm"))]
856        drop(db_file);
857    }
858}