A smart contract vulnerability detection method based on artificial intelligence
By using pre-trained models and vulnerability classification models, combining labelless and vulnerable label data, the problem of high false alarm rate in smart contract vulnerability detection is solved, achieving higher detection accuracy and lower false alarm rate.
Patent Information
- Application Number
- CN202210610675.2
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2022-05-31
- Publication Date
- 2025-03-21
- Estimated Expiration
- 2042-05-31
AI Technical Summary
The existing smart contract vulnerability detection methods based on machine learning and deep learning rely on traditional vulnerability analysis methods for data set annotation, resulting in high false positive rates and affecting the accuracy of the model.
Pre-trained models (such as Bert models) are used to combine vulnerability classification models (such as multi-layer perceptron models), pre-training is used to use labelless data, and fine-tuning is performed by a small amount of vulnerable label data to reduce the false positive rate of smart contract vulnerability detection.
It effectively reduces the false alarm rate of smart contract vulnerability detection, improves the accuracy of detection, and is suitable for smart contract vulnerability detection in real situations.
Smart Images
Figure CN115017513B_ABST
Abstract
Description
Technical Field
[0001] The present invention belongs to the field of smart contract security, and specifically relates to a smart contract vulnerability detection method based on artificial intelligence. Background Art
[0002] A smart contract is a program deployed and run on a blockchain platform, which is used to describe and automatically execute a contract. Due to the prominent financial attributes of blockchain, smart contracts are more likely to be targeted by attackers than ordinary programs, and attacks on smart contracts can bring higher economic value to attackers. There have been many smart contract attacks on the Ethereum platform, causing huge economic losses to users. Therefore, research on smart contract security vulnerabilities has become very urgent, and more and more researchers have begun to pay attention to the security vulnerabilities of smart contracts.
[0003] Traditional software analysis techniques including fuzz testing, symbolic execution, and formal verification have all been used by researchers to study smart contract security vulnerabilities, and have achieved good results. At the same time, some researchers have begun to focus on using machine learning technology to detect smart contract vulnerabilities. Compared with traditional software analysis techniques, machine learning technology can learn the pattern information of smart contract vulnerabilities without additional expert knowledge, thereby realizing the detection of smart contract vulnerabilities, and the detection time is significantly better than traditional software analysis techniques.
[0004] In the paper "VSCL: Automating Vulnerability Detection in Smart Contracts with Deep Learning", Feng Mi et al. parsed the smart contract bytecode into an opcode sequence, constructed a code control flow graph (CFG) based on the opcode, and used depth-first traversal to obtain a new opcode sequence, constructed a feature vector using the n-gram model, and then used a deep neural network to predict whether the smart contract has vulnerabilities. The paper "contractWard: Automated Vulnerability Detection Models for Ethereum Smart Contracts" also used the n-gram model to construct a feature vector based on the opcode sequence, and then combined it with a traditional machine learning model to predict smart contract vulnerabilities.
[0005] In the paper "Smart Contract Vulnerability Detection Using Graph Neural Networks", Yuan Zhuang et al. constructed data flow graphs and control flow graphs based on the smart contract source code to generate graph feature representations of smart contracts and used graph neural networks to predict smart contract vulnerabilities.
[0006] The datasets used by existing machine learning and deep learning-based methods are mostly classified and labeled using traditional vulnerability analysis methods. Although traditional vulnerability analysis methods have a high accuracy rate, they still have false positives, resulting in inaccurate dataset annotation and affecting the accuracy of the model. Summary of the invention
[0007] In view of the above-mentioned existing problems, the purpose of the present invention is to provide an artificial intelligence-based smart contract vulnerability detection method and system. This method uses a pre-trained model to train with a large amount of unlabeled data and a small amount of labeled data with vulnerabilities, which can effectively reduce the false alarm rate of smart contract vulnerability detection and is therefore more suitable for smart contract vulnerability detection in real situations.
[0008] In order to achieve the above object, the present invention adopts the following technical solution:
[0009] A smart contract vulnerability detection method based on artificial intelligence, the method comprising:
[0010] Obtain unlabeled smart contract dataset;
[0011] According to the smart contract vulnerability type, obtain a labeled smart contract dataset;
[0012] The smart contract bytecode in the smart contract dataset is converted into an opcode sequence, and the opcode sequence is used as a training sample, including unlabeled samples and labeled samples;
[0013] The smart contract vulnerability detection model includes a pre-training model and a vulnerability classification model; the pre-training model is the Bert model, and the vulnerability classification model is the multi-layer perceptron model;
[0014] Use unlabeled training samples to train the pre-trained model to obtain a trained pre-trained model;
[0015] Combine the trained pre-trained model with the vulnerability classification model, and use labeled training samples to train the smart contract vulnerability detection model to obtain a trained smart contract vulnerability detection model;
[0016] The smart contract to be tested is converted into an operation code sequence and input into the trained smart contract vulnerability detection model to obtain the vulnerability detection result corresponding to the smart contract.
[0017] Furthermore, converting the smart contract bytecode in the smart contract data set into an operation code sequence specifically includes:
[0018] According to the smart contract dataset, an unlabeled smart contract dataset and a smart contract dataset with vulnerability labels are obtained;
[0019] An operation code sequence is obtained according to the smart contracts in the unlabeled smart contract data set and the smart contract data set with vulnerability labels.
[0020] Furthermore, obtaining an unlabeled smart contract dataset and a smart contract dataset with vulnerability labels based on the smart contract dataset specifically includes:
[0021] Performing data cleaning on the smart contract data set to obtain a processed smart contract data set;
[0022] The processed smart contract dataset is compared with the reported smart contracts with vulnerabilities to obtain an unlabeled smart contract dataset and a smart contract dataset with vulnerability labels.
[0023] Furthermore, obtaining an operation code sequence according to the smart contracts in the smart contract data set without labels and the smart contract data set with vulnerability labels specifically includes:
[0024] The smart contract bytecodes in the unlabeled smart contract data set and the smart contract data set with vulnerability labels are converted into a smart contract operation code sequence containing operands through a disassembler;
[0025] The smart contract operation code sequence containing the operands is processed by removing the operands, and a custom token: [CLS] is added at the front position to obtain the operation code sequence.
[0026] Furthermore, the smart contract vulnerability detection model includes a pre-training model and a vulnerability classification model, wherein:
[0027] The pre-trained model is a Bert model, which is used to predict whether an opcode in an opcode sequence is replaced;
[0028] The pre-training model includes an input layer, an embedding layer, an encoder layer, a fully connected layer and a sigmoid layer;
[0029] The vulnerability classification model is a multi-layer perceptron model, which is used for vulnerability classification;
[0030] The vulnerability classification model includes a fully connected layer and a Softmax layer.
[0031] Furthermore, the input layer is used to receive an operation code sequence;
[0032] The Embedding layer is used to convert each opcode in the opcode sequence into a word vector, and map each opcode into a high-dimensional vector space as an input of the Encoder layer;
[0033] The Encoder layer uses the Attention mechanism to learn the context relationship of each opcode in the opcode sequence to obtain a feature vector of each opcode;
[0034] The fully connected layer is used to reduce the dimension of the feature vector of the operation code to obtain a vector with the same dimension as the classification label;
[0035] The Sigmoid layer is used to adjust the pre-trained model parameters and use the sigmoid function for back propagation;
[0036] The Softmax layer scales the vector output by the fully connected layer to obtain a predicted label vector for implementing a multi-classification task.
[0037] Furthermore, the use of unlabeled training samples to train the pre-trained model specifically includes:
[0038] Performing operation code replacement processing on the unlabeled training samples to obtain processed training samples;
[0039] The processed training samples are input into the input layer of the pre-trained model. The output of the pre-trained model predicts whether the opcode at each position is replaced, calculates the loss function and performs reverse gradient propagation to update the model parameters.
[0040] Furthermore, the operation code replacement process is performed on the unlabeled training sample, specifically including:
[0041] For the operation code sequence in the unlabeled training sample, except for the first position, randomly select positions that are 20% of the length of the operation code sequence, and randomly replace the operation codes at these positions with other operation codes to obtain a replaced operation code sequence;
[0042] The replaced operation code sequence is marked with a replacement label to indicate whether the operation code at the corresponding position is replaced, thereby obtaining a processed training sample.
[0043] Furthermore, the using of the Attention mechanism to learn the contextual relationship of each opcode in the opcode sequence specifically includes:
[0044] Perform a linear transformation on the word vector of each opcode in the input opcode sequence to obtain the query vector, key vector, and value vector of each opcode;
[0045] For each opcode, use the query vector of the opcode and the key vector of each opcode in the opcode sequence to make a vector product, and use the Softmax function to make the sum of the vector products equal to 1, which represents the contribution weight of each opcode to the calculation of the opcode feature vector;
[0046] The weight corresponding to each opcode is multiplied by the value vector and then summed to obtain the feature vector of the opcode. In this way, the feature vector of each opcode contains the context information in the opcode sequence.
[0047] Furthermore, the pre-trained model is combined with the vulnerability classification model, and the smart contract vulnerability detection model is trained using the training samples with vulnerability labels, specifically including:
[0048] Remove the fully connected layer and the Sigmoid layer of the trained pre-trained model;
[0049] The Encoder layer output vector corresponding to [CLS] is used as the input vector of the vulnerability classification model;
[0050] The vulnerability classification model uses the Softmax function to output multi-classification results, calculate the loss function, and perform reverse gradient propagation to update the model parameters.
[0051] An artificial intelligence-based smart contract vulnerability detection system, the system comprising:
[0052] Smart contract data set collection module, used to obtain smart contract data sets;
[0053] A training data processing module, used to convert the smart contracts in the smart contract data set into operation code sequences, and use the operation code sequences as training samples, including unlabeled training samples and vulnerability labeled training samples;
[0054] A pre-trained model training module, used to train the pre-trained model using the unlabeled training samples to obtain a trained pre-trained model; wherein the pre-trained model is a Bert model;
[0055] The smart contract vulnerability detection model training module combines the trained pre-trained model and the vulnerability classification model, and uses the training samples with vulnerability labels to train the smart contract vulnerability detection model to obtain a trained smart contract vulnerability detection model; wherein the vulnerability classification model is a multi-layer perceptron model.
[0056] A test data processing module, used to convert the smart contracts in the smart contract data set into an operation code sequence, and use the operation code sequence as the test operation code sequence;
[0057] The smart contract vulnerability detection model module is used to input the operation code sequence to be tested into the trained smart contract vulnerability detection model to obtain corresponding detection results.
[0058] Furthermore, the training data processing module specifically includes:
[0059] Clean the smart contract data set to obtain the processed smart contract data set;
[0060] Comparing the processed smart contract dataset with the reported smart contracts with vulnerabilities to obtain an unlabeled smart contract dataset and a smart contract dataset with vulnerability labels;
[0061] The smart contract bytecodes in the unlabeled smart contract data set and the smart contract data set with vulnerability labels are converted into a smart contract operation code sequence containing operands through a disassembler;
[0062] The smart contract operation code sequence containing the operands is processed by removing the operands, and a custom token: [CLS] is added at the front position to obtain the operation code sequence.
[0063] The present invention has the following beneficial effects compared with the prior art:
[0064] 1. The present invention detects vulnerabilities in smart contracts through opcode sequences, thus avoiding the problem that the smart contract source code cannot be obtained and the source code versions are different;
[0065] 2. The present invention uses an unsupervised learning method to train the pre-trained model, making full use of a large amount of unlabeled smart contract data to better learn the contextual relationship in the opcode sequence;
[0066] 3. The present invention uses the Encoder layer based on the Attention mechanism to extract the feature vector of the opcode sequence, avoiding the problems of gradient descent and gradient explosion caused by the long opcode sequence based on the long short-term memory network;
[0067] 4. The present invention adopts a solution of fine-tuning the pre-trained model using the labeled data with vulnerabilities. It does not require the use of a large amount of labeled data with vulnerabilities. Compared with the existing solution of labeling smart contracts using a method that already has false alarms, it effectively reduces the false alarm rate of smart contract vulnerability detection. BRIEF DESCRIPTION OF THE DRAWINGS
[0068] Figure 1is a flow chart of an artificial intelligence-based smart contract vulnerability detection method according to Embodiment 1 of the present invention;
[0069] Figure 2 This is a structural diagram of an artificial intelligence-based smart contract vulnerability detection system according to Embodiment 2 of the present invention. DETAILED DESCRIPTION
[0070] The present invention will be further described below in conjunction with the accompanying drawings and specific implementation methods.
[0071] Example 1
[0072] like Figure 1 As shown, the embodiment of the smart contract vulnerability detection method based on artificial intelligence provided by the present invention includes three stages, namely data preprocessing, model training and smart contract vulnerability detection, wherein:
[0073] Data preprocessing: Run the Ethereum client to obtain the smart contract data set, collect the publicly disclosed smart contracts with vulnerabilities in the data set from the vulnerability reporting platform and label them with vulnerabilities. Convert the smart contract bytecode into an opcode sequence and remove the operands and add token [CLS], finally obtaining an unlabeled data set and a labeled data set.
[0074] Model training: Establish a pre-trained model based on the Attention mechanism, use the unlabeled data set as training samples to train the model, and obtain a trained pre-trained model; combine the trained pre-trained model with the vulnerability classification model, use the labeled data set as training samples to train the model, and obtain a smart contract vulnerability detection model.
[0075] Smart contract vulnerability detection: Convert the smart contract to be detected into an operation code sequence, input it into the smart contract vulnerability detection model, and obtain the smart contract vulnerability detection result.
[0076] The AI-based smart contract vulnerability detection method provided in this example includes the following steps:
[0077] S1: Run the Ethereum client to obtain a smart contract dataset, where the smart contract dataset includes the smart contract bytecode.
[0078] S2: Collect the publicly disclosed smart contracts with vulnerabilities in the dataset from the vulnerability reporting platform and label them with vulnerability labels.
[0079] In this step, publicly disclosed smart contracts with vulnerabilities are collected from the vulnerability reporting platform according to the vulnerability type, ensuring the authenticity of the smart contract labels.
[0080] S3: Convert the smart contract bytecode in the smart contract dataset into an opcode sequence, and use the opcode sequence as training samples, including unlabeled samples and labeled samples;
[0081] Further, step S3 includes:
[0082] The smart contract bytecodes in the unlabeled smart contract dataset and the smart contract dataset with vulnerability labels are converted into smart contract opcode sequences containing operands through a disassembler;
[0083] Remove the operands from the smart contract opcode sequence containing operands, and add a custom token: [CLS] at the front to get the opcode sequence.
[0084] S4: Build a smart contract vulnerability detection model.
[0085] The smart contract vulnerability detection model includes a pre-training model and a vulnerability classification model. The pre-training model is a Bert model, which is used to predict whether the opcode in the opcode sequence is replaced, including the input layer, the Embedding layer, the Encoder layer, the fully connected layer and the Sigmoid layer; the vulnerability classification model is a multi-layer perceptron model, which is used for vulnerability classification, including the fully connected layer and the Softmax layer; where:
[0086] The input layer receives the opcode sequence;
[0087] The Embedding layer is used to convert each opcode in the opcode sequence into a word vector;
[0088] The Encoder layer uses the Attention mechanism to learn the contextual relationship of each opcode in the opcode sequence and obtain the feature vector of each opcode;
[0089] A fully connected layer, used for reducing the dimension of the feature vector of the operation code;
[0090] Sigmoid layer, used to adjust the parameters of the pre-trained model and use the sigmoid function for back propagation;
[0091] Softmax layer, used to implement multi-classification tasks.
[0092] S5: Use the unlabeled training samples to train the pre-trained model to obtain a trained pre-trained model.
[0093] Furthermore, step S5 specifically includes:
[0094] For the opcode sequence in the unlabeled training sample, except for the first position, randomly select positions that are 20% of the length of the opcode sequence, and randomly replace the opcodes at these positions with other opcodes to obtain the replaced opcode sequence;
[0095] The replaced operation code sequence is marked with a replacement label, and a label Y=(0, 0, 1, 0, 1..) is used to indicate whether the operation code at the corresponding position is replaced, thereby obtaining a processed training sample.
[0096] The processed training samples are input into the input layer of the pre-trained model, the result prediction output by the pre-trained model replaces the label, the loss function is calculated and the model parameters are updated by back-gradient propagation.
[0097] S6: Use the training samples with vulnerability labels to train the vulnerability classification model to obtain a trained smart contract vulnerability detection model.
[0098] Further, step S6 specifically includes:
[0099] Remove the fully connected layer and Sigmoid layer of the trained pre-trained model;
[0100] The Encoder layer output vector corresponding to [CLS] is used as the input vector of the vulnerability classification model;
[0101] The vulnerability classification model uses the Softmax function to output multi-classification results, calculate the loss function, and perform reverse gradient propagation to update the model parameters.
[0102] S7: Convert the smart contract to be tested into an operation code sequence, input the trained smart contract vulnerability detection model, and obtain the smart contract vulnerability detection result.
[0103] Example 2
[0104] like Figure 2 As shown, the artificial intelligence-based smart contract vulnerability detection system provided by the present invention includes a smart contract data set collection module, a training data processing module, a pre-trained model training module, a smart contract vulnerability detection model training module, a to-be-detected data processing module, and a smart contract vulnerability detection module. Among them:
[0105] The smart contract dataset collection module is used to obtain the smart contract dataset. It chooses whether to obtain the unlabeled smart contract dataset based on whether to retrain the pre-trained model. Otherwise, it only obtains the labeled smart contract dataset.
[0106] The training data processing module is used to convert the smart contract bytecode in the smart contract data into an operation code sequence and use the operation code sequence as a training sample;
[0107] Pre-trained model training module, use unlabeled training samples to train the pre-trained model. If there is a pre-trained model that has been trained, you can choose not to train it and go directly to the next step;
[0108] Smart contract vulnerability detection model training module, which uses labeled training samples to train the smart contract vulnerability detection model;
[0109] The data processing module to be tested is used to convert the smart contract to be tested into an operation code sequence;
[0110] The smart contract vulnerability detection module is used to input the operation code sequence to be detected into the trained smart contract vulnerability detection model to obtain the smart contract vulnerability detection result.
Claims
1. A smart contract vulnerability detection method based on artificial intelligence, characterized in that: The method comprises: Get smart contract dataset; The smart contract bytecode in the smart contract dataset is converted into an opcode sequence, and the opcode sequence is used as a training sample, including unlabeled samples and labeled samples; The smart contract vulnerability detection model includes a pre-trained model and a vulnerability classification model. The pre-trained model is trained using unlabeled training samples to obtain a trained pre-trained model. Combine the trained pre-trained model with the vulnerability classification model, and use labeled training samples to train the smart contract vulnerability detection model to obtain a trained smart contract vulnerability detection model; Convert the smart contract to be tested into an opcode sequence and input it into the trained smart contract vulnerability detection model to obtain the vulnerability detection result corresponding to the smart contract; The smart contract vulnerability detection model includes a pre-training model and a vulnerability classification model, wherein: The pre-trained model is a Bert model, which is used to predict whether an opcode in an opcode sequence is replaced; The pre-training model includes an input layer, an embedding layer, an encoder layer, a fully connected layer and a sigmoid layer; The vulnerability classification model is a multi-layer perceptron model, which is used for vulnerability classification; The vulnerability classification model includes a fully connected layer and a Softmax layer; The input layer is used to receive an operation code sequence; The Embedding layer is used to convert each opcode in the opcode sequence into a word vector, and map each opcode into a high-dimensional vector space as an input of the Encoder layer; The Encoder layer uses the Attention mechanism to learn the context relationship of each opcode in the opcode sequence to obtain a feature vector of each opcode; The fully connected layer is used to reduce the dimension of the feature vector of the operation code to obtain a vector with the same dimension as the classification label; The Sigmoid layer is used to adjust the pre-trained model parameters and use the sigmoid function for back propagation; The Softmax layer scales the vector output by the fully connected layer to obtain a predicted label vector for implementing a multi-classification task; The training of the pre-training model using unlabeled training samples specifically includes: Performing operation code replacement processing on the unlabeled training samples to obtain processed training samples; The processed training samples are input into the input layer of the pre-trained model. The output of the pre-trained model predicts whether the opcode at each position is replaced, calculates the loss function and performs reverse gradient propagation to update the model parameters. The pre-trained model is combined with the vulnerability classification model, and the smart contract vulnerability detection model is trained using the training samples with vulnerability labels, specifically including: Remove the fully connected layer and the Sigmoid layer of the trained pre-trained model; The Encoder layer output vector corresponding to [CLS] is used as the input vector of the vulnerability classification model; The vulnerability classification model uses the Softmax function to output multi-classification results, calculate the loss function, and perform reverse gradient propagation to update the model parameters.
2. The method for detecting smart contract vulnerabilities based on artificial intelligence according to claim 1, characterized in that: The smart contract bytecode in the smart contract dataset is converted into an opcode sequence, including: According to the smart contract dataset, an unlabeled smart contract dataset and a smart contract dataset with vulnerability labels are obtained; An operation code sequence is obtained according to the smart contracts in the unlabeled smart contract data set and the smart contract data set with vulnerability labels.
3. The method for detecting smart contract vulnerabilities based on artificial intelligence according to claim 2, characterized in that: The step of obtaining an unlabeled smart contract dataset and a smart contract dataset with vulnerability labels according to the smart contract dataset specifically includes: Performing data cleaning on the smart contract data set to obtain a processed smart contract data set; The processed smart contract dataset is compared with the reported smart contracts with vulnerabilities to obtain an unlabeled smart contract dataset and a smart contract dataset with vulnerability labels.
4. The method for detecting smart contract vulnerabilities based on artificial intelligence according to claim 2, characterized in that: The step of obtaining an operation code sequence based on the smart contracts in the smart contract dataset without labels and the smart contract dataset with vulnerability labels specifically includes: The smart contract bytecodes in the unlabeled smart contract data set and the smart contract data set with vulnerability labels are converted into a smart contract operation code sequence containing operands through a disassembler; The smart contract operation code sequence containing the operands is processed by removing the operands, and a custom token: [CLS] is added at the front position to obtain the operation code sequence.
5. The method for detecting smart contract vulnerabilities based on artificial intelligence according to claim 1, characterized in that: The performing operation code replacement processing on the unlabeled training sample specifically includes: For the operation code sequence in the unlabeled training sample, except for the first position, randomly select positions that are 20% of the length of the operation code sequence, and randomly replace the operation codes at these positions with other operation codes to obtain a replaced operation code sequence; The replaced operation code sequence is marked with a replacement label to indicate whether the operation code at the corresponding position is replaced, thereby obtaining a processed training sample.
6. The method for detecting smart contract vulnerabilities based on artificial intelligence according to claim 1, characterized in that: The using of the Attention mechanism to learn the contextual relationship of each opcode in the opcode sequence specifically includes: Perform a linear transformation on the word vector of each opcode in the input opcode sequence to obtain the query vector, key vector, and value vector of each opcode; For each opcode, use the query vector of the opcode and the key vector of each opcode in the opcode sequence to make a vector product, and use the Softmax function to make the sum of the vector products equal to 1, which represents the contribution weight of each opcode to the calculation of the opcode feature vector; The contribution weight corresponding to each opcode is multiplied by the value vector and then summed to obtain the feature vector of the opcode. In this way, the feature vector of each opcode contains the context information in the opcode sequence.