MindSpore DoS PoC — CWE-789: Unbounded Tensor Allocation in load_checkpoint()
Overview
This repository contains a proof-of-concept demonstrating a Denial of Service vulnerability in MindSpore (Huawei's open-source AI framework).
A crafted .ckpt checkpoint file (44 bytes) causes mindspore.load_checkpoint() to attempt a 37.3 GB memory allocation, crashing the process or exhausting system memory.
| Metric | Value |
|---|---|
| CWE | CWE-789 (Memory Allocation with Excessive Size Value) |
| CVSS | 7.5 High (AV:N/AC:L/PR:N/UI:R/S:U/C:N/I:N/A:H) |
| Affected | MindSpore 2.8.0 (latest) and all prior versions |
| File size | 44 bytes |
| Allocation | 40,000,000,000 bytes (~37.3 GB) |
| Amplification | 909,090,909x |
Vulnerability
MindSpore's .ckpt format uses Protocol Buffers. The TensorProto message contains a dims field (repeated int64) that specifies tensor shape and a tensor_content field with the actual data bytes.
When loading a checkpoint, _load_into_param_dict() in serialization.py reads dims from the protobuf and passes them to the C++ function ConvertBytesToTensor() in tensor_py.cc, which allocates a tensor based entirely on the attacker-controlled dims — without validating that dims match the actual data size.
serialization.py:1119 dims = element.tensor.dims # from file (attacker-controlled)
serialization.py:1127 Tensor_.convert_bytes_to_tensor(new_data, tuple(dims), ms_type)
tensor_py.cc:545 tensor = make_shared<Tensor>(data_type, shape) # unbounded allocation
Files
| File | Description |
|---|---|
craft_ckpt.py |
Generates malicious .ckpt files with arbitrary dimensions |
poc_demo.py |
Main PoC — triggers the vulnerability with safety limits |
poc_crash.py |
Shows SIGABRT crash variant with tight memory |
dos_moderate.ckpt |
44 bytes — dims [100000, 100000, 1] — 37.3 GB alloc |
dos_extreme.ckpt |
44 bytes — dims [100000, 100000, 100] — 3.7 TB alloc |
dos_minimal.ckpt |
43 bytes — dims [10000, 10000, 1000] — 372.5 GB alloc |
Reproduction
Requirements
- Python 3.10 (MindSpore does not support 3.12+)
- MindSpore 2.8.0:
pip install mindspore
Step 1: Generate malicious checkpoints
python3 craft_ckpt.py
Step 2: Trigger the vulnerability
python3 poc_demo.py
Expected output
[WARNING] CORE [tensor_data.h:565] Try to alloca a large memory, size is:40000000000
[CRITICAL] ME [serialization.py:1287] Failed to load the checkpoint file 'dos_moderate.ckpt'.
[+] VULNERABILITY CONFIRMED!
[+] Result: 44-byte file attempted ~37.3 GB allocation
[+] Amplification ratio: 909,090,909x
With tight memory limits (512 MB), the process crashes with SIGABRT (exit code 134) and corrupts the Python interpreter state:
Fatal Python error: PyThreadState_Get: the function must be called with the GIL held,
but the GIL is released (the current Python thread state is NULL)
Exit code: 134
Impact
- Denial of Service: A tiny file exhausts all available memory
- Process crash: SIGABRT on memory-constrained systems, unrecoverable Python interpreter corruption
- Supply chain vector: Malicious
.ckptfiles on model hubs (HuggingFace, ModelScope) affect any user who loads them
Disclaimer
This proof-of-concept is published for authorized security research purposes only, in accordance with responsible disclosure practices. The vulnerability has been reported to the vendor. Do not use this PoC to disrupt services or systems you do not own or have authorization to test.