Skip to main content

core_crypto/proteus/
prekey.rs

1use core_crypto_keystore::{Database, traits::FetchFromDatabase as _};
2use proteus_wasm::keys::PreKeyBundle;
3
4use super::ProteusCentral;
5use crate::{KeystoreError, ProteusError, Result};
6
7impl ProteusCentral {
8    /// Generates a new Proteus PreKey, stores it in the keystore and returns a serialized PreKeyBundle to be consumed
9    /// externally
10    pub(crate) async fn new_prekey(&self, id: u16, keystore: &Database) -> Result<Vec<u8>> {
11        use proteus_wasm::keys::{PreKey, PreKeyId};
12
13        let prekey_id = PreKeyId::new(id);
14        let prekey = PreKey::new(prekey_id);
15        let keystore_prekey = core_crypto_keystore::entities::ProteusPrekey::from_raw(
16            id,
17            prekey.serialise().map_err(ProteusError::wrap("serialising prekey"))?,
18        );
19        let bundle = PreKeyBundle::new(self.proteus_identity.as_ref().public_key.clone(), &prekey);
20        let bundle = bundle
21            .serialise()
22            .map_err(ProteusError::wrap("serialising prekey bundle"))?;
23        keystore
24            .save(keystore_prekey)
25            .await
26            .map_err(KeystoreError::wrap("saving keystore prekey"))?;
27        Ok(bundle)
28    }
29
30    /// Generates a new Proteus Prekey, with an automatically auto-incremented ID.
31    ///
32    /// See [ProteusCentral::new_prekey]
33    pub(crate) async fn new_prekey_auto(&self, keystore: &Database) -> Result<(u16, Vec<u8>)> {
34        let id = core_crypto_keystore::entities::ProteusPrekey::get_free_id(keystore)
35            .await
36            .map_err(KeystoreError::wrap("getting proteus prekey by id"))?;
37        Ok((id, self.new_prekey(id, keystore).await?))
38    }
39
40    /// Returns the Proteus last resort prekey ID (u16::MAX = 65535 = 0xFFFF)
41    pub fn last_resort_prekey_id() -> u16 {
42        proteus_wasm::keys::MAX_PREKEY_ID.value()
43    }
44
45    /// Returns the Proteus last resort prekey
46    /// If it cannot be found, one will be created.
47    pub(crate) async fn last_resort_prekey(&self, keystore: &Database) -> Result<Vec<u8>> {
48        let last_resort = if let Some(last_resort) = keystore
49            .get::<core_crypto_keystore::entities::ProteusPrekey>(&Self::last_resort_prekey_id())
50            .await
51            .map_err(KeystoreError::wrap("finding proteus prekey"))?
52        {
53            proteus_wasm::keys::PreKey::deserialise(&last_resort.prekey)
54                .map_err(ProteusError::wrap("deserialising proteus prekey"))?
55        } else {
56            let last_resort = proteus_wasm::keys::PreKey::last_resort();
57
58            use core_crypto_keystore::CryptoKeystoreProteus as _;
59            keystore
60                .proteus_store_prekey(
61                    Self::last_resort_prekey_id(),
62                    &last_resort
63                        .serialise()
64                        .map_err(ProteusError::wrap("serialising last resort prekey"))?,
65                )
66                .await
67                .map_err(KeystoreError::wrap("storing proteus prekey"))?;
68
69            last_resort
70        };
71
72        let bundle = PreKeyBundle::new(self.proteus_identity.as_ref().public_key.clone(), &last_resort);
73        let bundle = bundle
74            .serialise()
75            .map_err(ProteusError::wrap("serialising prekey bundle"))?;
76
77        Ok(bundle)
78    }
79
80    /// Hex-encoded fingerprint of the given prekey
81    ///
82    /// # Errors
83    /// If the prekey cannot be deserialized
84    pub fn fingerprint_prekeybundle(prekey: &[u8]) -> Result<String> {
85        let prekey = PreKeyBundle::deserialise(prekey).map_err(ProteusError::wrap("deserialising prekey bundle"))?;
86        Ok(prekey.identity_key.fingerprint())
87    }
88}
89
90#[cfg(test)]
91mod tests {
92    use core_crypto_keystore::{ConnectionType, DatabaseKey};
93
94    use super::*;
95    use crate::test_utils::{proteus_utils::*, *};
96
97    #[macro_rules_attribute::apply(smol_macros::test)]
98    async fn can_produce_proteus_consumed_prekeys() {
99        #[cfg(not(target_os = "unknown"))]
100        let (path, db_file) = tmp_db_file();
101        #[cfg(target_os = "unknown")]
102        let (path, _) = tmp_db_file();
103
104        let session_id = uuid::Uuid::new_v4().hyphenated().to_string();
105
106        let key = DatabaseKey::generate();
107        let mut keystore = core_crypto_keystore::Database::open(ConnectionType::Persistent(&path), &key)
108            .await
109            .unwrap();
110        keystore.new_transaction().await.unwrap();
111        let mut alice = ProteusCentral::try_new(&keystore).await.unwrap();
112
113        let mut bob = CryptoboxLike::init();
114
115        let alice_prekey_bundle_ser = alice.new_prekey(1, &keystore).await.unwrap();
116
117        bob.init_session_from_prekey_bundle(&session_id, &alice_prekey_bundle_ser);
118        let message = b"Hello world!";
119        let encrypted = bob.encrypt(&session_id, message);
120
121        let (_, decrypted) = alice
122            .session_from_message(&mut keystore, &session_id, &encrypted)
123            .await
124            .unwrap();
125
126        assert_eq!(message, decrypted.as_slice());
127
128        let encrypted = alice.encrypt(&mut keystore, &session_id, message).await.unwrap();
129        let decrypted = bob.decrypt(&session_id, &encrypted).await;
130
131        assert_eq!(message, decrypted.as_slice());
132        keystore.commit_transaction().await.unwrap();
133        keystore.wipe().await.unwrap();
134        #[cfg(not(target_os = "unknown"))]
135        drop(db_file);
136    }
137
138    #[macro_rules_attribute::apply(smol_macros::test)]
139    async fn auto_prekeys_are_sequential() {
140        use core_crypto_keystore::entities::ProteusPrekey;
141        const GAP_AMOUNT: u16 = 5;
142        const ID_TEST_RANGE: std::ops::RangeInclusive<u16> = 1..=30;
143
144        #[cfg(not(target_os = "unknown"))]
145        let (path, db_file) = tmp_db_file();
146        #[cfg(target_os = "unknown")]
147        let (path, _) = tmp_db_file();
148
149        let key = DatabaseKey::generate();
150        let keystore = core_crypto_keystore::Database::open(ConnectionType::Persistent(&path), &key)
151            .await
152            .unwrap();
153        keystore.new_transaction().await.unwrap();
154        let alice = ProteusCentral::try_new(&keystore).await.unwrap();
155
156        for i in ID_TEST_RANGE {
157            let (pk_id, pkb) = alice.new_prekey_auto(&keystore).await.unwrap();
158            assert_eq!(i, pk_id);
159            let prekey = proteus_wasm::keys::PreKeyBundle::deserialise(&pkb).unwrap();
160            assert_eq!(prekey.prekey_id.value(), pk_id);
161        }
162
163        use rand::Rng as _;
164        let mut rng = rand::thread_rng();
165        let mut gap_ids: Vec<u16> = (0..GAP_AMOUNT).map(|_| rng.gen_range(ID_TEST_RANGE)).collect();
166        gap_ids.sort();
167        gap_ids.dedup();
168        while gap_ids.len() < GAP_AMOUNT as usize {
169            gap_ids.push(rng.gen_range(ID_TEST_RANGE));
170            gap_ids.sort();
171            gap_ids.dedup();
172        }
173        for gap_id in gap_ids.iter() {
174            keystore.remove::<ProteusPrekey>(gap_id).await.unwrap();
175        }
176
177        gap_ids.sort();
178
179        for gap_id in gap_ids.iter() {
180            let (pk_id, pkb) = alice.new_prekey_auto(&keystore).await.unwrap();
181            assert_eq!(pk_id, *gap_id);
182            let prekey = proteus_wasm::keys::PreKeyBundle::deserialise(&pkb).unwrap();
183            assert_eq!(prekey.prekey_id.value(), *gap_id);
184        }
185
186        let mut gap_ids: Vec<u16> = (0..GAP_AMOUNT).map(|_| rng.gen_range(ID_TEST_RANGE)).collect();
187        gap_ids.sort();
188        gap_ids.dedup();
189        while gap_ids.len() < GAP_AMOUNT as usize {
190            gap_ids.push(rng.gen_range(ID_TEST_RANGE));
191            gap_ids.sort();
192            gap_ids.dedup();
193        }
194        for gap_id in gap_ids.iter() {
195            keystore.remove::<ProteusPrekey>(gap_id).await.unwrap();
196        }
197
198        let potential_range = *ID_TEST_RANGE.end()..=(*ID_TEST_RANGE.end() * 2);
199        let potential_range_check = potential_range.clone();
200        for _ in potential_range {
201            let (pk_id, pkb) = alice.new_prekey_auto(&keystore).await.unwrap();
202            assert!(gap_ids.contains(&pk_id) || potential_range_check.contains(&pk_id));
203            let prekey = proteus_wasm::keys::PreKeyBundle::deserialise(&pkb).unwrap();
204            assert_eq!(prekey.prekey_id.value(), pk_id);
205        }
206        keystore.commit_transaction().await.unwrap();
207        keystore.wipe().await.unwrap();
208        #[cfg(not(target_os = "unknown"))]
209        drop(db_file);
210    }
211}