gguf_rs_lib/format/
metadata.rs

1//! GGUF metadata structures and operations
2
3use crate::error::{GGUFError, Result};
4use crate::format::constants::*;
5use crate::format::types::GGUFValueType;
6
7#[cfg(feature = "std")]
8use crate::format::endian::{
9    read_f32, read_f64, read_i16, read_i32, read_i64, read_i8, read_u16, read_u32, read_u64,
10    read_u8, write_f32, write_f64, write_i16, write_i32, write_i64, write_i8, write_u16, write_u32,
11    write_u64, write_u8,
12};
13#[cfg(feature = "std")]
14use std::collections::HashMap;
15#[cfg(feature = "std")]
16use std::io::{Read, Write};
17
18#[cfg(feature = "std")]
19use serde::{Deserialize, Serialize};
20
21#[cfg(all(not(feature = "std"), feature = "alloc"))]
22extern crate alloc;
23#[cfg(all(not(feature = "std"), feature = "alloc"))]
24use alloc::{
25    boxed::Box,
26    format,
27    string::{String, ToString},
28    vec::Vec,
29};
30#[cfg(all(not(feature = "std"), feature = "alloc"))]
31use hashbrown::HashMap;
32
33/// Variable-sized metadata is grown in bounded increments rather than trusting
34/// attacker-controlled `u64` length fields as allocation sizes.
35#[cfg(feature = "std")]
36const METADATA_BYTE_READ_CHUNK: usize = 8 * 1024;
37#[cfg(feature = "std")]
38const METADATA_ARRAY_RESERVE_CHUNK: usize = 1024;
39
40#[cfg(feature = "std")]
41#[derive(Debug)]
42struct MetadataDecodeBudget {
43    remaining: usize,
44    limit: usize,
45}
46
47#[cfg(feature = "std")]
48impl MetadataDecodeBudget {
49    fn new(limit: usize) -> Self {
50        Self { remaining: limit, limit }
51    }
52
53    fn charge(&mut self, bytes: usize, description: &str) -> Result<()> {
54        self.remaining = self.remaining.checked_sub(bytes).ok_or_else(|| {
55            GGUFError::InvalidMetadata(format!(
56                "Decoded metadata allocation exceeds budget of {} bytes while reading {}",
57                self.limit, description
58            ))
59        })?;
60        Ok(())
61    }
62}
63
64// Import core modules for no_std compatibility
65#[cfg(not(feature = "std"))]
66use core::{fmt, slice};
67
68/// A metadata value in a GGUF file
69#[cfg_attr(feature = "std", derive(Serialize, Deserialize))]
70#[derive(Debug, Clone, PartialEq)]
71pub enum MetadataValue {
72    /// 8-bit unsigned integer
73    U8(u8),
74    /// 8-bit signed integer
75    I8(i8),
76    /// 16-bit unsigned integer
77    U16(u16),
78    /// 16-bit signed integer
79    I16(i16),
80    /// 32-bit unsigned integer
81    U32(u32),
82    /// 32-bit signed integer
83    I32(i32),
84    /// 32-bit floating point
85    F32(f32),
86    /// Boolean value
87    Bool(bool),
88    /// UTF-8 string
89    String(String),
90    /// Array of values
91    Array(Box<MetadataArray>),
92    /// 64-bit unsigned integer
93    U64(u64),
94    /// 64-bit signed integer
95    I64(i64),
96    /// 64-bit floating point
97    F64(f64),
98}
99
100/// An array of metadata values (all of the same type)
101#[cfg_attr(feature = "std", derive(Serialize, Deserialize))]
102#[derive(Debug, Clone, PartialEq)]
103pub struct MetadataArray {
104    /// Type of elements in the array
105    pub element_type: GGUFValueType,
106    /// Number of elements
107    pub length: u64,
108    /// Array elements
109    pub values: Vec<MetadataValue>,
110}
111
112/// Collection of metadata key-value pairs
113#[cfg_attr(feature = "std", derive(Serialize, Deserialize))]
114#[derive(Debug, Clone)]
115pub struct Metadata {
116    /// Metadata key-value pairs
117    #[cfg(any(feature = "std", feature = "alloc"))]
118    pub data: HashMap<String, MetadataValue>,
119    #[cfg(not(any(feature = "std", feature = "alloc")))]
120    pub data: (), // Placeholder for no_std + no_alloc builds
121}
122
123#[cfg(any(feature = "std", feature = "alloc"))]
124impl Default for Metadata {
125    fn default() -> Self {
126        Self::new()
127    }
128}
129
130#[cfg(not(any(feature = "std", feature = "alloc")))]
131impl Default for Metadata {
132    fn default() -> Self {
133        Self { data: () }
134    }
135}
136
137impl MetadataValue {
138    /// Get the type of this metadata value
139    pub fn value_type(&self) -> GGUFValueType {
140        match self {
141            MetadataValue::U8(_) => GGUFValueType::U8,
142            MetadataValue::I8(_) => GGUFValueType::I8,
143            MetadataValue::U16(_) => GGUFValueType::U16,
144            MetadataValue::I16(_) => GGUFValueType::I16,
145            MetadataValue::U32(_) => GGUFValueType::U32,
146            MetadataValue::I32(_) => GGUFValueType::I32,
147            MetadataValue::F32(_) => GGUFValueType::F32,
148            MetadataValue::Bool(_) => GGUFValueType::Bool,
149            MetadataValue::String(_) => GGUFValueType::String,
150            MetadataValue::Array(_) => GGUFValueType::Array,
151            MetadataValue::U64(_) => GGUFValueType::U64,
152            MetadataValue::I64(_) => GGUFValueType::I64,
153            MetadataValue::F64(_) => GGUFValueType::F64,
154        }
155    }
156
157    /// Calculate the serialized size of this value
158    pub fn serialized_size(&self) -> usize {
159        self.checked_serialized_size().unwrap_or(usize::MAX)
160    }
161
162    /// Calculate the serialized size without overflowing.
163    pub fn checked_serialized_size(&self) -> Option<usize> {
164        match self {
165            MetadataValue::U8(_) | MetadataValue::I8(_) | MetadataValue::Bool(_) => Some(1),
166            MetadataValue::U16(_) | MetadataValue::I16(_) => Some(2),
167            MetadataValue::U32(_) | MetadataValue::I32(_) | MetadataValue::F32(_) => Some(4),
168            MetadataValue::String(s) => 8usize.checked_add(s.len()),
169            MetadataValue::Array(arr) => arr.as_ref().checked_serialized_size(),
170            MetadataValue::U64(_) | MetadataValue::I64(_) | MetadataValue::F64(_) => Some(8),
171        }
172    }
173
174    /// Validate length, type, and nesting invariants before serialization.
175    pub fn validate(&self) -> Result<()> {
176        self.validate_at_depth(0)
177    }
178
179    fn validate_at_depth(&self, depth: usize) -> Result<()> {
180        match self {
181            MetadataValue::Array(array) => array.validate_at_depth(depth + 1),
182            _ => Ok(()),
183        }
184    }
185
186    /// Read a metadata value from a reader
187    #[cfg(feature = "std")]
188    pub fn read_from<R: Read>(reader: &mut R, value_type: GGUFValueType) -> Result<Self> {
189        let mut budget = MetadataDecodeBudget::new(GGUF_MAX_METADATA_DECODED_SIZE);
190        Self::read_from_at_depth(reader, value_type, 0, &mut budget)
191    }
192
193    #[cfg(feature = "std")]
194    fn read_from_at_depth<R: Read>(
195        reader: &mut R,
196        value_type: GGUFValueType,
197        depth: usize,
198        budget: &mut MetadataDecodeBudget,
199    ) -> Result<Self> {
200        let value = match value_type {
201            GGUFValueType::U8 => MetadataValue::U8(read_u8(reader)?),
202            GGUFValueType::I8 => MetadataValue::I8(read_i8(reader)?),
203            GGUFValueType::U16 => MetadataValue::U16(read_u16(reader)?),
204            GGUFValueType::I16 => MetadataValue::I16(read_i16(reader)?),
205            GGUFValueType::U32 => MetadataValue::U32(read_u32(reader)?),
206            GGUFValueType::I32 => MetadataValue::I32(read_i32(reader)?),
207            GGUFValueType::F32 => MetadataValue::F32(read_f32(reader)?),
208            GGUFValueType::Bool => match read_u8(reader)? {
209                0 => MetadataValue::Bool(false),
210                1 => MetadataValue::Bool(true),
211                value => {
212                    return Err(GGUFError::InvalidMetadata(format!(
213                        "Invalid boolean value: {} (expected 0 or 1)",
214                        value
215                    )));
216                }
217            },
218            GGUFValueType::String => {
219                let len = read_u64(reader)?;
220                let bytes = read_variable_bytes(reader, len, "string", budget)?;
221                let string = String::from_utf8(bytes)
222                    .map_err(|e| GGUFError::Format(format!("Invalid UTF-8 string: {}", e)))?;
223                MetadataValue::String(string)
224            }
225            GGUFValueType::Array => {
226                if depth >= GGUF_MAX_METADATA_NESTING_DEPTH {
227                    return Err(GGUFError::InvalidMetadata(format!(
228                        "Metadata array nesting exceeds limit of {}",
229                        GGUF_MAX_METADATA_NESTING_DEPTH
230                    )));
231                }
232                budget.charge(core::mem::size_of::<MetadataArray>(), "nested array")?;
233                let array = MetadataArray::read_from_at_depth(reader, depth + 1, budget)?;
234                MetadataValue::Array(Box::new(array))
235            }
236            GGUFValueType::U64 => MetadataValue::U64(read_u64(reader)?),
237            GGUFValueType::I64 => MetadataValue::I64(read_i64(reader)?),
238            GGUFValueType::F64 => MetadataValue::F64(read_f64(reader)?),
239        };
240
241        Ok(value)
242    }
243
244    /// Write a metadata value to a writer  
245    #[cfg(feature = "std")]
246    pub fn write_to<W: Write>(&self, writer: &mut W) -> Result<()> {
247        self.validate()?;
248        match self {
249            MetadataValue::U8(v) => write_u8(writer, *v)?,
250            MetadataValue::I8(v) => write_i8(writer, *v)?,
251            MetadataValue::U16(v) => write_u16(writer, *v)?,
252            MetadataValue::I16(v) => write_i16(writer, *v)?,
253            MetadataValue::U32(v) => write_u32(writer, *v)?,
254            MetadataValue::I32(v) => write_i32(writer, *v)?,
255            MetadataValue::F32(v) => write_f32(writer, *v)?,
256            MetadataValue::Bool(v) => write_u8(writer, if *v { 1 } else { 0 })?,
257            MetadataValue::String(s) => {
258                let length = u64::try_from(s.len()).map_err(|_| {
259                    GGUFError::InvalidMetadata(
260                        "String length does not fit the GGUF u64 length field".to_string(),
261                    )
262                })?;
263                write_u64(writer, length)?;
264                writer.write_all(s.as_bytes())?;
265            }
266            MetadataValue::Array(arr) => arr.as_ref().write_to(writer)?,
267            MetadataValue::U64(v) => write_u64(writer, *v)?,
268            MetadataValue::I64(v) => write_i64(writer, *v)?,
269            MetadataValue::F64(v) => write_f64(writer, *v)?,
270        }
271        Ok(())
272    }
273
274    /// Convert to a string representation for display
275    #[cfg(any(feature = "std", feature = "alloc"))]
276    pub fn to_string_representation(&self) -> String {
277        match self {
278            MetadataValue::U8(v) => v.to_string(),
279            MetadataValue::I8(v) => v.to_string(),
280            MetadataValue::U16(v) => v.to_string(),
281            MetadataValue::I16(v) => v.to_string(),
282            MetadataValue::U32(v) => v.to_string(),
283            MetadataValue::I32(v) => v.to_string(),
284            MetadataValue::F32(v) => v.to_string(),
285            MetadataValue::Bool(v) => v.to_string(),
286            MetadataValue::String(s) => s.clone(),
287            MetadataValue::Array(arr) => format!("Array[{}; {}]", arr.element_type, arr.length),
288            MetadataValue::U64(v) => v.to_string(),
289            MetadataValue::I64(v) => v.to_string(),
290            MetadataValue::F64(v) => v.to_string(),
291        }
292    }
293
294    /// Try to convert to a specific type
295    pub fn as_str(&self) -> Option<&str> {
296        match self {
297            MetadataValue::String(s) => Some(s),
298            _ => None,
299        }
300    }
301
302    pub fn as_u64(&self) -> Option<u64> {
303        match self {
304            MetadataValue::U8(v) => Some(*v as u64),
305            MetadataValue::U16(v) => Some(*v as u64),
306            MetadataValue::U32(v) => Some(*v as u64),
307            MetadataValue::U64(v) => Some(*v),
308            _ => None,
309        }
310    }
311
312    pub fn as_i64(&self) -> Option<i64> {
313        match self {
314            MetadataValue::I8(v) => Some(*v as i64),
315            MetadataValue::I16(v) => Some(*v as i64),
316            MetadataValue::I32(v) => Some(*v as i64),
317            MetadataValue::I64(v) => Some(*v),
318            _ => None,
319        }
320    }
321
322    pub fn as_f64(&self) -> Option<f64> {
323        match self {
324            MetadataValue::F32(v) => Some(*v as f64),
325            MetadataValue::F64(v) => Some(*v),
326            _ => None,
327        }
328    }
329
330    pub fn as_bool(&self) -> Option<bool> {
331        match self {
332            MetadataValue::Bool(v) => Some(*v),
333            _ => None,
334        }
335    }
336}
337
338impl MetadataArray {
339    /// Create a new metadata array
340    #[cfg(any(feature = "std", feature = "alloc"))]
341    pub fn new(element_type: GGUFValueType, values: Vec<MetadataValue>) -> Result<Self> {
342        // Validate that all values have the same type as specified
343        for value in &values {
344            if value.value_type() != element_type {
345                return Err(GGUFError::InvalidMetadata(format!(
346                    "Array element type mismatch: expected {}, got {}",
347                    element_type,
348                    value.value_type()
349                )));
350            }
351        }
352
353        let length = u64::try_from(values.len()).map_err(|_| {
354            GGUFError::InvalidMetadata(
355                "Array length does not fit the GGUF u64 length field".to_string(),
356            )
357        })?;
358
359        let array = Self { element_type, length, values };
360        array.validate()?;
361        Ok(array)
362    }
363
364    #[cfg(not(any(feature = "std", feature = "alloc")))]
365    pub fn new(_element_type: GGUFValueType, _values: &[MetadataValue]) -> Result<Self> {
366        Err(GGUFError::AllocationRequired)
367    }
368
369    /// Calculate the serialized size of this array
370    pub fn serialized_size(&self) -> usize {
371        self.checked_serialized_size().unwrap_or(usize::MAX)
372    }
373
374    /// Calculate the serialized array size without overflowing.
375    pub fn checked_serialized_size(&self) -> Option<usize> {
376        self.values
377            .iter()
378            .try_fold(12usize, |size, value| size.checked_add(value.checked_serialized_size()?))
379    }
380
381    /// Validate the array's length, homogeneous element type, and nesting depth.
382    pub fn validate(&self) -> Result<()> {
383        self.validate_at_depth(1)
384    }
385
386    fn validate_at_depth(&self, depth: usize) -> Result<()> {
387        if depth > GGUF_MAX_METADATA_NESTING_DEPTH {
388            return Err(GGUFError::InvalidMetadata(format!(
389                "Metadata array nesting exceeds limit of {}",
390                GGUF_MAX_METADATA_NESTING_DEPTH
391            )));
392        }
393        let actual_length = u64::try_from(self.values.len()).map_err(|_| {
394            GGUFError::InvalidMetadata(
395                "Array length does not fit the GGUF u64 length field".to_string(),
396            )
397        })?;
398        if self.length != actual_length {
399            return Err(GGUFError::InvalidMetadata(format!(
400                "Array length field {} does not match {} values",
401                self.length,
402                self.values.len()
403            )));
404        }
405        for value in &self.values {
406            if value.value_type() != self.element_type {
407                return Err(GGUFError::InvalidMetadata(format!(
408                    "Array element type mismatch: expected {}, got {}",
409                    self.element_type,
410                    value.value_type()
411                )));
412            }
413            value.validate_at_depth(depth)?;
414        }
415        Ok(())
416    }
417
418    /// Read a metadata array from a reader
419    #[cfg(feature = "std")]
420    pub fn read_from<R: Read>(reader: &mut R) -> Result<Self> {
421        let mut budget = MetadataDecodeBudget::new(GGUF_MAX_METADATA_DECODED_SIZE);
422        Self::read_from_at_depth(reader, 1, &mut budget)
423    }
424
425    #[cfg(feature = "std")]
426    fn read_from_at_depth<R: Read>(
427        reader: &mut R,
428        depth: usize,
429        budget: &mut MetadataDecodeBudget,
430    ) -> Result<Self> {
431        if depth > GGUF_MAX_METADATA_NESTING_DEPTH {
432            return Err(GGUFError::InvalidMetadata(format!(
433                "Metadata array nesting exceeds limit of {}",
434                GGUF_MAX_METADATA_NESTING_DEPTH
435            )));
436        }
437        let element_type_raw = read_u32(reader)?;
438        let element_type = GGUFValueType::from_u32(element_type_raw)?;
439        let length = read_u64(reader)?;
440
441        let platform_length = usize::try_from(length).map_err(|_| {
442            GGUFError::InvalidMetadata("Array length does not fit usize".to_string())
443        })?;
444        let decoded_slots = platform_length
445            .checked_mul(core::mem::size_of::<MetadataValue>())
446            .ok_or_else(|| {
447                GGUFError::InvalidMetadata(
448                    "Decoded metadata array allocation overflows usize".to_string(),
449                )
450            })?;
451        budget.charge(decoded_slots, "array elements")?;
452        let mut values = Vec::new();
453        while values.len() < platform_length {
454            if values.len() == values.capacity() {
455                let remaining = platform_length - values.len();
456                let additional = values.capacity().max(METADATA_ARRAY_RESERVE_CHUNK).min(remaining);
457                values.try_reserve_exact(additional).map_err(|_| {
458                    GGUFError::InvalidMetadata(format!(
459                        "Unable to allocate storage for metadata array of {} elements",
460                        length
461                    ))
462                })?;
463            }
464            let batch_end = values
465                .len()
466                .saturating_add(
467                    (values.capacity() - values.len()).min(METADATA_ARRAY_RESERVE_CHUNK),
468                )
469                .min(platform_length);
470            while values.len() < batch_end {
471                let value = MetadataValue::read_from_at_depth(reader, element_type, depth, budget)?;
472                values.push(value);
473            }
474        }
475
476        Ok(Self { element_type, length, values })
477    }
478
479    /// Write a metadata array to a writer
480    #[cfg(feature = "std")]
481    pub fn write_to<W: Write>(&self, writer: &mut W) -> Result<()> {
482        self.validate()?;
483        write_u32(writer, self.element_type as u32)?;
484        write_u64(writer, self.length)?;
485
486        for value in &self.values {
487            value.write_to(writer)?;
488        }
489
490        Ok(())
491    }
492
493    /// Get the length of the array
494    pub fn len(&self) -> usize {
495        self.values.len()
496    }
497
498    /// Check if the array is empty
499    pub fn is_empty(&self) -> bool {
500        self.values.is_empty()
501    }
502
503    /// Get an element by index
504    pub fn get(&self, index: usize) -> Option<&MetadataValue> {
505        self.values.get(index)
506    }
507
508    /// Iterate over the values
509    #[cfg(feature = "std")]
510    pub fn iter(&self) -> std::slice::Iter<'_, MetadataValue> {
511        self.values.iter()
512    }
513
514    #[cfg(not(feature = "std"))]
515    pub fn iter(&self) -> slice::Iter<'_, MetadataValue> {
516        self.values.iter()
517    }
518}
519
520#[cfg(feature = "std")]
521fn read_variable_bytes<R: Read>(
522    reader: &mut R,
523    length: u64,
524    kind: &str,
525    budget: &mut MetadataDecodeBudget,
526) -> Result<Vec<u8>> {
527    let length = usize::try_from(length)
528        .map_err(|_| GGUFError::Format(format!("{} length does not fit usize", kind)))?;
529    budget.charge(length, kind)?;
530    let mut bytes = Vec::new();
531    let mut chunk = [0u8; METADATA_BYTE_READ_CHUNK];
532    while bytes.len() < length {
533        if bytes.len() == bytes.capacity() {
534            let remaining = length - bytes.len();
535            let additional = bytes.capacity().max(METADATA_BYTE_READ_CHUNK).min(remaining);
536            bytes.try_reserve_exact(additional).map_err(|_| {
537                GGUFError::InvalidMetadata(format!(
538                    "Unable to allocate storage for metadata {} of {} bytes",
539                    kind, length
540                ))
541            })?;
542        }
543        let chunk_len = (length - bytes.len()).min(chunk.len());
544        reader.read_exact(&mut chunk[..chunk_len])?;
545        bytes.extend_from_slice(&chunk[..chunk_len]);
546    }
547    Ok(bytes)
548}
549
550#[cfg(any(feature = "std", feature = "alloc"))]
551impl Metadata {
552    /// Create a new empty metadata collection
553    pub fn new() -> Self {
554        Self { data: HashMap::new() }
555    }
556
557    /// Insert a key-value pair
558    pub fn insert(&mut self, key: String, value: MetadataValue) {
559        self.data.insert(key, value);
560    }
561
562    /// Get a value by key
563    pub fn get(&self, key: &str) -> Option<&MetadataValue> {
564        self.data.get(key)
565    }
566
567    /// Remove a key-value pair
568    pub fn remove(&mut self, key: &str) -> Option<MetadataValue> {
569        self.data.remove(key)
570    }
571
572    /// Check if a key exists
573    pub fn contains_key(&self, key: &str) -> bool {
574        self.data.contains_key(key)
575    }
576
577    /// Get the number of key-value pairs
578    pub fn len(&self) -> usize {
579        self.data.len()
580    }
581
582    /// Check if the metadata is empty
583    pub fn is_empty(&self) -> bool {
584        self.data.is_empty()
585    }
586
587    /// Iterate over key-value pairs
588    #[cfg(feature = "std")]
589    pub fn iter(&self) -> std::collections::hash_map::Iter<'_, String, MetadataValue> {
590        self.data.iter()
591    }
592
593    #[cfg(all(not(feature = "std"), feature = "alloc"))]
594    pub fn iter(&self) -> hashbrown::hash_map::Iter<'_, String, MetadataValue> {
595        self.data.iter()
596    }
597
598    /// Get all keys
599    #[cfg(feature = "std")]
600    pub fn keys(&self) -> std::collections::hash_map::Keys<'_, String, MetadataValue> {
601        self.data.keys()
602    }
603
604    #[cfg(all(not(feature = "std"), feature = "alloc"))]
605    pub fn keys(&self) -> hashbrown::hash_map::Keys<'_, String, MetadataValue> {
606        self.data.keys()
607    }
608
609    /// Get all values
610    #[cfg(feature = "std")]
611    pub fn values(&self) -> std::collections::hash_map::Values<'_, String, MetadataValue> {
612        self.data.values()
613    }
614
615    #[cfg(all(not(feature = "std"), feature = "alloc"))]
616    pub fn values(&self) -> hashbrown::hash_map::Values<'_, String, MetadataValue> {
617        self.data.values()
618    }
619
620    /// Validate metadata invariants before writing or using format-level settings.
621    pub fn validate(&self) -> Result<()> {
622        if self.data.len() as u64 > GGUF_MAX_METADATA_COUNT {
623            return Err(GGUFError::InvalidMetadata(format!(
624                "Too many metadata entries: {} exceeds limit of {}",
625                self.data.len(),
626                GGUF_MAX_METADATA_COUNT
627            )));
628        }
629        for (key, value) in &self.data {
630            if key.is_empty() {
631                return Err(GGUFError::InvalidMetadata("Metadata key cannot be empty".to_string()));
632            }
633            if key.len() > GGUF_MAX_METADATA_KEY_LENGTH {
634                return Err(GGUFError::InvalidMetadata(format!(
635                    "Metadata key is {} bytes; maximum is {}",
636                    key.len(),
637                    GGUF_MAX_METADATA_KEY_LENGTH
638                )));
639            }
640            if !is_valid_metadata_key(key) {
641                return Err(GGUFError::InvalidMetadata(format!(
642                    "Metadata key '{}' must be dot-separated lower_snake_case ASCII",
643                    key
644                )));
645            }
646            value.validate()?;
647        }
648        self.tensor_alignment()?;
649        Ok(())
650    }
651
652    /// Return the global tensor alignment declared by `general.alignment`.
653    pub fn tensor_alignment(&self) -> Result<usize> {
654        let alignment = match self.get("general.alignment") {
655            None => GGUF_DEFAULT_ALIGNMENT,
656            Some(MetadataValue::U32(value)) => usize::try_from(*value).map_err(|_| {
657                GGUFError::InvalidMetadata(
658                    "general.alignment does not fit this platform".to_string(),
659                )
660            })?,
661            Some(value) => {
662                return Err(GGUFError::InvalidMetadata(format!(
663                    "general.alignment must be u32, got {}",
664                    value.value_type()
665                )));
666            }
667        };
668        if alignment == 0 || alignment % 8 != 0 {
669            return Err(GGUFError::InvalidMetadata(format!(
670                "general.alignment must be a non-zero multiple of 8, got {}",
671                alignment
672            )));
673        }
674        Ok(alignment)
675    }
676}
677
678#[cfg(any(feature = "std", feature = "alloc"))]
679fn is_valid_metadata_key(key: &str) -> bool {
680    key.is_ascii()
681        && key.split('.').all(|component| {
682            !component.is_empty()
683                && component.split('_').all(|word| {
684                    !word.is_empty()
685                        && word
686                            .bytes()
687                            .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit())
688                })
689        })
690}
691
692#[cfg(not(any(feature = "std", feature = "alloc")))]
693impl Metadata {
694    /// Create a new empty metadata collection (no-op for no_std + no_alloc)
695    pub fn new() -> Self {
696        Self { data: () }
697    }
698
699    /// Insert a key-value pair (no-op for no_std + no_alloc)
700    pub fn insert(&mut self, _key: String, _value: MetadataValue) {
701        // No-op: can't store data without allocation
702    }
703
704    /// Get a value by key (always returns None for no_std + no_alloc)
705    pub fn get(&self, _key: &str) -> Option<&MetadataValue> {
706        None
707    }
708
709    /// Remove a key-value pair (always returns None for no_std + no_alloc)
710    pub fn remove(&mut self, _key: &str) -> Option<MetadataValue> {
711        None
712    }
713
714    /// Check if a key exists (always returns false for no_std + no_alloc)
715    pub fn contains_key(&self, _key: &str) -> bool {
716        false
717    }
718
719    /// Get the number of key-value pairs (always 0 for no_std + no_alloc)
720    pub fn len(&self) -> usize {
721        0
722    }
723
724    /// Check if the metadata is empty (always true for no_std + no_alloc)
725    pub fn is_empty(&self) -> bool {
726        true
727    }
728}
729
730#[cfg(any(feature = "std", feature = "alloc"))]
731impl Metadata {
732    /// Read metadata from a reader
733    #[cfg(feature = "std")]
734    pub fn read_from<R: Read>(reader: &mut R, count: u64) -> Result<Self> {
735        let mut budget = MetadataDecodeBudget::new(GGUF_MAX_METADATA_DECODED_SIZE);
736        Self::read_from_with_decode_budget(reader, count, &mut budget)
737    }
738
739    #[cfg(feature = "std")]
740    fn read_from_with_decode_budget<R: Read>(
741        reader: &mut R,
742        count: u64,
743        budget: &mut MetadataDecodeBudget,
744    ) -> Result<Self> {
745        if count > GGUF_MAX_METADATA_COUNT {
746            return Err(GGUFError::InvalidMetadata(format!(
747                "Metadata count {} exceeds limit of {}",
748                count, GGUF_MAX_METADATA_COUNT
749            )));
750        }
751        let capacity = usize::try_from(count).map_err(|_| {
752            GGUFError::InvalidMetadata("Metadata count does not fit usize".to_string())
753        })?;
754        let decoded_entries = capacity
755            .checked_mul(
756                core::mem::size_of::<(String, MetadataValue)>()
757                    .saturating_add(2 * core::mem::size_of::<usize>()),
758            )
759            .ok_or_else(|| {
760                GGUFError::InvalidMetadata(
761                    "Decoded metadata entry allocation overflows usize".to_string(),
762                )
763            })?;
764        budget.charge(decoded_entries, "metadata entries")?;
765        let mut data = HashMap::new();
766        data.try_reserve(capacity).map_err(|_| {
767            GGUFError::InvalidMetadata("Unable to allocate metadata map".to_string())
768        })?;
769
770        for _ in 0..count {
771            // Read key
772            let key_len = read_u64(reader)?;
773            if key_len > GGUF_MAX_METADATA_KEY_LENGTH as u64 {
774                return Err(GGUFError::Format(format!("Metadata key too long: {} bytes", key_len)));
775            }
776            let key_len = usize::try_from(key_len).map_err(|_| {
777                GGUFError::Format("Metadata key length does not fit usize".to_string())
778            })?;
779            budget.charge(key_len, "metadata key")?;
780
781            let mut key_bytes = vec![0u8; key_len];
782            reader.read_exact(&mut key_bytes)?;
783            let key = String::from_utf8(key_bytes)
784                .map_err(|e| GGUFError::Format(format!("Invalid UTF-8 in metadata key: {}", e)))?;
785            if key.is_empty() {
786                return Err(GGUFError::InvalidMetadata("Metadata key cannot be empty".to_string()));
787            }
788            if data.contains_key(&key) {
789                return Err(GGUFError::InvalidMetadata(format!("Duplicate metadata key: {}", key)));
790            }
791
792            // Read value type and value
793            let value_type_raw = read_u32(reader)?;
794            let value_type = GGUFValueType::from_u32(value_type_raw)?;
795            let value = MetadataValue::read_from_at_depth(reader, value_type, 0, budget)?;
796
797            data.insert(key, value);
798        }
799
800        let metadata = Self { data };
801        metadata.validate()?;
802        Ok(metadata)
803    }
804
805    /// Read metadata with a serialized-byte budget and the default decoded-allocation budget.
806    #[cfg(feature = "std")]
807    pub fn read_from_with_limit<R: Read>(
808        reader: &mut R,
809        count: u64,
810        max_bytes: usize,
811    ) -> Result<Self> {
812        Self::read_from_with_limits(reader, count, max_bytes, GGUF_MAX_METADATA_DECODED_SIZE)
813    }
814
815    /// Read metadata with independent serialized-byte and decoded-allocation budgets.
816    #[cfg(feature = "std")]
817    pub fn read_from_with_limits<R: Read>(
818        reader: &mut R,
819        count: u64,
820        max_bytes: usize,
821        max_decoded_bytes: usize,
822    ) -> Result<Self> {
823        let limit = u64::try_from(max_bytes).unwrap_or(u64::MAX);
824        let mut limited = reader.take(limit);
825        let mut budget = MetadataDecodeBudget::new(max_decoded_bytes);
826        match Self::read_from_with_decode_budget(&mut limited, count, &mut budget) {
827            Ok(metadata) => Ok(metadata),
828            Err(_) if limited.limit() == 0 => Err(GGUFError::InvalidMetadata(format!(
829                "Metadata exceeds byte limit of {}",
830                max_bytes
831            ))),
832            Err(error) => Err(error),
833        }
834    }
835
836    /// Write metadata to a writer
837    #[cfg(feature = "std")]
838    pub fn write_to<W: Write>(&self, writer: &mut W) -> Result<()> {
839        self.validate()?;
840        let mut entries = Vec::new();
841        entries.try_reserve_exact(self.data.len()).map_err(|_| {
842            GGUFError::InvalidMetadata(format!(
843                "Unable to allocate deterministic metadata index for {} entries",
844                self.data.len()
845            ))
846        })?;
847        entries.extend(self.data.iter());
848        entries.sort_unstable_by_key(|(key, _)| *key);
849        for (key, value) in entries {
850            // Write key
851            let key_len = u64::try_from(key.len()).map_err(|_| {
852                GGUFError::InvalidMetadata(format!(
853                    "Metadata key length does not fit the GGUF u64 length field: {} bytes",
854                    key.len()
855                ))
856            })?;
857            write_u64(writer, key_len)?;
858            writer.write_all(key.as_bytes())?;
859
860            // Write value type and value
861            write_u32(writer, value.value_type() as u32)?;
862            value.write_to(writer)?;
863        }
864
865        Ok(())
866    }
867
868    /// Calculate the serialized size of all metadata
869    pub fn serialized_size(&self) -> usize {
870        self.checked_serialized_size().unwrap_or(usize::MAX)
871    }
872
873    /// Calculate the serialized metadata size without overflowing.
874    pub fn checked_serialized_size(&self) -> Option<usize> {
875        self.data.iter().try_fold(0usize, |size, (key, value)| {
876            size.checked_add(8)?
877                .checked_add(key.len())?
878                .checked_add(4)?
879                .checked_add(value.checked_serialized_size()?)
880        })
881    }
882
883    /// Get a string value by key
884    pub fn get_string(&self, key: &str) -> Option<&str> {
885        self.get(key)?.as_str()
886    }
887
888    /// Get a u64 value by key
889    pub fn get_u64(&self, key: &str) -> Option<u64> {
890        self.get(key)?.as_u64()
891    }
892
893    /// Get an i64 value by key
894    pub fn get_i64(&self, key: &str) -> Option<i64> {
895        self.get(key)?.as_i64()
896    }
897
898    /// Get an f64 value by key
899    pub fn get_f64(&self, key: &str) -> Option<f64> {
900        self.get(key)?.as_f64()
901    }
902
903    /// Get a bool value by key
904    pub fn get_bool(&self, key: &str) -> Option<bool> {
905        self.get(key)?.as_bool()
906    }
907}
908
909#[cfg(feature = "std")]
910impl std::fmt::Display for MetadataValue {
911    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
912        write!(f, "{}", self.to_string_representation())
913    }
914}
915
916#[cfg(not(feature = "std"))]
917impl fmt::Display for MetadataValue {
918    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
919        match self {
920            MetadataValue::U8(v) => write!(f, "{}", v),
921            MetadataValue::I8(v) => write!(f, "{}", v),
922            MetadataValue::U16(v) => write!(f, "{}", v),
923            MetadataValue::I16(v) => write!(f, "{}", v),
924            MetadataValue::U32(v) => write!(f, "{}", v),
925            MetadataValue::I32(v) => write!(f, "{}", v),
926            MetadataValue::F32(v) => write!(f, "{}", v),
927            MetadataValue::Bool(v) => write!(f, "{}", v),
928            MetadataValue::String(v) => write!(f, "\"{}\"", v),
929            MetadataValue::Array(arr) => {
930                write!(f, "[")?;
931                for (i, value) in arr.iter().enumerate() {
932                    if i > 0 {
933                        write!(f, ", ")?;
934                    }
935                    write!(f, "{}", value)?;
936                }
937                write!(f, "]")
938            }
939            MetadataValue::U64(v) => write!(f, "{}", v),
940            MetadataValue::I64(v) => write!(f, "{}", v),
941            MetadataValue::F64(v) => write!(f, "{}", v),
942        }
943    }
944}
945
946#[cfg(all(test, feature = "std"))]
947mod tests {
948    use super::*;
949    use std::io::Cursor;
950
951    #[test]
952    fn test_metadata_value_types() {
953        let val = MetadataValue::U32(42);
954        assert_eq!(val.value_type(), GGUFValueType::U32);
955        assert_eq!(val.as_u64(), Some(42));
956
957        let val = MetadataValue::String("hello".to_string());
958        assert_eq!(val.value_type(), GGUFValueType::String);
959        assert_eq!(val.as_str(), Some("hello"));
960
961        let val = MetadataValue::Bool(true);
962        assert_eq!(val.value_type(), GGUFValueType::Bool);
963        assert_eq!(val.as_bool(), Some(true));
964    }
965
966    #[test]
967    fn test_metadata_value_io() {
968        let original = MetadataValue::F32(std::f32::consts::PI);
969
970        let mut buffer = Vec::new();
971        original.write_to(&mut buffer).unwrap();
972
973        let mut cursor = Cursor::new(buffer);
974        let read_value = MetadataValue::read_from(&mut cursor, GGUFValueType::F32).unwrap();
975
976        assert_eq!(original, read_value);
977    }
978
979    #[test]
980    fn test_metadata_array() {
981        let values = vec![MetadataValue::U32(1), MetadataValue::U32(2), MetadataValue::U32(3)];
982
983        let array = MetadataArray::new(GGUFValueType::U32, values).unwrap();
984        assert_eq!(array.len(), 3);
985        assert_eq!(array.element_type, GGUFValueType::U32);
986        assert_eq!(array.get(0), Some(&MetadataValue::U32(1)));
987    }
988
989    #[test]
990    fn test_metadata_array_type_mismatch() {
991        let values = vec![
992            MetadataValue::U32(1),
993            MetadataValue::String("hello".to_string()), // Wrong type!
994        ];
995
996        let result = MetadataArray::new(GGUFValueType::U32, values);
997        assert!(result.is_err());
998    }
999
1000    #[test]
1001    fn test_metadata_collection() {
1002        let mut metadata = Metadata::new();
1003        metadata.insert("name".to_string(), MetadataValue::String("test_model".to_string()));
1004        metadata.insert("version".to_string(), MetadataValue::U32(1));
1005
1006        assert_eq!(metadata.len(), 2);
1007        assert_eq!(metadata.get_string("name"), Some("test_model"));
1008        assert_eq!(metadata.get_u64("version"), Some(1));
1009        assert!(metadata.contains_key("name"));
1010        assert!(!metadata.contains_key("nonexistent"));
1011    }
1012
1013    #[test]
1014    fn test_metadata_io() {
1015        let mut original = Metadata::new();
1016        original.insert("name".to_string(), MetadataValue::String("test".to_string()));
1017        original.insert("count".to_string(), MetadataValue::U64(42));
1018
1019        let mut buffer = Vec::new();
1020        original.write_to(&mut buffer).unwrap();
1021
1022        let mut cursor = Cursor::new(buffer);
1023        let read_metadata = Metadata::read_from(&mut cursor, 2).unwrap();
1024
1025        assert_eq!(read_metadata.len(), 2);
1026        assert_eq!(read_metadata.get_string("name"), Some("test"));
1027        assert_eq!(read_metadata.get_u64("count"), Some(42));
1028    }
1029
1030    #[test]
1031    fn test_large_string_round_trip() {
1032        let long_value = MetadataValue::String("x".repeat(128 * 1024));
1033        let mut buffer = Vec::new();
1034        long_value.write_to(&mut buffer).unwrap();
1035
1036        let decoded =
1037            MetadataValue::read_from(&mut Cursor::new(buffer), GGUFValueType::String).unwrap();
1038        assert_eq!(decoded, long_value);
1039    }
1040
1041    #[test]
1042    fn test_array_io() {
1043        let values = vec![MetadataValue::I32(-1), MetadataValue::I32(0), MetadataValue::I32(42)];
1044
1045        let original = MetadataArray::new(GGUFValueType::I32, values).unwrap();
1046
1047        let mut buffer = Vec::new();
1048        original.write_to(&mut buffer).unwrap();
1049
1050        let mut cursor = Cursor::new(buffer);
1051        let read_array = MetadataArray::read_from(&mut cursor).unwrap();
1052
1053        assert_eq!(read_array.element_type, GGUFValueType::I32);
1054        assert_eq!(read_array.len(), 3);
1055        assert_eq!(read_array.get(2), Some(&MetadataValue::I32(42)));
1056    }
1057
1058    #[test]
1059    fn test_large_tokenizer_array_round_trip() {
1060        let values = (0..128_000).map(MetadataValue::U32).collect();
1061        let original = MetadataArray::new(GGUFValueType::U32, values).unwrap();
1062
1063        let mut buffer = Vec::new();
1064        original.write_to(&mut buffer).unwrap();
1065        let decoded = MetadataArray::read_from(&mut Cursor::new(buffer)).unwrap();
1066
1067        assert_eq!(decoded, original);
1068    }
1069
1070    #[test]
1071    fn test_truncated_huge_string_declaration_is_rejected_without_large_allocation() {
1072        let encoded_length = u64::MAX.to_le_bytes();
1073        let error =
1074            MetadataValue::read_from(&mut Cursor::new(encoded_length), GGUFValueType::String)
1075                .unwrap_err();
1076
1077        assert!(matches!(
1078            error,
1079            GGUFError::Io(_) | GGUFError::Format(_) | GGUFError::InvalidMetadata(_)
1080        ));
1081    }
1082
1083    #[test]
1084    fn test_truncated_huge_array_declaration_is_rejected_without_large_allocation() {
1085        let mut encoded = Vec::new();
1086        encoded.extend_from_slice(&(GGUFValueType::U8 as u32).to_le_bytes());
1087        encoded.extend_from_slice(&u64::MAX.to_le_bytes());
1088
1089        let error = MetadataArray::read_from(&mut Cursor::new(encoded)).unwrap_err();
1090        assert!(matches!(error, GGUFError::Io(_) | GGUFError::InvalidMetadata(_)));
1091    }
1092
1093    #[test]
1094    fn test_compact_array_is_rejected_before_decoded_memory_amplification() {
1095        let mut encoded = Vec::new();
1096        encoded.extend_from_slice(&6u64.to_le_bytes());
1097        encoded.extend_from_slice(b"tokens");
1098        encoded.extend_from_slice(&(GGUFValueType::Array as u32).to_le_bytes());
1099        encoded.extend_from_slice(&(GGUFValueType::U8 as u32).to_le_bytes());
1100        let oversized_elements =
1101            GGUF_MAX_METADATA_DECODED_SIZE / core::mem::size_of::<MetadataValue>() + 1;
1102        encoded.extend_from_slice(&(oversized_elements as u64).to_le_bytes());
1103
1104        let error =
1105            Metadata::read_from_with_limit(&mut Cursor::new(encoded), 1, GGUF_MAX_METADATA_SIZE)
1106                .unwrap_err();
1107        assert!(error.to_string().contains("Decoded metadata allocation exceeds budget"));
1108    }
1109
1110    #[test]
1111    fn test_metadata_key_validation() {
1112        for key in ["general.name", "tokenizer.ggml.model", "block_count", "layer.0.name"] {
1113            let mut metadata = Metadata::new();
1114            metadata.insert(key.to_string(), MetadataValue::U32(1));
1115            assert!(metadata.validate().is_ok(), "valid key rejected: {key}");
1116        }
1117
1118        for key in [
1119            "Uppercase",
1120            "has-dash",
1121            ".leading",
1122            "trailing.",
1123            "double..dot",
1124            "_leading",
1125            "trailing_",
1126            "double__underscore",
1127        ] {
1128            let mut metadata = Metadata::new();
1129            metadata.insert(key.to_string(), MetadataValue::U32(1));
1130            assert!(metadata.validate().is_err(), "invalid key accepted: {key}");
1131        }
1132    }
1133
1134    #[test]
1135    fn test_invalid_bool_is_rejected() {
1136        let mut cursor = Cursor::new([2u8]);
1137        assert!(MetadataValue::read_from(&mut cursor, GGUFValueType::Bool).is_err());
1138    }
1139
1140    #[test]
1141    fn test_array_length_mismatch_is_rejected() {
1142        let array = MetadataArray {
1143            element_type: GGUFValueType::U32,
1144            length: 2,
1145            values: vec![MetadataValue::U32(1)],
1146        };
1147        assert!(array.validate().is_err());
1148        assert!(array.write_to(&mut Vec::new()).is_err());
1149    }
1150}
gguf_rs_lib/format/metadata.rs

gguf_rs_lib/format/
metadata.rs
