core_crypto/
proteus.rs

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

core_crypto/
proteus.rs
