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    /// Auto prekeys fill the holes left by deleted prekeys before extending the id space.
130    ///
131    /// The order in which the holes get filled is deliberately not asserted: the keystore tracks
132    /// ids freed within a transaction in an unordered map, so all that is promised is that every
133    /// hole is filled exactly once before any fresh id is handed out.
134    #[macro_rules_attribute::apply(smol_macros::test)]
135    async fn auto_prekeys_fill_holes() {
136        use core_crypto_keystore::entities::ProteusPrekey;
137        const GAP_AMOUNT: usize = 5;
138        /// How far past the end of the filled id space to keep claiming ids
139        const EXTRA_AMOUNT: u16 = 10;
140        const ID_TEST_RANGE: std::ops::RangeInclusive<u16> = 1..=30;
141
142        /// Claim `count` auto prekeys, returning the assigned ids in the order they were handed out.
143        async fn claim(alice: &ProteusCentral, tx: &Transaction, count: usize) -> Vec<u16> {
144            let mut ids = Vec::with_capacity(count);
145            for _ in 0..count {
146                let (pk_id, pkb) = alice.new_prekey_auto(tx).await.unwrap();
147                let prekey = proteus_wasm::keys::PreKeyBundle::deserialise(&pkb).unwrap();
148                assert_eq!(prekey.prekey_id.value(), pk_id, "the bundle must carry the assigned id");
149                ids.push(pk_id);
150            }
151            ids
152        }
153
154        /// Pick `count` distinct ids from `ID_TEST_RANGE`, in ascending order.
155        fn pick_gap_ids(rng: &mut impl rand::Rng, count: usize) -> Vec<u16> {
156            let mut ids = Vec::with_capacity(count);
157            while ids.len() < count {
158                let id = rng.gen_range(ID_TEST_RANGE);
159                if !ids.contains(&id) {
160                    ids.push(id);
161                }
162            }
163            ids.sort();
164            ids
165        }
166
167        fn ascending(ids: &[u16]) -> Vec<u16> {
168            let mut ids = ids.to_owned();
169            ids.sort();
170            ids
171        }
172
173        #[cfg(not(target_os = "unknown"))]
174        let (path, db_file) = tmp_db_file();
175        #[cfg(target_os = "unknown")]
176        let (path, _) = tmp_db_file();
177
178        let key = DatabaseKey::generate();
179        let keystore = core_crypto_keystore::Database::open(&path, &key).await.unwrap();
180        let tx = keystore.new_transaction().await.unwrap();
181        let alice = ProteusCentral::try_new(&tx).await.unwrap();
182
183        // with no holes to fill, ids are simply assigned in ascending order from 1
184        let claimed = claim(&alice, &tx, ID_TEST_RANGE.count()).await;
185        assert_eq!(claimed, ID_TEST_RANGE.collect::<Vec<_>>());
186
187        let mut rng = rand::thread_rng();
188
189        // punch some holes; the next claims must fill exactly those
190        let gap_ids = pick_gap_ids(&mut rng, GAP_AMOUNT);
191        for gap_id in &gap_ids {
192            tx.remove::<ProteusPrekey>(gap_id).await.unwrap();
193        }
194        let claimed = claim(&alice, &tx, GAP_AMOUNT).await;
195        assert_eq!(
196            ascending(&claimed),
197            gap_ids,
198            "every deleted id must be reassigned exactly once"
199        );
200
201        // punch holes again, but this time keep claiming past them, to check both that holes take
202        // priority over fresh ids and that the fresh ids resume above the filled id space
203        let gap_ids = pick_gap_ids(&mut rng, GAP_AMOUNT);
204        for gap_id in &gap_ids {
205            tx.remove::<ProteusPrekey>(gap_id).await.unwrap();
206        }
207        let claimed = claim(&alice, &tx, GAP_AMOUNT + EXTRA_AMOUNT as usize).await;
208        let (filled, extended) = claimed.split_at(GAP_AMOUNT);
209        assert_eq!(
210            ascending(filled),
211            gap_ids,
212            "holes must all be filled before the id space is extended"
213        );
214        let high_water_mark = *ID_TEST_RANGE.end();
215        assert_eq!(
216            extended,
217            (high_water_mark + 1..=high_water_mark + EXTRA_AMOUNT).collect::<Vec<_>>(),
218            "once no holes remain, ids extend the id space in ascending order"
219        );
220
221        tx.commit().await.unwrap();
222        #[cfg(not(target_os = "unknown"))]
223        drop(db_file);
224    }
225
226    /// The last resort prekey is stored in the same table as ordinary prekeys, at `u16::MAX`.
227    /// Free-id selection must not treat it as though the id space were exhausted.
228    #[macro_rules_attribute::apply(smol_macros::test)]
229    async fn last_resort_prekey_does_not_exhaust_id_space() {
230        #[cfg(not(target_os = "unknown"))]
231        let (path, db_file) = tmp_db_file();
232        #[cfg(target_os = "unknown")]
233        let (path, _) = tmp_db_file();
234
235        let key = DatabaseKey::generate();
236        let keystore = core_crypto_keystore::Database::open(&path, &key).await.unwrap();
237
238        // in this transaction the last resort prekey is only cached, never persisted
239        let tx = keystore.new_transaction().await.unwrap();
240        let alice = ProteusCentral::try_new(&tx).await.unwrap();
241        alice.last_resort_prekey(&tx).await.unwrap();
242        let (pk_id, _) = alice.new_prekey_auto(&tx).await.unwrap();
243        assert_eq!(
244            pk_id, 1,
245            "an uncommitted last resort prekey must not block auto prekeys"
246        );
247        tx.commit().await.unwrap();
248
249        // now the last resort prekey is persisted, so it's the database query which has to skip it
250        let tx = keystore.new_transaction().await.unwrap();
251        let (pk_id, _) = alice.new_prekey_auto(&tx).await.unwrap();
252        assert_eq!(pk_id, 2, "a persisted last resort prekey must not block auto prekeys");
253        tx.commit().await.unwrap();
254
255        #[cfg(not(target_os = "unknown"))]
256        drop(db_file);
257    }
258}
core_crypto/proteus/prekey.rs

core_crypto/proteus/
prekey.rs
