#![allow(clippy::unwrap_used, clippy::expect_used)] #![allow(unsafe_code)]
use std::alloc::{GlobalAlloc, Layout, System};
use std::cell::Cell;
thread_local! {
static MAX_ALLOC: Cell<usize> = const { Cell::new(0) };
}
struct Watching;
unsafe impl GlobalAlloc for Watching {
unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
let _ = MAX_ALLOC.try_with(|m| m.set(m.get().max(layout.size())));
unsafe { System.alloc(layout) }
}
unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
unsafe { System.dealloc(ptr, layout) }
}
unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
let _ = MAX_ALLOC.try_with(|m| m.set(m.get().max(new_size)));
unsafe { System.realloc(ptr, layout, new_size) }
}
}
#[global_allocator]
static ALLOC: Watching = Watching;
fn largest_alloc_during<T>(f: impl FnOnce() -> T) -> usize {
MAX_ALLOC.with(|m| m.set(0));
let out = f();
drop(out);
MAX_ALLOC.with(Cell::get)
}
const TOLERANCE: usize = 4 << 20;
use chronix_encoding::{PcoDecoder, PcoEncoder};
fn sample_block() -> Vec<u8> {
let values: Vec<f64> = (0..1000)
.map(|i| 230.0 + f64::from(i % 400) / 100.0)
.collect();
PcoEncoder::encode_f64(&values).unwrap()
}
#[test]
fn our_count_header_bounds_the_allocation() {
let mut block = sample_block();
block[..4].copy_from_slice(&u32::MAX.to_le_bytes());
let peak = largest_alloc_during(|| {
assert!(PcoDecoder::decode_f64(&block).is_err());
});
assert!(
peak < TOLERANCE,
"a u32::MAX count allocated {peak} bytes before being rejected"
);
block[..4].copy_from_slice(&1_300_000u32.to_le_bytes());
let peak = largest_alloc_during(|| {
assert!(PcoDecoder::decode_f64(&block).is_err());
});
assert!(
peak <= 32 << 20,
"a count at the ceiling allocated {peak} bytes, above the 32 MiB bound"
);
}
#[test]
fn corrupting_pcos_own_metadata_cannot_allocate() {
let block = sample_block();
let mut worst = 0usize;
for i in 4..block.len().min(256) {
for patch in [0xFF_u8, 0x7F, 0x80, 0x00] {
let mut corrupt = block.clone();
corrupt[i] = patch;
let peak = largest_alloc_during(|| {
let _ = PcoDecoder::decode_f64(&corrupt);
let _ = PcoDecoder::decode_i64(&corrupt);
let _ = PcoDecoder::decode_u64(&corrupt);
});
worst = worst.max(peak);
assert!(
peak < TOLERANCE,
"byte {i} set to {patch:#04x} caused a {peak}-byte allocation"
);
}
}
println!("largest allocation across all single-byte corruptions: {worst} B");
}
#[test]
fn truncated_pco_blocks_cannot_allocate() {
let block = sample_block();
for cut in [4, 5, 8, 16, 32, 64, block.len() / 2, block.len() - 1] {
let peak = largest_alloc_during(|| {
let _ = PcoDecoder::decode_f64(&block[..cut]);
});
assert!(
peak < TOLERANCE,
"truncating to {cut} allocated {peak} bytes"
);
}
}