Smart contract vulnerability detection method based on code representation learning and graph neural network

Abstract

The invention discloses an intelligent contract vulnerability detection method based on code representation learning and a graph neural network, and the method comprises the steps: extracting the fine-grained semantic information of a capture function through Token-level features, constructing a graph structure based on a function call relation, and analyzing the global dependence. And a risk scoring mechanism is adopted, so that the interpretability of a vulnerability detection result is enhanced. Compared with an existing analysis method, the method provided by the invention has the advantages that the accuracy and comprehensiveness of vulnerability detection are remarkably improved, and how complex interaction relationships among a plurality of functions cause vulnerabilities can be explained. Experimental results show that the proposed method is superior to a traditional method in indexes such as accuracy, recall rate, F1 value and the like, particularly shows strong robustness and universality in a complex scene, and provides a brand new solution for intelligent contract vulnerability detection.

Description

Technical Field

[0001] The present invention belongs to the technical field of smart contract vulnerability detection, and specifically relates to a smart contract vulnerability detection method based on code representation learning and graph neural network. Background Art

[0002] Blockchain is a distributed ledger technology that organizes data into blocks and uses cryptographic algorithms to link each block in chronological order to form a secure and tamper-proof chain. The key features of blockchain include decentralization, transparency, and immutability, which have led to its widespread use in finance, supply chain management, and identity authentication. However, the rapid development of blockchain technology is accompanied by security threats. The complexity and insufficient testing of smart contracts make them vulnerable to attacks, which in turn lead to serious financial losses. The continuous malicious attacks on smart contracts have had a significant impact on the entire blockchain ecosystem.

[0003] At present, smart contract vulnerability detection mainly relies on static analysis, dynamic analysis, machine learning and artificial intelligence methods. Although static analysis can detect potential vulnerabilities through code structure and rules, it has limited detection effects on vulnerabilities that rely on runtime behavior or complex inputs. Dynamic analysis can capture actual runtime behavior, but it has limitations such as insufficient coverage, high operating costs, and difficulty in simulating external dependencies. Machine learning and deep learning models are often regarded as "black boxes" and their decision-making processes lack transparency, making it difficult for developers to understand and trust their detection results. In terms of interpretability, Chen et al. proposed applying LIME to smart contract vulnerability detection. By perturbing the code segments and observing the changes in the model output, the key features that the model focuses on can be identified. This method not only explains complex deep learning models, but also provides a basis for the credibility of the model. Jiang et al. proposed a model based on graph neural networks to identify reentrancy vulnerabilities in contracts. By analyzing the weights of the attention mechanism, they demonstrated how the model focuses on key nodes and explained the detection process.

[0004] The disadvantages of the above prior art are as follows:

[0005] (1) Static analysis lacks the ability to handle dynamic interactions and environmental dependencies, and is therefore unable to detect complex vulnerabilities that are dynamically generated or caused by external input. (2) Dynamic analysis can capture the actual behavior of the code at runtime, but it is clearly insufficient in covering all execution paths and may miss vulnerabilities on certain paths. (3) Machine learning and deep learning models perform well in detecting complex vulnerabilities, but because their decision-making process is opaque, it is difficult to explain why the model makes a certain detection conclusion. Summary of the invention

[0006] Purpose of the invention: This invention proposes a smart contract vulnerability detection method based on code representation learning and graph neural network, aiming to solve the security problems in smart contracts caused by the limitations and lack of interpretability of existing detection methods, including insufficient understanding of code semantics, insufficient modeling of complex interaction relationships between functions, and lack of interpretability of detection results. Improve the accuracy and transparency of vulnerability detection.

[0007] Technical solution: To achieve the purpose of the present invention, the technical solution adopted by the present invention is: a smart contract vulnerability detection method based on code representation learning and graph neural network, comprising the following steps:

[0008] Step 1: Learning smart contract code representation:

[0009] Tokenize each function in the smart contract, convert each function into a sequence of multiple tokens, use the pre-trained model to map each token into a high-dimensional space, and generate the corresponding token feature vector;

[0010] Step 2: Build the graph structure and update the nodes:

[0011] Build a graph structure based on the function call relationship in the smart contract, and use the token feature vector to initialize the features of each node in the graph. Define each function in the smart contract as a node in the graph, and the call relationship between functions as an edge. Different category attributes are assigned to the edge according to the return value type.

[0012] After the graph structure is built, message passing and feature aggregation are performed through the graph neural network. The node exchanges information with its neighboring nodes and aggregates features. Through message passing and attention weights, the information of the neighboring nodes is aggregated to the current node to achieve the update of node features.

[0013] Step 3: Calculate the risk score of each node and edge, quantify the vulnerability risk, obtain a comprehensive vulnerability score, and generate a vulnerability report for the smart contract based on the comprehensive vulnerability score, giving a panoramic view of the vulnerability.

[0014] Furthermore, in step 2, the graph structure is constructed and the nodes are updated, including:

[0015] Define each function in the smart contract as a node v in the graph i , e ij Represents the slave node v i To node v j The calling relationship, that is, the edge in the graph structure, where i and j are the nodes v i To node v j The identifier of the edge is given different category attributes according to the return value type;

[0016] The message passing mechanism is used to pass the information of neighbor nodes to the current node and aggregate the information. During the aggregation process, the information of neighbor nodes is weighted according to their relationship or importance with the current node to achieve information integration. The feature representation of each node is updated to capture the global structure and the relationship between nodes.

[0017] Use the attention mechanism to calculate the attention weight α of each edge ij , for node v i To neighbor node v j The sent messages are weighted, and the contribution of each node in the information aggregation is dynamically adjusted to capture the key features in the local structure.

[0018] Furthermore, the information aggregation process in step 2 is expressed as:

[0019]

[0020] in, is node v i Aggregate message at level l, v j is node v i Neighbors(v i ) is the node v i The set of neighbor nodes of is node v j Feature representation at layer l; In the graph neural network, the aggregation function AGGREGATE is used to aggregate the information of neighbor nodes, combine the information of neighbor nodes, i.e., feature representation, and update the feature representation of the current node; the AGGREGATE function integrates the features of neighbor nodes and uses the return value type of the edge as additional information;

[0021] Calculate the attention weights:

[0022]

[0023] Among them, the attention weight α ij Represents node v i and its neighbor node v j The relative importance between ij Measure at node v i During the feature update process, the neighbor node v j The influence on the current node; and They are node v i and node v j In the feature representation of the lth layer, softmax is the normalization function, LeakyRelu is the activation function, a is the learned attention vector, W is the linear transformation matrix, and e ij Represents the slave node vi To node v j The return value type characteristics of the edge where ;

[0024] For node v i To neighbor node v j The sent messages are weighted, and the message weighting process is expressed as:

[0025]

[0026] in, is node v i The weighted representation of the aggregated messages at layer l, α ij is the attention weight;

[0027] Use the weighted message to update the feature representation of the current node. The node feature update is expressed as:

[0028]

[0029] in, is node v i In the feature representation of the l+1th layer, the node update function UPDATED combines the node v i Current features and aggregate messages Generate new feature representations.

[0030] Furthermore, the calculation formula for the comprehensive vulnerability score in step 3 is:

[0031]

[0032] Among them, S i Represents each node v i Vulnerability score of node v i Score S self (v i ) is based on node v i The code representation of each node is calculated; the code representation of each node is generated by a pre-trained model, which captures the syntax and semantic information of the code, converts the code features of the node into a high-dimensional vector representation, and then processes the feature vector of the node through a multi-layer perceptron to output the vulnerability risk score S of the node. self (v i ); N(v i ) represents the node v i The set of connected neighbor nodes, |N(v i )| is the number of neighbor nodes, S neighbor (v j ) is the neighbor node v j Score, S edge (eij ) is the edge e ij The score of the connection node v i With v j The impact of the edge on the node score, W self , W neighbor and W edge They are the weights of the node itself, neighboring nodes, and edges;

[0033] By training three weights W self , W neighbor and W edge The node's own score is calculated and three weights are obtained through supervised learning training. During the training process, a labeled smart contract dataset is used as input, in which each node and edge is labeled with a vulnerability risk label. A loss function is defined to measure the difference between the model's predicted score and the actual label, and the weight W is adjusted through back propagation and gradient descent. self , W neighbor and W edge Optimize; each edge e ij Score S edge (e ij ) is composed of the scores of the two end nodes and the return value type weight W of the edge edge calculate.

[0034] Furthermore, the edge score is calculated as:

[0035] S edge (e ij )=W type i(W out iS token (v i )+W in iS token (v j )),

[0036] Among them, W type is the weight assigned according to the return value type of the edge, W out and W in are the weights of out-degree and in-degree, S token (v i ) and S token (v j ) are nodes v i and v j The code represents the score of the learning phase.

[0037] Beneficial effects: Compared with the prior art, the technical solution of the present invention has the following beneficial technical effects:

[0038] Unlike a single method that only relies on code representation learning or graph neural networks, the present invention makes up for the shortcomings between semantic modeling and global topological relationship analysis by deeply mining code semantic information and constructing a graph structure combined with the calling relationship between functions, thereby significantly improving the accuracy and explainability of vulnerability detection.

[0039] This paper proposes for the first time a smart contract vulnerability detection method based on code representation learning and graph neural network. First, a high-dimensional word vector for each function is generated using a pre-trained model to capture the grammatical and semantic information of the function, providing accurate initial features for subsequent vulnerability detection. Subsequently, a function call graph is constructed through a graph neural network to capture the global dependencies between functions. The message passing mechanism fuses the information of neighbor nodes with the features of the current node, which can effectively capture the fine-grained semantic information and global dependencies in the smart contract code, thereby significantly improving the accuracy of vulnerability detection.

[0040] At the same time, the present invention provides an explanatory basis for model decisions, clearly identifies functions that may cause vulnerabilities and their calling relationships, enhances the interpretability of the detection method, and enables developers to understand vulnerability risks and provide targeted security recommendations accordingly. In addition, the present invention not only identifies the vulnerability risks of a single function, but also analyzes how complex interactions between multiple functions contribute to the formation of vulnerabilities, thereby improving the comprehensiveness of detection and providing a panoramic perspective on the security of smart contracts. This invention helps to improve the overall security of smart contracts, reduce financial losses caused by security vulnerabilities, enhance users' trust in blockchain technology, and has significant application value in improving the security of smart contracts, reducing security risks, and improving vulnerability detection efficiency. BRIEF DESCRIPTION OF THE DRAWINGS

[0041] Figure 1 It is a flow chart of the method of the present invention.

[0042] Figure 2 It is a schematic diagram of the method of the present invention.

[0043] Figure 3 It is a schematic diagram of simulated vulnerability detection of this method. DETAILED DESCRIPTION

[0044] In order to make the purpose, technical scheme and advantages of the present invention clearer, the present invention is further described in detail below through the accompanying drawings and embodiments. However, it should be understood that the specific embodiments described herein are only used to explain the present invention and are not intended to limit the scope of the present invention. In addition, in the following description, the description of known structures and technologies is omitted to avoid unnecessary confusion of the concept of the present invention.

[0045] The overall process of the vulnerability detection and explanation method proposed in the present invention is as follows: Figure 1 and Figure 2As shown, first, by inputting the function-level word vector generated by code representation learning as the node feature into the graph neural network, and combining the function call relationship in the graph structure, a deep modeling method is proposed, which can simultaneously capture the fine-grained semantic information and global topological dependency of the code. This method realizes the organic combination of local semantic analysis and global structural modeling, and significantly improves the accuracy and comprehensiveness of the detection results. Secondly, the present invention designs an explainable risk assessment mechanism, which can provide a risk score for each function and its call relationship, and clearly points out the key path for the generation of vulnerabilities, enhances the transparency of the detection results, and provides developers with specific vulnerability repair suggestions, thereby improving the practicality of the detection method. The method comprises the following steps: 1. Code representation learning. 2. Graph structure construction and node update. 3. Explanatory analysis.

[0046] Step 1: Code representation learning.

[0047] Tokenize each function in the smart contract and convert it into a sequence of multiple tokens: i ={t1,t2,...,t k}, where T i represents the i-th function in the smart contract, as the unique identifier of the function, t k Represents the kth token in the function after tokenization. These tokens represent the basic elements of variables, keywords, and operators in the function code, and can capture the grammatical structure of the function and potential vulnerability clues. Use the pre-trained model to map each token to a high-dimensional space and generate the corresponding token vector. Specifically, for the function f i The token sequence T in i , each token t j are converted into a fixed-dimensional vector e through the embedding layer j The embedding is represented as: X i =[e1,e2,...,e k ], where X i Yes i The token embedding matrix contains the semantic features of the function.

[0048] Step 2: Graph structure construction and node update, including:

[0049] 2.1. Define each function in the smart contract as a node v in the graph i , where i represents the unique identifier of the function, and the calling relationship between functions i and j is edge e ij , and assign different category attributes to the edge according to the return value type.

[0050] 2.2. Pass the information of neighbor nodes to node v through the message passing mechanism i , and aggregate the information according to the node v i Aggregate messages at layer l The feature representation of each node is updated to capture the global structure and relationships between nodes.

[0051] 2.3. Use the attention mechanism to calculate the attention weight α of each edge ij , for node v i To neighbor node v j The messages sent are weighted to dynamically adjust the contribution of each node in the information aggregation and capture the key features in the local structure. On this basis, the information of neighbor nodes is aggregated to the current node through message passing and attention weights to update the node features. The updated features reflect the state of the node under the influence of its neighbors and gradually expand to the entire graph structure.

[0052] The information aggregation process is expressed as:

[0053]

[0054] in, is node v i Aggregate message at level l, v j is node v i Neighbors(v i ) is the node v i The set of neighbor nodes of is node v j Feature representation at the lth layer; In the graph neural network, the aggregation function AGGREGATE is used to aggregate the information of neighbor nodes, combine the information (feature representation) of neighbor nodes, and update the feature representation of the current node; The AGGREGATE function integrates the features of neighbor nodes and uses the return value type of the edge as additional information.

[0055] Calculation of attention weight:

[0056]

[0057] Among them, the attention weight α ij Represents node v i and its neighbor node v j The relative importance between ij Measured at node v i During the feature update process, the neighbor node v j The influence on the current node; and They are node v iand node v j Feature representation at the lth layer. Softmax is a normalization function that generates relative weights by normalizing the attention values ​​of all neighbor nodes so that the sum of these weights is 1. In the above formula, softmax ensures that the attention weight of each neighbor node is relative, and the sum of the weights of all neighbors is 1, so that the attention of each node to its neighbors is comparable. LeakyReLU (Leaky Rectified Linear Unit) is an activation function that aims to solve the "dead neuron" problem that may be caused when processing negative inputs. For input values ​​less than zero, LeakyReLU introduces a small negative slope to ensure that the output of the negative input is not zero, so that all neurons in the neural network can always participate in backpropagation during training. a is the learned attention vector, W is the linear transformation matrix, and e is the linear transformation matrix. ij Represents the slave node v i To node v j The return value type characteristics of the edge.

[0058] For node v i To neighbor node v j The messages sent are weighted, and the message weighting process is:

[0059]

[0060] in, is node v i The weighted representation of the aggregated messages at layer l, α ij is the attention weight.

[0061] Use the weighted message to update the feature representation of the current node. The node feature update is expressed as:

[0062]

[0063] in, is node v i In the feature representation of the l+1th layer, the node update function UPDATED combines the node v i Current features and aggregated messages Generate new feature representation, message Used to update node v i characteristics.

[0064] Step 3: Interpretive analysis, including:

[0065] 3.1 Node score: By training three weights: Node weight W self , the weight of the neighbor node W neighborand the edge weight W edge , making nodes and edges with potential problems score higher.

[0066] 3.2 Edge score calculation: Each edge e ij Score S edge (e ij ) is calculated by the scores of the two end nodes and the return value type weight of the edge.

[0067] 3.3 Comprehensive vulnerability scoring: Finally, each node v i The comprehensive vulnerability score is expressed as S i , the calculation expression is:

[0068]

[0069] Among them, S self (v i ) is the node v i Your own score, S self (v i ) is based on node v i Specifically, the code representation of each node is generated by a pre-trained model, which captures the syntax and semantic information of the code, converts the code features of the node into a high-dimensional vector representation, and then processes the feature vector of the node through a multi-layer perceptron to output the vulnerability risk score S of the node. self (v i ); N(v i ) represents the node v i The set of connected neighbor nodes, |N(v i )| is the number of neighbor nodes; S neighbor (v j ) is the neighbor node v j Score, S edge (e ij ) is the edge e ij The score of the connection node v i With v j The influence of the edge on the node score; W self , W neighbor and W edge They are the weights of the node itself, neighboring nodes, and edges.

[0070] By training three weights W self , W neighbor and W edgeThe node's own score is calculated and three weights are obtained through supervised learning training. During the training process, a labeled smart contract dataset is used as input, in which each node (function) and edge (call relationship between functions) is labeled with a vulnerability risk label. A loss function is defined to measure the difference between the model's predicted score and the actual label, and the weight W is adjusted through back propagation and gradient descent. self , W neighbor and W edge Optimize so that the model can more accurately predict the risk score of the node.

[0071] The calculation of node score combines three parts: node v i The score of S self (v i ), the score S of the neighbor node neighbor (v j ), and the edge score S edge (e ij ). Finally, the node score is calculated by the weighted sum of these three parts, where the weight W self , W neighbor and W edge Determines the relative importance of each part. Each edge e ij Score S edge (e ij ) is composed of the scores of the two end nodes and the return value type weight W of the edge edge Calculation. The edge score is calculated as:

[0072] S edge (e ij )=W type i(W out iS token (v i )+W in iS token (v j ))

[0073] Among them, W type is the weight assigned according to the return value type of the edge, W out and W in are the weights of out-degree and in-degree, S token (v i ) and S token (v j ) are nodes v i and v j The score during the code representation learning phase (token level step).

[0074] like Figure 3As shown in the figure, the code of the smart contract is input, tokenized, and the code is broken down into a series of tokens. Next, a word vector is generated for each token to capture the fine-grained features of each function in the code. Subsequently, a graph structure is constructed based on the function call relationship in the smart contract, and the features of each node in the graph are initialized using the token-level feature vector. After the graph structure is constructed, message passing and feature aggregation are performed through the graph neural network, and the node exchanges information with its neighbor nodes and aggregates features. Next, the attention mechanism is applied to dynamically adjust the weight of the message sent by the node to the neighbor to focus on the features that are critical to the task. Then, the risk score of each node and edge is calculated to quantify the potential vulnerability risk, and a vulnerability report for the smart contract is generated based on these risk scores, giving a panoramic view of the vulnerability generation.

[0075] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. For those skilled in the art, the present invention may have various modifications and variations. Any modification, equivalent replacement, improvement, etc. 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 code representation learning and graph neural network, characterized in that: The following steps are involved: Step 1: Learning smart contract code representation: Tokenize each function in the smart contract, convert each function into a sequence of multiple tokens, use the pre-trained model to map each token into a high-dimensional space, and generate the corresponding token feature vector; Step 2: Build the graph structure and update the nodes: Build a graph structure based on the function call relationship in the smart contract, and use the token feature vector to initialize the features of each node in the graph. Define each function in the smart contract as a node in the graph, and the call relationship between functions as an edge. Different category attributes are assigned to the edge according to the return value type. After the graph structure is built, message passing and feature aggregation are performed through the graph neural network. The node exchanges information with its neighboring nodes and aggregates features. Through message passing and attention weights, the information of the neighboring nodes is aggregated to the current node to achieve the update of node features. Step 3: Calculate the risk score of each node and edge, quantify the vulnerability risk, obtain a comprehensive vulnerability score, and generate a vulnerability report for the smart contract based on the comprehensive vulnerability score, giving a panoramic view of the vulnerability.

2. According to claim 1, a smart contract vulnerability detection method based on code representation learning and graph neural network is characterized in that: In step 2, the graph structure is constructed and nodes are updated, including: Define each function in the smart contract as a node v in the graph i , e ij Represents the slave node v i To node v j The calling relationship, that is, the edge in the graph structure, where i and j are the nodes v i To node v j The identifier of the edge is given different category attributes according to the return value type; The message passing mechanism is used to pass the information of neighbor nodes to the current node and aggregate the information. During the aggregation process, the information of neighbor nodes is weighted according to their relationship or importance with the current node to achieve information integration. The feature representation of each node is updated to capture the global structure and the relationship between nodes. Use the attention mechanism to calculate the attention weight α of each edge ij , for node v i To neighbor node v j The sent messages are weighted, and the contribution of each node in the information aggregation is dynamically adjusted to capture the key features in the local structure.

3. According to claim 2, a smart contract vulnerability detection method based on code representation learning and graph neural network is characterized in that: The information aggregation process in step 2 is expressed as: in, is node v i Aggregate message at level l, v j is node v i Neighbors(v i ) is the node v i The set of neighbor nodes of is node v j Feature representation at layer l; In the graph neural network, the aggregation function AGGREGATE is used to aggregate the information of neighbor nodes, combine the information of neighbor nodes, i.e., feature representation, and update the feature representation of the current node; the AGGREGATE function integrates the features of neighbor nodes and uses the return value type of the edge as additional information; Calculate the attention weights: Among them, the attention weight α ij Represents node v i and its neighbor node v j The relative importance between ij Measure at node v i During the feature update process, the neighbor node v j The influence on the current node; and They are node v i and node v j In the feature representation of the lth layer, softmax is the normalization function, LeakyRelu is the activation function, a is the learned attention vector, W is the linear transformation matrix, and e ij Represents the slave node v i To node v j The return value type characteristics of the edge where ; For node v i To neighbor node v j The sent messages are weighted, and the message weighting process is expressed as: in, is node v i The weighted representation of the aggregated messages at layer l, α ij is the attention weight; Use the weighted message to update the feature representation of the current node. The node feature update is expressed as: in, is node v i In the feature representation of the l+1th layer, the node update function UPDATED combines the node v i Current features and aggregate messages Generate new feature representations.

4. According to claim 1, a smart contract vulnerability detection method based on code representation learning and graph neural network is characterized in that: The calculation formula for the comprehensive vulnerability score in step 3 is: Among them, S i Represents each node v i Vulnerability score of node v i Score S self (v i ) is based on node v i The code representation of each node is calculated; the code representation of each node is generated by a pre-trained model, which captures the syntax and semantic information of the code, converts the code features of the node into a high-dimensional vector representation, and then processes the feature vector of the node through a multi-layer perceptron to output the vulnerability risk score S of the node. self (v i ); N(v i ) represents the node v i The set of connected neighbor nodes, |N(v i )| is the number of neighbor nodes, S neighbor (v j ) is the neighbor node v j Score, S edge (e ij ) is the edge e ij The score of the connected node v i With v j The influence of the edge on the node score, W self , W neighbor and W edge They are the weights of the node itself, neighboring nodes, and edges; By training three weights W self , W neighbor and W edge The node's own score is calculated and three weights are obtained through supervised learning training. During the training process, a labeled smart contract dataset is used as input, in which each node and edge is labeled with a vulnerability risk label. A loss function is defined to measure the difference between the model's predicted score and the actual label, and the weight W is adjusted through back propagation and gradient descent. self , W neighbor and W edge Optimize; each edge e ij Score S edge (e ij ) is composed of the scores of the two end nodes and the return value type weight W of the edge edge calculate.

5. According to claim 4, a smart contract vulnerability detection method based on code representation learning and graph neural network is characterized in that: The edge score is calculated as: S edge (e ij )=W type i(W out iS token (v i )+W in iS token (v j )), Among them, W type is the weight assigned according to the return value type of the edge, W out and W in are the weights of out-degree and in-degree, S token (v i ) and S token (v j ) are nodes v i and v j The code represents the score of the learning phase.

Images