trailofbits.python.pickles-in-torch-distributed.pickles-in-torch-distributed

profile photo of trailofbitstrailofbits
Author
unknown
Download Count*
CC BY-NC-SA 4.0
License

Functions reliant on pickle can result in arbitrary code execution

Run Locally

Run in CI

Semgrep CI docs

Defintion

Edit RuleView Source
Open in Playground 
rules:
  - id: pickles-in-torch-distributed
    message: Functions reliant on pickle can result in arbitrary code execution
    languages:
      - python
    severity: ERROR
    metadata:
      category: security
      cwe: "CWE-502: Deserialization of Untrusted Data"
      subcategory:
        - vuln
      confidence: MEDIUM
      likelihood: MEDIUM
      impact: HIGH
      references:
        - https://blog.trailofbits.com/2021/03/15/never-a-dill-moment-exploiting-machine-learning-pickle-files/
      license: CC-BY-NC-SA-4.0
    patterns:
      - pattern-either:
          - pattern: torch.distributed.broadcast_object_list(...)
          - pattern: torch.distributed.all_gather_object(...)
          - pattern: torch.distributed.gather_object(...)
          - pattern: torch.distributed.scatter_object_list(...)

Examples

pickles-in-torch-distributed.py

import torch.distributed as dist

if dist.get_rank() == 0:
    objects = ["f", 1] 
else:
    objects = [None, None]
    
# ruleid: pickles-in-torch-distributed
dist.broadcast_object_list(objects, src=0)

# ruleid: pickles-in-torch-distributed
dist.all_gather_object(output, gather_objects[dist.get_rank()])

# ruleid: pickles-in-torch-distributed
dist.gather_object(
        gather_objects[dist.get_rank()],
        output if dist.get_rank() == 0 else None,
        dst=0
    )

# ruleid: pickles-in-torch-distributed
dist.scatter_object_list(output_list, objects, src=0)

# ok: pickles-in-torch-distributed
dist.scatter(output_list, objects, src=0)