mls-storage: improve APIs of MLS storage traits

* Return `Result<Option<T>>` for query methods: the `Result` indicates if the backend successfully performed the query and the `Option` if the result has been found or not.
* Avoid returning the argument in saving methods
* Cleanup the implementation of some MLS trait methods

Discussion at #839 (comment) and #839 (comment)

Co-authored-by: Jeff Gardner <[email protected]>
Signed-off-by: Yuki Kishimoto <[email protected]>

Author: yukibtc

crates/nostr-mls-memory-storage/src/groups.rs

@@ -1,15 +1,15 @@
 //! Memory-based storage implementation of the NostrMlsStorageProvider trait for Nostr MLS groups
 
 use nostr::PublicKey;
-use nostr_mls_storage::groups::error::GroupError;
+use nostr_mls_storage::groups::error::{GroupError, InvalidGroupState};
 use nostr_mls_storage::groups::types::*;
 use nostr_mls_storage::groups::GroupStorage;
 use nostr_mls_storage::messages::types::Message;
 
 use crate::NostrMlsMemoryStorage;
 
 impl GroupStorage for NostrMlsMemoryStorage {
-    fn save_group(&self, group: Group) -> Result<Group, GroupError> {
+    fn save_group(&self, group: Group) -> Result<(), GroupError> {
         // Store in the MLS group ID cache
         {
             let mut cache = self.groups_cache.write();
@@ -19,10 +19,10 @@ impl GroupStorage for NostrMlsMemoryStorage {
         // Store in the Nostr group ID cache
         {
             let mut cache = self.groups_by_nostr_id_cache.write();
-            cache.put(group.nostr_group_id.clone(), group.clone());
+            cache.put(group.nostr_group_id.clone(), group);
         }
 
-        Ok(group)
+        Ok(())
     }
 
     fn all_groups(&self) -> Result<Vec<Group>, GroupError> {
@@ -32,92 +32,73 @@ impl GroupStorage for NostrMlsMemoryStorage {
         Ok(groups)
     }
 
-    fn find_group_by_mls_group_id(&self, mls_group_id: &[u8]) -> Result<Group, GroupError> {
+    fn find_group_by_mls_group_id(&self, mls_group_id: &[u8]) -> Result<Option<Group>, GroupError> {
         let cache = self.groups_cache.read();
-        if let Some(group) = cache.peek(mls_group_id) {
-            // Return a clone of the found group
-            return Ok(group.clone());
-        }
-
-        Err(GroupError::NotFound)
+        Ok(cache.peek(mls_group_id).cloned())
     }
 
-    fn find_group_by_nostr_group_id(&self, nostr_group_id: &str) -> Result<Group, GroupError> {
+    fn find_group_by_nostr_group_id(
+        &self,
+        nostr_group_id: &str,
+    ) -> Result<Option<Group>, GroupError> {
         let cache = self.groups_by_nostr_id_cache.read();
-        if let Some(group) = cache.peek(nostr_group_id) {
-            // Return a clone of the found group
-            return Ok(group.clone());
-        }
-
-        Err(GroupError::NotFound)
+        Ok(cache.peek(nostr_group_id).cloned())
     }
 
     fn messages(&self, mls_group_id: &[u8]) -> Result<Vec<Message>, GroupError> {
         // Check if the group exists first
         self.find_group_by_mls_group_id(mls_group_id)?;
 
         let cache = self.messages_by_group_cache.read();
-        if let Some(messages) = cache.peek(mls_group_id) {
-            return Ok(messages.clone());
+        match cache.peek(mls_group_id).cloned() {
+            Some(messages) => Ok(messages),
+            // If not in cache but group exists, return empty vector
+            None => Ok(Vec::new()),
         }
-
-        // If not in cache but group exists, return empty vector
-        Ok(Vec::new())
     }
 
     fn admins(&self, mls_group_id: &[u8]) -> Result<Vec<PublicKey>, GroupError> {
-        // Find the group first
-        if let Ok(group) = self.find_group_by_mls_group_id(mls_group_id) {
-            // Return the admin pubkeys from the group
-            return Ok(group.admin_pubkeys);
+        match self.find_group_by_mls_group_id(mls_group_id)? {
+            Some(group) => Ok(group.admin_pubkeys),
+            None => Err(GroupError::InvalidState(InvalidGroupState::NoAdmins)),
         }
-
-        Err(GroupError::NotFound)
     }
 
     fn group_relays(&self, mls_group_id: &[u8]) -> Result<Vec<GroupRelay>, GroupError> {
         // Check if the group exists first
         self.find_group_by_mls_group_id(mls_group_id)?;
 
         let cache = self.group_relays_cache.read();
-        if let Some(relays) = cache.peek(mls_group_id) {
-            return Ok(relays.clone());
+        match cache.peek(mls_group_id).cloned() {
+            Some(relays) => Ok(relays),
+            None => Err(GroupError::InvalidState(InvalidGroupState::NoRelays)),
         }
-
-        // If not in cache but group exists, return empty vector
-        Ok(Vec::new())
     }
 
-    fn save_group_relay(&self, group_relay: GroupRelay) -> Result<GroupRelay, GroupError> {
-        let mls_group_id = group_relay.mls_group_id.clone();
-
+    fn save_group_relay(&self, group_relay: GroupRelay) -> Result<(), GroupError> {
         // Check if the group exists first
-        self.find_group_by_mls_group_id(&mls_group_id)?;
-
-        let group_relay_clone = group_relay.clone();
+        self.find_group_by_mls_group_id(&group_relay.mls_group_id)?;
 
-        {
-            let mut cache = self.group_relays_cache.write();
-            // Get existing relays or create new vector
-            let relays = match cache.get(&mls_group_id) {
-                Some(existing_relays) => {
-                    let mut new_relays = existing_relays.clone();
-                    // Add the new relay if it doesn't already exist
-                    if !new_relays
-                        .iter()
-                        .any(|r| r.relay_url == group_relay.relay_url)
-                    {
-                        new_relays.push(group_relay_clone);
-                    }
-                    new_relays
-                }
-                None => vec![group_relay_clone],
-            };
+        let mut cache = self.group_relays_cache.write();
 
-            // Update the cache with the new vector
-            cache.put(mls_group_id, relays);
-        }
+        // Try to get the existing relays for the group
+        match cache.get_mut(&group_relay.mls_group_id) {
+            // If the group exists, add the new relay to the vector
+            Some(existing_relays) => {
+                // TODO: the time complexity here is O(n). The number of relays is likely to be small, but it's probably better to use a BTreeSet anyway.
 
-        Ok(group_relay)
+                // Add the new relay if it doesn't already exist
+                if !existing_relays.contains(&group_relay) {
+                    existing_relays.push(group_relay);
+                }
+            }
+            // If the group doesn't exist, create a new vector with the new relay
+            None => {
+                // Update the cache with the new vector
+                cache.put(group_relay.mls_group_id.clone(), vec![group_relay]);
+            }
+        };
+
+        Ok(())
     }
 }

crates/nostr-mls-memory-storage/src/lib.rs

@@ -245,7 +245,7 @@ mod tests {
         // Find the group by MLS group ID
         let found_group = nostr_storage.find_group_by_mls_group_id(&mls_group_id);
         assert!(found_group.is_ok());
-        let found_group = found_group.unwrap();
+        let found_group = found_group.unwrap().unwrap();
         assert_eq!(found_group.mls_group_id, mls_group_id);
         assert_eq!(found_group.nostr_group_id, "test_group_123");
 
@@ -371,9 +371,9 @@ mod tests {
         assert!(result.is_ok());
 
         // Find the welcome by event ID
-        let found_welcome = nostr_storage.find_welcome_by_event_id(event_id);
+        let found_welcome = nostr_storage.find_welcome_by_event_id(&event_id);
         assert!(found_welcome.is_ok());
-        let found_welcome = found_welcome.unwrap();
+        let found_welcome = found_welcome.unwrap().unwrap();
         assert_eq!(found_welcome.id, event_id);
         assert_eq!(found_welcome.mls_group_id, vec![9, 10, 11, 12]);
 
@@ -397,9 +397,9 @@ mod tests {
         assert!(result.is_ok());
 
         // Find the processed welcome by event ID
-        let found_processed_welcome = nostr_storage.find_processed_welcome_by_event_id(wrapper_id);
+        let found_processed_welcome = nostr_storage.find_processed_welcome_by_event_id(&wrapper_id);
         assert!(found_processed_welcome.is_ok());
-        let found_processed_welcome = found_processed_welcome.unwrap();
+        let found_processed_welcome = found_processed_welcome.unwrap().unwrap();
         assert_eq!(found_processed_welcome.wrapper_event_id, wrapper_id);
         assert_eq!(found_processed_welcome.welcome_event_id, Some(event_id));
 
@@ -466,9 +466,9 @@ mod tests {
         assert!(result.is_ok());
 
         // Find the message by event ID
-        let found_message = nostr_storage.find_message_by_event_id(event_id);
+        let found_message = nostr_storage.find_message_by_event_id(&event_id);
         assert!(found_message.is_ok());
-        let found_message = found_message.unwrap();
+        let found_message = found_message.unwrap().unwrap();
         assert_eq!(found_message.id, event_id);
         assert_eq!(found_message.mls_group_id, mls_group_id);
 
@@ -505,9 +505,9 @@ mod tests {
         assert!(result.is_ok());
 
         // Find the processed message by event ID
-        let found_processed_message = nostr_storage.find_processed_message_by_event_id(wrapper_id);
+        let found_processed_message = nostr_storage.find_processed_message_by_event_id(&wrapper_id);
         assert!(found_processed_message.is_ok());
-        let found_processed_message = found_processed_message.unwrap();
+        let found_processed_message = found_processed_message.unwrap().unwrap();
         assert_eq!(found_processed_message.wrapper_event_id, wrapper_id);
         assert_eq!(found_processed_message.message_event_id, Some(event_id));
 
@@ -545,7 +545,7 @@ mod tests {
         // Find the group by MLS group ID
         let found_group = nostr_storage.find_group_by_mls_group_id(&mls_group_id);
         assert!(found_group.is_ok());
-        let found_group = found_group.unwrap();
+        let found_group = found_group.unwrap().unwrap();
         assert_eq!(found_group.mls_group_id, mls_group_id);
     }
 
@@ -574,7 +574,7 @@ mod tests {
         // Find the group by MLS group ID
         let found_group = nostr_storage.find_group_by_mls_group_id(&mls_group_id);
         assert!(found_group.is_ok());
-        let found_group = found_group.unwrap();
+        let found_group = found_group.unwrap().unwrap();
         assert_eq!(found_group.mls_group_id, mls_group_id);
     }
 }

crates/nostr-mls-memory-storage/src/messages.rs

@@ -8,48 +8,35 @@ use nostr_mls_storage::messages::MessageStorage;
 use crate::NostrMlsMemoryStorage;
 
 impl MessageStorage for NostrMlsMemoryStorage {
-    fn save_message(&self, message: Message) -> Result<Message, MessageError> {
-        {
-            let mut cache = self.messages_cache.write();
-            cache.put(message.id, message.clone());
-        }
-
-        Ok(message)
+    fn save_message(&self, message: Message) -> Result<(), MessageError> {
+        let mut cache = self.messages_cache.write();
+        cache.put(message.id, message);
+        Ok(())
     }
 
-    fn find_message_by_event_id(&self, event_id: EventId) -> Result<Message, MessageError> {
+    fn find_message_by_event_id(
+        &self,
+        event_id: &EventId,
+    ) -> Result<Option<Message>, MessageError> {
         let cache = self.messages_cache.read();
-        if let Some(message) = cache.peek(&event_id) {
-            return Ok(message.clone());
-        }
-
-        Err(MessageError::NotFound)
+        Ok(cache.peek(event_id).cloned())
     }
 
     fn find_processed_message_by_event_id(
         &self,
-        event_id: EventId,
-    ) -> Result<ProcessedMessage, MessageError> {
+        event_id: &EventId,
+    ) -> Result<Option<ProcessedMessage>, MessageError> {
         let cache = self.processed_messages_cache.read();
-        if let Some(processed_message) = cache.peek(&event_id) {
-            return Ok(processed_message.clone());
-        }
-
-        Err(MessageError::NotFound)
+        Ok(cache.peek(event_id).cloned())
     }
 
     fn save_processed_message(
         &self,
         processed_message: ProcessedMessage,
-    ) -> Result<ProcessedMessage, MessageError> {
-        {
-            let mut cache = self.processed_messages_cache.write();
-            cache.put(
-                processed_message.wrapper_event_id,
-                processed_message.clone(),
-            );
-        }
+    ) -> Result<(), MessageError> {
+        let mut cache = self.processed_messages_cache.write();
+        cache.put(processed_message.wrapper_event_id, processed_message);
 
-        Ok(processed_message)
+        Ok(())
     }
 }

crates/nostr-mls-memory-storage/src/welcomes.rs

@@ -8,13 +8,11 @@ use nostr_mls_storage::welcomes::WelcomeStorage;
 use crate::NostrMlsMemoryStorage;
 
 impl WelcomeStorage for NostrMlsMemoryStorage {
-    fn save_welcome(&self, welcome: Welcome) -> Result<Welcome, WelcomeError> {
-        {
-            let mut cache = self.welcomes_cache.write();
-            cache.put(welcome.id, welcome.clone());
-        }
+    fn save_welcome(&self, welcome: Welcome) -> Result<(), WelcomeError> {
+        let mut cache = self.welcomes_cache.write();
+        cache.put(welcome.id, welcome);
 
-        Ok(welcome)
+        Ok(())
     }
 
     fn pending_welcomes(&self) -> Result<Vec<Welcome>, WelcomeError> {
@@ -28,39 +26,29 @@ impl WelcomeStorage for NostrMlsMemoryStorage {
         Ok(welcomes)
     }
 
-    fn find_welcome_by_event_id(&self, event_id: EventId) -> Result<Welcome, WelcomeError> {
+    fn find_welcome_by_event_id(
+        &self,
+        event_id: &EventId,
+    ) -> Result<Option<Welcome>, WelcomeError> {
         let cache = self.welcomes_cache.read();
-        if let Some(welcome) = cache.peek(&event_id) {
-            return Ok(welcome.clone());
-        }
-
-        Err(WelcomeError::NotFound)
+        Ok(cache.peek(event_id).cloned())
     }
 
-    fn find_processed_welcome_by_event_id(
+    fn save_processed_welcome(
         &self,
-        event_id: EventId,
-    ) -> Result<ProcessedWelcome, WelcomeError> {
-        let cache = self.processed_welcomes_cache.read();
-        if let Some(processed_welcome) = cache.peek(&event_id) {
-            return Ok(processed_welcome.clone());
-        }
+        processed_welcome: ProcessedWelcome,
+    ) -> Result<(), WelcomeError> {
+        let mut cache = self.processed_welcomes_cache.write();
+        cache.put(processed_welcome.wrapper_event_id, processed_welcome);
 
-        Err(WelcomeError::NotFound)
+        Ok(())
     }
 
-    fn save_processed_welcome(
+    fn find_processed_welcome_by_event_id(
         &self,
-        processed_welcome: ProcessedWelcome,
-    ) -> Result<ProcessedWelcome, WelcomeError> {
-        {
-            let mut cache = self.processed_welcomes_cache.write();
-            cache.put(
-                processed_welcome.wrapper_event_id,
-                processed_welcome.clone(),
-            );
-        }
-
-        Ok(processed_welcome)
+        event_id: &EventId,
+    ) -> Result<Option<ProcessedWelcome>, WelcomeError> {
+        let cache = self.processed_welcomes_cache.read();
+        Ok(cache.peek(event_id).cloned())
     }
 }