fix(format): bound what a VL read retains on a crafted global heap
VlResolver kept an owned copy of every object of every heap collection it parsed, for the whole read. Collections nested inside each other's object data, 32 bytes apart with each element pointing at a different one, made retained memory O(elements x file size): 1.58 GB for a 744 KB file (read_vl_strings did the same before VlResolver). Chaining every collection's objects into one shared run of tiny objects made parse time O(elements x objects) as well. libhdf5 refuses these files. - The cache records where each object lies (GlobalHeapCollection:: parse_index, new) instead of copying it, and is dropped past a 32 MiB budget. - A collection overlapping one already read is an error: libhdf5 gives every collection its own block, so only a crafted file has them. - parse and parse_index refuse a collection that runs past the end of the file and an object that runs past the end of its collection. tests/vl_heap_bounds.rs measures peak heap use with a counting allocator: 129 MB and 350 MB live before on its two crafted files (64 KB and 176 KB), 97 KB and 0.9 MB now. Conformance unchanged at 575 of 697. Co-Authored-By: Claude Opus 5.5 (1M context) <[email protected]>
This commit is contained in:
osobh
@@ -1,7 +1,7 @@
|
||||
//! HDF5 Global Heap collection parsing.
|
||||
|
||||
#[cfg(not(feature = "std"))]
|
||||
use alloc::vec::Vec;
|
||||
use alloc::{format, string::String, vec::Vec};
|
||||
|
||||
use crate::error::FormatError;
|
||||
|
||||
@@ -52,11 +52,42 @@ fn read_length(data: &[u8], offset: usize, length_size: u8) -> Result<u64, Forma
|
||||
})
|
||||
}
|
||||
|
||||
fn object_overrun_msg(index: u16, size: usize, collection_size: u64) -> String {
|
||||
format!(
|
||||
"global heap object {index} ({size} bytes) runs past the end of its \
|
||||
{collection_size}-byte collection"
|
||||
)
|
||||
}
|
||||
|
||||
/// Round up to next multiple of 8.
|
||||
fn pad8(x: usize) -> usize {
|
||||
(x + 7) & !7
|
||||
}
|
||||
|
||||
/// Where one object of a global heap collection lies in the file, without
|
||||
/// its data: see [`GlobalHeapCollection::parse_index`].
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub struct GlobalHeapObjectRef {
|
||||
/// Object index (1-based; 0 is the free space marker).
|
||||
pub index: u16,
|
||||
/// Reference count.
|
||||
pub reference_count: u16,
|
||||
/// Offset of the object's data in the file data the collection was
|
||||
/// parsed from.
|
||||
pub offset: usize,
|
||||
/// Size of the object's data in bytes.
|
||||
pub size: usize,
|
||||
}
|
||||
|
||||
/// A global heap collection's objects, located but not copied.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct GlobalHeapIndex {
|
||||
/// Total size of this collection including header.
|
||||
pub collection_size: u64,
|
||||
/// The objects, in file order.
|
||||
pub objects: Vec<GlobalHeapObjectRef>,
|
||||
}
|
||||
|
||||
impl GlobalHeapCollection {
|
||||
/// Parse a global heap collection at the given offset in the file data.
|
||||
pub fn parse(
|
||||
@@ -64,6 +95,33 @@ impl GlobalHeapCollection {
|
||||
offset: usize,
|
||||
length_size: u8,
|
||||
) -> Result<GlobalHeapCollection, FormatError> {
|
||||
let index = Self::parse_index(file_data, offset, length_size)?;
|
||||
Ok(GlobalHeapCollection {
|
||||
collection_size: index.collection_size,
|
||||
objects: index
|
||||
.objects
|
||||
.iter()
|
||||
.map(|o| GlobalHeapObject {
|
||||
index: o.index,
|
||||
reference_count: o.reference_count,
|
||||
data: file_data[o.offset..o.offset + o.size].to_vec(),
|
||||
})
|
||||
.collect(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Locate the objects of the global heap collection at `offset` without
|
||||
/// copying their data, so a caller can keep many collections indexed
|
||||
/// for the cost of their object headers.
|
||||
///
|
||||
/// The collection must lie inside `file_data`, and every object inside
|
||||
/// the collection, as libhdf5 lays them out; an object that runs past
|
||||
/// its collection is an error.
|
||||
pub fn parse_index(
|
||||
file_data: &[u8],
|
||||
offset: usize,
|
||||
length_size: u8,
|
||||
) -> Result<GlobalHeapIndex, FormatError> {
|
||||
// signature(4) + version(1) + reserved(3) + collection_size(length_size),
|
||||
// padded to a multiple of 8 as libhdf5 lays it out (`H5HG_SIZEOF_HDR`).
|
||||
// With 8-byte lengths the padding is 0; with 4-byte lengths it is 4,
|
||||
@@ -81,25 +139,25 @@ impl GlobalHeapCollection {
|
||||
}
|
||||
|
||||
let collection_size = read_length(file_data, offset + 8, length_size)?;
|
||||
let collection_size_usize =
|
||||
usize::try_from(collection_size).map_err(|_| FormatError::UnexpectedEof {
|
||||
expected: u64::MAX as usize,
|
||||
let collection_end = usize::try_from(collection_size)
|
||||
.ok()
|
||||
.and_then(|size| offset.checked_add(size))
|
||||
.ok_or(FormatError::UnexpectedEof {
|
||||
expected: usize::MAX,
|
||||
available: file_data.len(),
|
||||
})?;
|
||||
let collection_end =
|
||||
offset
|
||||
.checked_add(collection_size_usize)
|
||||
.ok_or(FormatError::UnexpectedEof {
|
||||
expected: usize::MAX,
|
||||
available: file_data.len(),
|
||||
})?;
|
||||
if collection_end > file_data.len() {
|
||||
return Err(FormatError::UnexpectedEof {
|
||||
expected: collection_end,
|
||||
available: file_data.len(),
|
||||
});
|
||||
}
|
||||
|
||||
let mut pos = offset + header_size;
|
||||
let mut objects = Vec::new();
|
||||
|
||||
// Parse objects until we hit index 0 (free space) or run out of space
|
||||
while pos + 2 <= collection_end {
|
||||
ensure_len(file_data, pos, 2)?;
|
||||
let object_index = u16::from_le_bytes([file_data[pos], file_data[pos + 1]]);
|
||||
|
||||
if object_index == 0 {
|
||||
@@ -110,26 +168,36 @@ impl GlobalHeapCollection {
|
||||
// object_index(2) + reference_count(2) + reserved(4) +
|
||||
// object_size(length_size), padded to 8 (`H5HG_SIZEOF_OBJHDR`).
|
||||
let obj_header_size = pad8(8 + length_size as usize);
|
||||
ensure_len(file_data, pos, obj_header_size)?;
|
||||
ensure_len(&file_data[..collection_end], pos, obj_header_size)?;
|
||||
|
||||
let reference_count = u16::from_le_bytes([file_data[pos + 2], file_data[pos + 3]]);
|
||||
let object_size = read_length(file_data, pos + 8, length_size)? as usize;
|
||||
let object_size = usize::try_from(read_length(file_data, pos + 8, length_size)?)
|
||||
.map_err(|_| FormatError::Overflow("global heap object size".into()))?;
|
||||
|
||||
pos += obj_header_size;
|
||||
ensure_len(file_data, pos, object_size)?;
|
||||
let data = file_data[pos..pos + object_size].to_vec();
|
||||
if pos
|
||||
.checked_add(object_size)
|
||||
.is_none_or(|end| end > collection_end)
|
||||
{
|
||||
return Err(FormatError::VlDataError(object_overrun_msg(
|
||||
object_index,
|
||||
object_size,
|
||||
collection_size,
|
||||
)));
|
||||
}
|
||||
|
||||
objects.push(GlobalHeapObject {
|
||||
objects.push(GlobalHeapObjectRef {
|
||||
index: object_index,
|
||||
reference_count,
|
||||
data,
|
||||
offset: pos,
|
||||
size: object_size,
|
||||
});
|
||||
|
||||
// Advance past data + padding to 8-byte boundary
|
||||
pos += pad8(object_size);
|
||||
pos = pos.saturating_add(pad8(object_size));
|
||||
}
|
||||
|
||||
Ok(GlobalHeapCollection {
|
||||
Ok(GlobalHeapIndex {
|
||||
collection_size,
|
||||
objects,
|
||||
})
|
||||
|
||||
@@ -10,7 +10,7 @@ use alloc::{collections::BTreeMap, format, string::String, vec, vec::Vec};
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
use crate::error::FormatError;
|
||||
use crate::global_heap::GlobalHeapCollection;
|
||||
use crate::global_heap::{GlobalHeapCollection, GlobalHeapIndex};
|
||||
|
||||
/// A parsed variable-length element reference (global heap ID).
|
||||
#[derive(Debug, Clone)]
|
||||
@@ -134,38 +134,43 @@ pub fn check_element_size(stored_size: u32, offset_size: u8) -> Result<(), Forma
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// A parsed collection, with its objects indexed for lookup.
|
||||
/// A collection's objects, located in the file data but not copied:
|
||||
/// `(index, offset, size)` of the first object with each index, sorted by
|
||||
/// index.
|
||||
struct CachedCollection {
|
||||
collection: GlobalHeapCollection,
|
||||
/// `slots[index]` is the position in `collection.objects` of the first
|
||||
/// object with that index.
|
||||
slots: Vec<Option<usize>>,
|
||||
objects: Vec<(u16, usize, usize)>,
|
||||
}
|
||||
|
||||
impl CachedCollection {
|
||||
fn new(collection: GlobalHeapCollection) -> Self {
|
||||
let max = collection
|
||||
fn new(index: GlobalHeapIndex) -> Self {
|
||||
let mut objects: Vec<(u16, usize, usize)> = index
|
||||
.objects
|
||||
.iter()
|
||||
.map(|o| o.index as usize)
|
||||
.max()
|
||||
.unwrap_or(0);
|
||||
let mut slots = vec![None; max + 1];
|
||||
for (pos, obj) in collection.objects.iter().enumerate() {
|
||||
let slot = &mut slots[obj.index as usize];
|
||||
if slot.is_none() {
|
||||
*slot = Some(pos);
|
||||
}
|
||||
}
|
||||
Self { collection, slots }
|
||||
.map(|o| (o.index, o.offset, o.size))
|
||||
.collect();
|
||||
// Stable, so the first object with a repeated index is kept.
|
||||
objects.sort_by_key(|o| o.0);
|
||||
objects.dedup_by_key(|o| o.0);
|
||||
Self { objects }
|
||||
}
|
||||
|
||||
fn get(&self, index: u32) -> Option<&[u8]> {
|
||||
let pos = (*self.slots.get(usize::try_from(index).ok()?)?)?;
|
||||
Some(&self.collection.objects[pos].data)
|
||||
/// What this entry costs to keep, in bytes (roughly).
|
||||
fn cost(&self) -> usize {
|
||||
64 + self.objects.len() * core::mem::size_of::<(u16, usize, usize)>()
|
||||
}
|
||||
|
||||
fn get(&self, index: u32) -> Option<(usize, usize)> {
|
||||
let index = u16::try_from(index).ok()?;
|
||||
let i = self.objects.binary_search_by_key(&index, |o| o.0).ok()?;
|
||||
Some((self.objects[i].1, self.objects[i].2))
|
||||
}
|
||||
}
|
||||
|
||||
/// How many bytes of collection indexes a [`VlResolver`] keeps before it
|
||||
/// drops them and starts again. Values are never copied into the cache, so
|
||||
/// this bounds what a read retains however many collections it visits.
|
||||
const CACHE_BUDGET: usize = 32 << 20;
|
||||
|
||||
/// Resolves variable-length elements against a file's global heap, parsing
|
||||
/// each heap collection once however many elements point into it.
|
||||
///
|
||||
@@ -173,11 +178,22 @@ impl CachedCollection {
|
||||
/// empty string or sequence), and an element whose heap object is not
|
||||
/// exactly `length × base size` bytes is an error ("Expected global heap
|
||||
/// object size does not match"), not a truncated or padded value.
|
||||
///
|
||||
/// Memory stays bounded on hostile files: the cache holds where each
|
||||
/// object lies, not a copy of it, up to a fixed budget; and collections
|
||||
/// that overlap one another are refused (libhdf5 never writes them), so a
|
||||
/// file cannot make the resolver parse the same bytes as the objects of
|
||||
/// many collections.
|
||||
pub struct VlResolver<'a> {
|
||||
file_data: &'a [u8],
|
||||
offset_size: u8,
|
||||
length_size: u8,
|
||||
cache: BTreeMap<u64, CachedCollection>,
|
||||
cached_bytes: usize,
|
||||
budget: usize,
|
||||
/// Start → end of every collection parsed so far (kept when the cache
|
||||
/// is dropped, to check overlaps).
|
||||
extents: BTreeMap<usize, usize>,
|
||||
}
|
||||
|
||||
impl<'a> VlResolver<'a> {
|
||||
@@ -189,6 +205,9 @@ impl<'a> VlResolver<'a> {
|
||||
offset_size,
|
||||
length_size,
|
||||
cache: BTreeMap::new(),
|
||||
cached_bytes: 0,
|
||||
budget: CACHE_BUDGET,
|
||||
extents: BTreeMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -210,11 +229,18 @@ impl<'a> VlResolver<'a> {
|
||||
}
|
||||
|
||||
/// The bytes of one element: `length × base_size` bytes from the heap,
|
||||
/// or empty for a null or zero-length element.
|
||||
fn resolve(&mut self, vl: &VlElement, base_size: usize) -> Result<&[u8], FormatError> {
|
||||
/// or `None` for a null element.
|
||||
fn resolve(
|
||||
&mut self,
|
||||
vl: &VlElement,
|
||||
base_size: usize,
|
||||
) -> Result<Option<&'a [u8]>, FormatError> {
|
||||
let addr = vl.collection_address;
|
||||
if addr == 0 || (vl.length == 0 && is_undefined_address(addr, self.offset_size)) {
|
||||
return Ok(&[]);
|
||||
if addr == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
if vl.length == 0 && is_undefined_address(addr, self.offset_size) {
|
||||
return Ok(Some(&[]));
|
||||
}
|
||||
let data = self.object(vl)?;
|
||||
let expected = (vl.length as usize)
|
||||
@@ -229,7 +255,27 @@ impl<'a> VlResolver<'a> {
|
||||
vl.length
|
||||
)));
|
||||
}
|
||||
Ok(data)
|
||||
Ok(Some(data))
|
||||
}
|
||||
|
||||
/// One element (the first [`element_size`](Self::element_size) bytes of
|
||||
/// `elem`) of a variable-length sequence whose base type is `base_size`
|
||||
/// bytes: its `length × base_size` bytes, or `None` for a null element
|
||||
/// (heap address 0).
|
||||
pub fn element(
|
||||
&mut self,
|
||||
elem: &[u8],
|
||||
base_size: usize,
|
||||
) -> Result<Option<&'a [u8]>, FormatError> {
|
||||
let vl = parse_vl_references(elem, 1, self.offset_size)?;
|
||||
self.resolve(&vl[0], base_size)
|
||||
}
|
||||
|
||||
/// One variable-length string element: its bytes up to the first NUL,
|
||||
/// or `None` for a null element (h5dump prints it as `NULL`, h5py
|
||||
/// returns it as empty).
|
||||
pub fn string_element(&mut self, elem: &[u8]) -> Result<Option<&'a [u8]>, FormatError> {
|
||||
Ok(self.element(elem, 1)?.map(cut_at_nul))
|
||||
}
|
||||
|
||||
/// The strings of the variable-length string elements in `raw`, as
|
||||
@@ -238,11 +284,7 @@ impl<'a> VlResolver<'a> {
|
||||
pub fn string_bytes(&mut self, raw: &[u8]) -> Result<Vec<Vec<u8>>, FormatError> {
|
||||
self.elements(raw)?
|
||||
.iter()
|
||||
.map(|vl| {
|
||||
let s = self.resolve(vl, 1)?;
|
||||
let end = s.iter().position(|&b| b == 0).unwrap_or(s.len());
|
||||
Ok(s[..end].to_vec())
|
||||
})
|
||||
.map(|vl| Ok(self.resolve(vl, 1)?.map(cut_at_nul).unwrap_or(&[]).to_vec()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
@@ -270,11 +312,16 @@ impl<'a> VlResolver<'a> {
|
||||
}
|
||||
self.elements(raw)?
|
||||
.iter()
|
||||
.map(|vl| self.resolve(vl, base_size).map(<[u8]>::to_vec))
|
||||
.map(|vl| Ok(self.resolve(vl, base_size)?.unwrap_or(&[]).to_vec()))
|
||||
.collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// A string's bytes up to its first NUL.
|
||||
fn cut_at_nul(s: &[u8]) -> &[u8] {
|
||||
&s[..s.iter().position(|&b| b == 0).unwrap_or(s.len())]
|
||||
}
|
||||
|
||||
/// Resolve VL strings from raw data by looking up each element in the global heap.
|
||||
///
|
||||
/// Reads the first `num_elements` elements of `raw`. Strings end at their
|
||||
@@ -342,25 +389,66 @@ pub fn read_vl_bytes(
|
||||
Ok(result)
|
||||
}
|
||||
|
||||
impl VlResolver<'_> {
|
||||
impl<'a> VlResolver<'a> {
|
||||
/// The heap object `vl` points to, whatever its size; its collection is
|
||||
/// parsed on first use.
|
||||
fn object(&mut self, vl: &VlElement) -> Result<&[u8], FormatError> {
|
||||
fn object(&mut self, vl: &VlElement) -> Result<&'a [u8], FormatError> {
|
||||
let addr = vl.collection_address;
|
||||
if !self.cache.contains_key(&addr) {
|
||||
let offset = usize::try_from(addr).map_err(|_| FormatError::UnexpectedEof {
|
||||
expected: usize::MAX,
|
||||
available: self.file_data.len(),
|
||||
})?;
|
||||
let coll = GlobalHeapCollection::parse(self.file_data, offset, self.length_size)?;
|
||||
self.cache.insert(addr, CachedCollection::new(coll));
|
||||
let index =
|
||||
GlobalHeapCollection::parse_index(self.file_data, offset, self.length_size)?;
|
||||
// parse_index checked that the collection lies in the file.
|
||||
let end = offset + index.collection_size as usize;
|
||||
self.check_overlap(offset, end)?;
|
||||
let coll = CachedCollection::new(index);
|
||||
if self.cached_bytes.saturating_add(coll.cost()) > self.budget {
|
||||
self.cache.clear();
|
||||
self.cached_bytes = 0;
|
||||
}
|
||||
self.cached_bytes += coll.cost();
|
||||
self.cache.insert(addr, coll);
|
||||
}
|
||||
self.cache[&addr]
|
||||
.get(vl.object_index)
|
||||
.ok_or(FormatError::GlobalHeapObjectNotFound {
|
||||
let (start, size) = self.cache[&addr].get(vl.object_index).ok_or(
|
||||
FormatError::GlobalHeapObjectNotFound {
|
||||
collection_address: addr,
|
||||
index: vl.object_index as u16,
|
||||
})
|
||||
},
|
||||
)?;
|
||||
Ok(&self.file_data[start..start + size])
|
||||
}
|
||||
|
||||
/// Record the collection at `start..end`, refusing one that overlaps a
|
||||
/// collection already read. libhdf5 allocates each collection its own
|
||||
/// block; overlapping ones only come from a crafted file, where they let
|
||||
/// every byte be parsed again as the objects of each collection.
|
||||
fn check_overlap(&mut self, start: usize, end: usize) -> Result<(), FormatError> {
|
||||
if let Some(&known) = self.extents.get(&start) {
|
||||
return if known == end {
|
||||
Ok(())
|
||||
} else {
|
||||
Err(FormatError::VlDataError(format!(
|
||||
"global heap collection at {start} changed size"
|
||||
)))
|
||||
};
|
||||
}
|
||||
let before = self.extents.range(..start).next_back();
|
||||
let after = self.extents.range(start..).next();
|
||||
let clash = match (before, after) {
|
||||
(Some((&s, &e)), _) if e > start => Some(s),
|
||||
(_, Some((&s, _))) if s < end => Some(s),
|
||||
_ => None,
|
||||
};
|
||||
if let Some(other) = clash {
|
||||
return Err(FormatError::VlDataError(format!(
|
||||
"global heap collection at {start} overlaps the one at {other}"
|
||||
)));
|
||||
}
|
||||
self.extents.insert(start, end);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -581,6 +669,43 @@ mod tests {
|
||||
assert!(r.strings(&raw[..30]).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn the_cache_stays_within_its_budget_and_rereads_what_it_dropped() {
|
||||
// Twenty collections of three objects each; a budget that holds
|
||||
// about two of them. Reading every element twice must still return
|
||||
// the right strings after the cache is dropped.
|
||||
let mut file_data = vec![0u8; 64];
|
||||
let mut raw = Vec::new();
|
||||
for c in 0..20u64 {
|
||||
let at = file_data.len();
|
||||
let names: Vec<String> = (0..3).map(|i| format!("c{c}o{i}")).collect();
|
||||
let objs: Vec<(u16, &[u8])> = names
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, n)| (i as u16 + 1, n.as_bytes()))
|
||||
.collect();
|
||||
build_gcol_at(&mut file_data, at, &objs);
|
||||
for (i, n) in names.iter().enumerate() {
|
||||
raw.extend(element(n.len() as u32, at as u64, i as u32 + 1, 8));
|
||||
}
|
||||
}
|
||||
raw.extend(raw.clone());
|
||||
let mut r = VlResolver::new(&file_data, 8, 8);
|
||||
let one = CachedCollection {
|
||||
objects: vec![(0, 0, 0); 3],
|
||||
}
|
||||
.cost();
|
||||
r.budget = 2 * one + 1;
|
||||
let want: Vec<String> = (0..2)
|
||||
.flat_map(|_| (0..20).flat_map(|c| (0..3).map(move |i| format!("c{c}o{i}"))))
|
||||
.collect();
|
||||
for (k, chunk) in raw.chunks(16).enumerate() {
|
||||
assert_eq!(r.strings(chunk).unwrap(), [want[k].clone()]);
|
||||
assert!(r.cached_bytes <= r.budget);
|
||||
assert!(r.cache.len() <= 2);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn element_size_is_checked_against_the_offset_size() {
|
||||
assert!(check_element_size(16, 8).is_ok());
|
||||
|
||||
@@ -0,0 +1,161 @@
|
||||
//! Crafted files cannot make variable-length reads retain memory, or take
|
||||
//! time, out of proportion to the file.
|
||||
//!
|
||||
//! `VlResolver` used to keep an owned copy of every object of every
|
||||
//! collection it parsed, for the whole read. A file whose global heap
|
||||
//! collections nest inside each other's object data — each element
|
||||
//! pointing at a different one — then made retained memory O(K × file
|
||||
//! size): a 744 KB file took 1.58 GB. The same nesting, with every
|
||||
//! collection's object chain jumping to one shared run of tiny objects,
|
||||
//! made the parse time O(K × M) as well. libhdf5 never writes overlapping
|
||||
//! collections; they are now refused, and the cache holds only where
|
||||
//! objects lie.
|
||||
//!
|
||||
//! Peak heap use is measured with a counting global allocator, so the
|
||||
//! cases run one after another in a single test.
|
||||
|
||||
use std::alloc::{GlobalAlloc, Layout, System};
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
use clawhdf5_format::vl_data::VlResolver;
|
||||
|
||||
struct Counting;
|
||||
|
||||
static CURRENT: AtomicUsize = AtomicUsize::new(0);
|
||||
static PEAK: AtomicUsize = AtomicUsize::new(0);
|
||||
|
||||
unsafe impl GlobalAlloc for Counting {
|
||||
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
|
||||
let p = unsafe { System.alloc(layout) };
|
||||
if !p.is_null() {
|
||||
let now = CURRENT.fetch_add(layout.size(), Ordering::Relaxed) + layout.size();
|
||||
PEAK.fetch_max(now, Ordering::Relaxed);
|
||||
}
|
||||
p
|
||||
}
|
||||
|
||||
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
|
||||
unsafe { System.dealloc(ptr, layout) };
|
||||
CURRENT.fetch_sub(layout.size(), Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
#[global_allocator]
|
||||
static ALLOC: Counting = Counting;
|
||||
|
||||
/// Bytes allocated at the peak of `f`, above what was live when it started.
|
||||
fn peak_during<T>(f: impl FnOnce() -> T) -> (T, usize) {
|
||||
let base = CURRENT.load(Ordering::Relaxed);
|
||||
PEAK.store(base, Ordering::Relaxed);
|
||||
let out = f();
|
||||
(out, PEAK.load(Ordering::Relaxed) - base)
|
||||
}
|
||||
|
||||
fn put_header(file: &mut [u8], at: usize, size: u64) {
|
||||
file[at..at + 4].copy_from_slice(b"GCOL");
|
||||
file[at + 4] = 1;
|
||||
file[at + 8..at + 16].copy_from_slice(&size.to_le_bytes());
|
||||
}
|
||||
|
||||
fn put_object(file: &mut [u8], at: usize, index: u16, size: u64) {
|
||||
file[at..at + 2].copy_from_slice(&index.to_le_bytes());
|
||||
file[at + 2..at + 4].copy_from_slice(&1u16.to_le_bytes());
|
||||
file[at + 8..at + 16].copy_from_slice(&size.to_le_bytes());
|
||||
}
|
||||
|
||||
fn element(length: u32, addr: u64, index: u32) -> Vec<u8> {
|
||||
let mut e = length.to_le_bytes().to_vec();
|
||||
e.extend_from_slice(&addr.to_le_bytes());
|
||||
e.extend_from_slice(&index.to_le_bytes());
|
||||
e
|
||||
}
|
||||
|
||||
/// K collections 32 bytes apart, each running to the end of the file with
|
||||
/// one object covering the rest of it (and so every later collection).
|
||||
/// Element i is that object of collection i.
|
||||
fn nested(k: usize) -> (Vec<u8>, Vec<u8>) {
|
||||
let base = 64;
|
||||
let end = base + 32 * k + 64;
|
||||
let mut file = vec![0u8; end];
|
||||
let mut raw = Vec::new();
|
||||
for i in 0..k {
|
||||
let at = base + 32 * i;
|
||||
put_header(&mut file, at, (end - at) as u64);
|
||||
let obj = (end - at - 32) as u64;
|
||||
put_object(&mut file, at + 16, 1, obj);
|
||||
raw.extend(element(obj as u32, at as u64, 1));
|
||||
}
|
||||
(file, raw)
|
||||
}
|
||||
|
||||
/// K collections 32 bytes apart, each with a first object that jumps over
|
||||
/// the later collections to one shared run of M empty objects, so parsing
|
||||
/// every collection walks all M.
|
||||
fn shared_tail(k: usize, m: usize) -> (Vec<u8>, Vec<u8>) {
|
||||
let base = 64;
|
||||
let tail = base + 32 * k + 32;
|
||||
let end = tail + 16 * m + 16;
|
||||
let mut file = vec![0u8; end];
|
||||
let mut raw = Vec::new();
|
||||
for i in 0..k {
|
||||
let at = base + 32 * i;
|
||||
put_header(&mut file, at, (end - at) as u64);
|
||||
let jump = (tail - at - 32) as u64;
|
||||
put_object(&mut file, at + 16, 1, jump);
|
||||
raw.extend(element(jump as u32, at as u64, 1));
|
||||
}
|
||||
for j in 0..m {
|
||||
put_object(&mut file, tail + 16 * j, (j % 65_000 + 2) as u16, 0);
|
||||
}
|
||||
(file, raw)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn overlapping_collections_are_refused_in_bounded_memory_and_time() {
|
||||
for (name, (file, raw)) in [
|
||||
("nested", nested(2000)),
|
||||
("shared tail", shared_tail(500, 10_000)),
|
||||
] {
|
||||
let start = Instant::now();
|
||||
let (result, peak) = peak_during(|| {
|
||||
let mut r = VlResolver::new(&file, 8, 8);
|
||||
(r.string_bytes(&raw), r.sequences(&raw, 1).map(|s| s.len()))
|
||||
});
|
||||
let took = start.elapsed();
|
||||
// libhdf5 never writes overlapping collections, and refuses these
|
||||
// files; so do we, rather than returning what they claim.
|
||||
let (strings, sequences) = result;
|
||||
let e = strings.expect_err(name).to_string();
|
||||
assert!(e.contains("overlaps"), "{name}: {e}");
|
||||
assert!(sequences.is_err(), "{name}");
|
||||
// Measured before the fix: 129 MB ("nested", 64 KB file) and 350 MB
|
||||
// ("shared tail", 176 KB file) live at the peak; after, 97 KB and
|
||||
// 0.9 MB.
|
||||
assert!(
|
||||
peak < 4 * file.len() + (1 << 20),
|
||||
"{name}: peak {peak} bytes for a {}-byte file",
|
||||
file.len()
|
||||
);
|
||||
assert!(took < Duration::from_secs(5), "{name}: took {took:?}");
|
||||
}
|
||||
}
|
||||
|
||||
/// Collections that do not overlap still read, however many elements point
|
||||
/// into them, and the first object of a collection is returned for its
|
||||
/// index (as before).
|
||||
#[test]
|
||||
fn separate_collections_still_read() {
|
||||
let mut file = vec![0u8; 64 + 3 * 64];
|
||||
let mut raw = Vec::new();
|
||||
for i in 0..3usize {
|
||||
let at = 64 + 64 * i;
|
||||
put_header(&mut file, at, 64);
|
||||
put_object(&mut file, at + 16, 1, 3);
|
||||
file[at + 32..at + 35].copy_from_slice(format!("s{i}!").as_bytes());
|
||||
raw.extend(element(3, at as u64, 1));
|
||||
}
|
||||
raw.extend(element(3, 64, 1));
|
||||
let mut r = VlResolver::new(&file, 8, 8);
|
||||
assert_eq!(r.strings(&raw).unwrap(), ["s0!", "s1!", "s2!", "s0!"]);
|
||||
}
|
||||
Reference in New Issue
Block a user