A smart contract vulnerability detection method based on multi-task learning

Abstract

The present invention discloses a smart contract vulnerability detection method based on multi-task learning, and implements smart contract vulnerability detection technology using a multi-task learning framework based on hard parameter sharing. First, in the data preparation stage, the smart contract data set is cleaned, and the vulnerability classification and labeling are performed through existing detection tools; in the data preprocessing stage, the smart contract sample source code is compiled into bytecode, and then the bytecode is cleaned and processed, and then decompiled and converted into an opcode sequence to form the input of the model; then in the model construction stage, a smart contract detection model based on multi-task learning is constructed; finally, in the training stage, the opcode sequence obtained in the data preprocessing stage is input into the model for training to realize the judgment and detection of vulnerabilities. The model shows good recognition function on the smart contract data set, the detection time is short and constant, and due to the characteristics of multi-task learning, the model has good generalization ability.

Description

Technical Field

[0001] The present invention belongs to the field of blockchain smart contract security technology, and specifically relates to a vulnerability detection method based on multi-task learning (MTL) for smart contracts. Background Art

[0002] Blockchain technology is a chain data structure formed by combining data blocks in a sequential manner in chronological order, and a distributed ledger that is cryptographically guaranteed to be tamper-proof and unforgeable. From a macro perspective, blockchain technology is a new distributed infrastructure and computing method that uses block chain data structures to verify and store data, uses distributed node consensus algorithms to generate and update data, uses cryptography to ensure the security of data transmission and access, and uses smart contracts to program and operate data.

[0003] Smart contracts are computer transaction protocols that execute contract terms. They are essentially executable codes. When users access the address, the smart contract will automatically run. With the continuous development of technology, smart contracts are widely used in various fields, especially in finance and insurance. Although the rapid development of smart contracts has brought many benefits, its security issues have also attracted the attention of many researchers. Since 2016, security incidents caused by smart contract vulnerabilities have emerged in an endless stream every year, causing serious economic losses of assets. Therefore, how to effectively and quickly detect smart contract vulnerabilities has become a key issue in academic research. At present, domestic vulnerability detection of smart contracts mainly relies on symbolic execution, formal verification, fuzz testing and other technologies. However, these methods have their own shortcomings. For example, vulnerability detection technology based on symbolic analysis is slow and time-consuming; vulnerability detection methods based on formal methods have high accuracy, but high cost and are limited to known vulnerabilities. They cannot analyze or detect the form of unknown vulnerabilities; vulnerability detection technology based on fuzz testing is trapped by the generation strategy of random tests, the global state of smart contracts and the call sequence.

[0004] Vulnerability detection methods implemented through machine learning have the characteristics of high detection efficiency, strong scalability, and short detection time, and have gradually attracted the attention of scholars at home and abroad. However, existing machine learning methods have characteristics such as weak generalization performance and the need for source code, which limits their applicability. In addition, in terms of vulnerability classification, many methods only implement the binary classification problem of smart contract samples, that is, whether there are vulnerabilities, but cannot specifically identify the type of vulnerability, nor can they be easily expanded to other vulnerability types. In order to solve some of the above-mentioned defects, the present invention provides a smart contract vulnerability detection method based on multi-task learning, which can use less detection time to complete the two tasks of vulnerability binary classification and multi-classification, thereby improving the classification accuracy of the model. Summary of the invention

[0005] The present invention uses a special form of transfer learning, multi-task learning, to solve the generalization performance of the model. Multi-task learning is a type of machine learning, which aims to jointly learn multiple related tasks so that the knowledge contained in the task can be used by other tasks, thereby improving the generalization performance of all tasks. Multi-task learning has been successfully applied to machine learning, from natural language processing, speech recognition to computer vision, etc., so the present invention introduces it into the field of smart contract vulnerability detection. In addition, using the easier to obtain contract operation code as the input of the model can not only reduce the difficulty of obtaining the contract source code, but also improve the stability of the detection model.

[0006] In order to solve the above technical problems, the present invention provides a smart contract vulnerability detection method based on multi-task learning, such as Figure 1 As shown, the process includes building a smart contract dataset, testing sample datasets to determine vulnerability labels, compiling smart contracts to generate and simplify opcode sequences, using oversampling methods to balance and divide datasets, training machine learning models, and using models for vulnerability detection.

[0007] Specifically, it includes:

[0008] In the data preparation stage, the smart contract sample set is classified and labeled with vulnerability labels;

[0009] In the data preprocessing stage, the smart contract source code is compiled into contract bytecode, invalid bytecode is deleted, and then the valid bytecode is decompiled into an opcode sequence as the input of the model;

[0010] In the model building phase, a multi-task learning framework based on hard parameter sharing is constructed. In the bottom shared layer of the framework, a feature network layer based on self-attention is constructed to extract and learn the features of the opcode sequence; and in the specific task layer, a binary classification model based on a convolutional neural network is established for the vulnerability judgment task; for the vulnerability identification task, a multi-label classification model based on a convolutional neural network is established as a vulnerability detection model for smart contracts;

[0011] In the smart contract vulnerability detection stage, the smart contract to be detected is processed into an operation code sequence according to the processing method of the data preprocessing stage, and then input into the multi-task-based smart contract vulnerability detection model to realize multi-task functions, including vulnerability judgment (binary classification), detection of multiple vulnerabilities (multi-label classification), etc.

[0012] In summary Figure 2 As shown, the steps of the vulnerability detection method of the present invention are as follows:

[0013] (1) The data preparation stage includes:

[0014] 1.1) Collect and obtain smart contract data sets based on the vulnerability types that need to be detected;

[0015] 1.2) Clean the smart contract data set to obtain valid smart contract data samples;

[0016] 1.3) Classify and label smart contract data samples: Use vulnerability detection tools to perform vulnerability detection and analysis on smart contract samples, and label each smart contract sample according to the vulnerability type, including labels for whether there is a vulnerability and the vulnerability type;

[0017] 1.4) The label is X = [x1, x2, ..., x n ], where n represents the multiple vulnerability types that need to be detected, and x i =0 means that the smart contract sample does not belong to the i-th vulnerability, x i =1 means that the contract sample belongs to the i-th vulnerability. (2) The data preprocessing stage includes:

[0018] 2.1) Compile the smart contract source code into contract bytecode and process it effectively;

[0019] 2.2) Decompile the bytecode to obtain the opcode sequence;

[0020] Bytecode is composed of a series of hexadecimal numbers, which represent a specific sequence of operations and parameters. However, it is often a long sequence, and its analysis and modeling requires a large amount of memory space, making it impractical to use bytecode directly as input. Therefore, it is necessary to further process the acquired contract bytecode, decompile and generate an opcode sequence to simplify the model input and improve the model training efficiency and training effect.

[0021] (3) The model construction stage includes:

[0022] 3.1) Construct a multi-task learning framework based on hard parameter sharing. At the bottom shared layer of the framework, a feature network layer based on self-attention is constructed to extract and learn the features of the opcode sequence; and at the specific task layer, a binary classification model based on a convolutional neural network is established for the vulnerability judgment task; for the vulnerability identification task, a multi-label classification model based on a convolutional neural network is established as a vulnerability detection model for smart contracts; the overall architecture of the model is as follows: Figure 3 shown.

[0023] 3.2) Bottom shared layer design

[0024] like Figure 4As shown, the bottom shared layer constructs a feature extraction network based on the self-attention layer. The word encoding (WordEmbedding) layer and the positional encoding (Positional Embedding) layer are used to describe the smart contract operation code sequence. Word encoding (Word Embedding) converts each input operation code into the form of a word vector, thereby mapping the operation code sequence into a multidimensional space. At the same time, in order to make this model understand the order of the operation code sequence, the position vector is added to the word encoding so that in the subsequent calculations, the distance between words can be better expressed to complete the description of the input sequence. The calculation of the positional encoding is shown in formula (1) and formula (2):

[0025] PE (pos,2i) =sin(pos / 10000 2i / d ) (1)

[0026] PE (pos,2i+1) =cos(pos / 10000 2i / d ) (2)

[0027] Position encoding using a sinusoidal function allows the model to easily learn to focus on relative position information, so for any fixed offset i, PE (pos+i) Both can be expressed as PE (pos) In this formula, pos represents the position of the word in the sequence, PE (pos,2i) PE (pos,2i+1) Represents the vector corresponding to the pos position, d represents the dimension of the vector, It represents the frequency w k Therefore, the positional encoding is a pair of sine and cosine containing each frequency. Positional Embedding uses the position information of the word to represent each word in the sequence twice, and combines the word order information with the word vector to form a new representation input to the model, so that the model has the ability to learn word order information.

[0028] After the contract operation code sequence is described by word encoding and position encoding, it will be input into the feature network based on the self-attention layer for feature learning and extraction. Self-attention, that is, when the model processes each word in the input sequence, Self-attention will focus on all the words in the entire input sequence, which can help the model to view other positions in the input sequence and find relevant clues to achieve better encoding effects. The specific calculation steps are as follows:

[0029] Generate three vectors from each encoder input vector. For each word, three vectors Q (query vector), K (key vector), and V (value vector) are formed by multiplying the word embedding with three weight matrices.

[0030] The word score is calculated by taking the dot product of the key vector (K) of the word in the input sequence and the query vector (Q) of the word at the current position, and then dividing the score by the square root of the dimension of the key vector to make the gradient more stable.

[0031] The scores are normalized using the softmax function so that they are all positive and sum to 1. The softmax score determines the contribution of each word to the encoding of the current position.

[0032] Multiply each value vector (V) by the softmax score and sum the weighted value vectors.

[0033] The calculation formula is shown in (3):

[0034]

[0035] By superimposing the multi-head attention mechanism, the self-attention layer is improved, the performance of the attention layer is improved, and the model's ability to focus on different positions is expanded. The multi-head self-attention layer performs the same self-attention calculation as above. Under the "multi-head" attention mechanism, an independent vector weight matrix is ​​maintained for each head, thereby generating unused vector matrices. The matrices generated by each head are then concatenated and multiplied with an additional weight matrix. The result is a matrix that integrates the information of all attention heads, which is then sent to the feed-forward neural network (Feed-Forward Network) to output the results of the feature network layer.

[0036] The specific calculation is shown in formula (4)(5)(6):

[0037]

[0038] MultiHead(Q,K,V)=Concat(head1,...,head h )W O (5)

[0039] Output=LayerNorm(x+FFN(x)) (6)

[0040] 3.3) Top-level task-specific layer design

[0041] The top specific task layer builds different branch network layers according to different tasks to achieve two tasks. In the judgment task branch, a binary classification convolutional neural network is constructed, which consists of a convolutional layer (Conv layer), a dropout layer (Dropout) and a fully connected layer (FC / Dense), such as Figure 5 As shown, the classification network uses binary cross entropy loss, and the calculation formula is shown in (7)(8):

[0042]

[0043]

[0044] In the recognition task branch, a multi-classification convolutional neural network is constructed. The classification network includes a convolutional layer, a dropout layer, and a fully connected layer (FC / Dense). Figure 6 As shown, the classification network uses cross entropy loss, and the calculation formula is shown in (9):

[0045]

[0046] Specifically, the classification network structure is composed as follows:

[0047] ① Convolution layer: This layer is composed of several convolution units, and the parameters of each convolution unit are optimized by the back propagation algorithm, which is mainly used to extract different features of the input. The present invention uses 256 convolution units to form a convolution layer.

[0048] ② Linear rectification layer: This layer uses the linear rectification function (ReLU) f(x) = max(0,x) as the excitation function, which is mainly used to enhance the nonlinear characteristics of the decision function and the entire neural network.

[0049] ③Dropout layer: This layer is used to retain a certain proportion of neuron connections and disconnect other neuron connections during each parameter update during model training, thereby preventing the model from overfitting.

[0050] ④FC layer: the fully connected layer, which is used to reduce the dimension of the operation code sequence features of the smart contract, thereby improving the accuracy of the model.

[0051] After going through two layers of FC layers, in order to complete their respective specific tasks, they will be divided into two branches. The judgment branch will be connected to a classification layer, and the classification operation will be performed through the sigmoid function to directly output the result; while the recognition branch will be connected to a FC layer, continue to reduce the dimension, and then connect to the classification layer, perform classification operations through the softmax function, and output the result. BRIEF DESCRIPTION OF THE DRAWINGS

[0052] Figure 1 Flow chart of the method of the present invention.

[0053] Figure 2 This is a step diagram of the smart contract detection method of the present invention.

[0054] Figure 3 It is a structural diagram of the vulnerability detection model of the present invention.

[0055] Figure 4A neural network for feature extraction based on self-attention layers.

[0056] Figure 5 Branching network for judgment tasks.

[0057] Figure 6 Branching networks for identification tasks. DETAILED DESCRIPTION

[0058] The present invention will be described in detail with reference to the accompanying drawings and embodiments.

[0059] In order to improve the scalability of the model, the present invention adopts a multi-task learning framework based on hard parameter sharing; in order to improve the accuracy of model vulnerability detection, the present invention adopts a feature extraction neural network based on a self-attention layer; in order to determine whether there is a vulnerability, a binary classification model based on a convolutional neural network is established and trained; in order to detect multiple vulnerability types, a multi-label classification idea is adopted, and a multi-classification model based on a convolutional neural network is established and trained, thereby realizing the detection of multiple vulnerability types.

[0060] like Figure 2 The step diagram of the vulnerability detection method of the present invention is divided into several stages, namely, a data preparation stage, a data preprocessing stage, a model construction stage and a model training stage.

[0061] (1) Data preparation stage

[0062] 1.1) The present invention obtains the Ethereum smart contract data set from the XBlock community, with a total of 149,363 contract files, which include 9 contract attributes, including address (open source contract address), ContractCode (contract bytecode), timestamp (contract creation timestamp), createValue (contract creation parameter), createBlockNumber (block height when the contract was created), createdTransactionHash (contract creation transaction), creationCode (contract creation bytecode), creator (contract creator address), and code (contract source code). The data set cleaning script is written to clean it, and redundant data, duplicate data, invalid data, and vacant data are removed. The cleaned data set has a total of 141,837 entries.

[0063] 1.2) Select multiple vulnerability types to be detected, and use existing vulnerability detection tools or manual analysis to mark the smart contract dataset, including vulnerability classification and labeling, to obtain a labeled smart contract dataset. The label includes whether it is a vulnerability (Flag), the specific type of vulnerability (Label), and the number of rows where the vulnerability occurs.

[0064] The present invention uses the detection tool SCStudio to detect and label all contract source code files. Since the extension program provided by SCStudio can only check one contract at a time, and the duration of each contract detection is uncertain, averaging 60 seconds, it takes a lot of time. Therefore, the program script is rewritten and the batch processing method is used to detect this data set. Each contract has 37 labels, and the labels are independent of each type of vulnerability. At the same time, the present invention stipulates that if a contract exceeds 200 seconds, it is defaulted to a timed-out contract and is discarded. 1.3) The present invention sets the number of vulnerability types that need to be classified to 3 categories, and constructs a label for each smart contract sample. The label consists of a 4-dimensional column, [x1, x2, x3, x4], and each element x i The value of x is either 0 or 1. i =1 means that the smart contract has the i-th vulnerability, x i =0 indicates that the smart contract does not have the i-th vulnerability. In particular, x1 indicates a label indicating whether the smart contract has a vulnerability. The vulnerability types detected by the present invention are shown in Table 1.

[0065] Table 1 Smart contract dataset

[0066]

[0067] After the above operations, the final smart contract dataset is converted into a labeled sample set, and each sample corresponds to a label.

[0068] (2) Data preprocessing stage

[0069] 2.1) For the smart contract dataset that has been marked in the above stage, obtain the smart contract source code from the dataset and compile it to obtain the contract bytecode.

[0070] 2.2) The compiled contract bytecode consists of three parts: deployment code, runtime code, and Auxdata code. The last 43 bytes in each contract are the Auxdata code, but the actual execution is the runtime code, so only the runtime code needs to be decompiled, so the Auxdata code needs to be cleared. So far, the bytecode cleaning method is defined, the bytecode cleaning script is written, and the invalid bytecode is deleted.

[0071] 2.3) Obtain the contract operation code sequence. The contract bytecode obtained by the above operation is converted into a contract operation code sequence containing operands through a disassembler. At this time, the sequence contains operands, but when used as input, the operands are not needed. Therefore, the sequence is converted into an operation code sequence represented by hexadecimal numbers without operands through hexadecimal encoding. Then, the operation code simplification method shown in Table 2 is used to simplify the operation code sequence, avoiding the dimensionality disaster caused by too many features.

[0072] The operation code simplification method is shown in Table 2.

[0073] Table 2 Operation code simplification method

[0074] Replace opcode Original opcode ARIT ADD,MUL,SUB,DIV,SDIV,SMOD,MOD,ADDMOD,MULMOD,EXP COMP LT,GT,SLT,SGT CONS1 BLOCKHASH,TIMESTAMP,NUMBER,DIFFICULTY,GASLIMIT,COINBASE CONS2 ADDRESS, ORIGIN, CALLER CONS3 GASPRICE,BALANCE,CALLVALUE,GAS LOGI AND,OR,XOR,NOT MOV SHL, SHR, SAR PUSH PUSH1-PUSH32 DUP DUP1-DUP16 SWAP SWAP1-SWAP16 LOG LOG1-LOG4

[0075] The dataset after simplifying the operation codes is shown in Table 3.

[0076] Table 3 Dataset after simplified operation code

[0077]

[0078] 2.4) As a typical anomaly detection problem, the number of smart contracts with vulnerabilities is far less than the number of normal contracts, and even the number of contracts with specified vulnerabilities is even rarer. In order to enable the model to be well trained and learn to collect enough feature information, an oversampling method (SMOTE) is used for the smart contract data set to increase the number of vulnerable contracts and the number of vulnerabilities of specified categories to achieve the purpose of balancing data. At the same time, it is divided into three data sets according to a certain ratio (7:1:2): training set, validation set and test set.

[0079] According to the above operations, write a data set sampling and division program script to process the smart contract data sample set as the input of the model. The training set is used to train the model's hyperparameters; the validation set is used to verify the model's hyperparameters and give the model feedback so that the model can be adjusted further; the test set is used to test the performance of the model.

[0080] (3) Model construction stage

[0081] Construct a smart contract vulnerability detection model based on multi-task learning as described above, and set the model parameters for the next stage of training.

[0082] Batch size: The number of samples selected for one training

[0083] The size of batch size affects the optimization degree and speed of the model, and directly affects the usage of GPU memory. Parallelization can improve the utilization of memory. Large batches can make the network converge faster, but it may lead to insufficient memory. The appropriate batch size can make the direction of gradient descent more accurate. In this invention, the batch size will be gradually increased from [16, 32, 64, 128, 256] to select the most appropriate batch size.

[0084] Learn rate: learning rate

[0085] The learning rate is one of the important hyperparameters in deep learning. It determines whether the objective function can converge to the local minimum and when it converges to the minimum. The appropriate learning rate can make the objective function converge to the local minimum in a suitable time. The lower the learning rate value, the slower the descent along the gradient. Using a smaller learning rate can avoid missing any local minimum points, but the convergence time is longer; however, if the learning rate value is higher, it is possible to ignore a local minimum point. The configuration of the learning rate is generally based on empirical methods and is configured by observing the loss value. Generally, when the model is initialized, a very small learning rate is set for training, and then the learning rate is increased each time until a better training effect is achieved.

[0086] Dropout: Random loss rate

[0087] Dropout refers to the process of temporarily discarding neural network units from the network according to a certain probability during the training process of the deep learning network. Its biggest role is to prevent overfitting and improve the performance of the model. The present invention sets the dropout value from 0.05, and increases the value each time training until a better training effect is achieved.

[0088] Optimizer: Optimizer

[0089] The essence of model training is to minimize the loss. After defining the loss function, an optimizer is used to optimize the gradient. The optimizers used in the present invention include SGD, Adam, etc. Adam is a fast-converging and commonly used optimizer; SGD converges slowly, but Momentum can be added to speed up convergence.

[0090] Epoch: number of iterations

[0091] Using all the data in the training set to train the model once is called an epoch. It is not enough to pass the complete data set once in the neural network, so multiple passes are required. As the number of epochs increases, the number of weight updates in the neural network also increases, and the model will change from underfitting to overfitting, so determining the appropriate epoch is important for the performance of the model. Experimental verification shows that when the difference between the test error rate and the training error rate is small, the current number of iterations can be considered appropriate. In the present invention, the epoch will be gradually increased from [20, 30, 40, 50, 60, 80, 100].

[0092] According to the above content, set the important hyperparameters of the model, carry out the model training phase, and continuously adjust the model hyperparameters according to the training results until the model achieves better results.

[0093] (4) Model training phase

[0094] According to the above operations, the training set is sent to the smart contract detection model to train the hyperparameters, the model hyperparameters are verified through the validation set, and the model is feedback optimized through the back propagation algorithm until the hyperparameters with optimal performance are obtained.

[0095] The algorithm pseudocode flow is as follows:

[0096]

[0097]

[0098] Experimental verification shows that when the smart contract vulnerability detection model based on multi-task learning adopts the settings shown in Table 4, better performance will be achieved.

[0099] Table 4 Hyperparameter settings

[0100] Hyperparameters value Epoch 50 Batchsize 128 Learnrate 0.0001 Dropout 0.4 Optimizer Adam

[0101] The detection experimental results of three types of contract vulnerabilities are shown in Table 5.

[0102] Table 5 Experimental results

[0103]

[0104] In summary, the vulnerability detection method proposed in the embodiment of the present invention realizes accurate and efficient automatic detection of three types of contract vulnerabilities by extracting semantic sequence features and utilizing deep learning algorithms. This method is suitable for application scenarios of large-scale contract vulnerability detection.

[0105] The above description is only an embodiment of the present invention and is not intended to limit the protection scope of the present invention. Any modification, equivalent replacement and improvement made within the spirit and principle of the present invention shall be included in the protection scope of the present invention.

Claims

1. A smart contract vulnerability detection method based on multi-task learning, characterized in that: include: (1) In the data preparation stage, the smart contract sample set is classified and labeled for vulnerabilities; (2) In the data preprocessing stage, the smart contract source code is compiled into contract bytecode, invalid bytecodes are deleted, and then the valid bytecodes are decompiled into operation code sequences as the input of the model; (3) In the model building stage, a multi-task learning framework based on hard parameter sharing is constructed; at the bottom shared layer of the framework, a feature network layer based on self-attention is constructed to extract learning operation code sequence features; At the specific task layer, a binary classification model based on a convolutional neural network is established for vulnerability judgment tasks; for vulnerability identification tasks, a multi-label classification model based on a convolutional neural network is established as a vulnerability detection model for smart contracts; (4) In the smart contract vulnerability detection stage, the smart contract to be detected is processed into an operation code sequence according to the processing method in the data preprocessing stage, and then input into the multi-task-based smart contract vulnerability detection model to realize multi-task functions, including vulnerability judgment and detection of multiple vulnerabilities; The data preparation stage includes: (1) Collect and obtain smart contract data sets based on the vulnerability types that need to be detected; (2) Clean the smart contract data set to obtain valid smart contract data samples; (3) Classify and label smart contract data samples: Use vulnerability detection tools to perform vulnerability detection and analysis on smart contract samples, and label each smart contract sample according to the vulnerability type, including labels for whether there is a vulnerability and the vulnerability type; (4) The label is N = [n1, n2, ... n i ,...n k ], where k represents the multiple vulnerability types that need to be detected, and n i =0 means that the smart contract sample does not belong to the nth vulnerability, n i =1 means that the contract sample belongs to the nth vulnerability; In the data preprocessing stage, the input construction of the model includes: (1) Compile the smart contract source code into contract bytecode; The bytecode compiled by the smart contract is divided into three parts: deployment code, runtime code, and auxdata code. Since the auxdata code at the end of each contract is invalid bytecode, the auxdata code is deleted; (2) Decompile the bytecode into a sequence of opcodes; The bytecode is composed of a series of hexadecimal numbers, which represent specific operation sequences and parameters. The acquired contract bytecode is processed and decompiled to generate an operation code sequence to simplify the model input and improve the model training efficiency and training effect. In the model building phase, the construction of the model includes: (1) Construct a multi-task learning framework based on hard parameter sharing; at the bottom shared layer of the framework, build a feature network layer based on self-attention to extract and learn the features of the opcode sequence; and at the specific task layer, build a binary classification model based on a convolutional neural network for the vulnerability judgment task; for the vulnerability identification task, build a multi-label classification model based on a convolutional neural network as a vulnerability detection model for smart contracts; (2) The feature network layer based on self-attention includes input layer, Word Embedding layer, Positional Embedding layer, Self-attention layer, and FFN layer in sequence; the classification network based on convolutional neural network includes convolution layer, Dropout layer, fully connected layer and classification layer in sequence; among them, the feature network layer is used to learn the features in the opcode sequence, and the classification network is used to judge and identify vulnerabilities.

Images