Skip to content

harden: collection count validation in persistence deserialization #124

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kacy merged 1 commit into from
Feb 14, 2026
Merged

harden: collection count validation in persistence deserialization #124

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions crates/ember-persistence/src/aof.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ fn read_string(r: &mut impl io::Read, field: &str) -> Result<String, FormatError
/// Used by SADD, SREM, HDEL, and ZREM deserialization.
fn read_string_list(r: &mut impl io::Read, field: &str) -> Result<Vec<String>, 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)?);
Expand Down Expand Up @@ -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)?));
Expand All @@ -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)?;
Expand Down Expand Up @@ -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")?;
Expand Down Expand Up @@ -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)?;
Expand All @@ -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)?;
Expand All @@ -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)?;
Expand All @@ -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)?;
Expand All @@ -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)?;
Expand Down
17 changes: 17 additions & 0 deletions crates/ember-persistence/src/format.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
8 changes: 8 additions & 0 deletions crates/ember-persistence/src/snapshot.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,7 @@ fn parse_snap_value(r: &mut impl io::Read) -> Result<SnapValue, FormatError> {
}
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)?));
Expand All @@ -76,6 +77,7 @@ fn parse_snap_value(r: &mut impl io::Read) -> Result<SnapValue, FormatError> {
}
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)?;
Expand All @@ -86,6 +88,7 @@ fn parse_snap_value(r: &mut impl io::Read) -> Result<SnapValue, FormatError> {
}
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")?;
Expand All @@ -96,6 +99,7 @@ fn parse_snap_value(r: &mut impl io::Read) -> Result<SnapValue, FormatError> {
}
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")?;
Expand Down Expand Up @@ -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 {
Expand All @@ -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 {
Expand All @@ -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 {
Expand All @@ -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 {
Expand Down