Learning Hub

How to Stop Pickle Deserialization Attacks in Python and Load Model Files Safely

Learn why unpickling runs code, then build an allowlist unpickler, a strict safetensors reader and a content-sniffing loader in Python, and prove them against hostile files with 31 tests.

October 15, 2026 34 Min Read
37

Loading a file feels harmless: you point a function at a path and get an object back. With pickle that assumption is wrong. A pickle file is not a record of data. It is a short list of instructions for a small stack-based interpreter inside Python, and those instructions can import any function they can name and call it with any arguments before your code ever sees the result. The Python documentation says it plainly: “The pickle module is not secure. Only unpickle data you trust.”

Table Of Content

That matters more every year, because machine learning work passes files around constantly. Scikit-learn models are saved with joblib, pandas frames with to_pickle, NumPy arrays with allow_pickle=True, and PyTorch checkpoints are zip archives with a pickle inside. A model downloaded from a hub, a checkpoint copied from a shared bucket, or a cache folder that every teammate can write to all lead to the same one-line mistake.

In this tutorial you will learn how that mistake works and then fix it properly, step by step. You will:

  • run a harmless hostile pickle and read its opcodes to see why it executes;
  • build an allowlist unpickler and see why a blocklist loses;
  • list the common loaders that unpickle behind your back;
  • write a strict reader for the safetensors format and attack it with ten broken files;
  • handle PyTorch checkpoints, add an audit-hook backstop, and combine everything into one loader that picks its parser from the file’s content;
  • teach your CI to find risky load calls, and prove all of it with 31 tests.

Terms used in this tutorial

Serialization turns an object into bytes so it can be stored or sent; deserialization (called unpickling in pickle’s case) turns the bytes back into an object. In a pickle, a global is a class or function the file asks the unpickler to import by module name and attribute name, such as builtins.print. An allowlist names the only things that are permitted and refuses everything else; a blocklist names the things that are forbidden and permits the rest. Safetensors is a file format for tensors that stores a length, a JSON header and raw bytes, with no way to name a function at all.

Prerequisites

  • Python 3.13 or newer. Everything here was run on Python 3.13.14 on Windows 11. Nothing in the code depends on Windows, but the outputs shown are from that machine.
  • pip and a terminal. You need about 850 MB of free disk space, almost all of it PyTorch (a 539 MB package in the virtual environment); no GPU is needed.
  • Comfort with Python functions, classes and bytes. You do not need to know how pickle works internally; that is the point of the tutorial.
  • A throwaway folder. Every hostile file in this tutorial only prints the line *** payload ran ***, but you should still treat the lab like a lab.

Step 1: Create the lab and a harmless marker

Make a folder, a virtual environment, and install the packages. Use whichever activation line matches your system (the first is for Windows PowerShell, the second for macOS and Linux) and skip the other. The package versions are the ones used to produce every output in this post; newer releases should behave the same, but error messages may differ in wording.

mkdir picklelab
cd picklelab
python -m venv venv
venv\Scripts\Activate.ps1
source venv/bin/activate
pip install numpy==2.5.4 pandas==3.0.6 joblib==1.6.0 safetensors==0.8.0 pyyaml==6.0.3 pytest==9.1.1 torch==2.14.1 bandit==1.9.4 ruff==0.17.0

Run every script from inside picklelab, because paths such as data/evil.pkl are relative. The first file is a tiny helper. It defines the marker text and a function that calls any loader with standard output captured, then reports whether the marker appeared. That gives us a yes or no answer to the question of whether attacker code ran, without reading output by eye.

# marker.py
"""Shared helpers for the lab: a harmless "payload ran" marker and a runner that reports it."""
import contextlib
import io

MARKER = "*** payload ran ***"


def run(fn, *args, **kwargs):
    """Call fn(*args, **kwargs) with stdout captured and report what happened.

    Returns (ran, outcome). ran is True if the payload printed MARKER while loading.
    outcome is 'loaded' or the name of the exception that was raised.
    """
    buf = io.StringIO()
    outcome = "loaded"
    try:
        with contextlib.redirect_stdout(buf):
            fn(*args, **kwargs)
    except Exception as exc:  # we want to report any failure, not only one type
        outcome = type(exc).__name__
    return MARKER in buf.getvalue(), outcome

Step 2: Watch a pickle run code

The attacker’s side

Pickle asks every object how it should be rebuilt. An object may answer with a tuple of a callable and its arguments, which means: to recreate me, call this function with these arguments and use the result. A class gives that answer by defining __reduce__. Our payload asks pickle to call print, which stands in for anything an attacker could call.

# attacker.py
"""The attacker's side: build a file that runs code when it is unpickled.

The payload only calls print(). An attacker could put any importable callable there.
"""
import pickle
from pathlib import Path

from marker import MARKER


class Payload:
    def __reduce__(self):
        # Tell pickle: "to rebuild this object, call print(MARKER)".
        return (print, (MARKER,))


if __name__ == "__main__":
    Path("data").mkdir(exist_ok=True)
    blob = pickle.dumps(Payload(), protocol=4)
    Path("data/evil.pkl").write_bytes(blob)
    print(f"wrote data/evil.pkl ({len(blob)} bytes); the payload did not run while writing it")
wrote data/evil.pkl (59 bytes); the payload did not run while writing it

Notice what did not happen: writing the file did not print the marker. pickle.dumps only records the instruction to call print with this text. The danger arrives on the other side.

The victim’s side

Here is the one-line mistake. The script first loads the example from the Python documentation, a hand-written pickle that imports os.system and applies it to a harmless echo command, and then loads our own file.

# victim_naive.py
"""The victim's side: the one-line mistake."""
import pickle

# 1. The example from the Python documentation: this pickle imports os.system
#    and applies it to a string, which runs a shell command.
docs_example = b"cos\nsystem\n(S'echo hello world'\ntR."
print("docs example returned:", pickle.loads(docs_example))

# 2. Our own payload from attacker.py.
with open("data/evil.pkl", "rb") as f:
    result = pickle.load(f)
print("pickle.load returned:", result)
hello world
docs example returned: 0
*** payload ran ***
pickle.load returned: None

Both loads executed code. hello world came from a shell command, and the second load printed our marker before returning None, which is simply what print returns. The loaded “object” is whatever the called function produced, so a victim that only checks the result may never notice anything. To confirm this step worked, you should see *** payload ran *** exactly once.

Step 3: Read the pickle as a program

The standard library can show you the instructions without running them. pickletools.dis is a disassembler: it reads bytes and prints opcodes, and it never imports anything the file names.

# dis_payload.py
"""Read the pickle as opcodes without executing anything."""
import pickletools
from pathlib import Path

pickletools.dis(Path("data/evil.pkl").read_bytes())
    0: \x80 PROTO      4
    2: \x95 FRAME      48
   11: \x8c SHORT_BINUNICODE 'builtins'
   21: \x94 MEMOIZE    (as 0)
   22: \x8c SHORT_BINUNICODE 'print'
   29: \x94 MEMOIZE    (as 1)
   30: \x93 STACK_GLOBAL
   31: \x94 MEMOIZE    (as 2)
   32: \x8c SHORT_BINUNICODE '*** payload ran ***'
   53: \x94 MEMOIZE    (as 3)
   54: \x85 TUPLE1
   55: \x94 MEMOIZE    (as 4)
   56: R    REDUCE
   57: \x94 MEMOIZE    (as 5)
   58: .    STOP
highest protocol among opcodes = 4

Five opcodes carry the whole attack. SHORT_BINUNICODE pushes a string onto the stack, here builtins and then print. STACK_GLOBAL pops those two strings, imports the module and fetches the attribute, so builtins.print is now on the stack. The next SHORT_BINUNICODE pushes the marker text and TUPLE1 wraps it in a one-item tuple. REDUCE pops the tuple and the callable and calls it. STOP returns whatever is on top. The MEMOIZE lines are bookkeeping.

The takeaway is that a pickle can only call something it can name, and naming a function or class means asking for a global. That is the lever the rest of the tutorial pulls. Be careful with the next idea, though: reading opcodes is useful for triage, but it is not a security decision. The Hugging Face documentation on its own pickle scanner, which reads files with pickletools.genops so it never executes them, still ends with a disclaimer: “this is not 100% foolproof. It is your responsibility as a user to check if something is safe or not.” The same page says the lists behind the scanner “are maintained in a best-effort manner.”

Step 4: Restrict what a pickle may import

Why find_class is the right place

Whenever an unpickler meets a global, it calls find_class(module, name) to fetch the object. The documentation explains the consequence: “Unlike its name suggests, Unpickler.find_class() is called whenever a global (i.e., a class or a function) is requested. Thus it is possible to either completely forbid globals or restrict them to a safe subset.” If you override that method, you decide what a pickle is allowed to reach.

The module below contains three unpicklers. AllowlistUnpickler is the one to use: it refuses every global unless the exact (module, name) pair is in the set you pass. RecordingUnpickler will help us build allowlists in the next step. BlocklistUnpickler is the version that looks sensible at first, and it is here so you can watch it fail.

# safe_unpickler.py
"""Three unpicklers: an allowlist (use this), a recorder (to build the allowlist), a blocklist (do not use)."""
import io
import pickle


class AllowlistUnpickler(pickle.Unpickler):
    """Refuses every global unless (module, name) is in `allowed`."""

    def __init__(self, file, *, allowed=frozenset()):
        super().__init__(file)
        self._allowed = frozenset(allowed)

    def find_class(self, module, name):
        if (module, name) in self._allowed:
            return super().find_class(module, name)
        raise pickle.UnpicklingError(f"forbidden global: {module}.{name}")


def safe_loads(data: bytes, allowed=frozenset()):
    return AllowlistUnpickler(io.BytesIO(data), allowed=allowed).load()


class RecordingUnpickler(pickle.Unpickler):
    """Lists the globals a pickle asks for. Run it only on data YOU produced."""

    def __init__(self, file):
        super().__init__(file)
        self.seen = []

    def find_class(self, module, name):
        self.seen.append((module, name))
        return super().find_class(module, name)


BLOCKED = {
    ("os", "system"),
    ("subprocess", "Popen"),
    ("builtins", "eval"),
    ("builtins", "exec"),
    ("builtins", "print"),
}


class BlocklistUnpickler(pickle.Unpickler):
    """A blocklist. It looks safe in a quick test. It is not."""

    def find_class(self, module, name):
        if (module, name) in BLOCKED:
            raise pickle.UnpicklingError(f"blocked: {module}.{name}")
        return super().find_class(module, name)

Try it against plain data and against hostile pickles

The demo script pickles a typical configuration dictionary with three protocol versions and loads it with an empty allowlist. Then it throws three hostile pickles at both the allowlist and the blocklist: our print payload, the documentation’s os.system example, and the same print call spelled the way Python 2 spelled it, __builtin__.print.

# step4_allowlist.py
import io
import pickle
from pathlib import Path

from marker import run
from safe_unpickler import BlocklistUnpickler, safe_loads

evil = Path("data/evil.pkl").read_bytes()
docs_example = b"cos\nsystem\n(S'echo hello world'\ntR."
# The same call as the payload, spelled with the old Python 2 name for builtins.
legacy_spelling = b"c__builtin__\nprint\n(S'*** payload ran ***'\ntR."

config = {
    "name": "demo", "layers": [64, 32], "lr": 0.001, "tags": ("a", "b"),
    "ids": {1, 2, 3}, "frozen": frozenset({1}), "raw": b"\x00\x01", "none": None, "ok": True,
}

print("-- plain data against an EMPTY allowlist --")
for proto in (2, 4, 5):
    blob = pickle.dumps(config, protocol=proto)
    try:
        print(f"protocol {proto}: round trip equal = {safe_loads(blob) == config}")
    except pickle.UnpicklingError as exc:
        print(f"protocol {proto}: refused ({exc})")

print("-- hostile pickles --")
print("allowlist vs print payload   :", run(safe_loads, evil))
print("allowlist vs docs example    :", run(safe_loads, docs_example))
print("allowlist vs legacy spelling :", run(safe_loads, legacy_spelling))


def blocklist_loads(data):
    return BlocklistUnpickler(io.BytesIO(data)).load()


print("blocklist vs print payload   :", run(blocklist_loads, evil))
print("blocklist vs docs example    :", run(blocklist_loads, docs_example))
print("blocklist vs legacy spelling :", run(blocklist_loads, legacy_spelling))
-- plain data against an EMPTY allowlist --
protocol 2: refused (forbidden global: __builtin__.set)
protocol 4: round trip equal = True
protocol 5: round trip equal = True
-- hostile pickles --
allowlist vs print payload   : (False, 'UnpicklingError')
allowlist vs docs example    : (False, 'UnpicklingError')
allowlist vs legacy spelling : (False, 'UnpicklingError')
blocklist vs print payload   : (False, 'UnpicklingError')
blocklist vs docs example    : (False, 'UnpicklingError')
blocklist vs legacy spelling : (True, 'loaded')

Read the output in two halves. First, plain data needs no globals at all on protocol 4 and 5: dictionaries, lists, tuples, sets, frozensets, bytes, None and booleans all have their own opcodes, so an empty allowlist loads them. Protocol 2 was refused with __builtin__.set, because that older protocol has no set opcode and spells built-in names the Python 2 way. The documentation confirms the spelling: “If fix_imports is true and protocol is less than 3, pickle will try to map the new Python 3 names to the old module names used in Python 2, so that the pickle data stream is readable with Python 2.”

Second, the allowlist stopped all three hostile pickles and none of them ran. The blocklist stopped two, and it even contained builtins.print, yet the third one ran (True). The legacy spelling slipped through because the unpickler translates old names to new ones inside its own find_class, after your override has already compared the raw names. You can see the order in the standard library source: the pure-Python find_class calls sys.audit first and only then applies _compat_pickle mappings when the protocol is below 3 and fix_imports is set. The rule to remember: you cannot list every bad name, because aliases, other modules and future callables keep appearing, but you can list the few names your legitimate data needs. Use an allowlist.

Step 5: Learn an allowlist from data you produced

Real data is rarely a plain dictionary. A NumPy array, for example, needs a few globals to rebuild itself. Rather than guessing, run RecordingUnpickler on a pickle you created yourself and record what it asks for. This is safe only because you trust the input: the recorder imports everything it sees.

# step5_numpy.py
import io
import pickle
from pathlib import Path

import numpy as np

from marker import run
from safe_unpickler import RecordingUnpickler, safe_loads

print("numpy", np.__version__)
arr = np.arange(6, dtype="float32").reshape(2, 3)

needed = {}
for proto in (2, 4, 5):
    blob = pickle.dumps(arr, protocol=proto)
    rec = RecordingUnpickler(io.BytesIO(blob))
    rec.load()
    needed[proto] = rec.seen
    print(f"protocol {proto}: {len(blob):3d} bytes, globals = {rec.seen}")

allowed = set(needed[4])
blob4 = pickle.dumps(arr, protocol=4)
print("allowlist built from protocol 4 loads protocol 4:", np.array_equal(safe_loads(blob4, allowed), arr))
for proto in (2, 5):
    try:
        safe_loads(pickle.dumps(arr, protocol=proto), allowed)
        print(f"allowlist built from protocol 4 loads protocol {proto}: yes")
    except pickle.UnpicklingError as exc:
        print(f"allowlist built from protocol 4 loads protocol {proto}: no ({exc})")

evil = Path("data/evil.pkl").read_bytes()
print("same allowlist vs the print payload:", run(safe_loads, evil, allowed))
numpy 2.5.4
protocol 2: 236 bytes, globals = [('numpy._core.multiarray', '_reconstruct'), ('numpy', 'ndarray'), ('_codecs', 'encode'), ('numpy', 'dtype')]
protocol 4: 174 bytes, globals = [('numpy._core.multiarray', '_reconstruct'), ('numpy', 'ndarray'), ('numpy', 'dtype')]
protocol 5: 153 bytes, globals = [('numpy._core.numeric', '_frombuffer'), ('numpy', 'dtype')]
allowlist built from protocol 4 loads protocol 4: True
allowlist built from protocol 4 loads protocol 2: no (forbidden global: _codecs.encode)
allowlist built from protocol 4 loads protocol 5: no (forbidden global: numpy._core.numeric._frombuffer)
same allowlist vs the print payload: (False, 'UnpicklingError')

Four things stand out. The array needs just three globals on protocol 4 in NumPy 2.5.4, and their names are implementation details. The NumPy 2.0 migration guide says “The np.core namespace is now officially private and has been renamed to np._core.” and adds that the layout of the private module might change without notice. Next, protocol 5 asks for different globals (_frombuffer instead of _reconstruct), and protocol 2 adds _codecs.encode. The list built from protocol 4 loaded protocol 4 and refused the other two. That matters because, according to the pickle documentation, Python 3.14 changed the default pickle protocol to 5, so a list recorded on Python 3.13 with default settings would refuse arrays pickled with protocol 5, which is what Python 3.14 and later use by default. Finally, the same allowlist still refused our hostile pickle.

The lesson is that allowlists work well for small, plain structures and badly for rich objects. For arrays and tensors, move to a format built for them, which is where Step 7 goes.

Step 6: Find every loader that unpickles behind your back

You rarely call pickle.load yourself. Many familiar functions do it for you. This script builds one hostile file per format and then calls the usual loaders, reporting whether the payload ran.

# step6_matrix.py
import pickle
import warnings
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
import yaml

from attacker import Payload
from marker import MARKER, run

warnings.simplefilter("ignore")
Path("data").mkdir(exist_ok=True)


def load_pickle(path):
    with open(path, "rb") as f:
        return pickle.load(f)


# Attacker side: one hostile file per format.
with open("data/m.pkl", "wb") as f:
    pickle.dump(Payload(), f)
joblib.dump(Payload(), "data/m.joblib")
np.save("data/m.npy", np.array([Payload()], dtype=object), allow_pickle=True)
yaml_doc = f"!!python/object/apply:builtins.print ['{MARKER}']"

# Victim side: every common way of loading them.
rows = [
    ("pickle.load", run(load_pickle, "data/m.pkl")),
    ("joblib.load", run(joblib.load, "data/m.joblib")),
    ("pandas.read_pickle", run(pd.read_pickle, "data/m.pkl")),
    ("numpy.load(allow_pickle=True)", run(np.load, "data/m.npy", allow_pickle=True)),
    ("numpy.load (default)", run(np.load, "data/m.npy")),
    ("yaml.unsafe_load", run(yaml.unsafe_load, yaml_doc)),
    ("yaml.full_load", run(yaml.full_load, yaml_doc)),
    ("yaml.safe_load", run(yaml.safe_load, yaml_doc)),
    ("yaml.load (no Loader)", run(yaml.load, yaml_doc)),
]
print(f"{'loader':<32}{'payload ran':<13}outcome")
for name, (ran, outcome) in rows:
    print(f"{name:<32}{str(ran):<13}{outcome}")
loader                          payload ran  outcome
pickle.load                     True         loaded
joblib.load                     True         loaded
pandas.read_pickle              True         loaded
numpy.load(allow_pickle=True)   True         loaded
numpy.load (default)            False        ValueError
yaml.unsafe_load                True         loaded
yaml.full_load                  False        ConstructorError
yaml.safe_load                  False        ConstructorError
yaml.load (no Loader)           False        TypeError

Five loaders ran the payload: pickle.load, joblib.load, pandas.read_pickle, numpy.load with allow_pickle=True, and yaml.unsafe_load. Each project warns about it in its own words. Joblib: “joblib.load relies on the pickle module and can therefore execute arbitrary Python code. It should therefore never be used to load files from untrusted sources.” Pandas: “Loading pickled data received from untrusted sources can be unsafe.” PyYAML: “yaml.load is as powerful as pickle.load and so may call any Python function.” The OWASP deserialization cheat sheet gives the same advice for Python in one line: “Unpickling can execute arbitrary code; never unpickle untrusted data.”

The other four rows are the good news. numpy.load with its default refused the object array with a ValueError, which matches the NumPy documentation: allow_pickle defaults to False, and “Reasons for disallowing pickles include security, as loading pickled data can execute arbitrary code.” For YAML, safe_load and full_load both refused this document with a ConstructorError, and yaml.load without a Loader argument raised a TypeError in PyYAML 6. Do not read the full_load row as a promise; the PyYAML documentation points untrusted input to safe_load, so use that.

Step 7: Move to formats that are data, not programs

What to use instead

The cleanest fix is to stop using a format that can carry code. Use JSON for configuration and plain NumPy .npy files (loaded with the default allow_pickle=False) for arrays. For tensors, use safetensors. Its README describes the entire layout in a few lines: 8 bytes holding an unsigned little-endian 64-bit integer N, then N bytes of UTF-8 JSON, then the rest of the file is the byte buffer. The JSON header maps each tensor name to a dtype, a shape and a pair of offsets into the buffer. There is no field that can name a function.

Write a strict reader so you understand the format

In production you should use the official safetensors library, which is fast and zero-copy. Writing a small reader teaches what “safe” actually means, and it gives you a second opinion to compare with. The reader below checks every field before touching tensor bytes: the length prefix must fit the file, the header must be UTF-8 JSON with no duplicate keys, each entry needs a known dtype, a list of non-negative integer dimensions and two integer offsets inside the buffer, the byte size must equal the item size times the number of elements, and the tensors sorted by start offset must tile the buffer exactly with no gap, overlap or trailing bytes.

# st_reader.py
"""A small, strict safetensors reader. Every field is checked before any tensor bytes are used."""
import json
import os
import struct
from math import prod

import numpy as np

DTYPES = {
    "F64": "<f8", "F32": "<f4", "F16": "<f2",
    "I64": "<i8", "I32": "<i4", "I16": "<i2", "I8": "i1", "U8": "u1", "BOOL": "?",
}
MAX_HEADER = 10_000_000  # our own cap; raise it if your real headers are larger


class BadSafetensors(ValueError):
    """The file is not a well-formed safetensors file."""


def _no_duplicates(pairs):
    out = {}
    for key, value in pairs:
        if key in out:
            raise BadSafetensors(f"duplicate key in header: {key!r}")
        out[key] = value
    return out


def read_safetensors(path):
    """Return (tensors, metadata). Raises BadSafetensors for anything unexpected."""
    size = os.path.getsize(path)
    with open(path, "rb") as f:
        raw = f.read(8)
        if len(raw) < 8:
            raise BadSafetensors("file is shorter than the 8-byte length prefix")
        (n,) = struct.unpack("<Q", raw)
        if n > MAX_HEADER or 8 + n > size:
            raise BadSafetensors(f"header length {n} does not fit this {size}-byte file")
        header_bytes = f.read(n)
        if not header_bytes.startswith(b"{"):
            raise BadSafetensors("header must start with '{'")
        try:
            header = json.loads(header_bytes.decode("utf-8"), object_pairs_hook=_no_duplicates)
        except (UnicodeDecodeError, json.JSONDecodeError) as exc:
            raise BadSafetensors(f"header is not valid UTF-8 JSON: {exc}") from exc
        data = f.read()

    meta = header.pop("__metadata__", {})
    if not isinstance(meta, dict) or not all(isinstance(v, str) for v in meta.values()):
        raise BadSafetensors("__metadata__ must map strings to strings")

    spans = []
    for name, info in header.items():
        try:
            dtype = DTYPES[info["dtype"]]
            shape = info["shape"]
            begin, end = info["data_offsets"]
        except (KeyError, TypeError, ValueError) as exc:
            raise BadSafetensors(f"tensor {name!r}: malformed entry ({type(exc).__name__}: {exc})") from exc
        if not (isinstance(shape, list) and all(type(d) is int and d >= 0 for d in shape)):
            raise BadSafetensors(f"tensor {name!r}: bad shape {shape!r}")
        if not (type(begin) is int and type(end) is int and 0 <= begin <= end <= len(data)):
            raise BadSafetensors(f"tensor {name!r}: offsets {begin}..{end} fall outside the {len(data)}-byte buffer")
        want = np.dtype(dtype).itemsize * prod(shape)
        if end - begin != want:
            raise BadSafetensors(f"tensor {name!r}: {end - begin} bytes cannot hold shape {shape} of {info['dtype']}")
        spans.append((begin, end, name, dtype, shape))

    spans.sort()
    cursor = 0
    for begin, end, name, _dtype, _shape in spans:
        if begin != cursor:
            raise BadSafetensors(f"tensor {name!r}: starts at byte {begin} but the previous tensor ended at {cursor}")
        cursor = end
    if cursor != len(data):
        raise BadSafetensors(f"{len(data) - cursor} bytes after the last tensor are not indexed")

    tensors = {
        name: np.frombuffer(data, dtype=dtype, count=prod(shape), offset=begin).reshape(shape)
        for begin, _end, name, dtype, shape in spans
    }
    return tensors, meta

The last check follows the format’s own rule: “The byte buffer needs to be entirely indexed, and cannot contain holes. This prevents the creation of polyglot files.” A polyglot file is one that is valid as two formats at once, which lets an attacker hide a second payload in bytes a loader never looks at.

Attack the reader with ten broken files

Next, build one good file and ten broken ones by hand, writing the length prefix, header and buffer ourselves so we can lie in each field.

# tamper_cases.py
"""Builders for one good safetensors file and nine broken ones, as raw bytes."""
import json
import struct

import numpy as np


def build(header=None, buffer=b"", *, claimed_len=None, raw_header=None):
    """Assemble 8-byte length + header + buffer by hand."""
    h = raw_header if raw_header is not None else json.dumps(header, separators=(",", ":")).encode()
    h += b" " * (-len(h) % 8)  # pad with spaces to a multiple of 8, as the format allows
    n = len(h) if claimed_len is None else claimed_len
    return struct.pack("<Q", n) + h + buffer


A = np.arange(4, dtype="<f4").tobytes()  # tensor "a": 16 bytes
B = bytes([1, 2, 3, 4])                  # tensor "b": 4 bytes
ENTRY_A = {"dtype": "F32", "shape": [2, 2], "data_offsets": [0, 16]}
ENTRY_B = {"dtype": "I8", "shape": [4], "data_offsets": [16, 20]}
GOOD = {"a": ENTRY_A, "b": ENTRY_B}

CASES = {
    "good file": build(GOOD, A + B),
    "length prefix lies": build(GOOD, A + B, claimed_len=2**40),
    "offsets run past the buffer": build({"a": ENTRY_A, "b": {**ENTRY_B, "data_offsets": [16, 100]}}, A + B),
    "two tensors overlap": build({"a": ENTRY_A, "b": {**ENTRY_B, "data_offsets": [12, 16]}}, A + B),
    "hole between tensors": build({"a": ENTRY_A, "b": {**ENTRY_B, "data_offsets": [20, 24]}}, A + b"\0" * 4 + B),
    "extra bytes at the end": build(GOOD, A + B + b"MZ-not-a-tensor"),
    "shape does not match size": build({"a": {**ENTRY_A, "data_offsets": [0, 8]}, "b": {**ENTRY_B, "data_offsets": [8, 12]}}, A[:8] + B),
    "unknown dtype": build({"a": {**ENTRY_A, "dtype": "F33"}, "b": ENTRY_B}, A + B),
    "duplicate tensor name": build(raw_header=(
        b'{"a":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]},'
        b'"a":{"dtype":"I8","shape":[4],"data_offsets":[16,20]}}'), buffer=A + B),
    "duplicate name, same entry twice": build(raw_header=(
        b'{"a":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]},'
        b'"a":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]},'
        b'"b":{"dtype":"I8","shape":[4],"data_offsets":[16,20]}}'), buffer=A + B),
    "header is not JSON": build(raw_header=b"this is not json", buffer=A + B),
}
# step7_tamper.py
from pathlib import Path

import numpy as np
from safetensors.numpy import load_file, save_file

from st_reader import read_safetensors
from tamper_cases import CASES

Path("data").mkdir(exist_ok=True)

# 1. Our reader and the official library agree on a file the library wrote.
tensors = {"w": np.arange(6, dtype="float32").reshape(2, 3), "idx": np.array([3, 1, 2], dtype="int64")}
save_file(tensors, "data/good.safetensors", metadata={"format": "np"})
mine, meta = read_safetensors("data/good.safetensors")
theirs = load_file("data/good.safetensors")
print("written by the library, read by ours:", {k: bool(np.array_equal(mine[k], theirs[k])) for k in tensors}, meta)
print("first 40 bytes of the file:", Path("data/good.safetensors").read_bytes()[:40])


def short(exc):
    return f"{type(exc).__name__}: {str(exc).splitlines()[0][:78]}"


def outcome(fn, path):
    try:
        fn(path)
        return "loaded"
    except Exception as exc:
        return short(exc)


# 2. Ten hand-built files: one good, nine broken.
for name, raw in CASES.items():
    path = Path("data") / "tamper.safetensors"
    path.write_bytes(raw)
    print(f"\ncase: {name}")
    print("  ours    :", outcome(read_safetensors, path))
    print("  library :", outcome(load_file, str(path)))
written by the library, read by ours: {'w': True, 'idx': True} {'format': 'np'}
first 40 bytes of the file: b'\x98\x00\x00\x00\x00\x00\x00\x00{"__metadata__":{"format":"np"},'

case: good file
  ours    : loaded
  library : loaded

case: length prefix lies
  ours    : BadSafetensors: header length 1099511627776 does not fit this 140-byte file
  library : SafetensorError: Error while deserializing header: header too large

case: offsets run past the buffer
  ours    : BadSafetensors: tensor 'b': offsets 16..100 fall outside the 20-byte buffer
  library : SafetensorError: Error while deserializing header: invalid shape, data type, or offset for tens

case: two tensors overlap
  ours    : BadSafetensors: tensor 'b': starts at byte 12 but the previous tensor ended at 16
  library : SafetensorError: Error while deserializing header: invalid offset for tensor `b`

case: hole between tensors
  ours    : BadSafetensors: tensor 'b': starts at byte 20 but the previous tensor ended at 16
  library : SafetensorError: Error while deserializing header: invalid offset for tensor `b`

case: extra bytes at the end
  ours    : BadSafetensors: 15 bytes after the last tensor are not indexed
  library : SafetensorError: Error while deserializing header: incomplete metadata, file not fully covered

case: shape does not match size
  ours    : BadSafetensors: tensor 'a': 8 bytes cannot hold shape [2, 2] of F32
  library : SafetensorError: Error while deserializing header: invalid shape, data type, or offset for tens

case: unknown dtype
  ours    : BadSafetensors: tensor 'a': malformed entry (KeyError: 'F33')
  library : SafetensorError: Error while deserializing header: invalid JSON in header: unknown variant `F33

case: duplicate tensor name
  ours    : BadSafetensors: duplicate key in header: 'a'
  library : SafetensorError: Error while deserializing header: invalid offset for tensor `a`

case: duplicate name, same entry twice
  ours    : BadSafetensors: duplicate key in header: 'a'
  library : loaded

case: header is not JSON
  ours    : BadSafetensors: header must start with '{'
  library : SafetensorError: Error while deserializing header: invalid JSON in header: expected ident at li

The first lines show our reader and the library agreeing on a file the library wrote; its first byte, \x98, is 152, the length of the JSON that follows. Both parsers then rejected nine of the ten broken files, each with its own message. There is one real difference. In the case duplicate name, same entry twice, our reader raised an error and the official library, version 0.8.0, loaded the file. The format description is explicit that “Duplicate keys are disallowed. Not all parsers may respect this.” So if your trust boundary needs duplicates to be an error, enforce it yourself, as the reader does with object_pairs_hook; Python’s json.loads would otherwise silently keep the last copy.

One more limit is worth knowing. The same README notes that “Tensor values are not checked against, in particular NaN and +/-Inf could be in the file.” A safe container tells you the file cannot run code. It does not tell you the numbers inside are trustworthy, which is why the PyTorch security guide adds that “The trustworthiness of a model is not binary.”

Step 8: Handle PyTorch checkpoints

What a .pt file really is

A checkpoint written by torch.save is a zip archive, and one of its members is a pickle called data.pkl. Loading it means unpickling that file, so everything above applies. PyTorch answers the problem with weights_only, a restricted unpickler with a built-in allowlist. In the version we installed, 2.14.1, that list holds 145 globals: 139 in the torch namespace and six general ones (collections.OrderedDict, collections.Counter, builtins.set, builtins.bytearray, builtins.complex and _codecs.encode). You can extend it with torch.serialization.add_safe_globals, the function named in the error message you will see below. PyTorch 2.6, released in January 2025, made it the default: “as an important security improvement measure we have changed the default value for weights_only parameter of torch.load.”

It is a filter, not a sandbox, and the history shows why. CVE-2025-32434 is rated Critical in NVD (CVSS 3.1 base score 9.8) and its description reads: “In version 2.5.1 and prior, a Remote Command Execution (RCE) vulnerability exists in PyTorch when loading a model using torch.load with weights_only=True. This issue has been patched in version 2.6.0.” So install PyTorch 2.6.0 or newer, and keep it updated. Even with the safe default, the torch.load documentation (checked in the docstring of the version we installed) opens with a warning: “torch.load() uses an unpickler under the hood. Never load data from an untrusted source.” The PyTorch security guide also advises to “Prefer to execute untrusted models within a secure, isolated environment such as a sandbox” and says “Safetensors gives the most safety but is the most restricted in what it supports.”

See it with a hostile checkpoint

First a small converter. It loads a checkpoint with weights_only=True, checks that the result is a flat dictionary of tensors, and writes a safetensors file. Run conversions in a throwaway environment, because the converter is the one place an untrusted checkpoint is still read.

# convert.py
"""Turn a PyTorch checkpoint into a safetensors file. Run it in a throwaway environment."""
import torch
from safetensors.torch import save_file


def pt_to_safetensors(src, dst):
    loaded = torch.load(src, map_location="cpu", weights_only=True)
    if not (isinstance(loaded, dict) and all(isinstance(k, str) and isinstance(v, torch.Tensor) for k, v in loaded.items())):
        raise TypeError("expected a flat dict of name -> tensor")
    save_file({k: v.contiguous() for k, v in loaded.items()}, dst)
# step8_torch.py
import pickletools
import re
import warnings
import zipfile
from pathlib import Path

import torch
from safetensors.torch import load_file

from attacker import Payload
from convert import pt_to_safetensors
from marker import run

warnings.simplefilter("ignore")
Path("data").mkdir(exist_ok=True)
print("torch", torch.__version__)

# 1. A checkpoint is a zip file with a pickle inside.
state = {"w": torch.arange(6, dtype=torch.float32).reshape(2, 3), "b": torch.zeros(3)}
torch.save(state, "data/ckpt.pt")
with zipfile.ZipFile("data/ckpt.pt") as z:
    names = z.namelist()
    pkl = z.read(next(n for n in names if n.endswith("data.pkl")))
print("zip members:", names)
print("globals in data.pkl:", [arg for op, arg, _ in pickletools.genops(pkl) if op.name == "GLOBAL"])

# 2. A hostile checkpoint, saved with the ordinary torch.save.
torch.save(Payload(), "data/evil.pt")
print("torch.load (default)         :", run(torch.load, "data/evil.pt"))
print("torch.load weights_only=False:", run(torch.load, "data/evil.pt", weights_only=False))
try:
    torch.load("data/evil.pt")
except Exception as exc:
    text = re.sub(r"\x1b\[[0-9;]*m", "", str(exc))  # strip the colour codes torch adds
    print("error type:", type(exc).__name__)
    print("error mentions:", [line.strip() for line in text.splitlines() if "Unsupported global" in line or "WeightsUnpickler error" in line])

# 3. The good checkpoint loads with the default.
back = torch.load("data/ckpt.pt")
print("good checkpoint loads:", {k: bool(torch.equal(back[k], state[k])) for k in state})

# 4. Convert inside a throwaway environment, then check the tensors survived.
pt_to_safetensors("data/ckpt.pt", "data/ckpt.safetensors")
round_trip = load_file("data/ckpt.safetensors")
print("converted file matches:", {k: bool(torch.equal(round_trip[k], state[k])) for k in state})
try:
    pt_to_safetensors("data/evil.pt", "data/never.safetensors")
except Exception as exc:
    print("converting the hostile file:", type(exc).__name__)
torch 2.14.1+cpu
zip members: ['ckpt/data.pkl', 'ckpt/.format_version', 'ckpt/.storage_alignment', 'ckpt/byteorder', 'ckpt/data/0', 'ckpt/data/1', 'ckpt/version', 'ckpt/.data/serialization_id']
globals in data.pkl: ['torch._utils _rebuild_tensor_v2', 'torch FloatStorage', 'collections OrderedDict']
torch.load (default)         : (False, 'UnpicklingError')
torch.load weights_only=False: (True, 'loaded')
error type: UnpicklingError
error mentions: ['WeightsUnpickler error: Unsupported global: GLOBAL print was not an allowed global by default. Please use `torch.serialization.add_safe_globals([print])` or the `torch.serialization.safe_globals([print])` context manager to allowlist this global if you trust this class/function.']
good checkpoint loads: {'w': True, 'b': True}
converted file matches: {'w': True, 'b': True}
converting the hostile file: UnpicklingError

The zip lists the pickle plus tensor storage files, and the pickle of a two-tensor state dictionary needs only three globals: collections.OrderedDict, torch._utils._rebuild_tensor_v2 and torch.FloatStorage. That short list is why an allowlist can work for checkpoints. Our hostile checkpoint, saved with an ordinary torch.save, was refused by the default load with an UnpicklingError that names the offending global, and it ran only when we passed weights_only=False. The good checkpoint loaded, the conversion round trip kept both tensors equal, and converting the hostile file failed before anything was written.

Step 9: Add a backstop with an audit hook

You do not control every call site. A dependency may call joblib.load or pickle.load deep inside its own code. Python raises an audit event every time an unpickler resolves a global; the audit events table lists pickle.find_class with the arguments module and name. A hook registered with sys.addaudithook can refuse anything not on your list, no matter who called.

# audit_guard.py
"""A process-wide backstop: refuse any pickle global that is not on the list.

sys.addaudithook() cannot be undone, and the hook sees every unpickling in the process,
so use it in a dedicated loader process, not in a general application.
"""
import sys


def install(allowed):
    allowed = frozenset(allowed)

    def hook(event, args):
        if event == "pickle.find_class":
            module, name = args
            if (module, name) not in allowed:
                raise RuntimeError(f"pickle global blocked by audit hook: {module}.{name}")

    sys.addaudithook(hook)

Because the hook never goes away (the sys module has no function to remove one), run this demo as its own process.

# audit_demo.py
"""Run this as its own process: the hook stays installed until the process ends."""
import pickle
from pathlib import Path

import joblib
import numpy as np

import audit_guard

Path("data").mkdir(exist_ok=True)
Path("data/good_array.pkl").write_bytes(pickle.dumps(np.arange(6, dtype="float32").reshape(2, 3), protocol=4))
audit_guard.install({("numpy._core.multiarray", "_reconstruct"), ("numpy", "ndarray"), ("numpy", "dtype")})


def load_pickle(path):
    with open(path, "rb") as f:
        return pickle.load(f)


LEGACY = b"c__builtin__\nprint\n(S'*** payload ran ***'\ntR."  # the old spelling of builtins.print

for label, fn in [
    ("good numpy pickle via pickle.load", lambda: load_pickle("data/good_array.pkl")),
    ("evil pickle via pickle.load", lambda: load_pickle("data/evil.pkl")),
    ("evil file via joblib.load", lambda: joblib.load("data/m.joblib")),
    ("legacy spelling via pickle.loads", lambda: pickle.loads(LEGACY)),
]:
    try:
        result = fn()
        print(f"{label}: loaded {type(result).__name__}")
    except Exception as exc:
        print(f"{label}: {type(exc).__name__}: {exc}")
good numpy pickle via pickle.load: loaded ndarray
evil pickle via pickle.load: RuntimeError: pickle global blocked by audit hook: builtins.print
evil file via joblib.load: RuntimeError: pickle global blocked by audit hook: builtins.print
legacy spelling via pickle.loads: RuntimeError: pickle global blocked by audit hook: __builtin__.print

The good array loaded, and the hostile pickles were stopped on every path, including the one through joblib.load. The fourth line is worth reading twice: the event carries the raw name from the file, __builtin__.print, so an allowlist of modern names refuses it automatically. A hook that tried to block names instead of allowing them would have the same blind spot as the blocklist in Step 4.

Keep your expectations realistic. The Python documentation says audit hooks are not suitable for implementing a sandbox, and that “malicious code can trivially disable or bypass hooks added using this function.” Use the hook as a tripwire for mistakes and for libraries you do not control, in a dedicated loader process, and never as the only defense. Because it sees every unpickling in the process, list every global your legitimate files need. For a wider tour of audit hooks, see our tutorial on catching a malicious Python package at import time with audit hooks.

Step 10: Put it together in one loader

Now combine the pieces into a single entry point, load_untrusted. It does five things. It optionally checks the file against a pinned SHA-256 hash. It identifies the format from the first bytes of the file, never from the extension. Then it routes to a parser that cannot run code: the strict safetensors reader, numpy.load with allow_pickle=False, or torch.load with weights_only=True. A pickle is refused by default, and even with allow_pickle=True it additionally requires a pinned hash and still goes through the allowlist unpickler. Anything unrecognized is refused too.

# safe_loader.py
"""load_untrusted(): choose the parser from the file's CONTENT and never fall back to plain pickle."""
import hashlib
import hmac
import zipfile
from pathlib import Path

import numpy as np

from safe_unpickler import AllowlistUnpickler
from st_reader import read_safetensors


class UnsafeFormat(Exception):
    """The file needs a loader that can run code, or it is not a format we recognise."""


def sha256_of(path) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def sniff(path) -> str:
    """Name the format from the first bytes, ignoring the file extension."""
    with open(path, "rb") as f:
        head = f.read(16)
    if head.startswith(b"\x93NUMPY"):
        return "npy"
    if head.startswith(b"PK\x03\x04"):
        with zipfile.ZipFile(path) as z:
            names = z.namelist()
        return "torch" if any(n == "data.pkl" or n.endswith("/data.pkl") for n in names) else "npz"
    if head[:1] == b"\x80" and head[1:2] in (b"\x02", b"\x03", b"\x04", b"\x05"):
        return "pickle"
    if len(head) > 8 and head[8:9] == b"{":
        return "safetensors"
    return "unknown"


def load_untrusted(path, *, sha256=None, allow_pickle=False, pickle_allowlist=frozenset()):
    path = Path(path)
    if sha256 is not None and not hmac.compare_digest(sha256_of(path), sha256.lower()):
        raise UnsafeFormat(f"{path.name}: sha256 does not match the pinned value")
    kind = sniff(path)
    if kind == "safetensors":
        tensors, _meta = read_safetensors(path)
        return tensors
    if kind in ("npy", "npz"):
        return np.load(path, allow_pickle=False)
    if kind == "torch":
        import torch  # imported late: it is slow and only needed for this branch

        return torch.load(path, map_location="cpu", weights_only=True)
    if kind == "pickle":
        if not allow_pickle:
            raise UnsafeFormat(f"{path.name}: pickle files are refused by default")
        if sha256 is None:
            raise UnsafeFormat(f"{path.name}: allow_pickle=True also needs a pinned sha256")
        with open(path, "rb") as f:
            return AllowlistUnpickler(f, allowed=pickle_allowlist).load()
    raise UnsafeFormat(f"{path.name}: unrecognised format")

The hash pin uses hmac.compare_digest for a constant-time comparison, a habit covered in our tutorial on verifying webhook signatures in Python. The Python documentation also suggests signing pickles you produce yourself: “Consider signing data with hmac if you need to ensure that it has not been tampered with.”

The demo script makes eleven load attempts on eight files, including hostile and disguised ones, and flags any attempt that prints the marker.

# step10_loader_demo.py
import contextlib
import io
import pickle
import re
import shutil
from pathlib import Path

import numpy as np

from marker import MARKER
from safe_loader import load_untrusted, sha256_of

D = Path("data")
D.mkdir(exist_ok=True)

# Files made earlier: good.safetensors (step 7), ckpt.pt and evil.pt (step 8), m.npy (step 6), evil.pkl (step 2).
np.save(D / "good.npy", np.arange(6, dtype="float32").reshape(2, 3))
(D / "good.pkl").write_bytes(pickle.dumps({"lr": 0.01, "layers": [64, 32]}, protocol=4))
shutil.copy(D / "evil.pkl", D / "evil_renamed.safetensors")


def attempt(name, note="", **kwargs):
    buf = io.StringIO()
    try:
        with contextlib.redirect_stdout(buf):
            result = load_untrusted(D / name, **kwargs)
        text = f"loaded {type(result).__name__}"
    except Exception as exc:
        first = re.sub(r"\x1b\[[0-9;]*m", "", str(exc)).splitlines()[0]
        text = f"{type(exc).__name__}: {first[:64]}"
    flag = "   <-- PAYLOAD RAN" if MARKER in buf.getvalue() else ""
    label = f"{name} {note}".strip()
    print(f"{label:<44}{text}{flag}")


attempt("good.safetensors")
attempt("good.npy")
attempt("ckpt.pt")
attempt("good.pkl")
attempt("good.pkl", "(allow_pickle)", allow_pickle=True)
attempt("good.pkl", "(allow, wrong pin)", allow_pickle=True, sha256="0" * 64)
attempt("good.pkl", "(allow, right pin)", allow_pickle=True, sha256=sha256_of(D / "good.pkl"))
attempt("evil.pkl", "(allow, right pin)", allow_pickle=True, sha256=sha256_of(D / "evil.pkl"))
attempt("evil_renamed.safetensors")
attempt("evil.pt")
attempt("m.npy")
good.safetensors                            loaded dict
good.npy                                    loaded ndarray
ckpt.pt                                     loaded dict
good.pkl                                    UnsafeFormat: good.pkl: pickle files are refused by default
good.pkl (allow_pickle)                     UnsafeFormat: good.pkl: allow_pickle=True also needs a pinned sha256
good.pkl (allow, wrong pin)                 UnsafeFormat: good.pkl: sha256 does not match the pinned value
good.pkl (allow, right pin)                 loaded dict
evil.pkl (allow, right pin)                 UnpicklingError: forbidden global: builtins.print
evil_renamed.safetensors                    UnsafeFormat: evil_renamed.safetensors: pickle files are refused by default
evil.pt                                     UnpicklingError: Weights only load failed. This file can still be loaded, to do s
m.npy                                       ValueError: Object arrays cannot be loaded when allow_pickle=False

No payload ran in any row. The safe formats loaded, the good pickle was refused until it had both the permission flag and a correct pin, and the hostile files were refused. Three rows deserve a closer look. A hostile pickle renamed to .safetensors was recognized as a pickle and refused, because the loader goes by content. The same hostile pickle with allow_pickle=True and a correct pin still failed in the allowlist, which shows what a pin means: it proves this is the file you reviewed or produced, not that the file is benign. And the object-array .npy was refused by NumPy’s default.

Step 11: Make your CI notice risky calls

The best loader does nothing if someone writes pickle.load next week. Static analysis can catch that, but each tool has gaps. This file contains one risky call per line and is only ever analyzed, never run.

# unsafe_calls.py
import pickle

import joblib
import numpy as np
import pandas as pd
import torch
import yaml

with open("model.pkl", "rb") as f:
    a = pickle.load(f)
b = pickle.loads(b"data")
c = joblib.load("model.joblib")
d = pd.read_pickle("frame.pkl")
e = np.load("arr.npy", allow_pickle=True)
g = yaml.load(open("cfg.yml"), Loader=yaml.UnsafeLoader)
h = torch.load("w.pt")
i = torch.load("w.pt", weights_only=False)

Run Bandit and Ruff against it, and then a small check of our own.

unsafe_calls.py:2: B403 Consider possible security implications associated with pickle module.
unsafe_calls.py:11: B301 Pickle and modules that wrap it can be unsafe when used to deserialize untrusted data, possible security issue.
unsafe_calls.py:12: B301 Pickle and modules that wrap it can be unsafe when used to deserialize untrusted data, possible security issue.
unsafe_calls.py:14: B301 Pickle and modules that wrap it can be unsafe when used to deserialize untrusted data, possible security issue.
unsafe_calls.py:16: B506 Use of unsafe yaml load. Allows instantiation of arbitrary objects. Consider yaml.safe_load().
unsafe_calls.py:17: B614 Use of unsafe PyTorch load
unsafe_calls.py:18: B614 Use of unsafe PyTorch load
unsafe_calls.py:11:9: S301 `pickle` and modules that wrap it can be unsafe when used to deserialize untrusted data, possible security issue
unsafe_calls.py:12:5: S301 `pickle` and modules that wrap it can be unsafe when used to deserialize untrusted data, possible security issue
unsafe_calls.py:14:5: S301 `pickle` and modules that wrap it can be unsafe when used to deserialize untrusted data, possible security issue
unsafe_calls.py:16:39: S506 Probable use of unsafe loader `UnsafeLoader` with `yaml.load`. Allows instantiation of arbitrary objects. Consider `yaml.safe_load`.
Found 4 errors.

Bandit 1.9.4 flagged pickle.load, pickle.loads, pandas.read_pickle, yaml.load with the unsafe loader and both torch.load calls, including the one that relies on the new default. Ruff 0.17.0 flagged the pickle calls (including pandas.read_pickle) and the YAML call. Neither flagged joblib.load (line 13) or numpy.load(..., allow_pickle=True) (line 15), which are the two spellings that hide the fact that pickle is involved. A short check built on the ast module closes those gaps. It resolves import aliases, so np.load becomes numpy.load, and it requires an explicit weights_only=True on torch.load, so the intent is visible in the code and does not depend on the default of whichever PyTorch version is installed.

# ci_check.py
"""Flag loader calls that can run code. Usage: python ci_check.py FILE_OR_DIR [...]  (exit 1 on findings)."""
import ast
import sys
from pathlib import Path

ALWAYS_UNSAFE = {
    "pickle.load", "pickle.loads", "pickle.Unpickler", "dill.load", "dill.loads",
    "cloudpickle.load", "cloudpickle.loads", "joblib.load", "pandas.read_pickle",
    "shelve.open", "marshal.load", "marshal.loads", "yaml.unsafe_load", "yaml.unsafe_load_all",
}


def qualified_name(node, aliases):
    """Turn np.load into 'numpy.load', following 'import ... as' and 'from ... import ... as'."""
    parts = []
    while isinstance(node, ast.Attribute):
        parts.append(node.attr)
        node = node.value
    if not isinstance(node, ast.Name):
        return None
    parts.append(aliases.get(node.id, node.id))
    return ".".join(reversed(parts))


def keyword(call, name):
    return next((k.value for k in call.keywords if k.arg == name), None)


def is_true(node):
    return isinstance(node, ast.Constant) and node.value is True


def check_source(source, filename="<string>"):
    tree = ast.parse(source, filename)
    aliases = {}
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for alias in node.names:
                aliases[alias.asname or alias.name.split(".")[0]] = alias.name if alias.asname else alias.name.split(".")[0]
        elif isinstance(node, ast.ImportFrom) and node.module:
            for alias in node.names:
                aliases[alias.asname or alias.name] = f"{node.module}.{alias.name}"
    findings = []
    for node in ast.walk(tree):
        if not isinstance(node, ast.Call):
            continue
        name = qualified_name(node.func, aliases)
        if name in ALWAYS_UNSAFE:
            findings.append((node.lineno, f"{name} can run code from the file"))
        elif name == "numpy.load" and is_true(keyword(node, "allow_pickle")):
            findings.append((node.lineno, "numpy.load(allow_pickle=True) unpickles object arrays"))
        elif name == "torch.load" and not is_true(keyword(node, "weights_only")):
            findings.append((node.lineno, "torch.load without weights_only=True"))
        elif name == "yaml.load":
            loader = keyword(node, "Loader") or (node.args[1] if len(node.args) > 1 else None)
            loader_name = qualified_name(loader, aliases) if loader is not None else None
            if not (loader_name and loader_name.endswith("SafeLoader")):
                findings.append((node.lineno, "yaml.load without a SafeLoader"))
    return sorted(findings)


def main(paths):
    total = 0
    for root in paths:
        files = sorted(Path(root).rglob("*.py")) if Path(root).is_dir() else [Path(root)]
        for file in files:
            for line, message in check_source(file.read_text(encoding="utf-8"), str(file)):
                print(f"{file}:{line}: {message}")
                total += 1
    return 1 if total else 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
unsafe_calls.py:11: pickle.load can run code from the file
unsafe_calls.py:12: pickle.loads can run code from the file
unsafe_calls.py:13: joblib.load can run code from the file
unsafe_calls.py:14: pandas.read_pickle can run code from the file
unsafe_calls.py:15: numpy.load(allow_pickle=True) unpickles object arrays
unsafe_calls.py:16: yaml.load without a SafeLoader
unsafe_calls.py:17: torch.load without weights_only=True
unsafe_calls.py:18: torch.load without weights_only=True

All eight lines are reported, and the command exits with status 1, so it fails a CI job. Wire it in next to the linters; our tutorial on catching vulnerable Python code with Bandit before it is committed shows the pre-commit side.

Step 12: Test everything

The test file covers the allowlist, the blocklist failure, the NumPy allowlist, all ten tampered safetensors files, the loader’s routing, pins, the audit hook in a child process and the PyTorch paths. Each hostile test asserts that the marker never appears.

# test_safe_loader.py
import io
import pickle
import subprocess
import sys
from pathlib import Path

import joblib
import numpy as np
import pytest
from safetensors.numpy import save_file

from attacker import Payload
from marker import MARKER, run
from safe_loader import UnsafeFormat, load_untrusted, sha256_of, sniff
from safe_unpickler import BlocklistUnpickler, RecordingUnpickler, safe_loads
from st_reader import BadSafetensors, read_safetensors
from tamper_cases import CASES

LAB = Path(__file__).parent
CONFIG = {"name": "demo", "layers": [64, 32], "lr": 0.001, "tags": ("a", "b"), "ids": {1, 2, 3}, "raw": b"\x00\x01", "none": None}


def evil_bytes():
    return pickle.dumps(Payload(), protocol=4)


@pytest.fixture(scope="session", autouse=True)
def lab_files():
    (LAB / "data").mkdir(exist_ok=True)
    (LAB / "data" / "evil.pkl").write_bytes(evil_bytes())
    joblib.dump(Payload(), LAB / "data" / "m.joblib")


# --- the allowlist ---------------------------------------------------------------
def test_plain_data_needs_no_globals_on_protocol_4():
    assert safe_loads(pickle.dumps(CONFIG, protocol=4)) == CONFIG


def test_protocol_2_data_is_refused_by_an_empty_allowlist():
    with pytest.raises(pickle.UnpicklingError, match="__builtin__"):
        safe_loads(pickle.dumps(CONFIG, protocol=2))


def test_allowlist_refuses_the_payload_and_it_never_runs():
    assert run(safe_loads, evil_bytes()) == (False, "UnpicklingError")


def test_blocklist_misses_a_legacy_spelling():
    legacy = b"c__builtin__\nprint\n(S'*** payload ran ***'\ntR."
    assert run(lambda b: BlocklistUnpickler(io.BytesIO(b)).load(), legacy) == (True, "loaded")


def test_numpy_allowlist_from_the_recorder():
    arr = np.arange(6, dtype="float32").reshape(2, 3)
    blob = pickle.dumps(arr, protocol=4)
    recorder = RecordingUnpickler(io.BytesIO(blob))
    recorder.load()
    assert np.array_equal(safe_loads(blob, recorder.seen), arr)
    assert run(safe_loads, evil_bytes(), recorder.seen) == (False, "UnpicklingError")


def test_numpy_load_refuses_object_arrays_by_default(tmp_path):
    path = tmp_path / "o.npy"
    np.save(path, np.array([Payload()], dtype=object), allow_pickle=True)
    assert run(np.load, path) == (False, "ValueError")


# --- safetensors -----------------------------------------------------------------
def test_reader_matches_the_library(tmp_path):
    tensors = {"w": np.arange(6, dtype="float32").reshape(2, 3), "i": np.array([3, 1, 2], dtype="int64")}
    path = tmp_path / "m.safetensors"
    save_file(tensors, str(path))
    mine, _meta = read_safetensors(path)
    assert all(np.array_equal(mine[k], v) for k, v in tensors.items())


def test_hand_built_good_file_loads(tmp_path):
    path = tmp_path / "g.safetensors"
    path.write_bytes(CASES["good file"])
    tensors, _meta = read_safetensors(path)
    assert tensors["a"].shape == (2, 2) and tensors["b"].tolist() == [1, 2, 3, 4]


@pytest.mark.parametrize("name", [n for n in CASES if n != "good file"])
def test_tampered_files_are_rejected(tmp_path, name):
    path = tmp_path / "t.safetensors"
    path.write_bytes(CASES[name])
    with pytest.raises(BadSafetensors):
        read_safetensors(path)


# --- the loader ------------------------------------------------------------------
def test_loader_goes_by_content_not_extension(tmp_path):
    path = tmp_path / "model.safetensors"
    path.write_bytes(evil_bytes())
    assert sniff(path) == "pickle"
    assert run(load_untrusted, path) == (False, "UnsafeFormat")


def test_pickle_is_refused_by_default(tmp_path):
    path = tmp_path / "x.pkl"
    path.write_bytes(pickle.dumps(CONFIG, protocol=4))
    with pytest.raises(UnsafeFormat, match="refused by default"):
        load_untrusted(path)


def test_allow_pickle_still_needs_a_pin(tmp_path):
    path = tmp_path / "x.pkl"
    path.write_bytes(pickle.dumps(CONFIG, protocol=4))
    with pytest.raises(UnsafeFormat, match="pinned sha256"):
        load_untrusted(path, allow_pickle=True)


def test_wrong_pin_is_refused(tmp_path):
    path = tmp_path / "x.pkl"
    path.write_bytes(pickle.dumps(CONFIG, protocol=4))
    with pytest.raises(UnsafeFormat, match="does not match"):
        load_untrusted(path, allow_pickle=True, sha256="0" * 64)


def test_pinned_good_pickle_loads(tmp_path):
    path = tmp_path / "x.pkl"
    path.write_bytes(pickle.dumps(CONFIG, protocol=4))
    assert load_untrusted(path, allow_pickle=True, sha256=sha256_of(path)) == CONFIG


def test_pinning_a_hostile_file_does_not_make_it_safe(tmp_path):
    path = tmp_path / "x.pkl"
    path.write_bytes(evil_bytes())
    assert run(load_untrusted, path, allow_pickle=True, sha256=sha256_of(path)) == (False, "UnpicklingError")


# --- the audit hook (needs its own process) ----------------------------------------
def test_audit_hook_blocks_hostile_loads_in_a_child_process():
    proc = subprocess.run([sys.executable, str(LAB / "audit_demo.py")], cwd=LAB, capture_output=True, text=True)
    assert "good numpy pickle via pickle.load: loaded ndarray" in proc.stdout
    assert proc.stdout.count("pickle global blocked by audit hook") == 3
    assert "blocked by audit hook: __builtin__.print" in proc.stdout
    assert MARKER not in proc.stdout


# --- PyTorch ---------------------------------------------------------------------
def test_torch_default_refuses_a_hostile_checkpoint(tmp_path):
    torch = pytest.importorskip("torch")
    path = tmp_path / "evil.pt"
    torch.save(Payload(), path)
    assert run(torch.load, path) == (False, "UnpicklingError")


def test_loader_reads_a_good_checkpoint(tmp_path):
    torch = pytest.importorskip("torch")
    path = tmp_path / "ok.pt"
    state = {"w": torch.arange(6.0)}
    torch.save(state, path)
    assert sniff(path) == "torch"
    assert torch.equal(load_untrusted(path)["w"], state["w"])


def test_checkpoint_converts_to_safetensors(tmp_path):
    torch = pytest.importorskip("torch")
    from convert import pt_to_safetensors
    from safetensors.torch import load_file

    src, dst = tmp_path / "ok.pt", tmp_path / "ok.safetensors"
    state = {"w": torch.arange(6.0).reshape(2, 3)}
    torch.save(state, src)
    pt_to_safetensors(src, dst)
    assert torch.equal(load_file(str(dst))["w"], state["w"])
# test_ci_check.py
from pathlib import Path

from ci_check import check_source

LAB = Path(__file__).parent


def test_every_unsafe_call_in_the_sample_is_flagged():
    findings = check_source((LAB / "unsafe_calls.py").read_text(encoding="utf-8"))
    assert [line for line, _ in findings] == [11, 12, 13, 14, 15, 16, 17, 18]


def test_safe_spellings_are_not_flagged():
    source = (
        "import numpy as np, torch, yaml\n"
        "np.load('a.npy')\n"
        "torch.load('w.pt', weights_only=True)\n"
        "yaml.safe_load('a: 1')\n"
        "yaml.load('a: 1', Loader=yaml.SafeLoader)\n"
    )
    assert check_source(source) == []


def test_import_aliases_are_followed():
    source = "from joblib import load as jl\nimport pickle as pk\njl('m.joblib')\npk.loads(b'')\n"
    assert [line for line, _ in check_source(source)] == [3, 4]
............................... [100%]
31 passed in 1.39s

All 31 tests pass in about a second and a half. If one of them ever fails, the marker line is the first thing to look for in the output.

Common mistakes and gotchas

Treating a scanner as the gate

Scanners that read opcodes are helpful, and the Hugging Face page quoted in Step 3 is honest about their limits. An allowlist does not try to recognize bad pickles; it only permits known-good ones.

Trusting the file extension

A name like model.safetensors proves nothing. Step 10 showed a pickle wearing that name, and the loader refused it because it read the bytes.

Pinning a hash and calling it safe

A pin tells you the file did not change since you recorded the hash. If the first version was hostile, the pin protects the hostile version. Keep the allowlist even when the hash matches.

Treating weights_only=True as a sandbox

CVE-2025-32434 existed in versions up to 2.5.1. Pin PyTorch 2.6.0 or newer, and still prefer safetensors and isolation for models from unknown sources.

Copying an allowlist between environments

Step 5 showed that the globals for one array depend on the pickle protocol and the NumPy release. Record the list in the environment that will load the files, and re-record it when you upgrade.

Forgetting the other loaders

Step 6 listed five that execute code. Add every loader your project uses to the CI check, and prefer yaml.safe_load for YAML.

Assuming a safe container means a safe model

Safetensors cannot run code, but it does not validate values, and a model can still be poisoned. Treat provenance, review and isolation as separate controls.

Verify the whole thing end to end

Run the scripts in this order from inside picklelab. After the last command you should see the final line 31 passed.

python attacker.py
python victim_naive.py
python dis_payload.py
python step4_allowlist.py
python step5_numpy.py
python step6_matrix.py
python step7_tamper.py
python step8_torch.py
python audit_demo.py
python step10_loader_demo.py
python -m bandit -q -f custom --msg-template "{relpath}:{line}: {test_id} {msg}" unsafe_calls.py
python -m ruff check --select S --output-format concise unsafe_calls.py
python ci_check.py unsafe_calls.py
python -m pytest -q

The payload should run only where this tutorial says it does: the naive victim in Step 2 (it prints the marker line), the blocklist row in Step 4 that fell for the legacy spelling and shows True, the rows in Step 6 that show True, and the weights_only=False row in Step 8. Everywhere else the payload must stay silent, and the loader table in Step 10 must never show PAYLOAD RAN. If it does, stop and investigate before trusting any loader you built.

Where to go next

  • Put the converter from Step 8 in a pipeline that downloads a model, converts it inside a container with no network access, and stores only the safetensors file and its hash.
  • If your own jobs save pickled checkpoints, as in our tutorial on checkpointing and resuming a Python training job, apply the same rules when they are loaded back from shared storage, or move them to a data format.
  • Extend ci_check.py with the loaders your team uses that this tutorial did not cover, and run it on every pull request.
  • Read the Restricting Globals section of the Python documentation and the PyTorch security guide in full; both are short and worth keeping close.

Tags:

AI SecurityApplication SecurityInsecure DeserializationMachine LearningPythonSafetensors

Share

No Comment! Be the first one.

Leave a Reply

Your email address will not be published. Required fields are marked *

Related Posts