GNNSCVulDetector

This repo is a python implementation of smart contract vulnerability detection using graph neural networks (TMP).

Citation

Please use this citation in your paper if you refer to our paper or code.

@inproceedings{zhuang2020smart,
  title={Smart Contract Vulnerability Detection using Graph Neural Network.},
  author={Zhuang, Yuan and Liu, Zhenguang and Qian, Peng and Liu, Qi and Wang, Xiang and He, Qinming},
  booktitle={IJCAI},
  pages={3283--3290},
  year={2020}
}

Requirements

Required Packages

python 3+
TensorFlow 1.14.0 (tf2.0 is not supported)
keras 2.2.4 with TensorFlow backend
sklearn 0.20.2
docopt as a command-line interface parser

Run the following script to install the required packages.

pip install --upgrade pip
pip install tensorflow==1.14.0
pip install keras==2.2.4
pip install scikit-learn==0.20.2
pip install docopt

Dataset

For each dataset, we randomly pick 80% contracts as the training set while the remainings are utilized for the testing set. In the comparison, metrics accuracy, recall, precision, and F1 score are all involved. In consideration of the distinct features of different platforms, experiments for reentrancy and timestamp dependence vulnerability are conducted on ESC (Ethereum smart contract) dataset, while infinite loop vulnerability is evaluated on VSC (Vntchain smart contract) dataset.

Here, we provide a tool for crawling the smart contract source code from Etherscan, which is developed in Aug 2018. If out of date, you can make the corresponding improvements.

For original dataset, please turn to the dataset repo.

Dataset structure in this project

All of the smart contract source code, graph data, and training data in these folders in the following structure respectively.

${GNNSCVulDetector}
├── data
│   ├── timestamp
│   │   └── source_code
│   │   └── graph_data
│   └── reentrancy
│       └── source_code
│       └── graph_data
├── train_data
    ├── timestamp
    │   └── train.json
    │   └── vaild.json
    └── reentrancy
        └── train.json
        └── vaild.json

data/reentrancy/source_code: This is the source code of smart contracts.
data/reentrancy/graph_data: This is the graph structure of smart contracts, consisting edges and nodes, which are extracted by our AutoExtractGraph.
graph_data/edge: It includes all edges and edge of each smart contract.
graph_data/node: It includes all nodes and node of each smart contract.
features/reentrancy: It includes all the reentrancy features of each smart contract extracted by our model.
train_data/reentrancy/train.json: This is the training data of all the smart contract for reentrancy.
train_data/reentrancy/valid.json: This is the testing data of all the smart contract for reentrancy.

Code Files

The tools for extracting graph features (vectors) are as follows:

${GNNSCVulDetector}
├── tools
│   ├── remove_comment.py
│   ├── construct_fragment.py
│   ├── reentrancy/AutoExtractGraph.py
│   └── reentrancy/graph2vec.py

AutoExtractGraph.py

All functions in the smart contract code are automatically split and stored.
Find the relationships between functions.
Extract all smart contracts source code into the corresponding contract graph consisting of nodes and edges.

python3 AutoExtractGraph.py

graph2vec.py

Feature ablation.
Convert contract graph into vectors.

python3 graph2vec.py

Running project

To run the program, please use this command: python3 GNNSCModel.py.

Examples:

python3 GNNSCModel.py --random_seed 9930 --thresholds 0.45

Note

We would like to point that the data processing code is available here. If any question, please email to [email protected]. And, the code is adapted from GGNN.

Reference

Li Y, Tarlow D, Brockschmidt M, et al. Gated graph sequence neural networks. ICLR, 2016. GGNN
Qian P, Liu Z, He Q, et al. Towards automated reentrancy detection for smart contracts based on sequential models. 2020. ReChecker

The code implementation seems to be inconsistent with the formulas in the paper.

In the Readout phase of your paper, formulas (7) to (10) are described as follows:

These formulas are implemented in lines 239 to 245 of GNNSCModel.py as follows:

        gate_input = tf.concat([last_h, self.placeholders['initial_node_representation']], axis=-1)  # [v x 2h]
        gated_outputs = tf.nn.sigmoid(regression_gate(gate_input)) * regression_transform(last_h)  # [v x 1] new_last_h

        # Sum up all nodes per graph
        graph_representations = tf.unsorted_segment_sum(data=gated_outputs,
                                                        segment_ids=self.placeholders['graph_nodes_list'],
                                                        num_segments=self.placeholders['num_graphs'])  # [g x 1]

According to the above,
the variable last_h corresponds to the symbol (h_i)^T in formulas (7),
the variable self.placeholders['initial_node_representation'] corresponds to the symbol (h_i)^0 in formulas (7),
the variable gate_input corresponds to the symbol s_i in formulas (7),
the variable regression_gate(gate_input) corresponds to the symbol g_i in formulas (8),
the variable regression_transform(last_h) corresponds to the symbol o_i in formulas (9),
the variable gated_ouputs should corresponds to Sigmoid(o_i * g_i) in formulas (10)

However, both formulas (8) and (9) take s_i as input, whereas the function regression_gate and regression_transform takes gate_input and last_h as input respectively, and the sigmod function has regression_gate(gate_input) ->o_i as input instead of o_i * g_i.

Is this a typo in the code implementation?

@Messi-Q

messi-q / gnnscvuldetector Goto Github PK