diff --git a/crates/segment/src/pt.rs b/crates/segment/src/pt.rs index 35ed57297..d10b32b20 100644 --- a/crates/segment/src/pt.rs +++ b/crates/segment/src/pt.rs @@ -108,10 +108,14 @@ fn read_entry<'a>(b: &'a [u8], e: &Entry) -> Result> match e.method { 0 => Ok(std::borrow::Cow::Borrowed(raw)), 8 => { - if e.size > MAX_ENTRY_BYTES { + // Deflate expands at most about 1032:1, so a stated size past that is not this entry's. + if e.size > MAX_ENTRY_BYTES || e.size > e.stored.saturating_mul(1032).saturating_add(64) { return Err("zip entry exceeds the size limit".into()); } - miniz_oxide::inflate::decompress_to_vec_with_limit(raw, e.size.max(1)).map(std::borrow::Cow::Owned).map_err(|e| format!("bad deflate data: {e:?}")) + // The status only: the error's partial output can be as large as the entry. + miniz_oxide::inflate::decompress_to_vec_with_limit(raw, e.size.max(1)) + .map(std::borrow::Cow::Owned) + .map_err(|e| format!("bad deflate data: {:?}", e.status)) } m => Err(format!("unsupported zip compression {m}")), } @@ -162,6 +166,9 @@ fn nodes(o: &Obj, budget: usize) -> usize { let kids: Vec<&Obj> = match o { Obj::Tuple(v) | Obj::List(v) => v.iter().collect(), Obj::Dict(v) => v.iter().flat_map(|(k, x)| [k, x]).collect(), + // A string counts one node per 64 bytes it holds: copies of a long one are charged for + // what they keep, not as one object. + Obj::Str(s) | Obj::Global(s, _) | Obj::Storage(s, _) => return 1 + s.len() / 64, _ => return 1, }; let mut n = 1; @@ -209,6 +216,17 @@ fn memo_get(memo: &HashMap, k: u64, cloned: &mut usize) -> Result Ok(o.clone()) } +/// A copy of the stack top to memoize, charged against the clone budget like a memo read: a +/// program memoizing one large object many times keeps a copy each time. +fn memo_put(stack: &[Obj], cloned: &mut usize) -> Result { + let o = stack.last().ok_or("pickle stack underflow")?; + *cloned = cloned.saturating_add(nodes(o, MAX_PICKLE_NODES)); + if *cloned > MAX_PICKLE_NODES { + return Err("pickle exceeds the object limit".into()); + } + Ok(o.clone()) +} + fn unpickle(p: &[u8]) -> Result { let mut stack: Vec = Vec::new(); let mut memo: HashMap = HashMap::new(); @@ -310,15 +328,15 @@ fn unpickle(p: &[u8]) -> Result { } b'q' => { let k = take(&mut i, 1)?[0] as u64; - memo.insert(k, stack.last().cloned().ok_or("pickle stack underflow")?); + memo.insert(k, memo_put(&stack, &mut cloned)?); } b'r' => { let s = take(&mut i, 4)?; - memo.insert(u32::from_le_bytes([s[0], s[1], s[2], s[3]]) as u64, stack.last().cloned().ok_or("pickle stack underflow")?); + memo.insert(u32::from_le_bytes([s[0], s[1], s[2], s[3]]) as u64, memo_put(&stack, &mut cloned)?); } 0x94 => { let k = memo.len() as u64; - memo.insert(k, stack.last().cloned().ok_or("pickle stack underflow")?); + memo.insert(k, memo_put(&stack, &mut cloned)?); } b'h' => { let k = take(&mut i, 1)?[0] as u64; @@ -609,4 +627,21 @@ mod tests { assert_eq!(half_to_f32(0x3c00), 1.0); assert_eq!(half_to_f32(0xc000), -2.0); } + + #[test] + fn memoizing_a_long_string_again_and_again_runs_out_of_the_budget() { + // A 1 MiB string memoized 300 times would keep 300 MiB of copies. + let mut p = vec![0x80, 2, b'X']; + p.extend_from_slice(&(1u32 << 20).to_le_bytes()); + p.extend(std::iter::repeat_n(b'a', 1 << 20)); + p.extend(std::iter::repeat_n(0x94, 300)); + p.push(b'.'); + assert!(unpickle(&p).is_err_and(|e| e.contains("object limit"))); + } + + #[test] + fn an_entry_stating_more_than_deflate_can_expand_to_is_refused() { + let e = Entry { data: 0, size: 1 << 29, stored: 4, method: 8 }; + assert!(read_entry(&[0u8; 4], &e).is_err_and(|m| m.contains("size limit"))); + } }