Skip to main content

core_crypto/proteus/
prekey.rs

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