MANDO-HGT: Heterogeneous Graph Transformers for Smart Contract Vulnerability Detection

Overview

This repository is the implementation of MANDO-HGT which should be applicable to either source code or bytecode/binary form of software programs to cater to different situations where the source code of the software may or may not be available.

Citation

Nguyen, H. H., Nguyen, N.M., Xie, C., Ahmadi, Z., Kudenko, D., Doan, T. N., & Jiang, L. (2023, May). MANDO-HGT: Heterogeneous Graph Transformers for Smart Contract Vulnerability Detection. In Proceedings of 20th International Conference on Mining Software Repositories (MSR' 23), Melbourne, Australia, 2023. Preprint

@inproceedings{nguyen2023msr,
  author = {Nguyen, Hoang H. and Nguyen, Nhat-Minh and Xie, Chunyao and Ahmadi, Zahra and Kudenko, Daniel and Doan, Thanh-Nam and Jiang, Lingxiao},
  title = {MANDO-HGT: Heterogeneous Graph Transformers for Smart Contract Vulnerability Detection},
  year = {2023},
  month = {5},
  booktitle = {Proceedings of the 20th International Conference on Mining Software Repositories},
  numpages = {13},
  keywords = {vulnerability detection, smart contracts, source code, bytecode, heterogeneous graph learning, graph transformer},
  location = {Melbourne, Australia},
  series = {MSR '23}
}

Table of contents

• MANDO-HGT
• Table of contents
• How to train the models?

- Dataset - Source code - Byte code - System Description - Install Environment - Inspection scripts - Graph Classification - Node Classification - Trainer - Graph Classification - Usage - Examples - Node Classification - Usage - Examples - Testing - Visuallization

How to train the models?

Dataset

Source code

• We prepared dataset for experiments.

Byte code

• We used Solc compilation of Crytic-compile to compile EVM creation/runtime bytecode from the source code of smart contracts. Then we used EtherSolve to generate CFGs for bytecode.

System Description

We run all experiments on

• Ubuntu 20.04
• CUDA 11.1
• NVIDA 3080

Install Environment

Install python required packages.

pip install -r requirements.txt -f https://download.pytorch.org/whl/lts/1.8/torch_lts.html -f https://data.pyg.org/whl/torch-1.8.0+cu111.html -f https://data.dgl.ai/wheels/repo.html

Inspection scripts

We provied inspection scripts for Graph Classification and Node Classification tasks as well as their required data.

Graph Classification

Training Phase

python -m experiments.graph_classification --epochs 50 --repeat 20

To show the result table

python -m experiments.graph_classification --result

Node Classification

Training Phase

python -m experiments.node_classification --epochs 50 --repeat 20

To show the result table

python -m experiments.node_classification --result

• We currently supported 7 types of bug: accesscontrol, arithmetic, denialofservice, frontrunning, reentrancy, timemanipulation, uncheckedlowlevelcalls.

• Run the inspection

Trainer

Graph Classification

Usage

usage: MANDO Graph Classifier [-h] [-s SEED] [-ld LOG_DIR]
                              [--outputmodels OUTPUTMODELS]
                              [--compressedgraph COMPRESSEDGRAPH]
                              [--dataset DATASET] [--testset TESTSET]
                              [--label LABEL] [--checkpoint CHECKPOINT]
                              [--featureextractor FEATUREEXTRACTOR]
                              [--nodefeature NODEFEATURE]
                              [--kfolds KFOLDS] [--test] [--non_visualize]
optional arguments:   -h, --help            show this help message and exit   -s SEED, --seed SEED  Random seed
Storage:   Directories for util results
-ld LOGDIR, --log-dir LOGDIR                         Directory for saving training logs and visualization   --outputmodels OUTPUTMODELS                         Where you want to save your models
Dataset:   Dataset paths
--compressedgraph COMPRESSEDGRAPH                         Compressed graphs of dataset which was extracted by                         graph helper tools   --dataset DATASET     Dicrectory of all souce code files which were used to                         extract the compressed graph   --testset TESTSET     Dicrectory of all souce code files which is a                         partition of the dataset for testing   --label LABEL         Label of sources in source code storage   --checkpoint CHECKPOINT                         Checkpoint of trained models
Node feature:   Define the way to get node features
--featureextractor FEATUREEXTRACTOR                         If "node_feature" is "GAE" or "LINE" or "Node2vec", we                         need a extracted features from those models   --nodefeature NODEFEATURE                         Kind of node features we want to use, here is one of                         "nodetype", "metapath2vec", "han", "gae", "line",                         "node2vec"
Optional configures:   Advanced options
--kfolds KFOLDS     Config for cross validate strategy   --test                Set true if you only want to run test phase   --non_visualize       Wheather you want to visualize the metrics

Examples

• We prepared some scripts for the custom MANDO structures bellow:

• Graph Classication for Heterogeous Control Flow Graphs (HCFGs) which detect vulnerabilites at the contract level.

- GAE as node features.

python graphclassifier.py -ld ./logs/graphclassification/cfg/gae/accesscontrol --outputmodels ./models/graphclassification/cfg/gae/accesscontrol --dataset ./experiments/ge-sc-data/sourcecode/accesscontrol/clean57buggycurated0/ --compressedgraph ./experiments/ge-sc-data/sourcecode/accesscontrol/clean57buggycurated0/cfgcompressedgraphs.gpickle --label ./experiments/ge-sc-data/sourcecode/accesscontrol/clean57buggycurated0/graphlabels.json --nodefeature gae --featureextractor ./experiments/ge-sc-data/sourcecode/gescmatricesnodeembedding/matrixgaedim128ofcoregraphofaccesscontrolcfgclean570.pkl --seed 1

• Graph Classication for Heterogeous Call Graphs (HCGs) which detect vulnerabilites at the contract level.

- LINE as node features.

python graphclassifier.py -ld ./logs/graphclassification/cg/line/accesscontrol --outputmodels ./models/graphclassification/cg/line/accesscontrol --dataset ./experiments/ge-sc-data/sourcecode/accesscontrol/clean57buggycurated0/ --compressedgraph ./experiments/ge-sc-data/sourcecode/accesscontrol/clean57buggycurated0/cgcompressedgraphs.gpickle --label ./experiments/ge-sc-data/sourcecode/accesscontrol/clean57buggycurated0/graphlabels.json --nodefeature line --featureextractor ./experiments/ge-sc-data/sourcecode/gescmatricesnodeembedding/matrixlinedim128ofcoregraphofaccesscontrolcgclean570.pkl --seed 1

• Graph Classication for combination of HCFGs and HCGs and which detect vulnerabilites at the contract level.

- node2vec as node features.

python graphclassifier.py -ld ./logs/graphclassification/cfgcg/node2vec/accesscontrol --outputmodels ./models/graphclassification/cfgcg/node2vec/accesscontrol --dataset ./experiments/ge-sc-data/sourcecode/accesscontrol/clean57buggycurated0/ --compressedgraph ./experiments/ge-sc-data/sourcecode/accesscontrol/clean57buggycurated0/cfgcgcompressedgraphs.gpickle --label ./experiments/ge-sc-data/sourcecode/accesscontrol/clean57buggycurated0/graphlabels.json --nodefeature node2vec --featureextractor ./experiments/ge-sc-data/sourcecode/gescmatricesnodeembedding/matrixnode2vecdim128ofcoregraphofaccesscontrolcfgcgclean570.pkl --seed 1

Node Classification

• We used node classification tasks to detect vulnerabilites at the line level and function level for Heterogeneous Control flow graph (HCFGs) and Call Graphs (HCGs) in corressponding.

Usage

usage: MANDO Node Classifier [-h] [-s SEED] [-ld LOG_DIR]
                             [--outputmodels OUTPUTMODELS]
                             [--compressedgraph COMPRESSEDGRAPH]
                             [--dataset DATASET] [--testset TESTSET]
                             [--label LABEL]
                             [--featurecompressedgraph FEATURECOMPRESSEDGRAPH]
                             [--cfgfeatureextractor CFGFEATUREEXTRACTOR]
                             [--featureextractor FEATUREEXTRACTOR]
                             [--nodefeature NODEFEATURE] [--kfolds KFOLDS]
                             [--test] [--non_visualize]
optional arguments:   -h, --help            show this help message and exit   -s SEED, --seed SEED  Random seed
Storage:   Directories \for util results
-ld LOGDIR, --log-dir LOGDIR                         Directory for saving training logs and visualization   --outputmodels OUTPUTMODELS                         Where you want to save your models
Dataset:   Dataset paths
--compressedgraph COMPRESSEDGRAPH                         Compressed graphs of dataset which was extracted by                         graph helper tools   --dataset DATASET     Dicrectory of all souce code files which were used to                         extract the compressed graph   --testset TESTSET     Dicrectory of all souce code files which is a                         partition of the dataset for testing   --label LABEL
Node feature:   Define the way to get node features
--featurecompressedgraph FEATURECOMPRESSEDGRAPH                         If "node_feature" is han, you mean use 2 HAN layers.                         The first one is HAN of CFGs as feature node for the                         second HAN of call graph, This is the compressed                         graphs were trained for the first HAN   --cfgfeatureextractor CFGFEATUREEXTRACTOR                         If "nodefeature" is han, featureextractor is a                         checkpoint of the first HAN layer   --featureextractor FEATUREEXTRACTOR                         If "node_feature" is "GAE" or "LINE" or "Node2vec", we                         need a extracted features from those models   --nodefeature NODEFEATURE                         Kind of node features we want to use, here is one of                         "nodetype", "metapath2vec", "han", "gae", "line",                         "node2vec"
Optional configures:   Advanced options
--kfolds KFOLDS     Config cross validate strategy   --test                If true you only want to run test phase   --non_visualize       Wheather you want to visualize the metrics

Examples

We prepared some scripts for the custom MANDO structures bellow:

• Node Classication for Heterogeous Control Flow Graphs (HCFGs) which detect vulnerabilites at the line level.

- GAE as node features for detection access_control bugs.

python nodeclassifier.py -ld ./logs/nodeclassification/cfg/gae/accesscontrol --outputmodels ./models/nodeclassification/cfg/gae/accesscontrol --dataset ./experiments/ge-sc-data/sourcecode/accesscontrol/buggycurated/ --compressedgraph ./experiments/ge-sc-data/sourcecode/accesscontrol/buggycurated/cfgcompressedgraphs.gpickle --nodefeature gae --featureextractor ./experiments/ge-sc-data/sourcecode/gescmatricesnodeembedding/matrixgaedim128ofcoregraphofaccesscontrolcfgbuggycurated.pkl --testset ./experiments/ge-sc-data/sourcecode/accesscontrol/curated --seed 1

• Node Classification for Heterogeous Call Graphs (HCGs) which detect vulnerabilites at the function level.
• The command lines are the same as CFG except the dataset.

- LINE as node features for detection access_control bugs.

python nodeclassifier.py -ld ./logs/nodeclassification/cg/line/accesscontrol --outputmodels ./models/nodeclassification/cg/line/accesscontrol --dataset ./experiments/ge-sc-data/sourcecode/accesscontrol/buggycurated --compressedgraph ./experiments/ge-sc-data/sourcecode/accesscontrol/buggycurated/cgcompressedgraphs.gpickle --nodefeature line --featureextractor ./experiments/ge-sc-data/sourcecode/gescmatricesnodeembedding/matrixlinedim128ofcoregraphofaccesscontrolcgbuggycurated.pkl --testset ./experiments/ge-sc-data/sourcecode/accesscontrol/curated --seed 1

• Node Classication for combination of HCFGs and HCGs and which detect vulnerabilites at the line level.

- node2vec as node features.

python nodeclassifier.py -ld ./logs/nodeclassification/cfgcg/node2vec/accesscontrol --outputmodels ./models/nodeclassification/cfgcg/node2vec/accesscontrol --dataset ./experiments/ge-sc-data/sourcecode/accesscontrol/buggycurated --compressedgraph ./experiments/ge-sc-data/sourcecode/accesscontrol/buggycurated/cfgcgcompressedgraphs.gpickle --nodefeature node2vec --featureextractor ./experiments/ge-sc-data/sourcecode/gescmatricesnodeembedding/matrixnode2vecdim128ofcoregraphofaccesscontrolcfgcgbuggycurated.pkl --testset ./experiments/ge-sc-data/sourcecode/accesscontrol/curated --seed 1

• We also stack 2 HAN layers for function-level detection. The first HAN layer is based on HCFGs used as feature for the second HAN layer based on HCGs (It will be deprecated in a future version).

python nodeclassifier.py -ld ./logs/nodeclassification/callgraph/node2vechan/accesscontrol --outputmodels ./models/nodeclassification/callgraph/node2vechan/accesscontrol --dataset ./ge-sc-data/nodeclassification/cg/accesscontrol/buggycurated --compressedgraph ./ge-sc-data/nodeclassification/cg/accesscontrol/buggycurated/compressedgraphs.gpickle --testset ./ge-sc-data/nodeclassification/cg/curated/accesscontrol --seed 1  --nodefeature han --featurecompressedgraph ./data/smartbugswild/binaryclasscfg/accesscontrol/buggycurated/compressedgraphs.gpickle --cfgfeatureextractor ./data/smartbugswild/embeddingsbuggycurratedmixed/cfgmixed/gescmatricesnodeembedding/matrixnode2vecdim128ofcoregraphofaccesscontrolcompressedgraphs.pkl --featureextractor ./models/nodeclassification/cfg/node2vec/accesscontrol/hanfold0.pth

Testing

• We automatically run testing after training phase for now.

Visuallization

• You also use tensorboard and take a look the trend of metrics for both training phase and testing phase.

tensorboard --logdir LOG_DIR