trailofbits.python.pytorch-package.pytorch-package

profile photo of trailofbitstrailofbits
Author
unknown
Download Count*

Avoid importing torch.package - it can result in arbitrary code execution via pickle

Run Locally

Run in CI

Defintion

rules:
  - id: pytorch-package
    message: Avoid importing torch.package - it can result in arbitrary code
      execution via pickle
    languages:
      - python
    severity: WARNING
    metadata:
      category: security
      cwe: "CWE-502: Deserialization of Untrusted Data"
      subcategory:
        - audit
      confidence: LOW
      likelihood: MEDIUM
      impact: HIGH
      technology:
        - pytorch
      description: Potential arbitrary code execution from `torch.package`
      references:
        - https://pytorch.org/docs/1.13/package.html#torch-package
        - https://blog.trailofbits.com/2021/03/15/never-a-dill-moment-exploiting-machine-learning-pickle-files/
      license: AGPL-3.0 license
      vulnerability_class:
        - "Insecure Deserialization "
    pattern: import torch.package

Examples

pytorch-package.py

# ruleid: pytorch-package
import torch.package

# ruleid: pytorch-package
from torch import package 

# ruleid: pytorch-package
import torch.package as tp 

# ok: pytorch-package
import torchx.package as tp