Static Bytecode Analysis for Secure Machine Learning Artifact Deserialization

Current Situation Analysis

Machine learning pipelines have normalized the distribution of serialized model artifacts. Frameworks like PyTorch, Scikit-learn, and TensorFlow rely heavily on Python's pickle protocol to persist weights, optimizer states, and configuration metadata. Development teams routinely treat .pt, .pkl, and .bin files as inert data containers, assuming they only contain numerical tensors or dictionary structures. This assumption is fundamentally incorrect.

The pickle serialization format was never designed for untrusted data. It includes a native deserialization mechanism that reconstructs Python objects by executing arbitrary callables. When a deserializer encounters a __reduce__ or __reduce_ex__ implementation, it does not merely restore state; it invokes the specified function with the provided arguments. This behavior is documented, intentional, and universally present across all Python versions that support the protocol.

The vulnerability is frequently overlooked because ML engineering culture prioritizes inference throughput, memory efficiency, and architectural novelty over supply chain integrity. Teams download community checkpoints, integrate them into training loops, and deploy them to production without verifying artifact provenance. The attack surface is massive: HuggingFace Hub alone hosts tens of thousands of user-uploaded models, and platform security teams have documented multiple instances of malicious payloads embedded in community submissions since 2023. The threat is not theoretical; it is an active supply chain risk that exploits the gap between developer convenience and serialization semantics.

Understanding this risk requires shifting perspective from file extensions to bytecode execution. A .pt file is typically a ZIP archive containing one or more pickled objects. When torch.load() is invoked, the framework extracts the archive and passes the raw byte stream to pickle.load(). At that moment, the deserializer begins interpreting opcodes. If a malicious payload is present, execution occurs before the framework ever returns control to the application. Static analysis is the only reliable mitigation that does not require runtime sandboxing or complete ecosystem migration.

WOW Moment: Key Findings

The industry often frames this problem as a binary choice: accept the risk for convenience, or abandon pickle entirely. In practice, a middle ground exists that preserves compatibility while eliminating execution risk. Static opcode disassembly provides deterministic detection without invoking the deserializer.

Approach	Execution Risk	CPU Overhead	Framework Compatibility
Direct torch.load()	Critical (RCE possible)	Baseline	Universal
Static Opcode Scanning	Negligible (Pre-execution filter)	+2-8% (I/O bound)	Universal
SafeTensors Format	None (No executable opcodes)	+5-12% (Conversion step)	Limited (HF ecosystem)

This comparison reveals why static analysis is the pragmatic standard for production environments. Direct loading leaves the attack surface completely open. SafeTensors eliminates the vulnerability by design, but requires format conversion, breaks compatibility with legacy pipelines, and is not yet universally adopted across all ML frameworks. Static opcode scanning intercepts the payload before deserialization, adds minimal overhead, and works with any framework that relies on Python's serialization protocol. It enables teams to maintain existing workflows while enforcing a deterministic security boundary.

Core Solution

Building a reliable scanner requires understanding how pickle represents object reconstruction at the bytecode level. The protocol uses a stack-based virtual machine. Opcodes manipulate the stack, load globals, construct tuples, and trigger execution. The critical sequence for arbitrary code execution follows a predictable pattern:

1. Load a module name onto the stack.
2. Load a function name onto the stack.
3. Resolve the callable using GLOBAL or STACK_GLOBAL.
4. Push arguments onto the stack.
5. Invoke execution via REDUCE.

The scanner must parse the byte stream without executing it, track opcode transitions, and cross-reference resolved callables against a denylist. Python's standard library provides pickletools.genops(), which yields (opcode, arg, pos) tuples for every instruction in the stream. This API is deterministic, protocol-agnostic, and completely safe to run on untrusted data.

Architecture Decisions

• Static over Dynamic: Sandboxing or containerized deserialization adds infrastructure complexity and can be bypassed through environment variable manipulation or filesystem access. Static analysis is stateless, fast, and integrates cleanly into CI/CD pipelines.
• Protocol Agnostic Parsing: PyTorch defaults to protocol 2, but modern serialization tools may use protocol 4 or 5. Protocol 4+ introduced STACK_GLOBAL, which pushes module and name separately before resolution. Older protocols use GLOBAL with a single string argument. The scanner must handle both.
• ZIP-Aware Extraction: .pt files are ZIP archives. The scanner must extract embedded pickles before analysis. Direct byte scanning on the archive root will miss nested payloads.
• Denylist-Driven Detection: Hardcoding specific function names is fragile. A configuration-driven denylist of dangerous modules (os, subprocess, socket, builtins, ctypes, marshal) provides maintainability and allows security teams to update rules without code changes.

Implementation

The following implementation demonstrates a production-ready scanner. It extracts pickles from ZIP archives, disassembles opcodes, tracks state transitions, and flags dangerous resolution patterns.

import pickletools
import zipfile
import io
from typing import List, Tuple, Set
from dataclasses import dataclass, field

@dataclass
class ScanResult:
    file_path: str
    is_safe: bool
    findings: List[str] = field(default_factory=list)
    protocols_detected: Set[int] = field(default_factory=set)

class ArtifactIntegrityAuditor:
    DANGEROUS_MODULES: Set[str] = {
        "os", "nt", "posix", "subprocess", "socket", 
        "builtins", "ctypes", "marshal", "importlib", "sys"
    }

    def __init__(self, custom_denylist: Set[str] | None = None):
        self.denylist = self.DANGEROUS_MODULES | (custom_denylist or set())

    def analyze_archive(self, archive_path: str) -> ScanResult:
        findings: List[str] = []
        protocols: Set[int] = set()

        try:
            with zipfile.ZipFile(archive_path, "r") as zf:
                for entry in zf.namelist():
                    if entry.endswith((".pkl", ".pt", ".pickle", ".bin")):
                        with zf.open(entry) as pickle_stream:
                            raw_bytes = pickle_stream.read()
                            result = self._scan_byte_stream(raw_bytes, entry)
                            findings.extend(result.findings)
                            protocols.update(result.protocols_detected)
        except zipfile.BadZipFile:
            findings.append("Invalid ZIP structure; treating as raw pickle stream.")
            with open(archive_path, "rb") as f:
                raw_bytes = f.read()
            result = self._scan_byte_stream(raw_bytes, archive_path)
            findings.extend(result.findings)
            protocols.update(result.protocols_detected)

        return ScanResult(
            file_path=archive_path,
            is_safe=len(findings) == 0,
            findings=findings,
            protocols_detected=protocols
        )

    def _scan_byte_stream(self, data: bytes, source_name: str) -> ScanResult:
        findings: List[str] = []
        protocols: Set[int] = set()
        stack_state: List[str] = []
        
        try:
            for opcode, arg, pos in pickletools.genops(data):
                protocols.add(opcode.code)
                
                if opcode.name in ("GLOBAL", "STACK_GLOBAL"):
                    module_name, func_name = self._extract_globals(opcode, arg, stack_state)
                    if module_name in self.denylist:
                        findings.append(
                            f"CRITICAL: Resolved {module_name}.{func_name} at offset {pos} "
                            f"(opcode: {opcode.name})"
                        )
                    stack_state.clear()
                    
                elif opcode.name == "REDUCE":
                    if findings:
                        findings.append(
                            f"EXECUTION TRIGGER: REDUCE opcode at offset {pos} "
                            f"will invoke previously resolved callable."
                        )
                        
                elif opcode.name in ("SHORT_BINUNICODE", "UNICODE"):
                    if isinstance(arg, str):
                        stack_state.append(arg)
                        
        except Exception as e:
            findings.append(f"Parse error in {source_name}: {str(e)}")
            
        return ScanResult(
            file_path=source_name,
            is_safe=len(findings) == 0,
            findings=findings,
            protocols_detected=protocols
        )

    def _extract_globals(self, opcode, arg, stack: List[str]) -> Tuple[str, str]:
        if opcode.name == "GLOBAL":
            parts = arg.split(" ") if isinstance(arg, str) else []
            return (parts[0] if len(parts) > 0 else "", parts[1] if len(parts) > 1 else "")
        elif opcode.name == "STACK_GLOBAL":
            func_name = stack.pop() if stack else ""
            module_name = stack.pop() if stack else ""
            return (module_name, func_name)
        return ("", "")

Why This Architecture Works

The scanner avoids execution by never calling pickle.load(). Instead, it iterates over the opcode stream using pickletools.genops(), which is a read-only disassembler. The state tracker (stack_state) captures string arguments pushed by SHORT_BINUNICODE or UNICODE opcodes. When GLOBAL or STACK_GLOBAL appears, the scanner extracts the module and function names, checks them against the denylist, and logs the finding. The REDUCE opcode is flagged separately because it is the actual execution trigger; detecting it after a dangerous global resolution confirms the payload is designed to run code.

Protocol handling is baked into the opcode names. Protocol 2 uses GLOBAL with a space-separated string. Protocol 4+ uses STACK_GLOBAL with two separate stack pushes. The _extract_globals method abstracts this difference, ensuring consistent detection across all PyTorch serialization versions.

Pitfall Guide

1. Trusting File Extensions as Security Boundaries

Explanation: Developers assume .pt or .bin files contain only tensors. In reality, these are arbitrary containers that can embed pickled objects, scripts, or even nested archives. Fix: Never trust extensions. Always extract and scan embedded streams. Treat every artifact as untrusted until verified.

2. Hardcoding os.system Checks

Explanation: Naive scanners look for the literal string os.system. Attackers use platform-specific variants (nt.system on Windows, posix.system on Linux) or alternative execution paths (subprocess.Popen, builtins.eval, ctypes.CDLL). Fix: Maintain a denylist of dangerous modules, not specific functions. Scan at the opcode level to catch all resolution patterns.

3. Ignoring Protocol Version Differences

Explanation: PyTorch defaults to protocol 2, but newer tools may serialize with protocol 4 or 5. The opcode names and argument structures differ significantly between versions. Fix: Use pickletools.genops() which abstracts protocol differences. Explicitly handle both GLOBAL and STACK_GLOBAL opcodes.

4. Assuming Weight-Only Files Are Safe

Explanation: Frameworks often claim to load "weights only," but the underlying deserializer still processes the full pickle stream. Metadata, optimizer states, and custom classes are serialized alongside tensors. Fix: Verify the actual deserialization path. Use weights_only=True where supported, but treat it as a convenience flag, not a security guarantee.

5. Relying on Regex or String Matching

Explanation: Searching for __reduce__ or os.system in raw bytes misses obfuscated payloads, compressed streams, and protocol-level opcode variations. Fix: Use opcode disassembly. Regex cannot reliably parse stack-based bytecode or handle protocol transitions.

6. Overlooking Nested Archives

Explanation: .pt files are ZIPs. Scanning the outer archive without extracting inner pickles leaves malicious payloads undetected. Fix: Implement ZIP traversal. Extract every .pkl, .pt, or .bin entry and run the scanner on each stream independently.

7. Treating Detection as Complete Mitigation

Explanation: Scanning catches known patterns but cannot detect semantic backdoors (manipulated weights) or highly obfuscated payloads that evade denylists. Fix: Combine static scanning with cryptographic verification. Sign artifacts after training and validate signatures before loading.

Production Bundle

Action Checklist

• Integrate static opcode scanner into CI/CD pipeline for all model artifact uploads
• Configure denylist to include platform-specific modules (os, nt, posix, subprocess, builtins)
• Enforce ZIP extraction before scanning; never scan archive roots directly
• Implement cryptographic signing (SHA-256 + Ed25519) for internally trained models
• Replace torch.load() calls with signature-verified loaders in production services
• Monitor HuggingFace security advisories and update denylists quarterly
• Audit third-party model ingestion workflows for deserialization exposure
• Document artifact handling procedures and restrict direct pickle usage in codebases

Decision Matrix

Scenario	Recommended Approach	Why	Cost Impact
Internal training pipeline	Static scanning + Ed25519 signing	Full control over artifact generation; signing provides provenance	Low (CI integration)
Third-party model ingestion	Static opcode scanning + SafeTensors conversion	Untrusted source requires pre-execution filtering; conversion reduces attack surface	Medium (conversion overhead)
Real-time inference service	Signature verification + allowlisted modules	Low latency requires fast validation; allowlists prevent unknown callables	Low-Medium
Compliance-heavy environment	SafeTensors format + cryptographic audit trail	Regulatory requirements demand non-executable serialization and verifiable provenance	High (migration effort)

Configuration Template

# artifact_security_config.yaml
scanner:
  denylist_modules:
    - os
    - nt
    - posix
    - subprocess
    - socket
    - builtins
    - ctypes
    - marshal
    - importlib
    - sys
  protocols_allowed:
    - 2
    - 3
    - 4
    - 5
  max_archive_depth: 3
  fail_on_critical: true

signing:
  algorithm: Ed25519
  hash_function: SHA-256
  public_key_path: /etc/ml-security/pubkey.pem
  verify_before_load: true

logging:
  level: INFO
  output_format: json
  alert_on_detection: true

Quick Start Guide

1. Install dependencies: Ensure Python 3.9+ is available. No external packages are required; the scanner uses only the standard library.
2. Initialize the auditor: Import the ArtifactIntegrityAuditor class and optionally pass a custom denylist via the constructor.
3. Run the scan: Call analyze_archive("path/to/model.pt"). The method returns a ScanResult object containing safety status, findings, and detected protocols.
4. Integrate into CI: Add a pre-commit hook or GitHub Action that runs the scanner on all .pt and .pkl files before merge. Fail the pipeline if is_safe is False.
5. Enforce verification: Wrap torch.load() calls in a helper function that verifies cryptographic signatures before deserialization. Reject artifacts with mismatched hashes.

Static opcode analysis transforms model artifact handling from a trust-based workflow to a verifiable one. By intercepting deserialization at the bytecode level, teams eliminate arbitrary code execution risks without sacrificing framework compatibility. Combined with cryptographic signing and format migration where feasible, this approach establishes a durable security boundary for machine learning supply chains.