From 8440ce47237e051a6a6a79238cfc836e7c30b58c Mon Sep 17 00:00:00 2001 From: Kacy Fortner Date: Sat, 14 Feb 2026 12:32:21 -0500 Subject: [PATCH] add collection count validation to persistence deserialization adds MAX_COLLECTION_COUNT (100M) and validate_collection_count() to reject corrupt count fields before entering iteration loops. previously capped_capacity only limited pre-allocation but the loop itself would iterate up to u32::MAX times from a crafted file. covers aof record deserialization, aof rewriter, and snapshot reader for all collection types (list, set, hash, sorted set). --- crates/ember-persistence/src/aof.rs | 9 +++++++++ crates/ember-persistence/src/format.rs | 17 +++++++++++++++++ crates/ember-persistence/src/snapshot.rs | 8 ++++++++ 3 files changed, 34 insertions(+) diff --git a/crates/ember-persistence/src/aof.rs b/crates/ember-persistence/src/aof.rs index 96e67a05..06f9327f 100644 --- a/crates/ember-persistence/src/aof.rs +++ b/crates/ember-persistence/src/aof.rs @@ -42,6 +42,7 @@ fn read_string(r: &mut impl io::Read, field: &str) -> Result Result, FormatError> { let count = format::read_u32(r)?; + format::validate_collection_count(count, field)?; let mut items = Vec::with_capacity(format::capped_capacity(count)); for _ in 0..count { items.push(read_string(r, field)?); @@ -505,6 +506,7 @@ impl AofRecord { TAG_LPUSH | TAG_RPUSH => { let key = read_string(&mut cursor, "key")?; let count = format::read_u32(&mut cursor)?; + format::validate_collection_count(count, "list")?; let mut values = Vec::with_capacity(format::capped_capacity(count)); for _ in 0..count { values.push(Bytes::from(format::read_bytes(&mut cursor)?)); @@ -526,6 +528,7 @@ impl AofRecord { TAG_ZADD => { let key = read_string(&mut cursor, "key")?; let count = format::read_u32(&mut cursor)?; + format::validate_collection_count(count, "sorted set")?; let mut members = Vec::with_capacity(format::capped_capacity(count)); for _ in 0..count { let score = format::read_f64(&mut cursor)?; @@ -559,6 +562,7 @@ impl AofRecord { TAG_HSET => { let key = read_string(&mut cursor, "key")?; let count = format::read_u32(&mut cursor)?; + format::validate_collection_count(count, "hash")?; let mut fields = Vec::with_capacity(format::capped_capacity(count)); for _ in 0..count { let field = read_string(&mut cursor, "field")?; @@ -992,6 +996,7 @@ impl AofReader { let key = format::read_bytes(&mut self.reader)?; format::write_bytes(&mut payload, &key)?; let count = format::read_u32(&mut self.reader)?; + format::validate_collection_count(count, "list")?; format::write_u32(&mut payload, count)?; for _ in 0..count { let val = format::read_bytes(&mut self.reader)?; @@ -1006,6 +1011,7 @@ impl AofReader { let key = format::read_bytes(&mut self.reader)?; format::write_bytes(&mut payload, &key)?; let count = format::read_u32(&mut self.reader)?; + format::validate_collection_count(count, "sorted set")?; format::write_u32(&mut payload, count)?; for _ in 0..count { let score = format::read_f64(&mut self.reader)?; @@ -1018,6 +1024,7 @@ impl AofReader { let key = format::read_bytes(&mut self.reader)?; format::write_bytes(&mut payload, &key)?; let count = format::read_u32(&mut self.reader)?; + format::validate_collection_count(count, "sorted set")?; format::write_u32(&mut payload, count)?; for _ in 0..count { let member = format::read_bytes(&mut self.reader)?; @@ -1042,6 +1049,7 @@ impl AofReader { let key = format::read_bytes(&mut self.reader)?; format::write_bytes(&mut payload, &key)?; let count = format::read_u32(&mut self.reader)?; + format::validate_collection_count(count, "hash")?; format::write_u32(&mut payload, count)?; for _ in 0..count { let field = format::read_bytes(&mut self.reader)?; @@ -1054,6 +1062,7 @@ impl AofReader { let key = format::read_bytes(&mut self.reader)?; format::write_bytes(&mut payload, &key)?; let count = format::read_u32(&mut self.reader)?; + format::validate_collection_count(count, "set")?; format::write_u32(&mut payload, count)?; for _ in 0..count { let item = format::read_bytes(&mut self.reader)?; diff --git a/crates/ember-persistence/src/format.rs b/crates/ember-persistence/src/format.rs index 2f818dbc..dc4ae977 100644 --- a/crates/ember-persistence/src/format.rs +++ b/crates/ember-persistence/src/format.rs @@ -250,6 +250,23 @@ pub fn capped_capacity(count: u32) -> usize { (count as usize).min(65_536) } +/// Maximum element count for collections (lists, sets, hashes, sorted sets) +/// in persistence formats. Prevents corrupt count fields from causing +/// unbounded iteration during deserialization. 100M is well beyond any +/// realistic collection while catching obviously corrupt u32 values. +pub const MAX_COLLECTION_COUNT: u32 = 100_000_000; + +/// Validates that a deserialized collection count is within bounds. +/// Returns `InvalidData` if the count exceeds `MAX_COLLECTION_COUNT`. +pub fn validate_collection_count(count: u32, label: &str) -> Result<(), FormatError> { + if count > MAX_COLLECTION_COUNT { + return Err(FormatError::InvalidData(format!( + "{label} count {count} exceeds max {MAX_COLLECTION_COUNT}" + ))); + } + Ok(()) +} + /// Maximum vector dimensions allowed in persistence formats. /// Matches the protocol-layer cap. Records exceeding this are rejected /// during deserialization to prevent OOM from corrupt files. diff --git a/crates/ember-persistence/src/snapshot.rs b/crates/ember-persistence/src/snapshot.rs index 41157349..a3abd64c 100644 --- a/crates/ember-persistence/src/snapshot.rs +++ b/crates/ember-persistence/src/snapshot.rs @@ -68,6 +68,7 @@ fn parse_snap_value(r: &mut impl io::Read) -> Result { } TYPE_LIST => { let count = format::read_u32(r)?; + format::validate_collection_count(count, "list")?; let mut deque = VecDeque::with_capacity(format::capped_capacity(count)); for _ in 0..count { deque.push_back(Bytes::from(format::read_bytes(r)?)); @@ -76,6 +77,7 @@ fn parse_snap_value(r: &mut impl io::Read) -> Result { } TYPE_SORTED_SET => { let count = format::read_u32(r)?; + format::validate_collection_count(count, "sorted set")?; let mut members = Vec::with_capacity(format::capped_capacity(count)); for _ in 0..count { let score = format::read_f64(r)?; @@ -86,6 +88,7 @@ fn parse_snap_value(r: &mut impl io::Read) -> Result { } TYPE_HASH => { let count = format::read_u32(r)?; + format::validate_collection_count(count, "hash")?; let mut map = HashMap::with_capacity(format::capped_capacity(count)); for _ in 0..count { let field = read_snap_string(r, "hash field")?; @@ -96,6 +99,7 @@ fn parse_snap_value(r: &mut impl io::Read) -> Result { } TYPE_SET => { let count = format::read_u32(r)?; + format::validate_collection_count(count, "set")?; let mut set = HashSet::with_capacity(format::capped_capacity(count)); for _ in 0..count { let member = read_snap_string(r, "set member")?; @@ -584,6 +588,7 @@ impl SnapshotReader { } TYPE_LIST => { let count = format::read_u32(&mut self.reader)?; + format::validate_collection_count(count, "list")?; format::write_u32(&mut buf, count)?; let mut deque = VecDeque::with_capacity(format::capped_capacity(count)); for _ in 0..count { @@ -595,6 +600,7 @@ impl SnapshotReader { } TYPE_SORTED_SET => { let count = format::read_u32(&mut self.reader)?; + format::validate_collection_count(count, "sorted set")?; format::write_u32(&mut buf, count)?; let mut members = Vec::with_capacity(format::capped_capacity(count)); for _ in 0..count { @@ -614,6 +620,7 @@ impl SnapshotReader { } TYPE_HASH => { let count = format::read_u32(&mut self.reader)?; + format::validate_collection_count(count, "hash")?; format::write_u32(&mut buf, count)?; let mut map = HashMap::with_capacity(format::capped_capacity(count)); for _ in 0..count { @@ -633,6 +640,7 @@ impl SnapshotReader { } TYPE_SET => { let count = format::read_u32(&mut self.reader)?; + format::validate_collection_count(count, "set")?; format::write_u32(&mut buf, count)?; let mut set = HashSet::with_capacity(format::capped_capacity(count)); for _ in 0..count {