libsignal_protocol/storage/
inmem.rs

1//
2// Copyright 2020-2022 Signal Messenger, LLC.
3// SPDX-License-Identifier: AGPL-3.0-only
4//
5
6//! Implementations for stores defined in [super::traits].
7//!
8//! These implementations are purely in-memory, and therefore most likely useful for testing.
9
10use std::borrow::Cow;
11use std::collections::HashMap;
12
13use async_trait::async_trait;
14use uuid::Uuid;
15
16use crate::storage::traits::{self, IdentityChange};
17use crate::{
18    CiphertextMessageType, IdentityKey, IdentityKeyPair, KyberPreKeyId, KyberPreKeyRecord,
19    PreKeyId, PreKeyRecord, ProtocolAddress, PublicKey, Result, SenderKeyRecord, SessionNotFound,
20    SessionRecord, SignalProtocolError, SignedPreKeyId, SignedPreKeyRecord,
21};
22
23/// Reference implementation of [traits::IdentityKeyStore].
24#[derive(Clone)]
25pub struct InMemIdentityKeyStore {
26    key_pair: IdentityKeyPair,
27    registration_id: u32,
28    known_keys: HashMap<ProtocolAddress, IdentityKey>,
29}
30
31impl InMemIdentityKeyStore {
32    /// Create a new instance.
33    ///
34    /// `key_pair` corresponds to [traits::IdentityKeyStore::get_identity_key_pair], and
35    /// `registration_id` corresponds to [traits::IdentityKeyStore::get_local_registration_id].
36    pub fn new(key_pair: IdentityKeyPair, registration_id: u32) -> Self {
37        Self {
38            key_pair,
39            registration_id,
40            known_keys: HashMap::new(),
41        }
42    }
43
44    /// Clear the mapping of known keys.
45    pub fn reset(&mut self) {
46        self.known_keys.clear();
47    }
48}
49
50#[async_trait(?Send)]
51impl traits::IdentityKeyStore for InMemIdentityKeyStore {
52    async fn get_identity_key_pair(&self) -> Result<IdentityKeyPair> {
53        Ok(self.key_pair)
54    }
55
56    async fn get_local_registration_id(&self) -> Result<u32> {
57        Ok(self.registration_id)
58    }
59
60    async fn save_identity(
61        &mut self,
62        address: &ProtocolAddress,
63        identity: &IdentityKey,
64    ) -> Result<IdentityChange> {
65        match self.known_keys.get(address) {
66            None => {
67                self.known_keys.insert(address.clone(), *identity);
68                Ok(IdentityChange::NewOrUnchanged)
69            }
70            Some(k) if k == identity => Ok(IdentityChange::NewOrUnchanged),
71            Some(_k) => {
72                self.known_keys.insert(address.clone(), *identity);
73                Ok(IdentityChange::ReplacedExisting)
74            }
75        }
76    }
77
78    async fn is_trusted_identity(
79        &self,
80        address: &ProtocolAddress,
81        identity: &IdentityKey,
82        _direction: traits::Direction,
83    ) -> Result<bool> {
84        match self.known_keys.get(address) {
85            None => {
86                Ok(true) // first use
87            }
88            Some(k) => Ok(k == identity),
89        }
90    }
91
92    async fn get_identity(&self, address: &ProtocolAddress) -> Result<Option<IdentityKey>> {
93        match self.known_keys.get(address) {
94            None => Ok(None),
95            Some(k) => Ok(Some(k.to_owned())),
96        }
97    }
98}
99
100/// Reference implementation of [traits::PreKeyStore].
101#[derive(Clone)]
102pub struct InMemPreKeyStore {
103    pre_keys: HashMap<PreKeyId, PreKeyRecord>,
104}
105
106impl InMemPreKeyStore {
107    /// Create an empty pre-key store.
108    pub fn new() -> Self {
109        Self {
110            pre_keys: HashMap::new(),
111        }
112    }
113
114    /// Returns all registered pre-key ids
115    pub fn all_pre_key_ids(&self) -> impl Iterator<Item = &PreKeyId> {
116        self.pre_keys.keys()
117    }
118}
119
120impl Default for InMemPreKeyStore {
121    fn default() -> Self {
122        Self::new()
123    }
124}
125
126#[async_trait(?Send)]
127impl traits::PreKeyStore for InMemPreKeyStore {
128    async fn get_pre_key(&self, id: PreKeyId) -> Result<PreKeyRecord> {
129        Ok(self
130            .pre_keys
131            .get(&id)
132            .ok_or(SignalProtocolError::InvalidPreKeyId)?
133            .clone())
134    }
135
136    async fn save_pre_key(&mut self, id: PreKeyId, record: &PreKeyRecord) -> Result<()> {
137        // This overwrites old values, which matches Java behavior, but is it correct?
138        self.pre_keys.insert(id, record.to_owned());
139        Ok(())
140    }
141
142    async fn remove_pre_key(&mut self, id: PreKeyId) -> Result<()> {
143        // If id does not exist this silently does nothing
144        self.pre_keys.remove(&id);
145        Ok(())
146    }
147}
148
149/// Reference implementation of [traits::SignedPreKeyStore].
150#[derive(Clone)]
151pub struct InMemSignedPreKeyStore {
152    signed_pre_keys: HashMap<SignedPreKeyId, SignedPreKeyRecord>,
153}
154
155impl InMemSignedPreKeyStore {
156    /// Create an empty signed pre-key store.
157    pub fn new() -> Self {
158        Self {
159            signed_pre_keys: HashMap::new(),
160        }
161    }
162
163    /// Returns all registered signed pre-key ids
164    pub fn all_signed_pre_key_ids(&self) -> impl Iterator<Item = &SignedPreKeyId> {
165        self.signed_pre_keys.keys()
166    }
167}
168
169impl Default for InMemSignedPreKeyStore {
170    fn default() -> Self {
171        Self::new()
172    }
173}
174
175#[async_trait(?Send)]
176impl traits::SignedPreKeyStore for InMemSignedPreKeyStore {
177    async fn get_signed_pre_key(&self, id: SignedPreKeyId) -> Result<SignedPreKeyRecord> {
178        Ok(self
179            .signed_pre_keys
180            .get(&id)
181            .ok_or(SignalProtocolError::InvalidSignedPreKeyId)?
182            .clone())
183    }
184
185    async fn save_signed_pre_key(
186        &mut self,
187        id: SignedPreKeyId,
188        record: &SignedPreKeyRecord,
189    ) -> Result<()> {
190        // This overwrites old values, which matches Java behavior, but is it correct?
191        self.signed_pre_keys.insert(id, record.to_owned());
192        Ok(())
193    }
194}
195
196/// Basic implementation of [traits::KyberPreKeyStore].
197///
198/// Note that this implementation does not clear any keys upon use! This is correct for last-resort
199/// keys, but a real client would normally have a set of one-time keys to use first.
200#[derive(Clone)]
201pub struct InMemKyberPreKeyStore {
202    kyber_pre_keys: HashMap<KyberPreKeyId, KyberPreKeyRecord>,
203    base_keys_seen: HashMap<(KyberPreKeyId, SignedPreKeyId), Vec<PublicKey>>,
204}
205
206impl InMemKyberPreKeyStore {
207    /// Create an empty kyber pre-key store.
208    pub fn new() -> Self {
209        Self {
210            kyber_pre_keys: HashMap::new(),
211            base_keys_seen: HashMap::new(),
212        }
213    }
214
215    /// Returns all registered Kyber pre-key ids
216    pub fn all_kyber_pre_key_ids(&self) -> impl Iterator<Item = &KyberPreKeyId> {
217        self.kyber_pre_keys.keys()
218    }
219}
220
221impl Default for InMemKyberPreKeyStore {
222    fn default() -> Self {
223        Self::new()
224    }
225}
226
227#[async_trait(?Send)]
228impl traits::KyberPreKeyStore for InMemKyberPreKeyStore {
229    async fn get_kyber_pre_key(&self, kyber_prekey_id: KyberPreKeyId) -> Result<KyberPreKeyRecord> {
230        Ok(self
231            .kyber_pre_keys
232            .get(&kyber_prekey_id)
233            .ok_or(SignalProtocolError::InvalidKyberPreKeyId)?
234            .clone())
235    }
236
237    async fn save_kyber_pre_key(
238        &mut self,
239        kyber_prekey_id: KyberPreKeyId,
240        record: &KyberPreKeyRecord,
241    ) -> Result<()> {
242        self.kyber_pre_keys
243            .insert(kyber_prekey_id, record.to_owned());
244        Ok(())
245    }
246
247    async fn mark_kyber_pre_key_used(
248        &mut self,
249        kyber_prekey_id: KyberPreKeyId,
250        ec_prekey_id: SignedPreKeyId,
251        base_key: &PublicKey,
252    ) -> Result<()> {
253        let base_keys_seen = self
254            .base_keys_seen
255            .entry((kyber_prekey_id, ec_prekey_id))
256            .or_default();
257        if base_keys_seen.contains(base_key) {
258            return Err(SignalProtocolError::InvalidMessage(
259                CiphertextMessageType::PreKey,
260                "reused base key".to_owned(),
261            ));
262        }
263        base_keys_seen.push(*base_key);
264        Ok(())
265    }
266}
267
268/// Reference implementation of [traits::SessionStore].
269#[derive(Clone)]
270pub struct InMemSessionStore {
271    sessions: HashMap<ProtocolAddress, SessionRecord>,
272}
273
274impl InMemSessionStore {
275    /// Create an empty session store.
276    pub fn new() -> Self {
277        Self {
278            sessions: HashMap::new(),
279        }
280    }
281
282    /// Bulk version of [`SessionStore::load_session`].
283    ///
284    /// Useful for [crate::sealed_sender_multi_recipient_encrypt].
285    ///
286    /// [`SessionStore::load_session`]: crate::SessionStore::load_session
287    pub fn load_existing_sessions(
288        &self,
289        addresses: &[&ProtocolAddress],
290    ) -> Result<Vec<&SessionRecord>> {
291        addresses
292            .iter()
293            .map(|&address| {
294                self.sessions.get(address).ok_or_else(|| {
295                    SignalProtocolError::SessionNotFound(SessionNotFound::new(
296                        address.clone(),
297                        "load_existing_sessions",
298                    ))
299                })
300            })
301            .collect()
302    }
303}
304
305impl Default for InMemSessionStore {
306    fn default() -> Self {
307        Self::new()
308    }
309}
310
311#[async_trait(?Send)]
312impl traits::SessionStore for InMemSessionStore {
313    async fn load_session(&self, address: &ProtocolAddress) -> Result<Option<SessionRecord>> {
314        match self.sessions.get(address) {
315            None => Ok(None),
316            Some(s) => Ok(Some(s.clone())),
317        }
318    }
319
320    async fn store_session(
321        &mut self,
322        address: &ProtocolAddress,
323        record: &SessionRecord,
324    ) -> Result<()> {
325        self.sessions.insert(address.clone(), record.clone());
326        Ok(())
327    }
328}
329
330/// Reference implementation of [traits::SenderKeyStore].
331#[derive(Clone)]
332pub struct InMemSenderKeyStore {
333    // We use Cow keys in order to store owned values but compare to referenced ones.
334    // See https://users.rust-lang.org/t/hashmap-with-tuple-keys/12711/6.
335    keys: HashMap<(Cow<'static, ProtocolAddress>, Uuid), SenderKeyRecord>,
336}
337
338impl InMemSenderKeyStore {
339    /// Create an empty sender key store.
340    pub fn new() -> Self {
341        Self {
342            keys: HashMap::new(),
343        }
344    }
345}
346
347impl Default for InMemSenderKeyStore {
348    fn default() -> Self {
349        Self::new()
350    }
351}
352
353#[async_trait(?Send)]
354impl traits::SenderKeyStore for InMemSenderKeyStore {
355    async fn store_sender_key(
356        &mut self,
357        sender: &ProtocolAddress,
358        distribution_id: Uuid,
359        record: &SenderKeyRecord,
360    ) -> Result<()> {
361        self.keys.insert(
362            (Cow::Owned(sender.clone()), distribution_id),
363            record.clone(),
364        );
365        Ok(())
366    }
367
368    async fn load_sender_key(
369        &mut self,
370        sender: &ProtocolAddress,
371        distribution_id: Uuid,
372    ) -> Result<Option<SenderKeyRecord>> {
373        Ok(self
374            .keys
375            .get(&(Cow::Borrowed(sender), distribution_id))
376            .cloned())
377    }
378}
379
380/// Reference implementation of [traits::ProtocolStore].
381#[allow(missing_docs)]
382#[derive(Clone)]
383pub struct InMemSignalProtocolStore {
384    pub session_store: InMemSessionStore,
385    pub pre_key_store: InMemPreKeyStore,
386    pub signed_pre_key_store: InMemSignedPreKeyStore,
387    pub kyber_pre_key_store: InMemKyberPreKeyStore,
388    pub identity_store: InMemIdentityKeyStore,
389    pub sender_key_store: InMemSenderKeyStore,
390}
391
392impl InMemSignalProtocolStore {
393    /// Create an object with the minimal implementation of [traits::ProtocolStore], representing
394    /// the given identity `key_pair` along with the separate randomly chosen `registration_id`.
395    pub fn new(key_pair: IdentityKeyPair, registration_id: u32) -> Result<Self> {
396        Ok(Self {
397            session_store: InMemSessionStore::new(),
398            pre_key_store: InMemPreKeyStore::new(),
399            signed_pre_key_store: InMemSignedPreKeyStore::new(),
400            kyber_pre_key_store: InMemKyberPreKeyStore::new(),
401            identity_store: InMemIdentityKeyStore::new(key_pair, registration_id),
402            sender_key_store: InMemSenderKeyStore::new(),
403        })
404    }
405
406    /// Returns all registered pre-key ids
407    pub fn all_pre_key_ids(&self) -> impl Iterator<Item = &PreKeyId> {
408        self.pre_key_store.all_pre_key_ids()
409    }
410
411    /// Returns all registered signed pre-key ids
412    pub fn all_signed_pre_key_ids(&self) -> impl Iterator<Item = &SignedPreKeyId> {
413        self.signed_pre_key_store.all_signed_pre_key_ids()
414    }
415
416    /// Returns all registered Kyber pre-key ids
417    pub fn all_kyber_pre_key_ids(&self) -> impl Iterator<Item = &KyberPreKeyId> {
418        self.kyber_pre_key_store.all_kyber_pre_key_ids()
419    }
420}
421
422#[async_trait(?Send)]
423impl traits::IdentityKeyStore for InMemSignalProtocolStore {
424    async fn get_identity_key_pair(&self) -> Result<IdentityKeyPair> {
425        self.identity_store.get_identity_key_pair().await
426    }
427
428    async fn get_local_registration_id(&self) -> Result<u32> {
429        self.identity_store.get_local_registration_id().await
430    }
431
432    async fn save_identity(
433        &mut self,
434        address: &ProtocolAddress,
435        identity: &IdentityKey,
436    ) -> Result<IdentityChange> {
437        self.identity_store.save_identity(address, identity).await
438    }
439
440    async fn is_trusted_identity(
441        &self,
442        address: &ProtocolAddress,
443        identity: &IdentityKey,
444        direction: traits::Direction,
445    ) -> Result<bool> {
446        self.identity_store
447            .is_trusted_identity(address, identity, direction)
448            .await
449    }
450
451    async fn get_identity(&self, address: &ProtocolAddress) -> Result<Option<IdentityKey>> {
452        self.identity_store.get_identity(address).await
453    }
454}
455
456#[async_trait(?Send)]
457impl traits::PreKeyStore for InMemSignalProtocolStore {
458    async fn get_pre_key(&self, id: PreKeyId) -> Result<PreKeyRecord> {
459        self.pre_key_store.get_pre_key(id).await
460    }
461
462    async fn save_pre_key(&mut self, id: PreKeyId, record: &PreKeyRecord) -> Result<()> {
463        self.pre_key_store.save_pre_key(id, record).await
464    }
465
466    async fn remove_pre_key(&mut self, id: PreKeyId) -> Result<()> {
467        self.pre_key_store.remove_pre_key(id).await
468    }
469}
470
471#[async_trait(?Send)]
472impl traits::SignedPreKeyStore for InMemSignalProtocolStore {
473    async fn get_signed_pre_key(&self, id: SignedPreKeyId) -> Result<SignedPreKeyRecord> {
474        self.signed_pre_key_store.get_signed_pre_key(id).await
475    }
476
477    async fn save_signed_pre_key(
478        &mut self,
479        id: SignedPreKeyId,
480        record: &SignedPreKeyRecord,
481    ) -> Result<()> {
482        self.signed_pre_key_store
483            .save_signed_pre_key(id, record)
484            .await
485    }
486}
487
488#[async_trait(?Send)]
489impl traits::KyberPreKeyStore for InMemSignalProtocolStore {
490    async fn get_kyber_pre_key(&self, kyber_prekey_id: KyberPreKeyId) -> Result<KyberPreKeyRecord> {
491        self.kyber_pre_key_store
492            .get_kyber_pre_key(kyber_prekey_id)
493            .await
494    }
495
496    async fn save_kyber_pre_key(
497        &mut self,
498        kyber_prekey_id: KyberPreKeyId,
499        record: &KyberPreKeyRecord,
500    ) -> Result<()> {
501        self.kyber_pre_key_store
502            .save_kyber_pre_key(kyber_prekey_id, record)
503            .await
504    }
505
506    async fn mark_kyber_pre_key_used(
507        &mut self,
508        kyber_prekey_id: KyberPreKeyId,
509        ec_prekey_id: SignedPreKeyId,
510        base_key: &PublicKey,
511    ) -> Result<()> {
512        self.kyber_pre_key_store
513            .mark_kyber_pre_key_used(kyber_prekey_id, ec_prekey_id, base_key)
514            .await
515    }
516}
517
518#[async_trait(?Send)]
519impl traits::SessionStore for InMemSignalProtocolStore {
520    async fn load_session(&self, address: &ProtocolAddress) -> Result<Option<SessionRecord>> {
521        self.session_store.load_session(address).await
522    }
523
524    async fn store_session(
525        &mut self,
526        address: &ProtocolAddress,
527        record: &SessionRecord,
528    ) -> Result<()> {
529        self.session_store.store_session(address, record).await
530    }
531}
532
533#[async_trait(?Send)]
534impl traits::SenderKeyStore for InMemSignalProtocolStore {
535    async fn store_sender_key(
536        &mut self,
537        sender: &ProtocolAddress,
538        distribution_id: Uuid,
539        record: &SenderKeyRecord,
540    ) -> Result<()> {
541        self.sender_key_store
542            .store_sender_key(sender, distribution_id, record)
543            .await
544    }
545
546    async fn load_sender_key(
547        &mut self,
548        sender: &ProtocolAddress,
549        distribution_id: Uuid,
550    ) -> Result<Option<SenderKeyRecord>> {
551        self.sender_key_store
552            .load_sender_key(sender, distribution_id)
553            .await
554    }
555}
556
557impl traits::ProtocolStore for InMemSignalProtocolStore {}
libsignal_protocol/storage/inmem.rs

libsignal_protocol/storage/
inmem.rs
