refactor(storage): Better message storage

This goes with the changes made in <https://github.com/whisperfish/presage/pull/67>. This will however break previous messages stored.

Additionally, a button in the primary menu was added to clear message
storage.

Cargo.lock

@@ -2108,7 +2108,7 @@ checksum = "eb9f9e6e233e5c4a35559a617bf40a4ec447db2e84c20b55a6f83167b7e57872"
 [[package]]
 name = "presage"
 version = "0.3.0"
-source = "git+https://github.com/Schmiddiii/presage?rev=0abbb46a92ec8f186ca43339cb109ad9132a02ed#0abbb46a92ec8f186ca43339cb109ad9132a02ed"
+source = "git+https://github.com/Schmiddiii/presage?rev=2ccbdb3fb14623625be9b51c76a8a0993203d535#2ccbdb3fb14623625be9b51c76a8a0993203d535"
 dependencies = [
  "async-trait",
  "base64 0.12.3",

Cargo.toml

@@ -18,7 +18,7 @@ futures = "0.3.24"
 tokio = { version = "1.20.1", features = ["rt", "macros", "rt-multi-thread", "time"] }
 async-recursion = "1.0.0"
 
-presage = { git = "https://github.com/Schmiddiii/presage", rev = "0abbb46a92ec8f186ca43339cb109ad9132a02ed" }
+presage = { git = "https://github.com/Schmiddiii/presage", rev = "2ccbdb3fb14623625be9b51c76a8a0993203d535" }
 libsignal-service = { git = "https://github.com/whisperfish/libsignal-service-rs", rev = "f9b54478c5413252e810089fd9648b6351bcc3bf" }
 
 async-trait = "0.1.57"

data/resources/ui/window.ui

@@ -125,6 +125,10 @@
         <attribute name="label" translatable="yes">Unlink</attribute>
         <attribute name="action">win.unlink</attribute>
       </item>
+      <item>
+        <attribute name="label" translatable="yes">Clear messages</attribute>
+        <attribute name="action">win.clear-messages</attribute>
+      </item>
       <item>
         <attribute name="label" translatable="yes">Keybindings</attribute>
         <attribute name="action">win.show-help-overlay</attribute>

src/backend/channel.rs

@@ -9,7 +9,7 @@ use gio::subclass::prelude::ObjectSubclassIsExt;
 use libsignal_service::{groups_v2::Group, prelude::Uuid, proto::DataMessage};
 use presage::{
     prelude::{GroupContextV2, GroupMasterKey, ServiceAddress},
-    MessageIdentity,
+    MessageIdentity, Thread,
 };
 
 use super::{Contact, Manager, Message};
@@ -48,20 +48,9 @@ impl Channel {
             if let Ok(group) = group {
                 return Self::from_group(group, group_context_v2, manager).await;
             } else {
-                if let Some(ref uuid) = contact.address().and_then(|a| a.uuid) {
-                    s.imp().unloaded_messages.swap(&RefCell::new(
-                        manager.messages_by_contact(uuid).unwrap_or_default(),
-                    ));
-                }
                 s.imp().contact.swap(&RefCell::new(Some(contact)));
             }
         } else {
-            if let Some(ref uuid) = contact.address().and_then(|a| a.uuid) {
-                s.imp().unloaded_messages.swap(&RefCell::new(
-                    manager.messages_by_contact(uuid).unwrap_or_default(),
-                ));
-            }
-
             s.imp().contact.swap(&RefCell::new(Some(contact)));
         }
         s
@@ -74,79 +63,65 @@ impl Channel {
     ) -> Self {
         let s: Self = Object::new(&[("manager", manager)]).expect("Failed to create `Channel`");
         s.imp().group.swap(&RefCell::new(Some(group)));
-        s.imp().unloaded_messages.swap(&RefCell::new(
-            manager
-                .messages_by_group(group_context_v2)
-                .unwrap_or_default(),
-        ));
         s.imp()
             .group_context
             .swap(&RefCell::new(Some(group_context_v2.clone())));
-        s.imp().unloaded_messages.swap(&RefCell::new(
-            manager
-                .messages_by_group(group_context_v2)
-                .unwrap_or_default(),
-        ));
         s
     }
 
+    fn thread(&self) -> Option<Thread> {
+        if let Some(key) = self
+            .group_context()
+            .and_then(|c| c.master_key)
+            .and_then(|k| <[u8; 32]>::try_from(k).ok())
+        {
+            Some(Thread::Group(key))
+        } else {
+            Some(Thread::Contact(self.uuid()?))
+        }
+    }
+
     #[async_recursion::async_recursion(?Send)]
     pub async fn load_last(&self, number: usize) -> Vec<Message> {
-        let to_load = {
-            let mut unloaded_msgs = self.imp().unloaded_messages.borrow_mut();
+        if let Some(thread) = self.thread() {
+            let mut results = vec![];
+            let manager = self.property::<Manager>("manager");
+            let messages_unborrowed = &self.imp().messages;
+            let first_timestamp = {
+                let msgs = messages_unborrowed.borrow();
+                msgs.get(0).and_then(|m| m.timestamp()).map(|t| t - 1)
+            };
             crate::trace!(
-                "Loading last messages for channel: {} ({:?}). Got {} total unloaded messages.",
+                "Loading {} last messages for channel: {} (Thread: {:?}). Needed earlier then {:?}",
+                number,
                 self.property::<String>("title"),
-                self.uuid(),
-                unloaded_msgs.len()
+                thread,
+                first_timestamp
             );
-            if unloaded_msgs.len() <= number {
-                unloaded_msgs.drain(..).collect::<Vec<_>>()
-            } else {
-                unloaded_msgs.drain(..number).collect::<Vec<_>>()
+            // TODO: Skip already loaded messages?
+            let iter = manager.messages_by_thread(&thread, first_timestamp);
+            if iter.is_err() {
+                return vec![];
             }
-        };
-
-        if to_load.is_empty() {
-            return vec![];
-        }
-
-        let mut empty = 0;
-        let mut results = vec![];
-        let manager = self.property::<Manager>("manager");
-        for msg_id in to_load {
-            if let Ok(Some(msg)) = manager.message_by_id(&msg_id).await {
+            for content in iter.unwrap() {
+                let msg = Message::from_content(content, &manager).await;
                 if !msg.is_empty() {
                     {
-                        let mut messages = self.imp().messages.borrow_mut();
+                        let mut messages = messages_unborrowed.borrow_mut();
                         messages.insert(0, msg.clone());
                     }
                     results.insert(0, msg.clone());
                     self.notify("last-message");
-                } else {
-                    log::trace!(
-                        "Channel {} got empty message, skipping.",
-                        self.property::<String>("title")
-                    );
-                    empty += 1
+                    let _ = self.do_new_message(&msg).await;
+                }
+                if results.len() == number {
+                    break;
                 }
-                // TODO: Error?
-                let _ = self.do_new_message(&msg).await;
-            } else {
-                log::warn!(
-                    "Channel {} got message that could not be found, skipping.",
-                    self.property::<String>("title")
-                );
             }
-        }
-
-        let mut previous = if empty != 0 {
-            self.load_last(empty).await
+            results
         } else {
             vec![]
-        };
-        previous.append(&mut results);
-        previous
+        }
     }
 
     pub(super) fn internal_hash(&self) -> u64 {
@@ -179,6 +154,7 @@ impl Channel {
                 body
             );
             if let Some(quote) = message.quote() {
+                log::trace!("Message claims to have a quote");
                 if let Ok(id) = MessageIdentity::try_from(&quote) {
                     if let Ok(Some(quoted_msg)) =
                         self.property::<Manager>("manager").message_by_id(&id).await
@@ -355,7 +331,6 @@ mod imp {
 
         pub(super) manager: RefCell<Option<Manager>>,
         pub(super) messages: RefCell<Vec<Message>>,
-        pub(super) unloaded_messages: RefCell<Vec<MessageIdentity>>,
         pub(super) pending_reactions: RefCell<HashMap<MessageIdentity, String>>,
     }
 

src/backend/manager.rs

@@ -26,7 +26,7 @@ use libsecret::{
 use crate::ApplicationError;
 use chacha20poly1305::ChaCha20Poly1305;
 use encrypted_sled::{CountingNonce, EncryptionCipher};
-use presage::{MessageIdentity, MessageStore};
+use presage::{MessageIdentity, MessageStore, Thread};
 
 const MESSAGE_BOUND: usize = 10;
 
@@ -137,6 +137,14 @@ impl Manager {
         Ok(())
     }
 
+    pub fn clear_messages(&self) -> Result<(), ApplicationError> {
+        log::trace!("Clearing messages from the manager");
+        if let Some(config_store) = self.imp().config_store.borrow().as_ref() {
+            config_store.clear_messages()?;
+        }
+        Ok(())
+    }
+
     pub fn save_message(
         &self,
         message: Content,
@@ -180,29 +188,18 @@ impl Manager {
         }
     }
 
-    pub fn messages_by_contact(
+    pub fn messages_by_thread(
         &self,
-        contact: &Uuid,
-    ) -> Result<Vec<MessageIdentity>, ApplicationError> {
-        crate::trace!("Querying message by contact: {:?}", contact);
+        thread: &Thread,
+        from: Option<u64>,
+    ) -> Result<impl Iterator<Item = Content>, ApplicationError> {
+        crate::trace!("Querying message by thread: {:?}", thread);
         if let Some(config_store) = self.imp().config_store.borrow().as_ref() {
-            Ok(config_store.messages_by_contact(contact)?)
+            Ok(config_store.messages_by_thread(thread, from)?)
         } else {
-            log::warn!("Query messages by contact without config store being set up");
-            Ok(vec![])
-        }
-    }
-
-    pub fn messages_by_group(
-        &self,
-        group: &GroupContextV2,
-    ) -> Result<Vec<MessageIdentity>, ApplicationError> {
-        crate::trace!("Querying message by group: {:?}", group.master_key);
-        if let Some(config_store) = self.imp().config_store.borrow().as_ref() {
-            Ok(config_store.messages_by_group(group)?)
-        } else {
-            log::warn!("Query messages by contact without config store being set up");
-            Ok(vec![])
+            log::error!("Query messages by contact without config store being set up");
+            // TODO: Error?
+            panic!("Query messages by contact without config store being set up");
         }
     }
 

src/gui/window.rs

@@ -24,7 +24,7 @@ impl Window {
                 &[&format!("<Control>{}", i)],
             );
         }
-        app.set_accels_for_action(&"channel-list.toggle-search", &[&"<Control>f"]);
+        app.set_accels_for_action("channel-list.toggle-search", &["<Control>f"]);
         Object::new(&[("application", app)]).expect("Failed to create Window")
     }
 
@@ -129,6 +129,35 @@ pub mod imp {
                 let settings = PreferencesWindow::new();
                 settings.show();
             });
+
+            let action_clear_messages = SimpleAction::new("clear-messages", None);
+            action_clear_messages.connect_activate(clone!(@weak obj => move |_, _| {
+                log::trace!("User requested to clear messages");
+                // TODO: AdwMessageDialog
+                let dialog = gtk::MessageDialog::builder()
+                    .buttons(gtk::ButtonsType::OkCancel)
+                    .message_type(gtk::MessageType::Warning)
+                    .secondary_text(&gettextrs::gettext(
+                        "This will clear all messages from the device. This will also close the application.",
+                    ))
+                    .text(&gettextrs::gettext("Are you sure?"))
+                    .transient_for(&obj)
+                    .build();
+                dialog.connect_response(clone!(@weak obj => move |dialog, response| {
+                    dialog.close();
+                    if response == gtk::ResponseType::Ok {
+                        log::info!("Clearing device");
+                        if let Some(man) = obj.imp().manager.borrow().as_ref() {
+                            if let Err(e)= man.clear_messages() {
+                                log::error!("Failed to clear db: {}", e);
+                            }
+                        }
+                        log::trace!("Closing the window after message clear");
+                        obj.close();
+                    }
+                }));
+                dialog.show();
+            }));
             log::trace!("Setting up unlink action");
             let action_unlink = SimpleAction::new("unlink", None);
             action_unlink.connect_activate(clone!(@weak obj => move |_, _| {
@@ -152,6 +181,7 @@ pub mod imp {
                                 log::error!("Failed to clear db: {}", e);
                             }
                         }
+                        log::trace!("Closing the window after unlink");
                         obj.close();
                     }
                 }));
@@ -201,6 +231,7 @@ pub mod imp {
             let actions = SimpleActionGroup::new();
             obj.insert_action_group("win", Some(&actions));
             actions.add_action(&action_settings);
+            actions.add_action(&action_clear_messages);
             actions.add_action(&action_unlink);
             actions.add_action(&action_show_help_overlay);
             actions.add_action(&action_about);

src/storage/encrypted_sled.rs

@@ -6,7 +6,7 @@ use std::{
 use async_trait::async_trait;
 use encrypted_sled::IVec;
 use libsignal_service::{
-    content::ContentBody,
+    content::Content,
     models::Contact,
     prelude::{
         protocol::{
@@ -16,14 +16,13 @@ use libsignal_service::{
         },
         ProtobufMessage, Uuid,
     },
-    proto::{sync_message::Sent, GroupContextV2, SyncMessage},
     ServiceAddress,
 };
 use log::{debug, trace, warn};
 
 use presage::{
     ConfigStore, ContactsStore, ContentProto, Error, MessageIdentity, MessageStore, Registered,
-    StateStore,
+    StateStore, Thread,
 };
 
 const SLED_KEY_REGISTRATION: &str = "registration";
@@ -32,8 +31,7 @@ const SLED_KEY_GROUPS: &str = "groups";
 
 const SLED_TREE_SESSIONS: &str = "sessions";
 const SLED_TREE_MESSAGES: &str = "messages";
-const SLED_TREE_CONTACTS_TO_MESSAGES: &str = "contacts-to-messages";
-const SLED_TREE_GROUPS_TO_MESSAGES: &str = "groups-to-messages";
+const SLED_TREE_THREAD_PREFIX: &str = "thread";
 
 #[derive(Debug, Clone)]
 pub struct EncryptedSledConfigStore<E> {
@@ -58,6 +56,23 @@ impl<E: encrypted_sled::Encryption + 'static> EncryptedSledConfigStore<E> {
         let tree_contacts = db.open_tree(SLED_KEY_CONTACTS)?;
         tree_contacts.clear()?;
         tree_contacts.flush()?;
+        drop(db);
+        self.clear_messages()?;
+        Ok(())
+    }
+
+    pub fn clear_messages(&self) -> Result<(), Error> {
+        let db = self.db.read().expect("poisoned mutex");
+        for name in db.tree_names().expect("Failed to get tree names") {
+            if name
+                .as_ref()
+                .starts_with(SLED_TREE_THREAD_PREFIX.as_bytes())
+            {
+                let tree_thread = db.open_tree(&name)?;
+                tree_thread.clear()?;
+                tree_thread.flush()?;
+            }
+        }
         Ok(())
     }
 
@@ -492,66 +507,44 @@ fn prefix_merge(_key: &[u8], old_value: Option<&[u8]>, merged_bytes: &[u8]) -> O
     Some(ret)
 }
 
+fn thread_key(t: &Thread) -> Vec<u8> {
+    let mut bytes = SLED_TREE_THREAD_PREFIX.as_bytes().to_owned();
+    bytes.append(&mut t.into());
+    bytes
+}
+
 impl<E: encrypted_sled::Encryption + 'static> MessageStore for EncryptedSledConfigStore<E> {
+    type MessagesIter = SledMessagesIter<E>;
     fn save_message(
         &mut self,
         message: libsignal_service::prelude::Content,
         receiver: Option<impl Into<ServiceAddress>>,
     ) -> Result<(), Error> {
+        let receiver = receiver.map(|s| s.into());
+
         let id = MessageIdentity::try_from(&message)?;
-        let store_uuid = if let Some(uuid) = receiver.map(|s| s.into()).and_then(|s| s.uuid) {
-            uuid
-        } else {
-            id.0
-        };
+        let thread = Thread::from_content_receiver(&message, receiver.as_ref())?;
         let timestamp = id.1;
 
-        crate::trace!("Storing a message with id: {:?}", id);
-
-        let group_master_key = match message.body {
-            ContentBody::DataMessage(ref msg)
-            | ContentBody::SynchronizeMessage(SyncMessage {
-                sent:
-                    Some(Sent {
-                        message: Some(ref msg),
-                        ..
-                    }),
-                ..
-            }) => msg.group_v2.clone().and_then(|g| g.master_key),
-            _ => None,
-        };
+        log::trace!("Storing a message with id: {:?}", id);
 
         let tree_messages = self
             .db
-            .try_read()
+            .read()
             .expect("poisoned mutex")
             .open_tree(SLED_TREE_MESSAGES)?;
-        let tree_contacts_to_messages = self
+        let tree_thread = self
             .db
-            .try_read()
+            .read()
             .expect("poisoned mutex")
-            .open_tree(SLED_TREE_CONTACTS_TO_MESSAGES)?;
-        let tree_groups_to_messages = self
-            .db
-            .try_read()
-            .expect("poisoned mutex")
-            .open_tree(SLED_TREE_GROUPS_TO_MESSAGES)?;
-        tree_contacts_to_messages.set_merge_operator(prefix_merge);
-        tree_groups_to_messages.set_merge_operator(prefix_merge);
+            .open_tree(thread_key(&thread))?;
 
-        let key_uuid = store_uuid.as_bytes();
-        let key_timestamp = timestamp.to_ne_bytes();
-        let key = [&key_uuid[..], &key_timestamp[..]].concat();
+        let key: [u8; 24] = id.into();
         let value = ContentProto::from_content(message);
+        let value_vec = value.encode_to_vec();
 
-        tree_messages.insert(key.clone(), value.encode_to_vec())?;
-        if let Some(group) = group_master_key {
-            crate::trace!("Storing message to group: {:?}", group);
-            tree_groups_to_messages.merge(group, key)?;
-        } else {
-            crate::trace!("Storing message to contact: {:?}", store_uuid);
-            tree_contacts_to_messages.merge(store_uuid.as_bytes(), key)?;
-        }
+        tree_messages.insert(key, value_vec.clone())?;
+        tree_thread.insert(timestamp.to_be_bytes(), value_vec)?;
         Ok(())
     }
 
@@ -572,17 +565,12 @@ impl<E: encrypted_sled::Encryption + 'static> MessageStore for EncryptedSledConf
         &self,
         id: &MessageIdentity,
     ) -> Result<Option<libsignal_service::prelude::Content>, Error> {
-        let sender = id.0;
-        let timestamp = id.1;
-
         let tree = self
             .db
             .read()
             .expect("poisoned mutex")
             .open_tree(SLED_TREE_MESSAGES)?;
-        let key_sender = sender.as_bytes();
-        let key_timestamp = timestamp.to_ne_bytes();
-        let key = [&key_sender[..], &key_timestamp[..]].concat();
+        let key: [u8; 24] = id.clone().into();
         let val = tree.get(key)?;
         if let Some(val) = val {
             let proto = ContentProto::decode(&*val)?;
@@ -592,31 +580,34 @@ impl<E: encrypted_sled::Encryption + 'static> MessageStore for EncryptedSledConf
         }
     }
 
-    fn messages_by_contact(&self, contact: &Uuid) -> Result<Vec<MessageIdentity>, Error> {
-        crate::trace!("Query messages by contact: {:?}", contact);
-        Ok(self
+    fn messages_by_thread(
+        &self,
+        thread: &Thread,
+        from: Option<u64>,
+    ) -> Result<Self::MessagesIter, Error> {
+        let tree_thread = self
             .db
             .read()
             .expect("poisoned mutex")
-            .open_tree(SLED_TREE_CONTACTS_TO_MESSAGES)?
-            .get(contact.as_bytes())?
-            .unwrap_or_default()
-            .chunks_exact(24)
-            .map(|c| MessageIdentity::from(<[u8; 24]>::try_from(c).unwrap()))
-            .collect())
+            .open_tree(thread_key(thread))?;
+        let iter = if let Some(from) = from {
+            tree_thread.range(..from.to_be_bytes())
+        } else {
+            tree_thread.range::<&[u8], std::ops::RangeFull>(..)
+        }?;
+        Ok(SledMessagesIter(iter.rev()))
     }
+}
 
-    fn messages_by_group(&self, group: &GroupContextV2) -> Result<Vec<MessageIdentity>, Error> {
-        crate::trace!("Query messages by group: {:?}", group.master_key());
-        Ok(self
-            .db
-            .read()
-            .expect("poisoned mutex")
-            .open_tree(SLED_TREE_GROUPS_TO_MESSAGES)?
-            .get(group.master_key())?
-            .unwrap_or_default()
-            .chunks_exact(24)
-            .map(|c| MessageIdentity::from(<[u8; 24]>::try_from(c).unwrap()))
-            .collect())
+pub struct SledMessagesIter<E>(std::iter::Rev<encrypted_sled::Iter<E>>);
+
+impl<E: encrypted_sled::Encryption + 'static> Iterator for SledMessagesIter<E> {
+    // TODO: If error, throw away the rest. Maybe return Result<Content, Error>?
+    type Item = Content;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let ivec = self.0.next()?.ok()?.1;
+        let proto = ContentProto::decode(&*ivec).ok()?;
+        Some(proto.into_content())
     }
 }