Fix: Add input validation to prevent memory exhaustion (#147)

* Add input validation to prevent memory exhaustion (Issue #82)

* fix: coordinated eviction from both caches and update test to trigger eviction

* feat: make validation limits configurable via ValidationLimits struct

* remove unnecessary deprecated constant aliases

* refactor: move cache_size into ValidationLimits

* fix: add zero-value validation to all ValidationLimits builder methods

Author: mubarakcoded

crates/mdk-core/src/groups.rs

@@ -3850,18 +3850,21 @@ mod tests {
             .expect("Failed to merge commit");
     }
 
-    /// Test group with very long name
+    /// Test group with long name within allowed limits
+    ///
+    /// Security fix (Issue #82): Group names are now limited to prevent memory exhaustion.
+    /// This test verifies that names within the limit work correctly.
     #[test]
     fn test_group_with_long_name() {
         let creator_mdk = create_test_mdk();
         let (creator, members, admins) = create_test_group_members();
         let group_id = create_test_group(&creator_mdk, &creator, &members, &admins);
 
-        // Update to very long name (1000 characters)
-        let long_name = "a".repeat(1000);
+        // Update to a long name within the allowed limit (256 bytes)
+        let long_name = "a".repeat(256);
         let update = NostrGroupDataUpdate::new().name(long_name);
         let result = creator_mdk.update_group_data(&group_id, update);
-        assert!(result.is_ok(), "Long group name should be valid");
+        assert!(result.is_ok(), "Group name at limit should be valid");
         creator_mdk
             .merge_pending_commit(&group_id)
             .expect("Failed to merge commit");

crates/mdk-memory-storage/CHANGELOG.md

@@ -40,6 +40,8 @@
 
 - Implemented pagination support using `Pagination` struct for group messages ([#111](https://github.com/marmot-protocol/mdk/pull/111))
 - Implemented pagination support using `Pagination` struct for pending welcomes ([#110](https://github.com/marmot-protocol/mdk/pull/110))
+- **Security (Audit Issue AM)**: Added input validation constants and enforcement to prevent memory exhaustion attacks. New public constants: `DEFAULT_MAX_RELAYS_PER_GROUP`, `DEFAULT_MAX_MESSAGES_PER_GROUP`, `DEFAULT_MAX_GROUP_NAME_LENGTH`, `DEFAULT_MAX_GROUP_DESCRIPTION_LENGTH`, `DEFAULT_MAX_ADMINS_PER_GROUP`, `DEFAULT_MAX_RELAYS_PER_WELCOME`, `DEFAULT_MAX_ADMINS_PER_WELCOME`, `DEFAULT_MAX_RELAY_URL_LENGTH`. Fixes [#82](https://github.com/marmot-protocol/mdk/issues/82) ([#147](https://github.com/marmot-protocol/mdk/pull/147))
+- Added `ValidationLimits` struct for configurable validation limits, allowing users to override default memory exhaustion protection limits via `MdkMemoryStorage::with_limits()` ([#147](https://github.com/marmot-protocol/mdk/pull/147))
 
 ### Fixed
 
@@ -52,7 +54,9 @@
 - **Security (Audit Issue AA)**: Added pagination to prevent memory exhaustion from unbounded loading of pending welcomes ([#110](https://github.com/marmot-protocol/mdk/pull/110))
 - **Security (Audit Issue AN)**: Fixed security issue where `save_message` would accept messages for non-existent groups, allowing cache pollution. Now verifies group existence before inserting messages into the cache. ([#113](https://github.com/marmot-protocol/mdk/pull/113))
 - **Security (Audit Issue AO)**: Removed MLS group identifiers from error messages to prevent metadata leakage in logs and telemetry. Error messages now use generic "Group not found" instead of including the sensitive 32-byte MLS group ID. ([#112](https://github.com/marmot-protocol/mdk/pull/112))
+- **Security (Audit Issue AM)**: Added input validation to prevent memory exhaustion from unbounded per-key values in LRU caches. Validation is enforced in `save_group`, `replace_group_relays`, `save_message`, and `save_welcome` to cap string lengths, collection sizes, and per-group message counts. Fixes [#82](https://github.com/marmot-protocol/mdk/issues/82) ([#147](https://github.com/marmot-protocol/mdk/pull/147))
 - Fix `admins()` to return `InvalidParameters` error when group not found, instead of incorrectly returning `NoAdmins` ([#104](https://github.com/marmot-protocol/mdk/pull/104))
+
 ### Removed
 
 ### Deprecated

crates/mdk-memory-storage/src/groups.rs

@@ -13,6 +13,33 @@ use crate::MdkMemoryStorage;
 
 impl GroupStorage for MdkMemoryStorage {
     fn save_group(&self, group: Group) -> Result<(), GroupError> {
+        // Validate group name length
+        if group.name.len() > self.limits.max_group_name_length {
+            return Err(GroupError::InvalidParameters(format!(
+                "Group name exceeds maximum length of {} bytes (got {} bytes)",
+                self.limits.max_group_name_length,
+                group.name.len()
+            )));
+        }
+
+        // Validate group description length
+        if group.description.len() > self.limits.max_group_description_length {
+            return Err(GroupError::InvalidParameters(format!(
+                "Group description exceeds maximum length of {} bytes (got {} bytes)",
+                self.limits.max_group_description_length,
+                group.description.len()
+            )));
+        }
+
+        // Validate admin pubkeys count
+        if group.admin_pubkeys.len() > self.limits.max_admins_per_group {
+            return Err(GroupError::InvalidParameters(format!(
+                "Group admin count exceeds maximum of {} (got {})",
+                self.limits.max_admins_per_group,
+                group.admin_pubkeys.len()
+            )));
+        }
+
         // Acquire both locks to ensure atomic update
         let mut groups_cache = self.groups_cache.write();
         let mut nostr_id_cache = self.groups_by_nostr_id_cache.write();
@@ -131,6 +158,25 @@ impl GroupStorage for MdkMemoryStorage {
         group_id: &GroupId,
         relays: BTreeSet<RelayUrl>,
     ) -> Result<(), GroupError> {
+        // Validate relay count to prevent memory exhaustion
+        if relays.len() > self.limits.max_relays_per_group {
+            return Err(GroupError::InvalidParameters(format!(
+                "Relay count exceeds maximum of {} (got {})",
+                self.limits.max_relays_per_group,
+                relays.len()
+            )));
+        }
+
+        // Validate individual relay URL lengths
+        for relay in &relays {
+            if relay.as_str().len() > self.limits.max_relay_url_length {
+                return Err(GroupError::InvalidParameters(format!(
+                    "Relay URL exceeds maximum length of {} bytes",
+                    self.limits.max_relay_url_length
+                )));
+            }
+        }
+
         // Check if the group exists first
         if self.find_group_by_mls_group_id(group_id)?.is_none() {
             return Err(GroupError::InvalidParameters("Group not found".to_string()));
@@ -195,12 +241,167 @@ impl GroupStorage for MdkMemoryStorage {
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::{
+        DEFAULT_MAX_ADMINS_PER_GROUP, DEFAULT_MAX_GROUP_DESCRIPTION_LENGTH,
+        DEFAULT_MAX_GROUP_NAME_LENGTH, DEFAULT_MAX_RELAY_URL_LENGTH, DEFAULT_MAX_RELAYS_PER_GROUP,
+    };
     use mdk_storage_traits::groups::types::GroupState;
     use mdk_storage_traits::messages::MessageStorage;
     use mdk_storage_traits::messages::types::{Message, MessageState};
-    use nostr::{EventId, Kind, PublicKey, Tags, Timestamp, UnsignedEvent};
+    use nostr::{EventId, Keys, Kind, Tags, Timestamp, UnsignedEvent};
     use openmls_memory_storage::MemoryStorage;
 
+    fn create_test_group(mls_group_id: GroupId, nostr_group_id: [u8; 32]) -> Group {
+        Group {
+            mls_group_id,
+            nostr_group_id,
+            name: "Test Group".to_string(),
+            description: "A test group".to_string(),
+            admin_pubkeys: BTreeSet::new(),
+            last_message_id: None,
+            last_message_at: None,
+            epoch: 0,
+            state: GroupState::Active,
+            image_hash: None,
+            image_key: None,
+            image_nonce: None,
+        }
+    }
+
+    #[test]
+    fn test_save_group_name_length_validation() {
+        let storage = MdkMemoryStorage::new(MemoryStorage::default());
+        let mls_group_id = GroupId::from_slice(&[1, 2, 3, 4]);
+
+        // Test with name at exactly the limit (should succeed)
+        let mut group = create_test_group(mls_group_id.clone(), [1u8; 32]);
+        group.name = "a".repeat(DEFAULT_MAX_GROUP_NAME_LENGTH);
+        assert!(storage.save_group(group).is_ok());
+
+        // Test with name exceeding the limit (should fail)
+        let mut group = create_test_group(GroupId::from_slice(&[2, 3, 4, 5]), [2u8; 32]);
+        group.name = "a".repeat(DEFAULT_MAX_GROUP_NAME_LENGTH + 1);
+        let result = storage.save_group(group);
+        assert!(result.is_err());
+        assert!(
+            result
+                .unwrap_err()
+                .to_string()
+                .contains("Group name exceeds maximum length")
+        );
+    }
+
+    #[test]
+    fn test_save_group_description_length_validation() {
+        let storage = MdkMemoryStorage::new(MemoryStorage::default());
+        let mls_group_id = GroupId::from_slice(&[1, 2, 3, 4]);
+
+        // Test with description at exactly the limit (should succeed)
+        let mut group = create_test_group(mls_group_id.clone(), [1u8; 32]);
+        group.description = "a".repeat(DEFAULT_MAX_GROUP_DESCRIPTION_LENGTH);
+        assert!(storage.save_group(group).is_ok());
+
+        // Test with description exceeding the limit (should fail)
+        let mut group = create_test_group(GroupId::from_slice(&[2, 3, 4, 5]), [2u8; 32]);
+        group.description = "a".repeat(DEFAULT_MAX_GROUP_DESCRIPTION_LENGTH + 1);
+        let result = storage.save_group(group);
+        assert!(result.is_err());
+        assert!(
+            result
+                .unwrap_err()
+                .to_string()
+                .contains("Group description exceeds maximum length")
+        );
+    }
+
+    #[test]
+    fn test_save_group_admin_count_validation() {
+        let storage = MdkMemoryStorage::new(MemoryStorage::default());
+        let mls_group_id = GroupId::from_slice(&[1, 2, 3, 4]);
+
+        // Test with admin count at exactly the limit (should succeed)
+        let mut group = create_test_group(mls_group_id.clone(), [1u8; 32]);
+        for _ in 0..DEFAULT_MAX_ADMINS_PER_GROUP {
+            group.admin_pubkeys.insert(Keys::generate().public_key());
+        }
+        assert!(storage.save_group(group).is_ok());
+
+        // Test with admin count exceeding the limit (should fail)
+        let mut group = create_test_group(GroupId::from_slice(&[2, 3, 4, 5]), [2u8; 32]);
+        for _ in 0..=DEFAULT_MAX_ADMINS_PER_GROUP {
+            group.admin_pubkeys.insert(Keys::generate().public_key());
+        }
+        let result = storage.save_group(group);
+        assert!(result.is_err());
+        assert!(
+            result
+                .unwrap_err()
+                .to_string()
+                .contains("Group admin count exceeds maximum")
+        );
+    }
+
+    #[test]
+    fn test_replace_group_relays_count_validation() {
+        let storage = MdkMemoryStorage::new(MemoryStorage::default());
+        let mls_group_id = GroupId::from_slice(&[1, 2, 3, 4]);
+
+        // Create a group first
+        let group = create_test_group(mls_group_id.clone(), [1u8; 32]);
+        storage.save_group(group).unwrap();
+
+        // Test with relay count at exactly the limit (should succeed)
+        let mut relays = BTreeSet::new();
+        for i in 0..DEFAULT_MAX_RELAYS_PER_GROUP {
+            relays.insert(RelayUrl::parse(&format!("wss://relay{}.example.com", i)).unwrap());
+        }
+        assert!(storage.replace_group_relays(&mls_group_id, relays).is_ok());
+
+        // Test with relay count exceeding the limit (should fail)
+        let mut relays = BTreeSet::new();
+        for i in 0..=DEFAULT_MAX_RELAYS_PER_GROUP {
+            relays.insert(RelayUrl::parse(&format!("wss://relay{}.example.com", i)).unwrap());
+        }
+        let result = storage.replace_group_relays(&mls_group_id, relays);
+        assert!(result.is_err());
+        assert!(
+            result
+                .unwrap_err()
+                .to_string()
+                .contains("Relay count exceeds maximum")
+        );
+    }
+
+    #[test]
+    fn test_replace_group_relays_url_length_validation() {
+        let storage = MdkMemoryStorage::new(MemoryStorage::default());
+        let mls_group_id = GroupId::from_slice(&[1, 2, 3, 4]);
+
+        // Create a group first
+        let group = create_test_group(mls_group_id.clone(), [1u8; 32]);
+        storage.save_group(group).unwrap();
+
+        // Test with URL at exactly the limit (should succeed)
+        // URL format: wss:// (6) + domain + .com (4) = need domain of DEFAULT_MAX_RELAY_URL_LENGTH - 10
+        let domain = "a".repeat(DEFAULT_MAX_RELAY_URL_LENGTH - 10);
+        let url = format!("wss://{}.com", domain);
+        let relays = BTreeSet::from([RelayUrl::parse(&url).unwrap()]);
+        assert!(storage.replace_group_relays(&mls_group_id, relays).is_ok());
+
+        // Test with URL exceeding the limit (should fail)
+        let domain = "a".repeat(DEFAULT_MAX_RELAY_URL_LENGTH); // This will exceed the limit
+        let url = format!("wss://{}.com", domain);
+        let relays = BTreeSet::from([RelayUrl::parse(&url).unwrap()]);
+        let result = storage.replace_group_relays(&mls_group_id, relays);
+        assert!(result.is_err());
+        assert!(
+            result
+                .unwrap_err()
+                .to_string()
+                .contains("Relay URL exceeds maximum length")
+        );
+    }
+
     #[test]
     fn test_messages_pagination_memory() {
         let storage = MdkMemoryStorage::new(MemoryStorage::default());
@@ -227,7 +428,7 @@ mod tests {
         storage.save_group(group).unwrap();
 
         // Create 25 test messages
-        let pubkey = PublicKey::from_slice(&[1u8; 32]).unwrap();
+        let pubkey = Keys::generate().public_key();
         for i in 0..25 {
             let event_id = EventId::from_slice(&[i as u8; 32]).unwrap();
             let wrapper_event_id = EventId::from_slice(&[100 + i as u8; 32]).unwrap();
@@ -360,6 +561,83 @@ mod tests {
         assert_eq!(result.unwrap().len(), 0); // No results at that offset
     }
 
+    /// Test that custom validation limits work correctly for groups
+    #[test]
+    fn test_custom_group_limits() {
+        use crate::ValidationLimits;
+
+        // Create storage with custom smaller limits
+        let limits = ValidationLimits::default()
+            .with_max_group_name_length(10)
+            .with_max_group_description_length(20)
+            .with_max_admins_per_group(2)
+            .with_max_relays_per_group(3);
+
+        let storage = MdkMemoryStorage::with_limits(MemoryStorage::default(), limits);
+
+        // Verify limits are accessible
+        assert_eq!(storage.limits().max_group_name_length, 10);
+        assert_eq!(storage.limits().max_group_description_length, 20);
+        assert_eq!(storage.limits().max_admins_per_group, 2);
+        assert_eq!(storage.limits().max_relays_per_group, 3);
+
+        let mls_group_id = GroupId::from_slice(&[1, 2, 3, 4]);
+
+        // Test name length with custom limit (10 chars should succeed)
+        let mut group = create_test_group(mls_group_id.clone(), [1u8; 32]);
+        group.name = "a".repeat(10);
+        assert!(storage.save_group(group).is_ok());
+
+        // Test name length exceeding custom limit (11 chars should fail)
+        let mut group = create_test_group(GroupId::from_slice(&[2, 3, 4, 5]), [2u8; 32]);
+        group.name = "a".repeat(11);
+        let result = storage.save_group(group);
+        assert!(result.is_err());
+        assert!(result.unwrap_err().to_string().contains("10 bytes"));
+
+        // Test admin count with custom limit (2 admins should succeed)
+        let mut group = create_test_group(GroupId::from_slice(&[3, 4, 5, 6]), [3u8; 32]);
+        group.admin_pubkeys.insert(Keys::generate().public_key());
+        group.admin_pubkeys.insert(Keys::generate().public_key());
+        assert!(storage.save_group(group).is_ok());
+
+        // Test admin count exceeding custom limit (3 admins should fail)
+        let mut group = create_test_group(GroupId::from_slice(&[4, 5, 6, 7]), [4u8; 32]);
+        for _ in 0..3 {
+            group.admin_pubkeys.insert(Keys::generate().public_key());
+        }
+        let result = storage.save_group(group);
+        assert!(result.is_err());
+        assert!(result.unwrap_err().to_string().contains("maximum of 2"));
+
+        // Test relay count with custom limit
+        let group = create_test_group(GroupId::from_slice(&[5, 6, 7, 8]), [5u8; 32]);
+        storage.save_group(group).unwrap();
+
+        // 3 relays should succeed
+        let relays = BTreeSet::from([
+            RelayUrl::parse("wss://r1.com").unwrap(),
+            RelayUrl::parse("wss://r2.com").unwrap(),
+            RelayUrl::parse("wss://r3.com").unwrap(),
+        ]);
+        assert!(
+            storage
+                .replace_group_relays(&GroupId::from_slice(&[5, 6, 7, 8]), relays)
+                .is_ok()
+        );
+
+        // 4 relays should fail
+        let relays = BTreeSet::from([
+            RelayUrl::parse("wss://r1.com").unwrap(),
+            RelayUrl::parse("wss://r2.com").unwrap(),
+            RelayUrl::parse("wss://r3.com").unwrap(),
+            RelayUrl::parse("wss://r4.com").unwrap(),
+        ]);
+        let result = storage.replace_group_relays(&GroupId::from_slice(&[5, 6, 7, 8]), relays);
+        assert!(result.is_err());
+        assert!(result.unwrap_err().to_string().contains("maximum of 3"));
+    }
+
     #[test]
     fn test_nostr_group_id_collision_rejected() {
         let storage = MdkMemoryStorage::new(MemoryStorage::default());