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

Author
unknown
Download Count*
License
Functions reliant on pickle can result in arbitrary code execution
Run Locally
Run in CI
Defintion
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)
Short Link: https://sg.run/kzGP