DeepWukong

中文文档

(TOSEM'21) DeepWukong: Statically Detecting Software Vulnerabilities Using Deep Graph Neural Network

Docker (Recommended)

Build

docker build -t deepwukong .

Train with preprocessed data

# Download and extract data first
# See "Data" section below for download link
wget -O Data.7z "https://bupteducn-my.sharepoint.com/:u:/g/personal/jackiecheng_bupt_edu_cn/EalnVAYC8zZDgwhPmGJ034cBYNZ8zB7-mNSNm-a7oYXkcw?e=eRUc50&download=1"
7z x Data.7z -odata/

# Train on GPU
docker run --gpus all -v $(pwd)/data:/workspace/data deepwukong

# Train with custom config
docker run --gpus all \
  -v $(pwd)/data:/workspace/data \
  -v $(pwd)/configs:/workspace/configs \
  deepwukong src/run.py -c configs/dwk.yaml

Train from source code (full pipeline)

Place your source code under data/<DATASET>/source-code/ with a manifest.xml (see Data Preparation Guide for format), then run each step:

# Step 1: Generate PDG from source code (joern)
docker run --gpus all -v $(pwd)/data:/workspace/data \
  deepwukong src/joern/joern-parse.py -c configs/dwk.yaml

# Step 2: Extract XFG from PDG
docker run --gpus all -v $(pwd)/data:/workspace/data \
  deepwukong src/data_generator.py -c configs/dwk.yaml

# Step 3: Symbolize and split into train/val/test
docker run --gpus all -v $(pwd)/data:/workspace/data \
  deepwukong src/preprocess/dataset_generator.py -c configs/dwk.yaml

# Step 4: Train word embeddings
docker run --gpus all -v $(pwd)/data:/workspace/data \
  deepwukong src/preprocess/word_embedding.py -c configs/dwk.yaml

# Step 5: Train model
docker run --gpus all -v $(pwd)/data:/workspace/data \
  deepwukong src/run.py -c configs/dwk.yaml

# Step 6 (optional): Evaluate a saved checkpoint
docker run --gpus all \
  -v $(pwd)/data:/workspace/data \
  -v $(pwd)/ts_logger:/workspace/ts_logger \
  deepwukong src/evaluate.py <path to checkpoint>

To use a different dataset, copy configs/dwk.yaml and change dataset.name:

docker run --gpus all \
  -v $(pwd)/data:/workspace/data \
  -v $(pwd)/configs:/workspace/configs \
  deepwukong src/run.py -c configs/my_dataset.yaml

GPU Requirements

• NVIDIA GPU with CUDA 12.4+ support
• NVIDIA Container Toolkit installed
• Use --gpus all to enable GPU access (CPU fallback is automatic)

Local Setup

• Requirements: Python >= 3.10, PyTorch >= 2.0 (CUDA recommended)

• Install dependencies

  bash env.sh

  Or manually:

  # Install PyTorch (use CUDA version for GPU support)
  # CPU only:
  pip install torch
  # CUDA 12.4:
  pip install torch --index-url https://download.pytorch.org/whl/cu124
  
  pip install torch-geometric
  pip install -r requirements.txt

Data

• Preprocessed CWE119 data: Download from data (7z, ~1.8GB), extract under data/:

  7z x Data.7z -odata/

• Custom datasets: See Data Preparation Guide (中文版) for the full pipeline from source code to training data.

• Expected structure:

  data/
  ├── sensiAPI.txt
  └── CWE119/
      ├── train.json, val.json, test.json
      ├── w2v.wv
      ├── XFG/          # XFG pickle files
      └── source-code/  # Original source (needed for preprocessing only)
  

---

One-Step Evaluation

• From Pretrained model

  • Download from pretrained model.
  • Note: Pretrained checkpoints saved with pytorch-lightning 1.x may not be compatible. Retraining is recommended.
  • PYTHONPATH="." python src/evaluate.py <path to the pretrained model>

• Training and Testing

  PYTHONPATH="." python src/run.py

  GPU is auto-detected. Training on CWE119 takes ~9 minutes on dual RTX A5000.

---

Run from scratch:

Data preparation

See Data Preparation Guide for detailed instructions on preparing custom datasets.

Use joern to Generate PDG

joern is included in joern/ directory (requires Java 8+).

PYTHONPATH="." python src/joern/joern-parse.py -c <config file>

Generate raw XFG

PYTHONPATH="." python src/data_generator.py -c <config file>

Symbolize and Split Dataset

PYTHONPATH="." python src/preprocess/dataset_generator.py -c <config file>

Word Embedding Pretraining

PYTHONPATH="." python src/preprocess/word_embedding.py -c <config file>

Evaluation

PYTHONPATH="." python src/run.py -c <config file>

Acknowledgements

The modernization of this project (dependency upgrades, Docker support, and data preparation documentation) was done with Claude Code.

Citation

Please kindly cite our paper if it benefits:

@article{xiao2021deepwukong,
author = {Cheng, Xiao and Wang, Haoyu and Hua, Jiayi and Xu, Guoai and Sui, Yulei},
title = {DeepWukong: Statically Detecting Software Vulnerabilities Using Deep Graph Neural Network},
year = {2021},
publisher = {ACM},
volume = {30},
number = {3},
url = {https://doi.org/10.1145/3436877},
doi = {10.1145/3436877},
journal = {ACM Trans. Softw. Eng. Methodol.},
articleno = {38},
numpages = {33}
}