CHEG:

This repository provides the implementation of CHEG, a vulnerability detection framework that integrates Pre-trained Language Models (PLMs) with heterogeneous graph neural networks (GNNs) for fine-grained vulnerability localization.

CHEG models the structural and semantic relationships in source code using heterogeneous graphs and leverages language model representations to enhance vulnerability detection performance. The framework also provides several variants to explore different graph modeling strategies and ablation settings.

---

Project Structure

.
├── main.py                # LM + GNN training and evaluation
├── main_meta_path.py      # Meta-path based heterogeneous graph model
├── main_withoutSSCG.py    # CPG + PLM version (without SSCG)
├── main_hom.py            # Homogeneous graph version
├── main_abulation.py      # Ablation experiments
├── models/                # Model implementations
├── data/                  # Dataset files
├── utils/                 # Utility functions
└── README.md

---

Installation

Requirements

• Python ≥ 3.8
• PyTorch
• DGL / PyG
• Transformers
• Numpy
• Scikit-learn

Install dependencies using:

pip install -r requirements.txt

---

Usage

1. LM + GNN Model (CHEG)

Run the main CHEG model that integrates language model representations with graph neural networks.

python main.py --dataset <dataset_name>

Example:

python main.py --dataset Libav

---

2. Test Only

Evaluate a trained model.

python main.py --dataset <dataset_name> --model_path <model_path> --test_only

Example:

python main.py --dataset Libav --model_path checkpoints/model.pt --test_only

---

3. Meta-path Based Model

Run the variant that incorporates meta-path based heterogeneous graph modeling.

python main_meta_path.py

---

4. Using CPG + PLM (without SSCG)

Run the variant that combines Code Property Graph (CPG) with Pre-trained Language Models, without using the SSCG structure.

python main_withoutSSCG.py

---

5. Homogeneous Graph Version

Run the model using homogeneous graph representation.

python main_hom.py

---

6. Ablation Study

Run ablation experiments used in the paper.

python main_abulation.py

---

Acknowledgements

This project builds upon several prior research efforts and open-source implementations in the field of vulnerability detection. We sincerely appreciate the contributions of the research community that made this work possible.

In particular, we thank the authors of SySeVR (Li et al.) for their pioneering work on vulnerability detection and their valuable insights into code representation and vulnerability analysis.

We also acknowledge the authors of FVD-DPM for their open research contributions, which inspired parts of this implementation.

---

References

@inproceedings{li2018sysevr,
  author = {Li, Zhen and Zou, Deqing and Xu, Shouhuai and Jin, Hai and Zhu, Yawei and Chen, Zhaoxuan and Wang, Sujuan},
  title = {SySeVR: A Framework for Using Deep Learning to Detect Software Vulnerabilities},
  booktitle = {Proceedings of the 2018 IEEE/ACM 40th International Conference on Software Engineering (ICSE)},
  year = {2018},
  pages = {224--234},
  publisher = {ACM}
}

@inproceedings{DBLP:conf/uss/ShaoD24,
  author       = {Miaomiao Shao and
                  Yuxin Ding},
  editor       = {Davide Balzarotti and
                  Wenyuan Xu},
  title        = {{FVD-DPM:} Fine-grained Vulnerability Detection via Conditional Diffusion
                  Probabilistic Models},
  booktitle    = {33rd USENIX Security Symposium (USENIX Security 2024)},
  address      = {Philadelphia, PA, USA},
  publisher    = {USENIX Association},
  year         = {2024},
  url          = {https://www.usenix.org/conference/usenixsecurity24/presentation/shao}
}

---

Citation

If you find this repository useful for your research, please consider citing the related papers above.