From dd325c7ab8556e7ff33d64391c78e2316561da55 Mon Sep 17 00:00:00 2001 From: "Dylan M. Taylor" Date: Fri, 23 Oct 2026 10:44:42 -0400 Subject: [PATCH] segment: bound checkpoint tensors and pickle work (#2192) * segment: bound checkpoint tensors and pickle work * segment: bound entry size, memo clones, and total tensor elements * segment: cap cumulative decompressed storage bytes --- crates/segment/src/pt.rs | 139 ++++++++++++++++++++++++++++++++++----- 1 file changed, 123 insertions(+), 16 deletions(-) diff --git a/crates/segment/src/pt.rs b/crates/segment/src/pt.rs index d85ff80c..35ed5729 100644 --- a/crates/segment/src/pt.rs +++ b/crates/segment/src/pt.rs @@ -6,6 +6,15 @@ use std::collections::HashMap; +const MAX_ARCHIVE_BYTES: usize = 1 << 30; +const MAX_PICKLE_OPS: usize = 4_000_000; +const MAX_TENSORS: usize = 1_000_000; +const MAX_TENSOR_ELEMENTS: usize = 1 << 28; +const MAX_ENTRY_BYTES: usize = 1 << 30; +const MAX_PICKLE_NODES: usize = 4_000_000; +const MAX_TOTAL_ELEMENTS: usize = 1 << 29; +const MAX_TENSOR_DIMS: usize = 64; + /// A tensor read from a checkpoint: its shape and its values as f32 (row-major, contiguous). #[derive(Clone, Debug, PartialEq)] pub struct Tensor { @@ -38,6 +47,9 @@ struct Entry { /// The entries of a zip archive by name. fn zip_entries(b: &[u8]) -> Result> { + if b.len() > MAX_ARCHIVE_BYTES { + return Err("checkpoint exceeds the 1024 MiB archive limit".into()); + } // End of central directory: the last `PK\x05\x06` (a comment may follow it). let eocd = (0..b.len().saturating_sub(21)).rev().find(|&i| b.get(i..i + 4) == Some(b"PK\x05\x06")).ok_or("not a zip archive (no end of central directory)")?; @@ -96,6 +108,9 @@ 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 { + 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:?}")) } m => Err(format!("unsupported zip compression {m}")), @@ -142,7 +157,58 @@ fn int_list(o: &Obj) -> Option> { } } +/// Nodes in `o`, stopping once the count passes `budget`. +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(), + _ => return 1, + }; + let mut n = 1; + for k in kids { + if n > budget { + break; + } + n += nodes(k, budget - n); + } + n +} + +fn tensor_element_count(name: &str, offset: usize, size: &[usize], stride: &[usize], values_len: usize) -> Result { + if size.len() != stride.len() || size.len() > MAX_TENSOR_DIMS { + return Err(format!("{name}: invalid tensor dimensionality")); + } + let n = size + .iter() + .try_fold(1usize, |n, &dim| n.checked_mul(dim)) + .filter(|n| *n <= MAX_TENSOR_ELEMENTS) + .ok_or_else(|| format!("{name}: tensor exceeds the element limit"))?; + if n == 0 { + return Ok(0); + } + let mut max_at = offset; + for (&dim, &step) in size.iter().zip(stride) { + max_at = max_at + .checked_add((dim - 1).checked_mul(step).ok_or_else(|| format!("{name}: tensor stride overflow"))?) + .ok_or_else(|| format!("{name}: tensor offset overflow"))?; + } + if max_at >= values_len { + return Err(format!("{name}: data out of range")); + } + Ok(n) +} + /// Run the pickle program and return its result. +/// A memoized object, charged against the clone budget (memo reuse can otherwise grow trees exponentially). +fn memo_get(memo: &HashMap, k: u64, cloned: &mut usize) -> Result { + let o = memo.get(&k).ok_or("pickle: bad memo")?; + *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(); @@ -166,7 +232,10 @@ fn unpickle(p: &[u8]) -> Result { *i += n + 1; Ok(s) }; + let mut op_count = 0usize; + let mut cloned = 0usize; loop { + op_count = op_count.checked_add(1).filter(|n| *n <= MAX_PICKLE_OPS).ok_or("pickle exceeds the opcode limit")?; let op = *p.get(i).ok_or("truncated pickle")?; i += 1; match op { @@ -253,11 +322,11 @@ fn unpickle(p: &[u8]) -> Result { } b'h' => { let k = take(&mut i, 1)?[0] as u64; - stack.push(memo.get(&k).cloned().ok_or("pickle: bad memo")?); + stack.push(memo_get(&memo, k, &mut cloned)?); } b'j' => { let s = take(&mut i, 4)?; - stack.push(memo.get(&(u32::from_le_bytes([s[0], s[1], s[2], s[3]]) as u64)).cloned().ok_or("pickle: bad memo")?); + stack.push(memo_get(&memo, u32::from_le_bytes([s[0], s[1], s[2], s[3]]) as u64, &mut cloned)?); } b't' => { let items = pop_mark(&mut stack)?; @@ -290,9 +359,10 @@ fn unpickle(p: &[u8]) -> Result { stack.push(match f { Obj::Global(m, n) if m == "torch._utils" && n == "_rebuild_tensor_v2" => { match (args.first(), args.get(1), args.get(2).and_then(int_list), args.get(3).and_then(int_list)) { - (Some(Obj::Storage(dtype, key)), Some(Obj::Int(off)), Some(size), Some(stride)) => { - Obj::Tensor { dtype: dtype.clone(), key: key.clone(), offset: usize::try_from(*off).unwrap_or(0), size, stride } - } + (Some(Obj::Storage(dtype, key)), Some(Obj::Int(off)), Some(size), Some(stride)) => match usize::try_from(*off) { + Ok(offset) => Obj::Tensor { dtype: dtype.clone(), key: key.clone(), offset, size, stride }, + Err(_) => Obj::Other, + }, _ => Obj::Other, } } @@ -337,16 +407,24 @@ fn unpickle(p: &[u8]) -> Result { } } -/// Element bytes and conversion of a storage dtype. -fn to_f32(dtype: &str, bytes: &[u8]) -> Option> { +/// Bytes per element of a float storage dtype. +fn elem_size(dtype: &str) -> Option { match dtype { - "FloatStorage" => Some(bytes.as_chunks::<4>().0.iter().map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect()), - "HalfStorage" => Some(bytes.as_chunks::<2>().0.iter().map(|c| half_to_f32(u16::from_le_bytes([c[0], c[1]]))).collect()), - "BFloat16Storage" => Some(bytes.as_chunks::<2>().0.iter().map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16)).collect()), + "FloatStorage" => Some(4), + "HalfStorage" | "BFloat16Storage" => Some(2), _ => None, } } +/// Element `i` of a storage (`i` must be below `bytes.len() / elem_size`). +fn elem_f32(dtype: &str, bytes: &[u8], i: usize) -> f32 { + match dtype { + "FloatStorage" => f32::from_le_bytes([bytes[4 * i], bytes[4 * i + 1], bytes[4 * i + 2], bytes[4 * i + 3]]), + "HalfStorage" => half_to_f32(u16::from_le_bytes([bytes[2 * i], bytes[2 * i + 1]])), + _ => f32::from_bits((u16::from_le_bytes([bytes[2 * i], bytes[2 * i + 1]]) as u32) << 16), + } +} + pub(crate) fn half_to_f32(h: u16) -> f32 { let (s, e, m) = ((h >> 15) as u32, ((h >> 10) & 0x1f) as u32, (h & 0x3ff) as u32); let bits = match (e, m) { @@ -382,22 +460,42 @@ pub fn read(bytes: &[u8]) -> Result { Some((_, Obj::Dict(inner))) => inner.clone(), _ => dict, }; + if dict.len() > MAX_TENSORS { + return Err("checkpoint exceeds the tensor count limit".into()); + } let mut storages: HashMap> = HashMap::new(); let mut out = Weights::new(); + let mut total = 0usize; + let mut inflated = 0usize; for (k, v) in dict { let (Obj::Str(name), Obj::Tensor { dtype, key, offset, size, stride }) = (k, v) else { continue }; if !storages.contains_key(&key) { let e = entries.get(&format!("{prefix}data/{key}")).ok_or_else(|| format!("missing storage {key}"))?; + if e.method == 8 { + inflated = inflated.saturating_add(e.size); + if inflated > MAX_ENTRY_BYTES { + return Err("checkpoint exceeds the decompressed storage limit".into()); + } + } storages.insert(key.clone(), read_entry(bytes, e)?); } - let Some(values) = storages.get(&key).and_then(|b| to_f32(&dtype, b)) else { continue }; - let n: usize = size.iter().product(); - // Gather (any strides) into a contiguous row-major tensor. - let mut data = Vec::with_capacity(n); + let (Some(esz), Some(values)) = (elem_size(&dtype), storages.get(&key)) else { continue }; + // Check hostile stride math before allocating or indexing the backing storage. + let n = tensor_element_count(&name, offset, &size, &stride, values.len() / esz)?; + total = total.saturating_add(n); + if total > MAX_TOTAL_ELEMENTS { + return Err("checkpoint exceeds the total element limit".into()); + } + let mut data = Vec::new(); + data.try_reserve_exact(n).map_err(|_| format!("{name}: tensor allocation failed"))?; let mut idx = vec![0usize; size.len()]; for _ in 0..n { - let at = offset + idx.iter().zip(&stride).map(|(i, s)| i * s).sum::(); - data.push(*values.get(at).ok_or_else(|| format!("{name}: data out of range"))?); + let at = idx + .iter() + .zip(&stride) + .try_fold(offset, |at, (i, s)| i.checked_mul(*s).and_then(|v| at.checked_add(v))) + .ok_or_else(|| format!("{name}: tensor offset overflow"))?; + data.push(elem_f32(&dtype, values, at)); for d in (0..idx.len()).rev() { idx[d] += 1; if idx[d] < size[d] { @@ -490,6 +588,15 @@ mod tests { assert!(!w.contains_key("n"), "integer tensors are skipped"); } + #[test] + fn rejects_tensor_shape_overflow_and_invalid_strides() { + assert!(tensor_element_count("bad", 0, &[usize::MAX, 2], &[1, 1], 16).is_err()); + assert!(tensor_element_count("bad", 0, &[2, 2], &[usize::MAX, 1], 16).is_err()); + assert!(tensor_element_count("bad", 0, &[2, 2], &[1], 16).is_err()); + assert!(tensor_element_count("bad", usize::MAX, &[2], &[1], 16).is_err()); + assert_eq!(tensor_element_count("empty", 0, &[0, usize::MAX], &[0, 0], 0), Ok(0)); + } + #[test] fn damaged_files_are_errors() { assert!(read(b"not a zip").is_err());