Add storage traits and in-mem implementations

src/lib.rs

@@ -8,6 +8,10 @@ mod proto;
 mod protocol;
 mod ratchet;
 mod state;
+mod storage;
 
 pub use error::SignalProtocolError;
+pub use address::ProtocolAddress;
 pub use identity_key::{IdentityKey, IdentityKeyPair};
+
+pub use storage::{IdentityKeyStore, PreKeyStore, SignedPreKeyStore, SessionStore};

src/storage.rs

@@ -0,0 +1,9 @@
+mod traits;
+mod inmem;
+
+pub use traits::{IdentityKeyStore,
+                 PreKeyStore,
+                 SignedPreKeyStore,
+                 SessionStore};
+
+pub use inmem::InMemIdentityKeyStore;

src/storage/inmem.rs

@@ -0,0 +1,188 @@
+use crate::error::{SignalProtocolError, Result};
+use crate::storage::traits;
+use crate::state::{PreKeyRecord, PreKeyId, SignedPreKeyRecord, SignedPreKeyId, SessionRecord};
+use crate::{IdentityKeyPair, IdentityKey, ProtocolAddress};
+
+use std::collections::HashMap;
+
+pub struct InMemIdentityKeyStore {
+    key_pair: IdentityKeyPair,
+    id: u32,
+    known_keys: HashMap<ProtocolAddress, IdentityKey>,
+}
+
+impl InMemIdentityKeyStore {
+    fn new(key_pair: IdentityKeyPair, id: u32) -> Self {
+        Self {
+            key_pair,
+            id,
+            known_keys: HashMap::new()
+        }
+    }
+}
+
+impl traits::IdentityKeyStore for InMemIdentityKeyStore {
+    fn get_identity_key_pair(&self) -> Result<IdentityKeyPair> {
+        Ok(self.key_pair.clone())
+    }
+
+    fn get_local_registration_id(&self) -> Result<u32> {
+        Ok(self.id)
+    }
+
+    fn save_identity(&mut self, address: &ProtocolAddress, identity: &IdentityKey) -> Result<bool> {
+        match self.known_keys.get(address) {
+            None => {
+                self.known_keys.insert(address.clone(), identity.clone());
+                Ok(false) // new key
+            }
+            Some(k) if k == identity => {
+                Ok(false) // same key
+            }
+            Some(k) => {
+                self.known_keys.insert(address.clone(), identity.clone());
+                Ok(true) // overwrite
+            }
+        }
+    }
+
+    fn is_trusted_identity(&self, address: &ProtocolAddress, identity: &IdentityKey, direction: traits::Direction) -> Result<bool> {
+        match self.known_keys.get(address) {
+            None => {
+                Ok(true) // first use
+            }
+            Some(k) => {
+                Ok(k == identity)
+            }
+        }
+    }
+
+    fn get_identity(&self, address: &ProtocolAddress) -> Result<Option<IdentityKey>> {
+        match self.known_keys.get(address) {
+            None => Ok(None),
+            Some(k) => Ok(Some(k.to_owned()))
+        }
+    }
+}
+
+pub struct InMemPreKeyStore {
+    pre_keys: HashMap<PreKeyId, PreKeyRecord>,
+}
+
+impl InMemPreKeyStore {
+    fn new() -> Self {
+        Self { pre_keys: HashMap::new() }
+    }
+}
+
+impl traits::PreKeyStore for InMemPreKeyStore {
+    fn get_pre_key(&self, id: PreKeyId) -> Result<PreKeyRecord> {
+        Ok(self.pre_keys.get(&id).ok_or(SignalProtocolError::InvalidPreKeyId)?.clone())
+    }
+
+    fn save_pre_key(&mut self, id: PreKeyId, record: &PreKeyRecord) -> Result<()> {
+        // This overwrites old values, which matches Java behavior, but is it correct?
+        self.pre_keys.insert(id, record.to_owned());
+        Ok(())
+    }
+
+    fn has_pre_key(&self, id: PreKeyId) -> Result<bool> {
+        Ok(self.pre_keys.get(&id).is_some())
+    }
+
+    fn remove_pre_key(&mut self, id: PreKeyId) -> Result<()> {
+        // If id does not exist this silently does nothing
+        self.pre_keys.remove(&id);
+        Ok(())
+    }
+}
+
+pub struct InMemSignedPreKeyStore {
+    signed_pre_keys: HashMap<SignedPreKeyId, SignedPreKeyRecord>,
+}
+
+impl InMemSignedPreKeyStore {
+    fn new() -> Self {
+        Self { signed_pre_keys: HashMap::new() }
+    }
+}
+
+impl traits::SignedPreKeyStore for InMemSignedPreKeyStore {
+    fn get_signed_pre_key(&self, id: SignedPreKeyId) -> Result<SignedPreKeyRecord> {
+        Ok(self.signed_pre_keys.get(&id).ok_or(SignalProtocolError::InvalidSignedPreKeyId)?.clone())
+    }
+
+    fn get_all_signed_prekeys(&self) -> Result<Vec<SignedPreKeyRecord>> {
+        let mut result = Vec::with_capacity(self.signed_pre_keys.len());
+        for v in self.signed_pre_keys.values() {
+            result.push(v.clone());
+        }
+        Ok(result)
+    }
+
+    fn save_signed_pre_key(&mut self, id: SignedPreKeyId, record: &SignedPreKeyRecord) -> Result<()> {
+        // This overwrites old values, which matches Java behavior, but is it correct?
+        self.signed_pre_keys.insert(id, record.to_owned());
+        Ok(())
+    }
+
+    fn has_signed_pre_key(&self, id: SignedPreKeyId) -> Result<bool> {
+        Ok(self.signed_pre_keys.get(&id).is_some())
+    }
+
+    fn remove_pre_key(&mut self, id: SignedPreKeyId) -> Result<()> {
+        // If id does not exist this silently does nothing
+        self.signed_pre_keys.remove(&id);
+        Ok(())
+    }
+}
+
+pub struct InMemSessionStore {
+    sessions: HashMap<ProtocolAddress, SessionRecord>,
+}
+
+impl InMemSessionStore {
+    fn new() -> Self {
+        Self { sessions: HashMap::new() }
+    }
+}
+
+impl traits::SessionStore for InMemSessionStore {
+    fn load_session(&self, address: &ProtocolAddress) -> Result<Option<SessionRecord>> {
+        match self.sessions.get(address) {
+            None => Ok(None),
+            Some(s) => Ok(Some(s.clone()))
+        }
+    }
+
+    fn get_sub_device_sessions(&self, name: &str) -> Result<Vec<u32>> {
+        let mut result = vec![];
+
+        for address in self.sessions.keys() {
+            if address.name() == name && address.device_id() != 1 {
+                result.push(address.device_id());
+            }
+        }
+
+        Ok(result)
+    }
+
+    fn store_session(&mut self, address: &ProtocolAddress, record: &SessionRecord) -> Result<()> {
+        self.sessions.insert(address.clone(), record.clone());
+        Ok(())
+    }
+
+    fn contains_session(&self, address: &ProtocolAddress) -> Result<bool> {
+        Ok(self.sessions.get(address).is_some())
+    }
+
+    fn delete_session(&mut self, address: &ProtocolAddress) -> Result<()> {
+        self.sessions.remove(address);
+        Ok(())
+    }
+
+    fn delete_all_sessions(&mut self, name: &str) -> Result<()> {
+        self.sessions.retain(|a,_| a.name() != name);
+        Ok(())
+    }
+}

src/storage/traits.rs

@@ -0,0 +1,58 @@
+
+use crate::{IdentityKey, IdentityKeyPair, ProtocolAddress};
+use crate::state::{PreKeyRecord, PreKeyId, SignedPreKeyRecord, SignedPreKeyId, SessionRecord};
+use crate::error::Result;
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum Direction {
+    Sending,
+    Receiving
+}
+
+pub trait IdentityKeyStore {
+    fn get_identity_key_pair(&self) -> Result<IdentityKeyPair>;
+
+    fn get_local_registration_id(&self) -> Result<u32>;
+
+    fn save_identity(&mut self, address: &ProtocolAddress, identity: &IdentityKey) -> Result<bool>;
+
+    fn is_trusted_identity(&self, address: &ProtocolAddress, identity: &IdentityKey, direction: Direction) -> Result<bool>;
+
+    fn get_identity(&self, address: &ProtocolAddress) -> Result<Option<IdentityKey>>;
+}
+
+pub trait PreKeyStore {
+    fn get_pre_key(&self, prekey_id: PreKeyId) -> Result<PreKeyRecord>;
+
+    fn save_pre_key(&mut self, prekey_id: PreKeyId, record: &PreKeyRecord) -> Result<()>;
+
+    fn has_pre_key(&self, prekey_id: PreKeyId) -> Result<bool>;
+
+    fn remove_pre_key(&mut self, prekey_id: PreKeyId) -> Result<()>;
+}
+
+pub trait SignedPreKeyStore {
+    fn get_signed_pre_key(&self, signed_prekey_id: SignedPreKeyId) -> Result<SignedPreKeyRecord>;
+
+    fn get_all_signed_prekeys(&self) -> Result<Vec<SignedPreKeyRecord>>;
+
+    fn save_signed_pre_key(&mut self, signed_prekey_id: SignedPreKeyId, record: &SignedPreKeyRecord) -> Result<()>;
+
+    fn has_signed_pre_key(&self, signed_prekey_id: SignedPreKeyId) -> Result<bool>;
+
+    fn remove_pre_key(&mut self, signed_prekey_id: SignedPreKeyId) -> Result<()>;
+}
+
+pub trait SessionStore {
+    fn load_session(&self, address: &ProtocolAddress) -> Result<Option<SessionRecord>>;
+
+    fn get_sub_device_sessions(&self, name: &str) -> Result<Vec<u32>>;
+
+    fn store_session(&mut self, address: &ProtocolAddress, record: &SessionRecord) -> Result<()>;
+
+    fn contains_session(&self, address: &ProtocolAddress) -> Result<bool>;
+
+    fn delete_session(&mut self, address: &ProtocolAddress) -> Result<()>;
+
+    fn delete_all_sessions(&mut self, name: &str) -> Result<()>;
+}