Talson/poc-torch-weights-only-sparse-oob

PoC: PyTorch .pt OOB memory corruption bypassing weights_only=True

Security research proof-of-concept for the huntr Model File Vulnerability program (target: PyTorch Package .pt). The model file is non-destructive: loading it with the documented-safe loader succeeds with no error, and a routine matmul then crashes the process (SIGSEGV) by reading out of bounds. It writes nothing and runs no shell commands.

What this demonstrates

malicious_sparse.pt is a sparse CSR tensor of declared shape (1, 4) whose crow_indices claim 1,048,576 nonzeros while its col_indices/values arrays hold a single element. PyTorch's weights_only unpickler allow-lists torch._utils._rebuild_sparse_tensor, which builds it with check_invariants=False, and the deferred validator _validate_loaded_sparse_tensors skips validation by default (check_sparse_tensor_invariants is off). So:

import torch
t = torch.load("malicious_sparse.pt", weights_only=True)   # SAFE loader: NO error
torch.mm(t, torch.ones(t.size(-1), 2))                      # OOB read -> SIGSEGV

The sparse matmul kernel walks 2^20 entries off the end of the one-element buffers → out-of-bounds read → segfault (DoS); a tuned, smaller count instead folds adjacent heap memory into the matmul output (information disclosure).

Both standard defenses are bypassed:

• torch.load(weights_only=True) — the mode documented as safe for untrusted files.
• ProtectAI ModelScan → No issues found! (the file has no unsafe pickle opcodes; the danger is in tensor data).

Reproduce

pip install torch==2.12.0
python3 poc.py                # build + load(weights_only=True) + matmul -> SIGSEGV (exit 139)
# to use THIS repo's bundled file instead of rebuilding:
MAL_PT=./malicious_sparse.pt python3 poc.py trigger
modelscan -p malicious_sparse.pt   # -> "No issues found!"

Files

• malicious_sparse.pt — the malicious (non-destructive) PyTorch model file.
• poc.py — builds the file (build) and demonstrates the load+matmul crash (trigger).

Fix

Validate sparse-tensor invariants when loading untrusted data regardless of the opt-in check_sparse_tensor_invariants flag, or bounds-check crow_indices[-1] against the col_indices/values lengths in _rebuild_sparse_tensor.

Tested: PyTorch 2.12.0, ModelScan main @ 61fcec9.