safetensors-dos-poc

SafeTensors DoS Proof-of-Concept Files

Overview

These proof-of-concept files demonstrate potential Denial-of-Service (DoS) vectors in SafeTensors file parsing. While the reference Rust implementation (safetensors crate) validates and rejects these malformed files, third-party parsers, custom loaders, or older versions that trust header metadata without validation may be vulnerable.

SafeTensors Binary Format

[8 bytes: header_size (u64 LE)] [header_size bytes: JSON header] [remaining: raw tensor data]

PoC Files

1. poc_oom_shape.safetensors (168 bytes)

Vector: Enormous tensor shape causing OOM on allocation.

The header declares a tensor with shape [1000000, 1000000] and dtype F32. A naive parser that allocates shape_product * dtype_size before reading data will attempt a 4 TB allocation, causing OOM or system instability.

2. poc_huge_header.safetensors (27 bytes)

Vector: Header size field set to 4 GB.

The first 8 bytes declare header_size = 0xFFFFFFFF (4,294,967,295 bytes). A parser that does malloc(header_size) followed by read(header_size) will attempt to allocate 4 GB for the header alone.

3. poc_overlapping.safetensors (243 bytes)

Vector: Overlapping tensor data offsets.

Two tensors (tensor_a at offsets [0, 16] and tensor_b at offsets [8, 24]) share overlapping byte ranges. This can cause undefined behavior in parsers that create independent mutable views of the underlying data buffer.

4. benign.safetensors (149 bytes)

A valid, well-formed SafeTensors file containing a single [4] float32 tensor for comparison.

Reference Implementation Results

The Rust-based safetensors library (v0.7.0) correctly rejects all three malicious files:

Attack Surface

These PoCs target implementations that:

• Allocate memory based on declared tensor dimensions without checking against actual file size
• Read header_size bytes without an upper bound check
• Allow overlapping data_offsets without validation
• Examples: custom Python parsers, C/C++ loaders, web-based model inspectors

Previous Finding: Integer Overflow in TensorView::new()

CWE-190 / CWE-125: Integer overflow in TensorView::new() bypasses size validation. In safetensors/src/tensor.rs, n_elements: usize = shape.iter().product() uses unchecked multiplication (unlike validate() which uses checked_mul). See overflow_poc.rs for details.

Reproduction

pip install safetensors
python generate_pocs.py