MANDO: Multi-Level Heterogeneous Graph Embeddings for Fine-Grained Detection of Smart Contract Vulnerabilities

Multi-Level Heterogeneous Graph Embeddings

This repository is an implementation of MANDO: Multi-Level Heterogeneous Graph Embeddings for Fine-Grained Detection of Smart Contract Vulnerabilities. The source code is based on the implementation of HAN and GAT model using Deep Graph Library.

Citation

Nguyen, H. H., Nguyen, N. M., Xie, C., Ahmadi, Z., Kudendo, D., Doan, T. N., & Jiang, L. (2022, October). MANDO: Multi-Level Heterogeneous Graph Embeddings for Fine-Grained Detection of Smart Contract Vulnerabilities, 2022 IEEE 9th International Conference on Data Science and Advanced Analytics (DSAA' 22), Shenzhen, China, 2022, pp. 1-10. Preprint

@inproceedings{nguyen2022dsaa,
  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: Multi-Level Heterogeneous Graph Embeddings for Fine-Grained Detection of Smart Contract Vulnerabilities},
  year = {2022},
  month = {10},
  booktitle = {Proceedings of the 9th IEEE International Conference on Data Science and Advanced Analytics},
  pages = {1-10},
  numpages = {10},
  keywords = {heterogeneous graphs, graph embedding, graph neural networks, vulnerability detection, smart contracts, Ethereum blockchain},
  location = {Shenzhen, China},
  doi = {10.1109/DSAA54385.2022.10032337},
  series = {DSAA '22}
}

Table of contents

• Smart Contract Vulnerabilities
• Multi-Level Graph Embeddings
• Table of contents
• How to train the models?

- Dataset - System Description - Install Environment - Inspection scripts - Graph Classification - Node Classification - Trainer - Graph Classification - Usage - Examples - Node Classification - Usage - Examples - Testing - Visuallization - Results - Combine HCFGs and HCGs in Form-A Fusion. - Coarse-Grained Contract-Level Detection - Fine-Grained Line-Level Detection - HCFGs only - Coarse-Grained Contract-Level Detection - Combine CFGs and CGs in Form-B Fusion. - Fine-Grained Function-Level Detection

How to train the models?

Dataset

• We prepared dataset for experiments.

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

Results

Combine HCFGs and HCGs in Form-A Fusion.

Coarse-Grained Contract-Level Detection

Fine-Grained Line-Level Detection

HCFGs only

Coarse-Grained Contract-Level Detection

Combine CFGs and CGs in Form-B Fusion.

Fine-Grained Function-Level Detection