From 4578766bb2b8adea6d7939c749399c7070a2c7d3 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Tue, 15 Sep 2026 05:17:20 +0000 Subject: [PATCH 1/2] security(gguf): add ParseLimits budgets and checked host-size conversions Fail closed on untrusted GGUF headers before allocation or loops proportional to declared counts. Typed LimitExceeded and HostSizeOverflow errors name the budget or field. Default and mmap readers share the same policy; trusted callers can override without changing defaults. Linear: RM-1358 Co-authored-by: Raul Cardenas Montoya --- CHANGELOG.md | 9 + README.md | 9 +- src/error.rs | 209 ++++++++++++++++++- src/gguf/cursor.rs | 257 ++++++++++++++++++++--- src/gguf/layout.rs | 110 +++++----- src/gguf/limits.rs | 212 +++++++++++++++++++ src/gguf/map.rs | 14 +- src/gguf/mod.rs | 31 ++- src/gguf/tensor.rs | 42 ++-- src/lib.rs | 14 +- tests/common/mod.rs | 29 +++ tests/gguf_limits.rs | 475 +++++++++++++++++++++++++++++++++++++++++++ 12 files changed, 1311 insertions(+), 100 deletions(-) create mode 100644 src/gguf/limits.rs create mode 100644 tests/gguf_limits.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 63b4541..20a4e46 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,15 @@ All notable changes to this project are documented in this file. crates: `[dependencies]` stays empty for this feature; the upstream `safetensors` crate, `serde_json`, and `corinth-canal` are not used. Does not include payload mmap or Hugging Face `config.json` policy. +- **GGUF `ParseLimits` (RM-1358):** documented resource budgets for untrusted + GGUF headers — KV count, tensor count, string bytes, array work items, + tensor rank, and metadata bytes. File-declared `u64` sizes convert with + typed [`HostSizeField`] errors; exhausted budgets return + [`ParserError::LimitExceeded`] naming the limit. Alignment, tensor offsets, + element counts, and packed byte sizes use checked arithmetic. Default and + `mmap` readers share the same policy. Trusted callers override via + `load_gguf_with_limits` / `parse_bytes_with_limits` / + `load_gguf_mmap_with_limits` without changing default safety. - **Optional `mmap` feature (#45 option 1):** `load_gguf_mmap` maps a GGUF with `memmap2` 0.9.11 instead of `fs::read` into a `Vec`. Default builds stay zero-dep (`default = []`). Packed CPU dequant for **Q8_0**, **Q5_K**, diff --git a/README.md b/README.md index 1fe1d85..d8d5caa 100644 --- a/README.md +++ b/README.md @@ -16,6 +16,7 @@ Today, `engram-parser` ships GGUF v3 deserialization, per-expert raw-weight extr ### Shipped now — GGUF - Parses GGUF v3 magic, header, KV metadata, and tensor directory into an in-memory [`GgufLayout`]. +- Applies a documented [`ParseLimits`] budget (KV/tensor counts, string sizes, array work, tensor rank, metadata bytes) before allocation or loops proportional to file-declared values. Defaults are generous; trusted callers can override via `load_gguf_with_limits` / `parse_bytes_with_limits` without weakening the default path. - Enumerates MoE experts discovered in a checkpoint. - Extracts the raw byte buffers for one expert's `gate`, `up`, and `down` projections. - Supports stacked (`blk.{B}.ffn_{role}_exps.weight`) and per-expert (`blk.{B}.ffn_{role}.{E}.weight`) conventions. @@ -218,13 +219,13 @@ Numeric helpers: `dequantize_f16`, `dequantize_q8_0`, `dequantize_q5_k`, `dequan Current GGUF surface includes: -- `load_gguf`, `parse_bytes`; -- `#[cfg(feature = "mmap")] load_gguf_mmap` → `GgufLayoutMmap` (page-aligned tensor slices via `tensor_page_aligned_bytes`); -- `GgufLayout`, `GgufMetadata`, `Tensor`, `DType`; +- `load_gguf`, `parse_bytes` (and `*_with_limits` for an explicit [`ParseLimits`] policy); +- `#[cfg(feature = "mmap")] load_gguf_mmap` / `load_gguf_mmap_with_limits` → `GgufLayoutMmap` (page-aligned tensor slices via `tensor_page_aligned_bytes`); +- `GgufLayout`, `GgufMetadata`, `Tensor`, `DType`, `ParseLimits`; - `dequantize_f16` (on `Tensor`), `dequantize_q8_0`, `dequantize_q5_k`, `dequantize_q6_k`, `dequantize_iq3_m`; - `extract_expert`, `list_experts`; - `MoeExpertWeights`, `RawTensor`; -- `ParserError`, `Result`; +- `ParserError`, `ParseLimitKind`, `HostSizeField`, `Result`; - public `GGML_TYPE_*` / `GGUF_VALUE_TYPE_*` constants and the `ggml_type_label` label function. Safetensors surface (`--features safetensors`) includes: diff --git a/src/error.rs b/src/error.rs index 1d24feb..1f63a0e 100644 --- a/src/error.rs +++ b/src/error.rs @@ -9,6 +9,86 @@ use std::fmt; use std::io; +/// Which [`crate::ParseLimits`] field was exhausted. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum ParseLimitKind { + /// [`crate::ParseLimits::max_kv_count`]. + KvCount, + /// [`crate::ParseLimits::max_tensor_count`]. + TensorCount, + /// [`crate::ParseLimits::max_string_bytes`]. + StringBytes, + /// [`crate::ParseLimits::max_array_work_items`]. + ArrayWorkItems, + /// [`crate::ParseLimits::max_tensor_rank`]. + TensorRank, + /// [`crate::ParseLimits::max_metadata_bytes`]. + MetadataBytes, +} + +impl ParseLimitKind { + /// Stable name used in parser errors and tests. + pub const fn as_str(self) -> &'static str { + match self { + Self::KvCount => "max_kv_count", + Self::TensorCount => "max_tensor_count", + Self::StringBytes => "max_string_bytes", + Self::ArrayWorkItems => "max_array_work_items", + Self::TensorRank => "max_tensor_rank", + Self::MetadataBytes => "max_metadata_bytes", + } + } +} + +impl fmt::Display for ParseLimitKind { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +/// File-declared field that failed `u64 -> usize` conversion. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum HostSizeField { + /// Length prefix of a GGUF string (KV key, string value, or tensor name). + StringLen, + /// Header `kv_count`. + KvCount, + /// Header `tensor_count`. + TensorCount, + /// Tensor directory `n_dims`. + TensorRank, + /// One tensor dimension. + TensorDim, + /// Tensor `relative_offset` in the data region. + RelativeOffset, + /// `general.alignment` layout field. + Alignment, + /// GGUF array element count. + ArrayLen, +} + +impl HostSizeField { + /// Stable name used in parser errors and tests. + pub const fn as_str(self) -> &'static str { + match self { + Self::StringLen => "string_len", + Self::KvCount => "kv_count", + Self::TensorCount => "tensor_count", + Self::TensorRank => "tensor_rank", + Self::TensorDim => "tensor_dim", + Self::RelativeOffset => "relative_offset", + Self::Alignment => "general.alignment", + Self::ArrayLen => "array_len", + } + } +} + +impl fmt::Display for HostSizeField { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + /// Unified error type for checkpoint parsing and MoE weight extraction. #[derive(Debug)] pub enum ParserError { @@ -52,6 +132,50 @@ pub enum ParserError { /// Number of experts actually available. available: usize, }, + /// A [`crate::ParseLimits`] budget was exhausted by a file-declared value. + LimitExceeded { + /// Path of the checkpoint. + path: String, + /// Budget that was exhausted. + limit: ParseLimitKind, + /// Declared or accumulated value that exceeded the budget. + declared: u64, + /// Configured budget for `limit`. + budget: u64, + }, + /// A file-declared `u64` could not be represented as a host `usize`. + HostSizeOverflow { + /// Path of the checkpoint. + path: String, + /// On-wire field that overflowed. + field: HostSizeField, + /// Declared value that does not fit in `usize`. + value: u64, + }, +} + +impl ParserError { + pub(crate) fn limit_exceeded( + path: impl Into, + limit: ParseLimitKind, + declared: u64, + budget: u64, + ) -> Self { + Self::LimitExceeded { + path: path.into(), + limit, + declared, + budget, + } + } + + pub(crate) fn host_size(path: impl Into, field: HostSizeField, value: u64) -> Self { + Self::HostSizeOverflow { + path: path.into(), + field, + value, + } + } } impl fmt::Display for ParserError { @@ -75,6 +199,19 @@ impl fmt::Display for ParserError { f, "expert index out of range: block={block}, expert={expert}, available={available}" ), + Self::LimitExceeded { + path, + limit, + declared, + budget, + } => write!( + f, + "parse limit {limit} exceeded in '{path}': declared {declared}, budget {budget}" + ), + Self::HostSizeOverflow { path, field, value } => write!( + f, + "host-size conversion failed in '{path}': field {field} value {value} does not fit usize" + ), } } } @@ -83,10 +220,80 @@ impl std::error::Error for ParserError { fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { match self { Self::Io { source, .. } => Some(source), - _ => None, + Self::UnsupportedFormat { .. } + | Self::MissingTensor { .. } + | Self::InvalidLayout { .. } + | Self::ExpertOutOfRange { .. } + | Self::LimitExceeded { .. } + | Self::HostSizeOverflow { .. } => None, } } } /// Convenience alias used throughout the crate. pub type Result = std::result::Result; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn limit_and_host_errors_name_the_field() { + let limit = ParserError::limit_exceeded("mem://e", ParseLimitKind::StringBytes, 17, 16); + let host = ParserError::host_size("mem://e", HostSizeField::RelativeOffset, u64::MAX); + assert!(limit.to_string().contains("max_string_bytes")); + assert!(host.to_string().contains("relative_offset")); + assert!(std::error::Error::source(&limit).is_none()); + assert!(std::error::Error::source(&host).is_none()); + } + + #[test] + fn limit_kind_match_is_exhaustive() { + let kinds = [ + ParseLimitKind::KvCount, + ParseLimitKind::TensorCount, + ParseLimitKind::StringBytes, + ParseLimitKind::ArrayWorkItems, + ParseLimitKind::TensorRank, + ParseLimitKind::MetadataBytes, + ]; + for kind in kinds { + let name = match kind { + ParseLimitKind::KvCount => "max_kv_count", + ParseLimitKind::TensorCount => "max_tensor_count", + ParseLimitKind::StringBytes => "max_string_bytes", + ParseLimitKind::ArrayWorkItems => "max_array_work_items", + ParseLimitKind::TensorRank => "max_tensor_rank", + ParseLimitKind::MetadataBytes => "max_metadata_bytes", + }; + assert_eq!(kind.as_str(), name); + } + } + + #[test] + fn host_field_match_is_exhaustive() { + let fields = [ + HostSizeField::StringLen, + HostSizeField::KvCount, + HostSizeField::TensorCount, + HostSizeField::TensorRank, + HostSizeField::TensorDim, + HostSizeField::RelativeOffset, + HostSizeField::Alignment, + HostSizeField::ArrayLen, + ]; + for field in fields { + let name = match field { + HostSizeField::StringLen => "string_len", + HostSizeField::KvCount => "kv_count", + HostSizeField::TensorCount => "tensor_count", + HostSizeField::TensorRank => "tensor_rank", + HostSizeField::TensorDim => "tensor_dim", + HostSizeField::RelativeOffset => "relative_offset", + HostSizeField::Alignment => "general.alignment", + HostSizeField::ArrayLen => "array_len", + }; + assert_eq!(field.as_str(), name); + } + } +} diff --git a/src/gguf/cursor.rs b/src/gguf/cursor.rs index 4c9087b..435a72a 100644 --- a/src/gguf/cursor.rs +++ b/src/gguf/cursor.rs @@ -5,7 +5,8 @@ //! Only what the parser needs: little-endian scalar reads, length-prefixed //! strings, and the ability to skip (or stringify) unknown KV values. -use crate::error::{ParserError, Result}; +use super::limits::{ParseLimits, u64_to_usize}; +use crate::error::{HostSizeField, ParseLimitKind, ParserError, Result}; pub(crate) const GGUF_MAGIC: [u8; 4] = *b"GGUF"; pub(crate) const GGUF_VERSION: u32 = 3; @@ -81,18 +82,39 @@ fn nonneg_signed(path: &str, v: i64) -> Result { }) } +fn fixed_scalar_size(value_type: u32) -> Option { + match value_type { + GGUF_VALUE_TYPE_UINT8 | GGUF_VALUE_TYPE_INT8 | GGUF_VALUE_TYPE_BOOL => Some(1), + GGUF_VALUE_TYPE_UINT16 | GGUF_VALUE_TYPE_INT16 => Some(2), + GGUF_VALUE_TYPE_UINT32 | GGUF_VALUE_TYPE_INT32 | GGUF_VALUE_TYPE_FLOAT32 => Some(4), + GGUF_VALUE_TYPE_UINT64 | GGUF_VALUE_TYPE_INT64 | GGUF_VALUE_TYPE_FLOAT64 => Some(8), + _ => None, + } +} + pub(crate) struct GgufCursor<'a> { bytes: &'a [u8], offset: usize, path: &'a str, + limits: ParseLimits, + array_work_used: u64, + metadata_origin: Option, } impl<'a> GgufCursor<'a> { + #[cfg(test)] pub(crate) fn new(bytes: &'a [u8], path: &'a str) -> Self { + Self::with_limits(bytes, path, ParseLimits::default()) + } + + pub(crate) fn with_limits(bytes: &'a [u8], path: &'a str, limits: ParseLimits) -> Self { Self { bytes, offset: 0, path, + limits, + array_work_used: 0, + metadata_origin: None, } } @@ -100,6 +122,18 @@ impl<'a> GgufCursor<'a> { self.offset } + pub(crate) fn limits(&self) -> ParseLimits { + self.limits + } + + pub(crate) fn begin_metadata_section(&mut self) { + self.metadata_origin = Some(self.offset); + } + + pub(crate) fn end_metadata_section(&mut self) { + self.metadata_origin = None; + } + fn unsupported(&self, reason: String) -> ParserError { ParserError::UnsupportedFormat { path: self.path.to_owned(), @@ -107,7 +141,75 @@ impl<'a> GgufCursor<'a> { } } + fn remaining_bytes(&self) -> u64 { + self.bytes.len().saturating_sub(self.offset) as u64 + } + + fn metadata_used(&self) -> Option { + let origin = self.metadata_origin?; + Some(self.offset.saturating_sub(origin) as u64) + } + + fn ensure_metadata_room(&self, extra: u64) -> Result<()> { + let Some(used) = self.metadata_used() else { + return Ok(()); + }; + let Some(projected) = used.checked_add(extra) else { + return Err(ParserError::limit_exceeded( + self.path, + ParseLimitKind::MetadataBytes, + u64::MAX, + self.limits.max_metadata_bytes, + )); + }; + self.limits + .reject(self.path, ParseLimitKind::MetadataBytes, projected) + } + + fn charge_metadata_used(&self) -> Result<()> { + let Some(used) = self.metadata_used() else { + return Ok(()); + }; + self.limits + .reject(self.path, ParseLimitKind::MetadataBytes, used) + } + + fn consume_array_work(&mut self, items: u64) -> Result<()> { + let Some(declared) = self.array_work_used.checked_add(items) else { + return Err(ParserError::limit_exceeded( + self.path, + ParseLimitKind::ArrayWorkItems, + u64::MAX, + self.limits.max_array_work_items, + )); + }; + self.limits + .reject(self.path, ParseLimitKind::ArrayWorkItems, declared)?; + self.array_work_used = declared; + Ok(()) + } + + fn prepare_array_elements(&mut self, element_type: u32, len: u64) -> Result<()> { + self.consume_array_work(len)?; + if let Some(size) = fixed_scalar_size(element_type) { + let need = len.saturating_mul(size); + self.ensure_metadata_room(need)?; + if need > self.remaining_bytes() { + return Err( + self.unsupported("GGUF array length exceeds remaining metadata bytes".into()) + ); + } + } else if len > self.remaining_bytes() { + return Err( + self.unsupported("GGUF array length exceeds remaining metadata bytes".into()) + ); + } + Ok(()) + } + pub(crate) fn read_exact(&mut self, len: usize) -> Result<&'a [u8]> { + let extra = u64::try_from(len).unwrap_or(u64::MAX); + self.ensure_metadata_room(extra)?; let end = self .offset .checked_add(len) @@ -117,6 +219,7 @@ impl<'a> GgufCursor<'a> { } let slice = &self.bytes[self.offset..end]; self.offset = end; + self.charge_metadata_used()?; Ok(slice) } @@ -166,7 +269,11 @@ impl<'a> GgufCursor<'a> { } pub(crate) fn read_string(&mut self) -> Result { - let len = self.read_u64()? as usize; + let len = self.read_u64()?; + self.limits + .reject(self.path, ParseLimitKind::StringBytes, len)?; + self.ensure_metadata_room(len)?; + let len = u64_to_usize(len, HostSizeField::StringLen, self.path)?; let bytes = self.read_exact(len)?; String::from_utf8(bytes.to_vec()) .map_err(|e| self.unsupported(format!("invalid UTF-8 in GGUF string: {e}"))) @@ -242,7 +349,7 @@ impl<'a> GgufCursor<'a> { "expected integer GGUF value for layout field, got type {value_type}" ))); }; - Ok(v as usize) + u64_to_usize(v, HostSizeField::Alignment, self.path) } /// Render a scalar GGUF value as a string (used for metadata KV). @@ -300,18 +407,12 @@ impl<'a> GgufCursor<'a> { /// Skip a GGUF array value using an explicit stack instead of recursion, /// so deep-but-valid metadata arrays are not rejected by an arbitrary - /// depth limit. Total work is still bounded by the remaining byte range. + /// depth limit. Total work is bounded by [`ParseLimits::max_array_work_items`] + /// and, when parsing KV metadata, [`ParseLimits::max_metadata_bytes`]. fn skip_array_value(&mut self) -> Result<()> { let nested = self.read_u32()?; let len = self.read_u64()?; - - // Reject lengths that cannot possibly fit in the remaining buffer. - let remaining = self.bytes.len().saturating_sub(self.offset) as u64; - if len > remaining { - return Err( - self.unsupported("GGUF array length exceeds remaining metadata bytes".into()) - ); - } + self.prepare_array_elements(nested, len)?; // Stack of (element_type, elements_remaining) pairs. Depth is bounded // only by nesting of arrays, not by a hard-coded recursion limit. @@ -323,22 +424,17 @@ impl<'a> GgufCursor<'a> { if count == 0 { continue; } - // Each element is an independent sub-array; read one header. let sub_ty = self.read_u32()?; let sub_len = self.read_u64()?; - let bytes_left = self.bytes.len().saturating_sub(self.offset) as u64; - if sub_len > bytes_left { - return Err(self.unsupported( - "GGUF nested array length exceeds remaining metadata bytes".into(), - )); - } + self.prepare_array_elements(sub_ty, sub_len)?; count -= 1; if count > 0 { stack.push((ty, count)); } stack.push((sub_ty, sub_len)); } else { - for _ in 0..count { + let n = u64_to_usize(count, HostSizeField::ArrayLen, self.path)?; + for _ in 0..n { self.skip_scalar_value(ty)?; } } @@ -350,6 +446,7 @@ impl<'a> GgufCursor<'a> { #[cfg(test)] mod tests { use super::*; + use crate::error::ParseLimitKind; fn push_u32(out: &mut Vec, v: u32) { out.extend_from_slice(&v.to_le_bytes()); @@ -359,20 +456,23 @@ mod tests { out.extend_from_slice(&v.to_le_bytes()); } - #[test] - fn skip_array_value_handles_deep_nesting() { - // Build a 32-level nested array of arrays ending in an empty UINT8 array. - const DEPTH: usize = 32; - let mut bytes = Vec::with_capacity(DEPTH * 12); - for level in 0..DEPTH { - if level == DEPTH - 1 { + fn nested_array_bytes(depth: usize) -> Vec { + let mut bytes = Vec::with_capacity(depth * 12); + for level in 0..depth { + if level == depth - 1 { push_u32(&mut bytes, GGUF_VALUE_TYPE_UINT8); } else { push_u32(&mut bytes, GGUF_VALUE_TYPE_ARRAY); } - push_u64(&mut bytes, if level == DEPTH - 1 { 0 } else { 1 }); + push_u64(&mut bytes, if level == depth - 1 { 0 } else { 1 }); } + bytes + } + #[test] + fn skip_array_value_handles_deep_nesting() { + const DEPTH: usize = 32; + let bytes = nested_array_bytes(DEPTH); let mut cursor = GgufCursor::new(&bytes, "mem://deep-array"); cursor .skip_value(GGUF_VALUE_TYPE_ARRAY) @@ -383,4 +483,105 @@ mod tests { "did not consume entire nested array" ); } + + #[test] + fn skip_array_value_rejects_work_budget_before_loop() { + let mut bytes = Vec::new(); + push_u32(&mut bytes, GGUF_VALUE_TYPE_UINT8); + push_u64(&mut bytes, 4); + bytes.extend_from_slice(&[1, 2, 3, 4]); + let limits = ParseLimits { + max_array_work_items: 3, + ..ParseLimits::default() + }; + let mut cursor = GgufCursor::with_limits(&bytes, "mem://array-work", limits); + let err = cursor.skip_value(GGUF_VALUE_TYPE_ARRAY).unwrap_err(); + match err { + ParserError::LimitExceeded { + limit, + declared, + budget, + .. + } => { + assert_eq!(limit, ParseLimitKind::ArrayWorkItems); + assert_eq!(declared, 4); + assert_eq!(budget, 3); + } + other => panic!("expected LimitExceeded, got {other}"), + } + assert_eq!(cursor.offset, 12, "must reject before iterating elements"); + } + + #[test] + fn skip_array_value_accepts_exact_work_budget() { + let mut bytes = Vec::new(); + push_u32(&mut bytes, GGUF_VALUE_TYPE_UINT8); + push_u64(&mut bytes, 3); + bytes.extend_from_slice(&[1, 2, 3]); + let limits = ParseLimits { + max_array_work_items: 3, + ..ParseLimits::default() + }; + let mut cursor = GgufCursor::with_limits(&bytes, "mem://array-work-ok", limits); + cursor + .skip_value(GGUF_VALUE_TYPE_ARRAY) + .expect("exact work budget"); + } + + #[test] + fn skip_array_value_rejects_u64_max_count() { + let mut bytes = Vec::new(); + push_u32(&mut bytes, GGUF_VALUE_TYPE_UINT8); + push_u64(&mut bytes, u64::MAX); + let mut cursor = GgufCursor::new(&bytes, "mem://array-max"); + let err = cursor.skip_value(GGUF_VALUE_TYPE_ARRAY).unwrap_err(); + match err { + ParserError::LimitExceeded { + limit, declared, .. + } => { + assert_eq!(limit, ParseLimitKind::ArrayWorkItems); + assert_eq!(declared, u64::MAX); + } + other => panic!("expected LimitExceeded, got {other}"), + } + } + + #[test] + fn read_string_rejects_oversize_before_allocating() { + let mut bytes = Vec::new(); + push_u64(&mut bytes, 5); + bytes.extend_from_slice(b"hello"); + let limits = ParseLimits { + max_string_bytes: 4, + ..ParseLimits::default() + }; + let mut cursor = GgufCursor::with_limits(&bytes, "mem://str", limits); + let err = cursor.read_string().unwrap_err(); + match err { + ParserError::LimitExceeded { + limit, + declared, + budget, + .. + } => { + assert_eq!(limit, ParseLimitKind::StringBytes); + assert_eq!(declared, 5); + assert_eq!(budget, 4); + } + other => panic!("expected LimitExceeded, got {other}"), + } + } + + #[test] + fn read_string_accepts_exact_budget() { + let mut bytes = Vec::new(); + push_u64(&mut bytes, 4); + bytes.extend_from_slice(b"abcd"); + let limits = ParseLimits { + max_string_bytes: 4, + ..ParseLimits::default() + }; + let mut cursor = GgufCursor::with_limits(&bytes, "mem://str-ok", limits); + assert_eq!(cursor.read_string().unwrap(), "abcd"); + } } diff --git a/src/gguf/layout.rs b/src/gguf/layout.rs index 154df9f..3195b26 100644 --- a/src/gguf/layout.rs +++ b/src/gguf/layout.rs @@ -10,16 +10,11 @@ use std::collections::HashMap; use super::cursor::{ - GGUF_MAGIC, GGUF_VALUE_TYPE_STRING, GGUF_VERSION, GgufCursor, invalid_layout, - is_signed_layout_type, unsupported, + GGUF_MAGIC, GGUF_VERSION, GgufCursor, invalid_layout, is_signed_layout_type, unsupported, }; +use super::limits::{ParseLimits, align_up_checked, u64_to_usize}; use super::tensor::{DType, Tensor}; -use crate::error::{ParserError, Result}; - -/// Upper bounds to prevent OOM from malformed inputs. -const MAX_TENSOR_COUNT: u64 = 1_000_000; -const MAX_KV_COUNT: u64 = 1_000_000; -const MAX_TENSOR_DIMS: usize = 8; +use crate::error::{HostSizeField, ParseLimitKind, ParserError, Result}; /// Parsed GGUF metadata key-value store. /// @@ -245,19 +240,32 @@ struct LayoutHeader { pub(crate) fn parse_layout( bytes: &[u8], path: &str, + limits: ParseLimits, ) -> Result<(GgufMetadata, HashMap, usize, usize)> { - let mut cursor = GgufCursor::new(bytes, path); + let mut cursor = GgufCursor::with_limits(bytes, path, limits); let header = read_layout_header(&mut cursor, path)?; let (alignment, metadata) = read_metadata_section(&mut cursor, path, header.kv_count)?; let mut tensors = read_tensor_directory(&mut cursor, path, header.tensor_count)?; - let tensor_data_offset = finalize_tensor_offsets(&mut tensors, cursor.offset(), alignment); + let tensor_data_offset = + finalize_tensor_offsets(&mut tensors, cursor.offset(), alignment, path)?; Ok((metadata, tensors, alignment, tensor_data_offset)) } fn read_layout_header(cursor: &mut GgufCursor<'_>, path: &str) -> Result { validate_gguf_header(cursor, path)?; - let tensor_count = bounded_count(cursor.read_u64()?, MAX_TENSOR_COUNT, "tensor_count", path)?; - let kv_count = bounded_count(cursor.read_u64()?, MAX_KV_COUNT, "kv_count", path)?; + let limits = cursor.limits(); + let tensor_count = limits.bounded_usize( + path, + ParseLimitKind::TensorCount, + HostSizeField::TensorCount, + cursor.read_u64()?, + )?; + let kv_count = limits.bounded_usize( + path, + ParseLimitKind::KvCount, + HostSizeField::KvCount, + cursor.read_u64()?, + )?; Ok(LayoutHeader { tensor_count, kv_count, @@ -284,16 +292,6 @@ fn validate_gguf_header(cursor: &mut GgufCursor<'_>, path: &str) -> Result<()> { Ok(()) } -fn bounded_count(raw: u64, limit: u64, label: &str, path: &str) -> Result { - if raw > limit { - return Err(unsupported( - path, - format!("{label} {raw} exceeds sanity limit {limit}"), - )); - } - Ok(raw as usize) -} - fn read_metadata_section( cursor: &mut GgufCursor<'_>, path: &str, @@ -301,6 +299,7 @@ fn read_metadata_section( ) -> Result<(usize, GgufMetadata)> { let mut alignment: usize = 32; let mut metadata = GgufMetadata::default(); + cursor.begin_metadata_section(); for _ in 0..kv_count { let key = cursor.read_string()?; @@ -320,6 +319,7 @@ fn read_metadata_section( } } + cursor.end_metadata_section(); Ok((alignment, metadata)) } @@ -340,7 +340,7 @@ fn read_tensor_entry(cursor: &mut GgufCursor<'_>, path: &str) -> Result let name = cursor.read_string()?; let dims = read_tensor_dims(cursor, path, &name)?; let ggml_type = cursor.read_u32()?; - let relative_offset = cursor.read_u64()? as usize; + let relative_offset = u64_to_usize(cursor.read_u64()?, HostSizeField::RelativeOffset, path)?; let dtype = DType::from_ggml_type(ggml_type); let n_elements = tensor_element_count(&dims, &name, path)?; validate_blocked_inner_dim(dtype, &dims, &name, path)?; @@ -381,18 +381,22 @@ fn validate_blocked_inner_dim(dtype: DType, dims: &[usize], name: &str, path: &s )) } -fn read_tensor_dims(cursor: &mut GgufCursor<'_>, path: &str, name: &str) -> Result> { - let n_dims_raw = cursor.read_u32()? as usize; - if n_dims_raw > MAX_TENSOR_DIMS { - return Err(unsupported( +fn read_tensor_dims(cursor: &mut GgufCursor<'_>, path: &str, _name: &str) -> Result> { + let n_dims_raw = u64::from(cursor.read_u32()?); + let n_dims = cursor.limits().bounded_usize( + path, + ParseLimitKind::TensorRank, + HostSizeField::TensorRank, + n_dims_raw, + )?; + + let mut dims = Vec::with_capacity(n_dims); + for _ in 0..n_dims { + dims.push(u64_to_usize( + cursor.read_u64()?, + HostSizeField::TensorDim, path, - format!("tensor '{name}' has {n_dims_raw} dims; max {MAX_TENSOR_DIMS}"), - )); - } - - let mut dims = Vec::with_capacity(n_dims_raw); - for _ in 0..n_dims_raw { - dims.push(cursor.read_u64()? as usize); + )?); } Ok(dims) } @@ -424,12 +428,30 @@ fn finalize_tensor_offsets( tensors: &mut HashMap, cursor_offset: usize, alignment: usize, -) -> usize { - let tensor_data_offset = align_up(cursor_offset, alignment); + path: &str, +) -> Result { + let tensor_data_offset = align_up_checked(cursor_offset, alignment, path)?; for tensor in tensors.values_mut() { - tensor.absolute_offset = tensor_data_offset + tensor.relative_offset; + tensor.absolute_offset = tensor_data_offset + .checked_add(tensor.relative_offset) + .ok_or_else(|| { + invalid_layout( + path, + format!("tensor '{}' absolute offset overflow", tensor.name), + ) + })?; + if tensor + .absolute_offset + .checked_add(tensor.byte_len) + .is_none() + { + return Err(invalid_layout( + path, + format!("tensor '{}' byte-length overflow", tensor.name), + )); + } } - tensor_data_offset + Ok(tensor_data_offset) } fn capture_kv( @@ -441,8 +463,8 @@ fn capture_kv( use super::cursor::{ GGUF_VALUE_TYPE_BOOL, GGUF_VALUE_TYPE_FLOAT32, GGUF_VALUE_TYPE_FLOAT64, GGUF_VALUE_TYPE_INT8, GGUF_VALUE_TYPE_INT16, GGUF_VALUE_TYPE_INT32, GGUF_VALUE_TYPE_INT64, - GGUF_VALUE_TYPE_UINT8, GGUF_VALUE_TYPE_UINT16, GGUF_VALUE_TYPE_UINT32, - GGUF_VALUE_TYPE_UINT64, + GGUF_VALUE_TYPE_STRING, GGUF_VALUE_TYPE_UINT8, GGUF_VALUE_TYPE_UINT16, + GGUF_VALUE_TYPE_UINT32, GGUF_VALUE_TYPE_UINT64, }; match value_type { GGUF_VALUE_TYPE_UINT8 @@ -509,14 +531,6 @@ fn capture_skipped_kv(cursor: &mut GgufCursor<'_>, value_type: u32) -> Result<() cursor.skip_value(value_type) } -fn align_up(value: usize, alignment: usize) -> usize { - if alignment <= 1 { - value - } else { - value.div_ceil(alignment) * alignment - } -} - fn tensor_block_sort_key(name: &str) -> (usize, String) { let block = name .strip_prefix("blk.") diff --git a/src/gguf/limits.rs b/src/gguf/limits.rs new file mode 100644 index 0000000..9a43764 --- /dev/null +++ b/src/gguf/limits.rs @@ -0,0 +1,212 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +//! Explicit resource budgets for untrusted GGUF inputs. +//! +//! GGUF headers declare counts, string lengths, array lengths, tensor ranks, +//! and offsets as `u64`. Those values must not drive unbounded allocation or +//! iteration. [`ParseLimits`] is the single policy object for those checks. +//! Default values are generous enough for production checkpoints and every +//! committed fixture; trusted callers may raise or disable budgets without +//! changing the default path. + +use crate::error::{HostSizeField, ParseLimitKind, ParserError, Result}; + +/// Resource budgets applied before allocation or loops proportional to +/// attacker-declared GGUF values. +/// +/// Defaults accept large real checkpoints. Override individual fields for +/// tighter sandboxes, or use [`ParseLimits::trusted`] when the input is +/// already authenticated. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ParseLimits { + /// Maximum number of KV metadata pairs. + pub max_kv_count: u64, + /// Maximum number of tensor directory entries. + pub max_tensor_count: u64, + /// Maximum UTF-8 bytes in a single GGUF string. + pub max_string_bytes: u64, + /// Maximum nested-array elements visited while skipping metadata arrays. + pub max_array_work_items: u64, + /// Maximum tensor rank (`n_dims`). + pub max_tensor_rank: u64, + /// Maximum bytes consumed by the KV metadata section. + pub max_metadata_bytes: u64, +} + +impl ParseLimits { + /// Generous production defaults. Same values as [`Default`]. + pub const DEFAULT: Self = Self { + max_kv_count: 1_000_000, + max_tensor_count: 1_000_000, + max_string_bytes: 16 * 1024 * 1024, + max_array_work_items: 16_000_000, + max_tensor_rank: 8, + max_metadata_bytes: 256 * 1024 * 1024, + }; + + /// Disable every budget (`u64::MAX`). Host-size conversions, alignment + /// overflow, and EOF checks still apply. + pub const TRUSTED: Self = Self { + max_kv_count: u64::MAX, + max_tensor_count: u64::MAX, + max_string_bytes: u64::MAX, + max_array_work_items: u64::MAX, + max_tensor_rank: u64::MAX, + max_metadata_bytes: u64::MAX, + }; + + /// Limits for already-authenticated inputs. Does not change [`Default`]. + pub const fn trusted() -> Self { + Self::TRUSTED + } + + /// Budget associated with `kind`. + pub const fn budget(self, kind: ParseLimitKind) -> u64 { + match kind { + ParseLimitKind::KvCount => self.max_kv_count, + ParseLimitKind::TensorCount => self.max_tensor_count, + ParseLimitKind::StringBytes => self.max_string_bytes, + ParseLimitKind::ArrayWorkItems => self.max_array_work_items, + ParseLimitKind::TensorRank => self.max_tensor_rank, + ParseLimitKind::MetadataBytes => self.max_metadata_bytes, + } + } + + /// Fail closed when `declared` is strictly greater than the named budget. + pub(crate) fn reject(self, path: &str, kind: ParseLimitKind, declared: u64) -> Result<()> { + let budget = self.budget(kind); + if declared > budget { + Err(ParserError::limit_exceeded(path, kind, declared, budget)) + } else { + Ok(()) + } + } + + /// Reject `declared` against `kind`, then convert it to a host `usize`. + pub(crate) fn bounded_usize( + self, + path: &str, + kind: ParseLimitKind, + field: HostSizeField, + declared: u64, + ) -> Result { + self.reject(path, kind, declared)?; + u64_to_usize(declared, field, path) + } +} + +impl Default for ParseLimits { + fn default() -> Self { + Self::DEFAULT + } +} + +/// Convert a file-declared `u64` to `usize`, naming the field on failure. +pub(crate) fn u64_to_usize(value: u64, field: HostSizeField, path: &str) -> Result { + usize::try_from(value).map_err(|_| ParserError::host_size(path, field, value)) +} + +/// Round `value` up to a multiple of `alignment` without wrapping. +pub(crate) fn align_up_checked(value: usize, alignment: usize, path: &str) -> Result { + if alignment <= 1 { + return Ok(value); + } + let padding = alignment - 1; + let padded = value + .checked_add(padding) + .ok_or_else(|| ParserError::InvalidLayout { + path: path.to_owned(), + reason: "aligned tensor data offset overflow".into(), + })?; + Ok(padded - (padded % alignment)) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn default_limits_are_generous_and_named() { + let limits = ParseLimits::default(); + assert_eq!(limits, ParseLimits::DEFAULT); + assert!(limits.max_kv_count >= 1_000_000); + assert!(limits.max_tensor_count >= 1_000_000); + assert!(limits.max_string_bytes >= 1024); + assert!(limits.max_array_work_items >= 32); + assert!(limits.max_tensor_rank >= 8); + assert!(limits.max_metadata_bytes >= 1024); + } + + #[test] + fn trusted_does_not_change_default() { + assert_ne!(ParseLimits::trusted(), ParseLimits::default()); + assert_eq!(ParseLimits::trusted().max_kv_count, u64::MAX); + assert_eq!(ParseLimits::default().max_kv_count, 1_000_000); + } + + #[test] + fn reject_is_inclusive_at_the_budget() { + let limits = ParseLimits { + max_kv_count: 4, + ..ParseLimits::default() + }; + limits + .reject("mem://lim", ParseLimitKind::KvCount, 4) + .expect("exact budget must pass"); + let err = limits + .reject("mem://lim", ParseLimitKind::KvCount, 5) + .unwrap_err(); + match err { + ParserError::LimitExceeded { + limit, + declared, + budget, + .. + } => { + assert_eq!(limit, ParseLimitKind::KvCount); + assert_eq!(declared, 5); + assert_eq!(budget, 4); + } + other => panic!("expected LimitExceeded, got {other}"), + } + } + + #[test] + fn u64_to_usize_reports_field_when_unrepresentable() { + #[cfg(target_pointer_width = "32")] + { + let err = u64_to_usize(u64::MAX, HostSizeField::StringLen, "mem://host").unwrap_err(); + match err { + ParserError::HostSizeOverflow { field, value, .. } => { + assert_eq!(field, HostSizeField::StringLen); + assert_eq!(value, u64::MAX); + } + other => panic!("expected HostSizeOverflow, got {other}"), + } + } + #[cfg(not(target_pointer_width = "32"))] + { + assert_eq!( + u64_to_usize(u64::MAX, HostSizeField::StringLen, "mem://host").unwrap(), + usize::MAX + ); + } + } + + #[test] + fn align_up_checked_rejects_overflow() { + let err = align_up_checked(usize::MAX, 32, "mem://align").unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("aligned tensor data offset overflow"), + "got: {msg}" + ); + } + + #[test] + fn align_up_checked_rounds_up() { + assert_eq!(align_up_checked(1, 32, "mem://align").unwrap(), 32); + assert_eq!(align_up_checked(32, 32, "mem://align").unwrap(), 32); + assert_eq!(align_up_checked(7, 1, "mem://align").unwrap(), 7); + } +} diff --git a/src/gguf/map.rs b/src/gguf/map.rs index e72e830..9968d68 100644 --- a/src/gguf/map.rs +++ b/src/gguf/map.rs @@ -15,6 +15,7 @@ use std::path::Path; use memmap2::{Mmap, MmapOptions}; use super::layout::{GgufLayout, GgufMetadata, parse_layout, tensor_payload_bytes}; +use super::limits::ParseLimits; use super::tensor::Tensor; use crate::error::{ParserError, Result}; @@ -80,6 +81,16 @@ impl<'a> PageAlignedTensorBytes<'a> { /// Requires the `mmap` cargo feature (`memmap2`). The default /// [`super::load_gguf`] path is unchanged. pub fn load_gguf_mmap>(path: P) -> Result { + load_gguf_mmap_with_limits(path, ParseLimits::default()) +} + +/// Memory-map a `.gguf` file with an explicit [`ParseLimits`] policy. +/// +/// Default and mmap readers apply equivalent limits for the same input. +pub fn load_gguf_mmap_with_limits>( + path: P, + limits: ParseLimits, +) -> Result { let path_ref = path.as_ref(); let path_str = path_ref.display().to_string(); let file = File::open(path_ref).map_err(|e| ParserError::Io { @@ -94,7 +105,8 @@ pub fn load_gguf_mmap>(path: P) -> Result { path: path_str.clone(), source: e, })?; - let (metadata, tensors, alignment, tensor_data_offset) = parse_layout(&mmap, &path_str)?; + let (metadata, tensors, alignment, tensor_data_offset) = + parse_layout(&mmap, &path_str, limits)?; Ok(GgufLayoutMmap { path: path_str, metadata, diff --git a/src/gguf/mod.rs b/src/gguf/mod.rs index 1079501..8d63411 100644 --- a/src/gguf/mod.rs +++ b/src/gguf/mod.rs @@ -10,6 +10,7 @@ mod cursor; mod dequant; mod layout; +mod limits; #[cfg(feature = "mmap")] mod map; mod tensor; @@ -22,8 +23,12 @@ pub use dequant::{ packed_row_size, }; pub use layout::{GgufLayout, GgufMetadata}; +pub use limits::ParseLimits; #[cfg(feature = "mmap")] -pub use map::{GgufLayoutMmap, PageAlignedTensorBytes, load_gguf_mmap, os_page_size}; +pub use map::{ + GgufLayoutMmap, PageAlignedTensorBytes, load_gguf_mmap, load_gguf_mmap_with_limits, + os_page_size, +}; pub use tensor::{ DType, GGML_TYPE_BF16, GGML_TYPE_F16, GGML_TYPE_F32, GGML_TYPE_F64, GGML_TYPE_I8, GGML_TYPE_I16, GGML_TYPE_I32, GGML_TYPE_I64, GGML_TYPE_IQ1_M, GGML_TYPE_IQ1_S, GGML_TYPE_IQ2_S, @@ -51,20 +56,38 @@ use crate::error::{ParserError, Result}; /// default builds stay zero-dep). For multi-GB checkpoints enable the /// `mmap` feature and use `load_gguf_mmap`. Tensor payloads remain /// available as raw byte slices via [`GgufLayout::tensor_bytes`]. +/// +/// Uses [`ParseLimits::default`]. Trusted callers can raise budgets with +/// [`load_gguf_with_limits`]. pub fn load_gguf>(path: P) -> Result { + load_gguf_with_limits(path, ParseLimits::default()) +} + +/// Load a `.gguf` checkpoint with an explicit [`ParseLimits`] policy. +pub fn load_gguf_with_limits>(path: P, limits: ParseLimits) -> Result { let path_ref = path.as_ref(); let path_str = path_ref.display().to_string(); let bytes = fs::read(path_ref).map_err(|e| ParserError::Io { path: path_str.clone(), source: e, })?; - parse_bytes(bytes, path_str) + parse_bytes_with_limits(bytes, path_str, limits) } /// Parse an already-loaded byte buffer as a GGUF checkpoint. Useful for -/// unit tests and in-memory round-trips. +/// unit tests and in-memory round-trips. Uses [`ParseLimits::default`]. pub fn parse_bytes(bytes: Vec, path: String) -> Result { - let (metadata, tensors, alignment, tensor_data_offset) = layout::parse_layout(&bytes, &path)?; + parse_bytes_with_limits(bytes, path, ParseLimits::default()) +} + +/// Parse an already-loaded byte buffer with an explicit [`ParseLimits`] policy. +pub fn parse_bytes_with_limits( + bytes: Vec, + path: String, + limits: ParseLimits, +) -> Result { + let (metadata, tensors, alignment, tensor_data_offset) = + layout::parse_layout(&bytes, &path, limits)?; Ok(GgufLayout { path, metadata, diff --git a/src/gguf/tensor.rs b/src/gguf/tensor.rs index f13f0c0..069d600 100644 --- a/src/gguf/tensor.rs +++ b/src/gguf/tensor.rs @@ -490,13 +490,19 @@ impl Tensor { reason: format!("read_f32_values called on dtype {:?}", self.dtype), }); } - if bytes.len() != self.n_elements * 4 { + let expected = + self.n_elements + .checked_mul(4) + .ok_or_else(|| ParserError::InvalidLayout { + path: self.name.clone(), + reason: "f32 element-count overflow".into(), + })?; + if bytes.len() != expected { return Err(ParserError::InvalidLayout { path: self.name.clone(), reason: format!( - "f32 byte-length mismatch: bytes={}, expected={}", - bytes.len(), - self.n_elements * 4 + "f32 byte-length mismatch: bytes={}, expected={expected}", + bytes.len() ), }); } @@ -517,13 +523,19 @@ impl Tensor { reason: format!("read_u16_values called on dtype {:?}", self.dtype), }); } - if bytes.len() != self.n_elements * 2 { + let expected = + self.n_elements + .checked_mul(2) + .ok_or_else(|| ParserError::InvalidLayout { + path: self.name.clone(), + reason: "16-bit element-count overflow".into(), + })?; + if bytes.len() != expected { return Err(ParserError::InvalidLayout { path: self.name.clone(), reason: format!( - "16-bit byte-length mismatch: bytes={}, expected={}", - bytes.len(), - self.n_elements * 2 + "16-bit byte-length mismatch: bytes={}, expected={expected}", + bytes.len() ), }); } @@ -545,13 +557,19 @@ impl Tensor { reason: format!("dequantize_f16 called on dtype {:?}", self.dtype), }); } - if bytes.len() != self.n_elements * 2 { + let expected = + self.n_elements + .checked_mul(2) + .ok_or_else(|| ParserError::InvalidLayout { + path: self.name.clone(), + reason: "f16 element-count overflow".into(), + })?; + if bytes.len() != expected { return Err(ParserError::InvalidLayout { path: self.name.clone(), reason: format!( - "f16 byte-length mismatch: bytes={}, expected={}", - bytes.len(), - self.n_elements * 2 + "f16 byte-length mismatch: bytes={}, expected={expected}", + bytes.len() ), }); } diff --git a/src/lib.rs b/src/lib.rs index 89da17c..0afc9ff 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -10,6 +10,10 @@ //! //! - **Zero dependencies by default**: the default path is pure Rust with an //! empty `[dependencies]`. The optional `mmap` feature adds `memmap2`. +//! - **Parse limits**: [`ParseLimits`] is a documented budget for KV/tensor +//! counts, string sizes, array work, tensor rank, and metadata bytes. +//! File-declared `u64` sizes convert with [`HostSizeField`] errors. +//! Defaults stay generous; trusted callers override explicitly. //! - **GGUF v3 support**: Full parsing of headers, metadata, and tensor directories //! - **GGUF wire-type metadata**: labels + packed `byte_len` for known quant //! codes (F32/F16/BF16, Q*/IQ*, integers, historical wire 31 = `Q4_0_4_4`). @@ -50,7 +54,7 @@ pub mod moe; pub mod safetensors; // Re-export commonly used types at the crate root for convenience. -pub use error::{ParserError, Result}; +pub use error::{HostSizeField, ParseLimitKind, ParserError, Result}; pub use gguf::{ DType, // GGML type constants @@ -101,6 +105,7 @@ pub use gguf::{ GGUF_VALUE_TYPE_UINT64, GgufLayout, GgufMetadata, + ParseLimits, Tensor, dequantize_iq3_m, dequantize_packed, @@ -110,9 +115,14 @@ pub use gguf::{ f16_bits_to_f32, ggml_type_label, load_gguf, + load_gguf_with_limits, packed_row_size, parse_bytes, + parse_bytes_with_limits, }; #[cfg(feature = "mmap")] -pub use gguf::{GgufLayoutMmap, PageAlignedTensorBytes, load_gguf_mmap, os_page_size}; +pub use gguf::{ + GgufLayoutMmap, PageAlignedTensorBytes, load_gguf_mmap, load_gguf_mmap_with_limits, + os_page_size, +}; pub use moe::{MoeExpertWeights, RawTensor, extract_expert, list_experts}; diff --git a/tests/common/mod.rs b/tests/common/mod.rs index d24d0be..426e16c 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -9,8 +9,10 @@ pub const GGUF_VERSION: u32 = 3; pub const ALIGNMENT: u32 = 32; // Value types. +pub const VT_UINT8: u32 = 0; pub const VT_UINT32: u32 = 4; pub const VT_STRING: u32 = 8; +pub const VT_ARRAY: u32 = 9; // Dtypes (GGUF wire type ids). pub const GGML_F32: u32 = 0; @@ -56,6 +58,33 @@ pub fn push_kv_string(out: &mut Vec, key: &str, v: &str) { push_string(out, v); } +pub fn gguf_header(tensor_count: u64, kv_count: u64) -> Vec { + let mut out = Vec::new(); + out.extend_from_slice(&GGUF_MAGIC); + push_u32(&mut out, GGUF_VERSION); + push_u64(&mut out, tensor_count); + push_u64(&mut out, kv_count); + out +} + +pub fn push_array_uint8_header(out: &mut Vec, key: &str, len: u64) { + push_string(out, key); + push_u32(out, VT_ARRAY); + push_u32(out, VT_UINT8); + push_u64(out, len); +} + +pub fn push_nested_array_payload(out: &mut Vec, depth: usize) { + for level in 0..depth { + if level == depth - 1 { + push_u32(out, VT_UINT8); + } else { + push_u32(out, VT_ARRAY); + } + push_u64(out, if level == depth - 1 { 0 } else { 1 }); + } +} + pub fn build_gguf(kv: &[(&str, KvValue)], tensors: &[TensorSpec]) -> Vec { let mut out = Vec::new(); out.extend_from_slice(&GGUF_MAGIC); diff --git a/tests/gguf_limits.rs b/tests/gguf_limits.rs new file mode 100644 index 0000000..402529e --- /dev/null +++ b/tests/gguf_limits.rs @@ -0,0 +1,475 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +//! Parser resource budgets, host-size conversions, and mmap error parity. + +mod common; +use common::*; + +use engram_parser::{ + HostSizeField, ParseLimitKind, ParseLimits, ParserError, parse_bytes, parse_bytes_with_limits, +}; + +fn expect_limit(err: ParserError, kind: ParseLimitKind, declared: u64, budget: u64) { + let msg = err.to_string(); + match err { + ParserError::LimitExceeded { + limit, + declared: got_declared, + budget: got_budget, + .. + } => { + assert_eq!(limit, kind, "limit kind"); + assert_eq!(got_declared, declared, "declared"); + assert_eq!(got_budget, budget, "budget"); + assert!( + msg.contains(kind.as_str()), + "Display must name the limit: {msg}" + ); + } + other => panic!("expected LimitExceeded({kind}), got {other}"), + } +} + +fn error_payload(err: &ParserError) -> String { + match err { + ParserError::Io { source, .. } => format!("io:{source}"), + ParserError::UnsupportedFormat { reason, .. } => format!("unsupported:{reason}"), + ParserError::MissingTensor { name, .. } => format!("missing:{name}"), + ParserError::InvalidLayout { reason, .. } => format!("invalid:{reason}"), + ParserError::ExpertOutOfRange { + block, + expert, + available, + } => format!("oor:{block}:{expert}:{available}"), + ParserError::LimitExceeded { + limit, + declared, + budget, + .. + } => format!("limit:{limit}:{declared}:{budget}"), + ParserError::HostSizeOverflow { field, value, .. } => { + format!("host:{field}:{value}") + } + } +} + +fn parse_ok(bytes: Vec, limits: ParseLimits) { + parse_bytes_with_limits(bytes, "mem://ok".into(), limits).expect("parse should succeed"); +} + +fn parse_err(bytes: Vec, limits: ParseLimits) -> ParserError { + parse_bytes_with_limits(bytes, "mem://err".into(), limits).expect_err("parse should fail") +} + +fn kv_only(kv: &[(&str, KvValue)]) -> Vec { + build_gguf(kv, &[]) +} + +fn one_tensor(name: &'static str, dims: Vec) -> Vec { + let n: usize = dims.iter().product(); + let tensors = [TensorSpec { + name, + dims, + ggml_type: GGML_F32, + payload: vec![0u8; n * 4], + }]; + build_gguf(&[("general.architecture", KvValue::Str("test"))], &tensors) +} + +#[test] +fn default_limits_accept_tiny_fixture() { + let bytes = one_tensor("token_embd.weight", vec![4, 2]); + parse_bytes(bytes, "mem://tiny".into()).expect("default limits must accept tiny fixtures"); +} + +#[test] +fn kv_count_boundary_and_plus_one() { + let two = kv_only(&[ + ("general.architecture", KvValue::Str("a")), + ("k2", KvValue::U32(1)), + ]); + let three = kv_only(&[ + ("general.architecture", KvValue::Str("a")), + ("k2", KvValue::U32(1)), + ("k3", KvValue::U32(2)), + ]); + let limits = ParseLimits { + max_kv_count: 2, + ..ParseLimits::default() + }; + parse_ok(two, limits); + expect_limit(parse_err(three, limits), ParseLimitKind::KvCount, 3, 2); +} + +#[test] +fn tensor_count_boundary_and_plus_one() { + let one = one_tensor("a", vec![2, 2]); + let two = { + let payload = vec![0u8; 16]; + build_gguf( + &[("general.architecture", KvValue::Str("test"))], + &[ + TensorSpec { + name: "a", + dims: vec![2, 2], + ggml_type: GGML_F32, + payload: payload.clone(), + }, + TensorSpec { + name: "b", + dims: vec![2, 2], + ggml_type: GGML_F32, + payload, + }, + ], + ) + }; + let limits = ParseLimits { + max_tensor_count: 1, + ..ParseLimits::default() + }; + parse_ok(one, limits); + expect_limit(parse_err(two, limits), ParseLimitKind::TensorCount, 2, 1); +} + +#[test] +fn string_bytes_boundary_and_plus_one() { + let exact = kv_only(&[("abcd", KvValue::U32(1))]); + let over = kv_only(&[("abcde", KvValue::U32(1))]); + let limits = ParseLimits { + max_string_bytes: 4, + ..ParseLimits::default() + }; + parse_ok(exact, limits); + expect_limit(parse_err(over, limits), ParseLimitKind::StringBytes, 5, 4); +} + +#[test] +fn tensor_rank_boundary_and_plus_one() { + let rank2 = one_tensor("t", vec![2, 2]); + let rank3 = one_tensor("t", vec![2, 2, 1]); + let limits = ParseLimits { + max_tensor_rank: 2, + ..ParseLimits::default() + }; + parse_ok(rank2, limits); + expect_limit(parse_err(rank3, limits), ParseLimitKind::TensorRank, 3, 2); +} + +#[test] +fn array_work_items_boundary_and_plus_one() { + fn file_with_uint8_array(n: u64) -> Vec { + let mut out = gguf_header(0, 1); + push_array_uint8_header(&mut out, "arr", n); + for i in 0..n { + out.push(i as u8); + } + out + } + let limits = ParseLimits { + max_array_work_items: 4, + ..ParseLimits::default() + }; + parse_ok(file_with_uint8_array(4), limits); + expect_limit( + parse_err(file_with_uint8_array(5), limits), + ParseLimitKind::ArrayWorkItems, + 5, + 4, + ); +} + +#[test] +fn metadata_bytes_boundary_and_plus_one() { + let mut bytes = gguf_header(0, 1); + push_kv_u32(&mut bytes, "a", 1); + let used = (bytes.len() - 24) as u64; + let ok_limits = ParseLimits { + max_metadata_bytes: used, + ..ParseLimits::default() + }; + parse_ok(bytes.clone(), ok_limits); + let tight = ParseLimits { + max_metadata_bytes: used - 1, + ..ParseLimits::default() + }; + expect_limit( + parse_err(bytes, tight), + ParseLimitKind::MetadataBytes, + used, + used - 1, + ); +} + +#[test] +fn u64_max_counts_and_strings_name_the_limit() { + let tensor_max = gguf_header(u64::MAX, 0); + expect_limit( + parse_err(tensor_max, ParseLimits::default()), + ParseLimitKind::TensorCount, + u64::MAX, + ParseLimits::DEFAULT.max_tensor_count, + ); + + let kv_max = gguf_header(0, u64::MAX); + expect_limit( + parse_err(kv_max, ParseLimits::default()), + ParseLimitKind::KvCount, + u64::MAX, + ParseLimits::DEFAULT.max_kv_count, + ); + + let mut huge_string = gguf_header(0, 1); + push_u64(&mut huge_string, u64::MAX); + expect_limit( + parse_err(huge_string, ParseLimits::default()), + ParseLimitKind::StringBytes, + u64::MAX, + ParseLimits::DEFAULT.max_string_bytes, + ); + + let mut huge_array = gguf_header(0, 1); + push_array_uint8_header(&mut huge_array, "arr", u64::MAX); + expect_limit( + parse_err(huge_array, ParseLimits::default()), + ParseLimitKind::ArrayWorkItems, + u64::MAX, + ParseLimits::DEFAULT.max_array_work_items, + ); +} + +#[test] +fn u64_max_relative_offset_cannot_wrap() { + let mut out = gguf_header(1, 0); + push_string(&mut out, "t"); + push_u32(&mut out, 1); + push_u64(&mut out, 1); + push_u32(&mut out, GGML_F32); + push_u64(&mut out, u64::MAX); + let err = parse_err(out, ParseLimits::default()); + match err { + ParserError::InvalidLayout { reason, .. } => { + assert!( + reason.contains("absolute offset overflow") || reason.contains("overflow"), + "got: {reason}" + ); + } + ParserError::HostSizeOverflow { field, value, .. } => { + assert_eq!(field, HostSizeField::RelativeOffset); + assert_eq!(value, u64::MAX); + } + other => panic!("expected wrap/host overflow, got {other}"), + } +} + +#[test] +fn deep_nested_arrays_are_bounded_by_total_work() { + const DEPTH: usize = 32; + let mut ok = gguf_header(0, 1); + push_string(&mut ok, "nested"); + push_u32(&mut ok, VT_ARRAY); + push_nested_array_payload(&mut ok, DEPTH); + + parse_ok(ok.clone(), ParseLimits::default()); + + // 32-level nest of length-1 arrays charges 31 work items (see cursor tests). + let exact = ParseLimits { + max_array_work_items: 31, + ..ParseLimits::default() + }; + parse_ok(ok.clone(), exact); + + let too_small = ParseLimits { + max_array_work_items: 30, + ..ParseLimits::default() + }; + let err = parse_err(ok, too_small); + match err { + ParserError::LimitExceeded { limit, .. } => { + assert_eq!(limit, ParseLimitKind::ArrayWorkItems); + } + other => panic!("expected array work limit, got {other}"), + } +} + +#[test] +fn pathological_rank_is_rejected_before_dim_iteration() { + let mut out = gguf_header(1, 0); + push_string(&mut out, "t"); + push_u32(&mut out, u32::MAX); + expect_limit( + parse_err(out, ParseLimits::default()), + ParseLimitKind::TensorRank, + u64::from(u32::MAX), + ParseLimits::DEFAULT.max_tensor_rank, + ); +} + +#[test] +fn element_count_multiplication_overflow() { + let half = (usize::MAX / 2) as u64 + 1; + let mut out = gguf_header(1, 0); + push_string(&mut out, "t"); + push_u32(&mut out, 2); + push_u64(&mut out, half); + push_u64(&mut out, half); + push_u32(&mut out, GGML_F32); + push_u64(&mut out, 0); + let err = parse_err(out, ParseLimits::default()); + match err { + ParserError::InvalidLayout { reason, .. } => { + assert!(reason.contains("element count overflow"), "got: {reason}"); + } + ParserError::HostSizeOverflow { field, .. } => { + assert_eq!(field, HostSizeField::TensorDim); + } + other => panic!("expected overflow, got {other}"), + } +} + +#[test] +fn truncated_buffers_at_each_header_directory_stage() { + let mut kv_and_tensor = gguf_header(1, 1); + push_kv_string(&mut kv_and_tensor, "general.architecture", "t"); + push_string(&mut kv_and_tensor, "weight"); + push_u32(&mut kv_and_tensor, 1); + push_u64(&mut kv_and_tensor, 1); + push_u32(&mut kv_and_tensor, GGML_F32); + push_u64(&mut kv_and_tensor, 0); + + let stages: [(&[u8], &str); 8] = [ + (&[][..], "empty"), + (b"GGU", "partial-magic"), + (b"GGUF", "magic-only"), + (&kv_and_tensor[..8], "magic-version"), + (&kv_and_tensor[..16], "tensor-count-only"), + (&kv_and_tensor[..24], "counts-no-kv"), + (&kv_and_tensor[..28], "partial-kv-key-len"), + ( + &kv_and_tensor[..kv_and_tensor.len() - 4], + "partial-tensor-offset", + ), + ]; + for (bytes, label) in stages { + let err = parse_bytes(bytes.to_vec(), format!("mem://trunc-{label}")).expect_err(label); + let msg = err.to_string(); + assert!( + msg.contains("EOF") + || msg.contains("overflow") + || msg.contains("unsupported") + || msg.contains("invalid") + || msg.contains("parse limit") + || msg.contains("host-size"), + "{label}: {msg}" + ); + } +} + +#[test] +fn trusted_override_does_not_change_default_safety() { + let three = kv_only(&[ + ("a", KvValue::U32(1)), + ("b", KvValue::U32(2)), + ("c", KvValue::U32(3)), + ]); + let tight = ParseLimits { + max_kv_count: 2, + ..ParseLimits::default() + }; + expect_limit( + parse_err(three.clone(), tight), + ParseLimitKind::KvCount, + 3, + 2, + ); + parse_ok(three, ParseLimits::trusted()); + assert_eq!(ParseLimits::default().max_kv_count, 1_000_000); +} + +#[cfg(feature = "mmap")] +fn write_temp_gguf(bytes: &[u8]) -> (std::path::PathBuf, TempGguf) { + use std::fs; + use std::process; + use std::time::{SystemTime, UNIX_EPOCH}; + + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("clock") + .as_nanos(); + let path = std::env::temp_dir().join(format!( + "engram-parser-limits-{}-{nanos}.gguf", + process::id() + )); + fs::write(&path, bytes).expect("write temp gguf"); + (path.clone(), TempGguf(path)) +} + +#[cfg(feature = "mmap")] +struct TempGguf(std::path::PathBuf); +#[cfg(feature = "mmap")] +impl Drop for TempGguf { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.0); + } +} + +#[cfg(feature = "mmap")] +#[test] +fn default_and_mmap_error_parity_on_limit_failures() { + use engram_parser::load_gguf_mmap_with_limits; + + let cases: Vec<(Vec, ParseLimits)> = vec![ + (gguf_header(u64::MAX, 0), ParseLimits::default()), + (gguf_header(0, u64::MAX), ParseLimits::default()), + ( + { + let mut b = gguf_header(0, 1); + push_u64(&mut b, u64::MAX); + b + }, + ParseLimits::default(), + ), + ( + kv_only(&[ + ("a", KvValue::U32(1)), + ("b", KvValue::U32(2)), + ("c", KvValue::U32(3)), + ]), + ParseLimits { + max_kv_count: 2, + ..ParseLimits::default() + }, + ), + ( + one_tensor("t", vec![2, 2, 1]), + ParseLimits { + max_tensor_rank: 2, + ..ParseLimits::default() + }, + ), + ]; + + for (i, (bytes, limits)) in cases.into_iter().enumerate() { + let owned = parse_bytes_with_limits(bytes.clone(), format!("mem://mmap-{i}"), limits) + .expect_err("owned"); + let (path, _guard) = write_temp_gguf(&bytes); + let mapped = load_gguf_mmap_with_limits(&path, limits).expect_err("mmap"); + assert_eq!( + error_payload(&owned), + error_payload(&mapped), + "default vs mmap error payload mismatch\n owned={owned}\n mmap={mapped}" + ); + } +} + +#[cfg(feature = "mmap")] +#[test] +fn default_mmap_success_parity_on_tiny_file() { + use engram_parser::load_gguf_mmap; + + let bytes = one_tensor("token_embd.weight", vec![4, 2]); + let (path, _guard) = write_temp_gguf(&bytes); + let owned = parse_bytes(bytes, "mem://mmap-ok".into()).expect("owned"); + let mapped = load_gguf_mmap(&path).expect("mmap"); + assert!(mapped.directory_matches(&owned)); +} From 2833e8f1f21c3c932da9743f36ee4fb472b8e605 Mon Sep 17 00:00:00 2001 From: Raul Cardenas Montoya Date: Fri, 18 Sep 2026 16:08:21 -0500 Subject: [PATCH 2/2] test(gguf): lock relative-offset metadata isolation Relative offset 0 is the GGUF data-section start. Checked addition cannot place a tensor in the header, and overlapping offsets still stay there. Co-authored-by: Cursor --- tests/gguf_limits.rs | 66 ++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/tests/gguf_limits.rs b/tests/gguf_limits.rs index c4b767d..81146cb 100644 --- a/tests/gguf_limits.rs +++ b/tests/gguf_limits.rs @@ -242,6 +242,72 @@ fn u64_max_counts_and_strings_name_the_limit() { ); } +fn two_f32_tensors_at_relative_offset(offset: u64) -> Vec { + let mut out = gguf_header(2, 0); + for name in ["left", "right"] { + push_string(&mut out, name); + push_u32(&mut out, 1); + push_u64(&mut out, 1); + push_u32(&mut out, GGML_F32); + push_u64(&mut out, offset); + } + while !out.len().is_multiple_of(ALIGNMENT as usize) { + out.push(0); + } + out.extend_from_slice(&1.0f32.to_le_bytes()); + out +} + +#[test] +fn relative_zero_cannot_expose_header_bytes() { + let bytes = one_tensor("w", vec![2]); + let layout = parse_bytes(bytes, "mem://rel0".into()).expect("parse"); + let tensor = layout.tensor("w").expect("tensor"); + let payload = layout.tensor_bytes(tensor).expect("payload"); + assert_all(&[ + (tensor.relative_offset == 0, "relative-zero"), + ( + tensor.absolute_offset == layout.tensor_data_offset, + "absolute-at-data-section", + ), + ( + tensor.absolute_offset >= layout.tensor_data_offset, + "not-before-data-section", + ), + (layout.bytes[..4] == GGUF_MAGIC, "header-magic-present"), + (!payload.starts_with(&GGUF_MAGIC), "payload-not-header"), + ( + layout.bytes[..layout.tensor_data_offset].starts_with(&GGUF_MAGIC), + "metadata-keeps-header", + ), + ]); +} + +#[test] +fn overlapping_relative_offsets_stay_in_tensor_data() { + let bytes = two_f32_tensors_at_relative_offset(0); + let layout = parse_bytes(bytes, "mem://overlap".into()).expect("parse overlap"); + let left = layout.tensor("left").expect("left"); + let right = layout.tensor("right").expect("right"); + let left_bytes = layout.tensor_bytes(left).expect("left payload"); + let right_bytes = layout.tensor_bytes(right).expect("right payload"); + assert_all(&[ + (left.relative_offset == 0, "left-rel"), + (right.relative_offset == 0, "right-rel"), + ( + left.absolute_offset == layout.tensor_data_offset, + "left-abs", + ), + ( + right.absolute_offset == layout.tensor_data_offset, + "right-abs", + ), + (!left_bytes.starts_with(&GGUF_MAGIC), "left-not-header"), + (!right_bytes.starts_with(&GGUF_MAGIC), "right-not-header"), + (left_bytes == right_bytes, "shared-data-section-bytes"), + ]); +} + #[test] fn u64_max_relative_offset_cannot_wrap() { let mut out = gguf_header(1, 0);