diff --git a/CHANGELOG.md b/CHANGELOG.md index 2fcd89c..998112d 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. - **Safetensors shard-index invariants ([RM-1360](https://linear.app/rpd-34/issue/RM-1360)):** index shard paths resolve relative to the checkpoint root; absolute paths and `..` traversal that would leave the root are rejected; existing paths diff --git a/README.md b/README.md index f5ebd6f..aee6d1a 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. @@ -220,13 +221,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 644ea3c..9a270d6 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,26 @@ 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, + }, /// A tensor name is claimed by more than one Safetensors shard. DuplicateTensorOwnership { /// Tensor name with conflicting owners. @@ -70,6 +170,30 @@ pub enum ParserError { }, } +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 { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { @@ -91,6 +215,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" + ), Self::DuplicateTensorOwnership { name, shards, path } => { write!( f, @@ -113,6 +250,8 @@ impl std::error::Error for ParserError { | Self::MissingTensor { .. } | Self::InvalidLayout { .. } | Self::ExpertOutOfRange { .. } + | Self::LimitExceeded { .. } + | Self::HostSizeOverflow { .. } | Self::DuplicateTensorOwnership { .. } | Self::MissingShard { .. } => None, } @@ -121,3 +260,68 @@ impl std::error::Error for ParserError { /// 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..81146cb --- /dev/null +++ b/tests/gguf_limits.rs @@ -0,0 +1,545 @@ +// 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}") + } + ParserError::DuplicateTensorOwnership { name, shards, .. } => { + format!("dup:{name}:{}", shards.join(",")) + } + ParserError::MissingShard { shard, .. } => format!("missing-shard:{shard}"), + } +} + +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, + ); +} + +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); + 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)); +}