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            keystore
59                .proteus_store_prekey(
60                    Self::last_resort_prekey_id(),
61                    &last_resort
62                        .serialise()
63                        .map_err(ProteusError::wrap("serialising last resort prekey"))?,
64                )
65                .await
66                .map_err(KeystoreError::wrap("storing proteus prekey"))?;
67
68            last_resort
69        };
70
71        let bundle = PreKeyBundle::new(self.proteus_identity.as_ref().public_key.clone(), &last_resort);
72        let bundle = bundle
73            .serialise()
74            .map_err(ProteusError::wrap("serialising prekey bundle"))?;
75
76        Ok(bundle)
77    }
78
79    /// Hex-encoded fingerprint of the given prekey
80    ///
81    /// # Errors
82    /// If the prekey cannot be deserialized
83    pub fn fingerprint_prekeybundle(prekey: &[u8]) -> Result<String> {
84        let prekey = PreKeyBundle::deserialise(prekey).map_err(ProteusError::wrap("deserialising prekey bundle"))?;
85        Ok(prekey.identity_key.fingerprint())
86    }
87}
88
89#[cfg(test)]
90mod tests {
91    use core_crypto_keystore::DatabaseKey;
92
93    use super::*;
94    use crate::test_utils::{proteus_utils::*, *};
95
96    #[macro_rules_attribute::apply(smol_macros::test)]
97    async fn can_produce_proteus_consumed_prekeys() {
98        #[cfg(not(target_os = "unknown"))]
99        let (path, db_file) = tmp_db_file();
100        #[cfg(target_os = "unknown")]
101        let (path, _) = tmp_db_file();
102
103        let session_id = uuid::Uuid::new_v4().hyphenated().to_string();
104
105        let key = DatabaseKey::generate();
106        let keystore = core_crypto_keystore::Database::open(&path, &key).await.unwrap();
107        keystore.new_transaction().await.unwrap();
108        let mut alice = ProteusCentral::try_new(&keystore).await.unwrap();
109
110        let mut bob = CryptoboxLike::init();
111
112        let alice_prekey_bundle_ser = alice.new_prekey(1, &keystore).await.unwrap();
113
114        bob.init_session_from_prekey_bundle(&session_id, &alice_prekey_bundle_ser);
115        let message = b"Hello world!";
116        let encrypted = bob.encrypt(&session_id, message);
117
118        let (_, decrypted) = alice
119            .session_from_message(&keystore, &session_id, &encrypted)
120            .await
121            .unwrap();
122
123        assert_eq!(message, decrypted.as_slice());
124
125        let encrypted = alice.encrypt(&keystore, &session_id, message).await.unwrap();
126        let decrypted = bob.decrypt(&session_id, &encrypted).await;
127
128        assert_eq!(message, decrypted.as_slice());
129        keystore.commit_transaction().await.unwrap();
130        #[cfg(not(target_os = "unknown"))]
131        drop(db_file);
132    }
133
134    #[macro_rules_attribute::apply(smol_macros::test)]
135    async fn auto_prekeys_are_sequential() {
136        use core_crypto_keystore::entities::ProteusPrekey;
137        const GAP_AMOUNT: u16 = 5;
138        const ID_TEST_RANGE: std::ops::RangeInclusive<u16> = 1..=30;
139
140        #[cfg(not(target_os = "unknown"))]
141        let (path, db_file) = tmp_db_file();
142        #[cfg(target_os = "unknown")]
143        let (path, _) = tmp_db_file();
144
145        let key = DatabaseKey::generate();
146        let keystore = core_crypto_keystore::Database::open(&path, &key).await.unwrap();
147        keystore.new_transaction().await.unwrap();
148        let alice = ProteusCentral::try_new(&keystore).await.unwrap();
149
150        for i in ID_TEST_RANGE {
151            let (pk_id, pkb) = alice.new_prekey_auto(&keystore).await.unwrap();
152            assert_eq!(i, pk_id);
153            let prekey = proteus_wasm::keys::PreKeyBundle::deserialise(&pkb).unwrap();
154            assert_eq!(prekey.prekey_id.value(), pk_id);
155        }
156
157        use rand::Rng as _;
158        let mut rng = rand::thread_rng();
159        let mut gap_ids: Vec<u16> = (0..GAP_AMOUNT).map(|_| rng.gen_range(ID_TEST_RANGE)).collect();
160        gap_ids.sort();
161        gap_ids.dedup();
162        while gap_ids.len() < GAP_AMOUNT as usize {
163            gap_ids.push(rng.gen_range(ID_TEST_RANGE));
164            gap_ids.sort();
165            gap_ids.dedup();
166        }
167        for gap_id in gap_ids.iter() {
168            keystore.remove::<ProteusPrekey>(gap_id).await.unwrap();
169        }
170
171        gap_ids.sort();
172
173        for gap_id in gap_ids.iter() {
174            let (pk_id, pkb) = alice.new_prekey_auto(&keystore).await.unwrap();
175            assert_eq!(pk_id, *gap_id);
176            let prekey = proteus_wasm::keys::PreKeyBundle::deserialise(&pkb).unwrap();
177            assert_eq!(prekey.prekey_id.value(), *gap_id);
178        }
179
180        let mut gap_ids: Vec<u16> = (0..GAP_AMOUNT).map(|_| rng.gen_range(ID_TEST_RANGE)).collect();
181        gap_ids.sort();
182        gap_ids.dedup();
183        while gap_ids.len() < GAP_AMOUNT as usize {
184            gap_ids.push(rng.gen_range(ID_TEST_RANGE));
185            gap_ids.sort();
186            gap_ids.dedup();
187        }
188        for gap_id in gap_ids.iter() {
189            keystore.remove::<ProteusPrekey>(gap_id).await.unwrap();
190        }
191
192        let potential_range = *ID_TEST_RANGE.end()..=(*ID_TEST_RANGE.end() * 2);
193        let potential_range_check = potential_range.clone();
194        for _ in potential_range {
195            let (pk_id, pkb) = alice.new_prekey_auto(&keystore).await.unwrap();
196            assert!(gap_ids.contains(&pk_id) || potential_range_check.contains(&pk_id));
197            let prekey = proteus_wasm::keys::PreKeyBundle::deserialise(&pkb).unwrap();
198            assert_eq!(prekey.prekey_id.value(), pk_id);
199        }
200        keystore.commit_transaction().await.unwrap();
201        #[cfg(not(target_os = "unknown"))]
202        drop(db_file);
203    }
204}