mls-storage: add exporter secret methods to GroupStorage trait

Signed-off-by: Yuki Kishimoto <[email protected]>

Author: erskingardner

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

@@ -10,6 +10,15 @@ use nostr_mls_storage::messages::types::Message;
 
 use crate::NostrMlsMemoryStorage;
 
+/// Creates a compound key from an MLS group ID and epoch
+///
+/// The key is created by concatenating the MLS group ID and the epoch as bytes
+fn create_compound_key(mls_group_id: &[u8], epoch: u64) -> Vec<u8> {
+    let mut key = mls_group_id.to_vec();
+    key.extend_from_slice(&epoch.to_be_bytes());
+    key
+}
+
 impl GroupStorage for NostrMlsMemoryStorage {
     fn save_group(&self, group: Group) -> Result<(), GroupError> {
         // Store in the MLS group ID cache
@@ -102,4 +111,36 @@ impl GroupStorage for NostrMlsMemoryStorage {
 
         Ok(())
     }
+
+    fn get_group_exporter_secret(
+        &self,
+        mls_group_id: &[u8],
+        epoch: u64,
+    ) -> Result<Option<GroupExporterSecret>, GroupError> {
+        // Check if the group exists first
+        self.find_group_by_mls_group_id(mls_group_id)?;
+
+        let cache = self.group_exporter_secrets_cache.read();
+        // Create a compound key from mls_group_id and epoch
+        let key = create_compound_key(mls_group_id, epoch);
+        Ok(cache.peek(&key).cloned())
+    }
+
+    fn save_group_exporter_secret(
+        &self,
+        group_exporter_secret: GroupExporterSecret,
+    ) -> Result<(), GroupError> {
+        // Check if the group exists first
+        self.find_group_by_mls_group_id(&group_exporter_secret.mls_group_id)?;
+
+        let mut cache = self.group_exporter_secrets_cache.write();
+        // Create a compound key from mls_group_id and epoch
+        let key = create_compound_key(
+            &group_exporter_secret.mls_group_id,
+            group_exporter_secret.epoch,
+        );
+        cache.put(key, group_exporter_secret);
+
+        Ok(())
+    }
 }

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

@@ -15,7 +15,7 @@ use std::num::NonZeroUsize;
 
 use lru::LruCache;
 use nostr::EventId;
-use nostr_mls_storage::groups::types::{Group, GroupRelay};
+use nostr_mls_storage::groups::types::{Group, GroupExporterSecret, GroupRelay};
 use nostr_mls_storage::messages::types::{Message, ProcessedMessage};
 use nostr_mls_storage::welcomes::types::{ProcessedWelcome, Welcome};
 use nostr_mls_storage::{Backend, NostrMlsStorageProvider};
@@ -70,6 +70,8 @@ pub struct NostrMlsMemoryStorage {
     messages_by_group_cache: RwLock<LruCache<Vec<u8>, Vec<Message>>>,
     /// LRU Cache for ProcessedMessage objects, keyed by Event ID
     processed_messages_cache: RwLock<LruCache<EventId, ProcessedMessage>>,
+    /// LRU Cache for GroupExporterSecret objects, keyed by a compound key of MLS group ID and epoch
+    group_exporter_secrets_cache: RwLock<LruCache<Vec<u8>, GroupExporterSecret>>,
 }
 
 impl Default for NostrMlsMemoryStorage {
@@ -121,6 +123,7 @@ impl NostrMlsMemoryStorage {
             messages_cache: RwLock::new(LruCache::new(cache_size)),
             messages_by_group_cache: RwLock::new(LruCache::new(cache_size)),
             processed_messages_cache: RwLock::new(LruCache::new(cache_size)),
+            group_exporter_secrets_cache: RwLock::new(LruCache::new(cache_size)),
         }
     }
 }
@@ -168,7 +171,7 @@ mod tests {
     use std::collections::BTreeSet;
 
     use nostr::{EventId, Kind, PublicKey, RelayUrl, Tags, Timestamp, UnsignedEvent};
-    use nostr_mls_storage::groups::types::{Group, GroupState, GroupType};
+    use nostr_mls_storage::groups::types::{Group, GroupExporterSecret, GroupState, GroupType};
     use nostr_mls_storage::groups::GroupStorage;
     use nostr_mls_storage::messages::types::{Message, ProcessedMessageState};
     use nostr_mls_storage::messages::MessageStorage;
@@ -333,6 +336,74 @@ mod tests {
         assert_eq!(found_relays.len(), 2);
     }
 
+    #[test]
+    fn test_group_exporter_secret_cache() {
+        let storage = MemoryStorage::default();
+        let nostr_storage = NostrMlsMemoryStorage::new(storage);
+
+        // Create a test group
+        let mls_group_id = vec![1, 2, 3, 4];
+        let group = Group {
+            mls_group_id: mls_group_id.clone(),
+            nostr_group_id: "test_group_123".to_string(),
+            name: "Test Group".to_string(),
+            description: "A test group".to_string(),
+            admin_pubkeys: BTreeSet::new(),
+            last_message_id: None,
+            last_message_at: None,
+            group_type: GroupType::Group,
+            epoch: 0,
+            state: GroupState::Active,
+        };
+
+        // Save the group
+        nostr_storage.save_group(group.clone()).unwrap();
+
+        // Create a test group exporter secret for epoch 0
+        let group_exporter_secret_0 = GroupExporterSecret {
+            mls_group_id: mls_group_id.clone(),
+            epoch: 0,
+            secret: vec![1, 2, 3, 4],
+        };
+
+        // Create a test group exporter secret for epoch 1
+        let group_exporter_secret_1 = GroupExporterSecret {
+            mls_group_id: mls_group_id.clone(),
+            epoch: 1,
+            secret: vec![5, 6, 7, 8],
+        };
+
+        // Save the group exporter secrets
+        nostr_storage
+            .save_group_exporter_secret(group_exporter_secret_0.clone())
+            .unwrap();
+        nostr_storage
+            .save_group_exporter_secret(group_exporter_secret_1.clone())
+            .unwrap();
+
+        // Get the group exporter secret for epoch 0
+        let found_group_exporter_secret_0 = nostr_storage
+            .get_group_exporter_secret(&mls_group_id, 0)
+            .unwrap();
+        assert!(found_group_exporter_secret_0.is_some());
+        let found_group_exporter_secret_0 = found_group_exporter_secret_0.unwrap();
+        assert_eq!(found_group_exporter_secret_0, group_exporter_secret_0);
+
+        // Get the group exporter secret for epoch 1
+        let found_group_exporter_secret_1 = nostr_storage
+            .get_group_exporter_secret(&mls_group_id, 1)
+            .unwrap();
+        assert!(found_group_exporter_secret_1.is_some());
+        let found_group_exporter_secret_1 = found_group_exporter_secret_1.unwrap();
+        assert_eq!(found_group_exporter_secret_1, group_exporter_secret_1);
+
+        // Verify we can't find an exporter secret for a non-existing epoch
+        let non_existent_secret = nostr_storage
+            .get_group_exporter_secret(&mls_group_id, 999)
+            .unwrap();
+        assert!(non_existent_secret.is_none());
+    }
+
     #[test]
     fn test_welcome_cache() {
         let storage = MemoryStorage::default();

crates/nostr-mls-sqlite-storage/migrations/V100__initial.sql

@@ -29,6 +29,17 @@ CREATE TABLE IF NOT EXISTS group_relays (
 -- Create index on mls_group_id for faster lookups
 CREATE INDEX IF NOT EXISTS idx_group_relays_mls_group_id ON group_relays(mls_group_id);
 
+-- Group Exporter Secrets table
+CREATE TABLE IF NOT EXISTS group_exporter_secrets (
+    mls_group_id BLOB NOT NULL,
+    epoch INTEGER NOT NULL,
+    secret BLOB NOT NULL,
+    PRIMARY KEY (mls_group_id, epoch)
+);
+
+-- Create index on mls_group_id for faster lookups
+CREATE INDEX IF NOT EXISTS idx_group_exporter_secrets_mls_group_id ON group_exporter_secrets(mls_group_id);
+
 -- Messages table
 CREATE TABLE IF NOT EXISTS messages (
     id BLOB PRIMARY KEY,  -- Event ID as byte array

crates/nostr-mls-sqlite-storage/src/db.rs

@@ -5,7 +5,9 @@ use std::io::{Error as IoError, ErrorKind};
 use std::str::FromStr;
 
 use nostr::{EventId, JsonUtil, Kind, PublicKey, RelayUrl, Tags, Timestamp, UnsignedEvent};
-use nostr_mls_storage::groups::types::{Group, GroupRelay, GroupState, GroupType};
+use nostr_mls_storage::groups::types::{
+    Group, GroupExporterSecret, GroupRelay, GroupState, GroupType,
+};
 use nostr_mls_storage::messages::types::{Message, ProcessedMessage, ProcessedMessageState};
 use nostr_mls_storage::welcomes::types::{
     ProcessedWelcome, ProcessedWelcomeState, Welcome, WelcomeState,
@@ -97,6 +99,19 @@ pub fn row_to_group_relay(row: &Row) -> SqliteResult<GroupRelay> {
     })
 }
 
+/// Convert a row to a GroupExporterSecret struct
+pub fn row_to_group_exporter_secret(row: &Row) -> SqliteResult<GroupExporterSecret> {
+    let mls_group_id: Vec<u8> = row.get("mls_group_id")?;
+    let epoch: u64 = row.get("epoch")?;
+    let secret: Vec<u8> = row.get("secret")?;
+
+    Ok(GroupExporterSecret {
+        mls_group_id,
+        epoch,
+        secret,
+    })
+}
+
 /// Convert a row to a Message struct
 pub fn row_to_message(row: &Row) -> SqliteResult<Message> {
     let id_blob: &[u8] = row.get_ref("id")?.as_blob()?;

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

@@ -4,7 +4,7 @@ use std::collections::BTreeSet;
 
 use nostr::PublicKey;
 use nostr_mls_storage::groups::error::GroupError;
-use nostr_mls_storage::groups::types::{Group, GroupRelay};
+use nostr_mls_storage::groups::types::{Group, GroupExporterSecret, GroupRelay};
 use nostr_mls_storage::groups::GroupStorage;
 use nostr_mls_storage::messages::types::Message;
 use rusqlite::{params, OptionalExtension};
@@ -197,6 +197,58 @@ impl GroupStorage for NostrMlsSqliteStorage {
 
         Ok(())
     }
+
+    fn get_group_exporter_secret(
+        &self,
+        mls_group_id: &[u8],
+        epoch: u64,
+    ) -> Result<Option<GroupExporterSecret>, GroupError> {
+        // First verify the group exists
+        if self.find_group_by_mls_group_id(mls_group_id)?.is_none() {
+            return Err(GroupError::InvalidParameters(format!(
+                "Group with MLS ID {:?} not found",
+                mls_group_id
+            )));
+        }
+
+        let conn_guard = self.db_connection.lock().map_err(into_group_err)?;
+
+        let mut stmt = conn_guard
+            .prepare("SELECT * FROM group_exporter_secrets WHERE mls_group_id = ? AND epoch = ?")
+            .map_err(into_group_err)?;
+
+        stmt.query_row(
+            params![mls_group_id, epoch],
+            db::row_to_group_exporter_secret,
+        )
+        .optional()
+        .map_err(into_group_err)
+    }
+
+    fn save_group_exporter_secret(
+        &self,
+        group_exporter_secret: GroupExporterSecret,
+    ) -> Result<(), GroupError> {
+        if self
+            .find_group_by_mls_group_id(&group_exporter_secret.mls_group_id)?
+            .is_none()
+        {
+            return Err(GroupError::InvalidParameters(format!(
+                "Group with MLS ID {:?} not found",
+                group_exporter_secret.mls_group_id
+            )));
+        }
+
+        let conn_guard = self.db_connection.lock().map_err(into_group_err)?;
+
+        conn_guard.execute(
+            "INSERT OR REPLACE INTO group_exporter_secrets (mls_group_id, epoch, secret) VALUES (?, ?, ?)",
+            params![&group_exporter_secret.mls_group_id, &group_exporter_secret.epoch, &group_exporter_secret.secret],
+        )
+        .map_err(into_group_err)?;
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -290,4 +342,83 @@ mod tests {
             "wss://relay.example.com"
         );
     }
+
+    #[test]
+    fn test_group_exporter_secret() {
+        let storage = NostrMlsSqliteStorage::new_in_memory().unwrap();
+
+        // Create a test group
+        let mls_group_id = vec![1, 2, 3, 4];
+        let group = Group {
+            mls_group_id: mls_group_id.clone(),
+            nostr_group_id: "test_group_123".to_string(),
+            name: "Test Group".to_string(),
+            description: "A test group".to_string(),
+            admin_pubkeys: BTreeSet::new(),
+            last_message_id: None,
+            last_message_at: None,
+            group_type: GroupType::Group,
+            epoch: 0,
+            state: GroupState::Active,
+        };
+
+        // Save the group
+        storage.save_group(group.clone()).unwrap();
+
+        // Create a group exporter secret
+        let secret1 = GroupExporterSecret {
+            mls_group_id: mls_group_id.clone(),
+            epoch: 1,
+            secret: vec![5, 6, 7, 8],
+        };
+
+        // Save the secret
+        storage.save_group_exporter_secret(secret1.clone()).unwrap();
+
+        // Get the secret and verify it was saved correctly
+        let retrieved_secret = storage
+            .get_group_exporter_secret(&mls_group_id, 1)
+            .unwrap()
+            .unwrap();
+        assert_eq!(retrieved_secret.secret, vec![5, 6, 7, 8]);
+
+        // Create a second secret with same group_id and epoch but different secret value
+        let secret2 = GroupExporterSecret {
+            mls_group_id: mls_group_id.clone(),
+            epoch: 1,
+            secret: vec![9, 10, 11, 12],
+        };
+
+        // Save the second secret - this should replace the first one due to the "OR REPLACE" in the SQL
+        storage.save_group_exporter_secret(secret2.clone()).unwrap();
+
+        // Get the secret again and verify it was updated
+        let retrieved_secret = storage
+            .get_group_exporter_secret(&mls_group_id, 1)
+            .unwrap()
+            .unwrap();
+        assert_eq!(retrieved_secret.secret, vec![9, 10, 11, 12]);
+
+        // Verify we can still save a different epoch
+        let secret3 = GroupExporterSecret {
+            mls_group_id: mls_group_id.clone(),
+            epoch: 2,
+            secret: vec![13, 14, 15, 16],
+        };
+
+        storage.save_group_exporter_secret(secret3.clone()).unwrap();
+
+        // Verify both epochs exist
+        let retrieved_secret1 = storage
+            .get_group_exporter_secret(&mls_group_id, 1)
+            .unwrap()
+            .unwrap();
+        let retrieved_secret2 = storage
+            .get_group_exporter_secret(&mls_group_id, 2)
+            .unwrap()
+            .unwrap();
+
+        assert_eq!(retrieved_secret1.secret, vec![9, 10, 11, 12]);
+        assert_eq!(retrieved_secret2.secret, vec![13, 14, 15, 16]);
+    }
 }