Hybrid model driven reentrant vulnerability maintenance method and device and computer equipment
Through a hybrid model-driven method, we obtain the slice set in the code, perform feature classification and attention score calculation, and combine it with a large language model to generate a repair template, which solves the problem of low accuracy of code vulnerability repair and achieves high-accuracy vulnerability repair.
Patent Information
- Application Number
- CN202411746778.7
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-11-29
- Publication Date
- 2025-03-25
AI Technical Summary
In the existing technology, the accuracy of code vulnerability repair is low, and it is impossible to accurately detect reentrant vulnerabilities, and the accuracy of vulnerability repair is insufficient.
Using a hybrid model-driven method, we obtain the slice set of the objective functions associated with the preset characters in the code to be repaired, perform feature classification training, calculate the attention score of the code line, and use the pre-trained large language model to generate a repair template, and output the repaired code.
Improves the accuracy of vulnerability repair, accurately locates the location of vulnerabilities, and enhances the detection and repair capabilities of reentrant vulnerabilities.
Smart Images

Figure CN119691747A_ABST
Abstract
Description
Technical Field
[0001] The present application relates to the field of vulnerability repair, and in particular to a hybrid model-driven reentrant vulnerability repair method, apparatus, and computer device. Background Art
[0002] In the related technologies, the methods for reentrancy vulnerability detection and repair are mainly divided into three categories: rule-based methods, deep learning-based methods, and large language model-based methods. These methods have their own characteristics and limitations, and representative applications of research tools include Oyente, Slither, Securify, SmartCheck, SliSE, Mythril, etc. For example, the smart contract vulnerability detection method based on graph neural network constructs key nodes by extracting the features of functions, variables and callback functions, thereby realizing the detection of reentrancy vulnerabilities, time-dependent vulnerabilities and infinite loop vulnerabilities in smart contracts; the new smart contract vulnerability detection method based on expert knowledge and graph neural network integrates global graph features and local expert features to further optimize the three vulnerability detection capabilities of smart contracts. Although many methods have been deeply discussed for reentrancy vulnerability detection, they still face problems such as insufficient vulnerability features and lack of domain knowledge, which leads to the inability to accurately detect reentrancy vulnerabilities and low accuracy of vulnerability repair.
[0003] There is currently no effective solution to the problem of low accuracy in repairing code vulnerabilities in related technologies. Summary of the invention
[0004] Based on this, it is necessary to provide a hybrid model-driven reentrant vulnerability repair method, device and computer equipment that can solve the problem of low accuracy in code vulnerability repair in response to the above technical problems.
[0005] In a first aspect, a hybrid model-driven reentrant vulnerability troubleshooting method is provided in this embodiment, wherein the method comprises:
[0006] Obtaining a slice set of a target function corresponding to a preset character in the code to be repaired, wherein the preset character is associated with the vulnerability to be repaired in the code;
[0007] Performing classification training on the features of the slice set to obtain a code snippet containing the vulnerability to be fixed;
[0008] Calculating the attention score of each code line in the code snippet, and selecting a specified number of target code lines with high attention scores from the code lines according to the order of the attention scores;
[0009] A repair template of a pre-trained large language model is obtained according to the target code line, and the code to be repaired is input into the large language model, and the repaired code is output.
[0010] In some embodiments, obtaining a slice set of a target function corresponding to a preset character in the code to be repaired includes:
[0011] Obtaining, in the code to be repaired, a first function including a preset character and a second function having a dependency relationship with the first function;
[0012] A slice set of the objective function is obtained according to the first function and the second function.
[0013] In some embodiments, before performing classification training on the features of the slice set to obtain the code fragment containing the vulnerability to be repaired, the method further includes:
[0014] inserting a code line randomly selected from the slice set into the end of the slice set; and / or,
[0015] A character randomly selected from the slice set is inserted at the end of the slice set.
[0016] In some embodiments, the classifying and training the features of the slice set to obtain the code fragment containing the vulnerability to be repaired includes:
[0017] vectorizing the set of slices;
[0018] Calculate the correlation between the vectors corresponding to the slice set based on the self-attention mechanism to obtain a matrix containing the attention weights of each tag in the slice set;
[0019] Extracting features from the matrix containing the attention weights;
[0020] According to the features, classification training is performed to obtain a code snippet containing the vulnerability to be repaired.
[0021] In some embodiments, calculating the attention score of each code line in the code snippet includes:
[0022] Obtaining the attention weight of each marker in the slice set;
[0023] Mapping the attention weight to the tag of the code snippet to obtain an attention score for each code line in the code snippet;
[0024] In some embodiments, obtaining a repair template of a pre-trained large language model according to the target code line, inputting the code to be repaired into the large language model, and outputting the repaired code includes:
[0025] Obtaining a preset code including the vulnerability to be repaired, a preset repair strategy, and an example of the vulnerability repaired based on the repair strategy;
[0026] The target code line, the preset code, the repair strategy and the vulnerability example are used as a repair template, and the code to be repaired is input into the large language model to obtain the repaired code.
[0027] In some embodiments, after obtaining a repair template of a pre-trained large language model according to the target code line, inputting the code to be repaired into the large language model, and outputting the repaired code, the method further includes:
[0028] The repaired code is detected to obtain the large language model repair detection report.
[0029] In some of the embodiments, the code to be repaired is the source code of the smart contract.
[0030] In a second aspect, a hybrid model-driven reentrant vulnerability repair device is provided in this embodiment, and the device includes:
[0031] A feature extraction module, used to obtain a slice set of the target function corresponding to a preset character in the code to be repaired, wherein the preset character is associated with the vulnerability to be repaired in the code;
[0032] A deep learning model is used to classify and train the features of the slice set to obtain a code snippet containing the vulnerability to be fixed;
[0033] A template construction module, used to calculate the attention score of each code line in the code snippet, and select a specified number of target code lines with high attention scores from the code lines according to the order of the attention scores;
[0034] A large language model is used to obtain a repair template of a pre-trained large language model according to the target code line, and the code to be repaired is input into the large language model, and the repaired code is output.
[0035] In a third aspect, a computer device is provided in this embodiment, including a memory and a processor, wherein the memory stores a computer program, and when the processor executes the computer program, the hybrid model-driven reentrant vulnerability repair method described in the first aspect is implemented.
[0036] In a fourth aspect, a computer-readable storage medium is provided in this embodiment, on which a computer program is stored. When the computer program is executed by a processor, the hybrid model-driven reentrant vulnerability repair method described in the first aspect is implemented.
[0037] In a fifth aspect, a computer program product is provided in this embodiment, including a computer program, which, when executed by a processor, implements the hybrid model-driven reentrant vulnerability repair method described in the first aspect above.
[0038] The above-mentioned hybrid model-driven reentrant vulnerability repair method, apparatus and computer device obtains a code snippet containing the vulnerability to be repaired in a slice set associated with the vulnerability to be repaired in the code, and determines the target code line containing the potential vulnerability in the code snippet based on the attention score. The target code line is used as a guide for the large language model to accurately locate the location of the vulnerability when performing the vulnerability repair task, and the accuracy of vulnerability repair is improved. BRIEF DESCRIPTION OF THE DRAWINGS
[0039] Figure 1 An application environment diagram of a hybrid model-driven reentrant vulnerability troubleshooting method in one embodiment;
[0040] Figure 2 A flow chart of a hybrid model-driven reentrant vulnerability troubleshooting method in one embodiment;
[0041] Figure 3 A schematic diagram of a hybrid model-driven reentrancy vulnerability detection and repair method in one embodiment;
[0042] Figure 4 is a flow chart of a feature extraction method in one embodiment;
[0043] Figure 5 A schematic diagram of deep learning model training in one embodiment;
[0044] Figure 6 A schematic diagram of a vulnerability repair method in an embodiment;
[0045] Figure 7 A structural block diagram of a hybrid model-driven reentrant vulnerability troubleshooting method in one embodiment;
[0046] Figure 8 FIG. 4 is a diagram showing the internal structure of a computer device in one embodiment. DETAILED DESCRIPTION
[0047] In order to make the purpose, technical solution and advantages of the present application more clearly understood, the present application is further described in detail below in conjunction with the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain the present application and are not used to limit the present application.
[0048] The hybrid model-driven reentrant vulnerability troubleshooting method provided in the embodiment of the present application can be applied to Figure 1In the application environment shown. Among them, the terminal 102 communicates with the server 104 through the network. The data storage system can store the data that the server 104 needs to process. The data storage system can be integrated on the server 104, or it can be placed on the cloud or other network servers. Based on the terminal, the code to be repaired is processed and a repair module is obtained, and the large language model in the terminal outputs the repaired code based on the repair module. Among them, the terminal 102 can be but not limited to various personal computers, laptops, smart phones, tablets, etc. The server 104 can be implemented with an independent server or a server cluster consisting of multiple servers.
[0049] In one embodiment, Figure 2 As shown in the figure, a hybrid model-driven reentrant vulnerability troubleshooting method is provided. Figure 1 The terminal in is used as an example to illustrate, including the following steps:
[0050] Step S202: obtaining a slice set of a target function corresponding to a preset character in the code to be repaired, wherein the preset character is associated with the vulnerability to be repaired in the code.
[0051] Among them, the target function corresponding to the preset character can be a function containing the preset character, or a function that has a dependency relationship with the function containing the preset character. The preset character can be determined based on the vulnerability to be repaired. For example, through the analysis of repeatable vulnerabilities, it is determined that the generation of reentrant vulnerabilities is closely related to the DAO (Distributed Autonomous Organization) attack. The smart contract attacked by DAO needs to contain the call.value keyword. According to the analysis results, call.value is used as the preset character. It is understandable that when the vulnerability to be repaired is other vulnerabilities, the preset character can also be set according to the vulnerability characteristics of the corresponding attack.
[0052] Optionally, the code of the smart contract is obtained as the code to be repaired, and after determining the target function associated with the preset character in the code, the function slices associated with the target function in the code are extracted based on a predefined slicing standard to obtain a slice set. In the case of multiple target functions, the function slices of each target function are integrated to obtain a slice set.
[0053] Step S204: Classify and train the features of the slice set to obtain code snippets containing vulnerabilities to be repaired.
[0054] The features of the slice set are the vulnerability features contained in the code corresponding to the target function. Optionally, a pre-trained model corresponding to the programming language of the code to be repaired can be obtained, for example, a code representation model such as CodeBERT, Roberta, and CodeBERTa, and after extracting the semantic features of the code corresponding to the slice set, the features of the slice set are extracted through a neural network. The features of the slice set can also be extracted based on other existing static code analysis tools.
[0055] After extracting the features of the slice set, a classifier may perform classification training based on the features of the slice set, and obtain a code snippet in the slice set containing the vulnerability to be fixed based on the classification result.
[0056] Optionally, a deep learning model including a classifier is designed, the slice set is input into the deep learning model, and after the slice set is vectorized based on the deep learning model, features of the slice set are extracted through an encoder and / or a neural network convolutional layer, and classification training is implemented through the classifier in the deep learning model, so that the deep model distinguishes between code snippets containing vulnerabilities to be fixed and code snippets that do not contain vulnerabilities to be fixed.
[0057] Step S206, calculating the attention score of each code line in the code snippet, and selecting a specified number of target code lines with high attention scores from the code lines according to the order of the attention scores.
[0058] The code snippet is the code containing the vulnerability to be fixed determined based on the classification result, and the specified number of target code lines can be preset or modified according to the requirements. Optionally, the code lines are sorted according to the attention score; based on the sorting, a specified number of code lines with higher attention scores are selected as target code lines, and the target code lines can help the large language model locate the vulnerability more accurately.
[0059] Step S208, obtaining a repair template of a pre-trained large language model according to the target code line, inputting the code to be repaired into the large language model, and outputting the repaired code.
[0060] Among them, the repair template is obtained based on the code line with a higher attention score, so that when the large language model repairs the code, it can refer to the target code line in the repair template to obtain a prompt of the location of the vulnerability, guide the large model to focus on these target code lines, and output the repaired code based on the repair template. Optionally, after the repair template is constructed according to the target code line, the repair template and the code to be repaired are input into the large language model. Under the guidance of the repair template, the pre-trained large language model identifies the vulnerability in the code to be repaired, repairs the vulnerability, and outputs the repaired code.
[0061] In the above hybrid model-driven reentrant vulnerability repair method, after obtaining the slice set associated with the vulnerability to be repaired in the code, the code snippet containing the vulnerability to be repaired in the slice set is determined, and the target code line containing the potential vulnerability is determined in the code snippet based on the attention score. The target code line is used as a repair template to help the large language model accurately locate the location of the vulnerability when performing the vulnerability repair task, thereby achieving more accurate vulnerability location and improving the accuracy of vulnerability repair.
[0062] In one embodiment, obtaining a slice set of a target function corresponding to a preset character in the code to be repaired includes: obtaining a first function containing the preset character and a second function having a dependency relationship with the first function in the code to be repaired; and obtaining a slice set of the target function according to the first function and the second function.
[0063] Optionally, a structured analysis is performed on the functions of the code to determine a first function containing preset characters, a function slice of the first function is extracted, and the function slice of the first function is used as the core feature of the vulnerability to be fixed. At the same time, by determining to perform a structured analysis on the functions of the code, a function that does not contain preset characters is determined, and the slice of the function that does not contain preset characters is analyzed: it is determined whether there is a function dependency relationship between the function that does not contain preset characters and the first function. If there is a function dependency relationship, the function that does not contain preset characters is used as the second function, and a function slice of the second function is extracted. The function slice of the first function and the function slice of the second function are used as output to obtain a slice set of the target function.
[0064] The code extraction method of the traditional method is highly dependent on the rules defined by experts, resulting in insufficient mining of the deep features of reentrant vulnerabilities and complex function call chains, resulting in insufficient vulnerability features contained in the extracted code. In this embodiment, by analyzing the code to obtain a first function and a second function that has a dependency relationship with the first function, a slice set of the target function corresponding to the vulnerability to be repaired is obtained, and the vulnerability features contained in the extracted code are increased, so that the slice set based on the first function and the second function can effectively extract and locate the features of the vulnerability to be repaired.
[0065] After obtaining the slice set, if the slice set belongs to an unbalanced data set, the output is prone to overfitting when using the unbalanced data set for classification training, and there is a problem of low robustness of the experimental results. Based on this, in one embodiment, before performing classification training on the features of the slice set to obtain the code snippet containing the vulnerability to be repaired, the method also includes: inserting the code line randomly selected from the slice set to the end of the slice set; and / or inserting the character randomly selected from the slice set to the end of the slice set.
[0066] Among them, the randomly selected characters in the slice set include a word or symbol of the target function. By inserting characters at the end of the slice set, the noise in the slice set can be increased from the token level without destroying the function structure; by inserting code lines at the end of the slice set, the noise in the slice set can be increased from the Line level without destroying the function structure; data enhancement is performed on the slice set of the target function, which improves the accuracy of reentrant vulnerability detection and repair.
[0067] In one embodiment, classification training is performed on features of a slice set to obtain a code snippet containing a vulnerability to be fixed, including: vectorizing the slice set; calculating the correlation between vectors corresponding to the slice set based on a self-attention mechanism to obtain a matrix containing attention weights of each tag in the slice set; extracting features from the matrix containing the attention weights; and classifying and training according to the features to obtain a code snippet containing the vulnerability to be fixed.
[0068] Among them, the code slices can be converted into numerical vector representations based on a pre-trained model corresponding to the language into which the code is converted or a cross-language pre-trained code representation model. After obtaining the vectorized slice set, optionally, based on the self-attention mechanism, the correlation between the query, key and value is calculated by converting the vectorized slice set into a vector of query, key and value. Based on the calculated similarity between the query and the key, the self-attention weight of the input vector is obtained, that is, the attention weight of each tag in the code; and a matrix containing the attention weights is output.
[0069] Optionally, a deep learning model can be designed to implement classification training for the slice set: the deep learning model includes: a pre-training layer, an encoder layer, a convolution layer, and a classification layer. The slice set is vectorized through the pre-training layer; a matrix containing attention weights of each tag in the slice set is calculated through the encoder layer; features in the matrix containing attention weights are extracted through the convolution layer; and classification training is performed on the features extracted by the convolution layer based on the classification layer.
[0070] In this embodiment, the attention mechanism enables the model to focus on the key parts of the code and improve the ability to identify vulnerability patterns. Code snippets containing vulnerabilities to be fixed are obtained through code feature classification, thereby helping the large language model to locate vulnerabilities more accurately.
[0071] In one embodiment, the attention score of each code line in the code snippet is calculated, including: obtaining the attention weight of each tag in the slice set; mapping the attention weight to the tag of the code snippet to obtain the attention score of each code line in the code snippet.
[0072] Among them, the different attention scores of each code line reflect the different recognition capabilities of the deep learning network for vulnerability features: a high code line attention score indicates that when the deep learning network identifies the current code line, it is more likely that the code line contains code containing vulnerabilities; conversely, it indicates that when the deep learning network identifies the current code line, it is less likely that the code line contains code containing vulnerabilities.
[0073] Optionally, the attention weights can be calculated based on a self-attention mechanism, for example, a Transformer encoder. If the attention weights have been calculated using a self-attention mechanism when the code is classified and trained, the attention weights of each tag can be directly obtained at this time. The attention weights are used to indicate the importance of each tag in the vulnerability detection process. By mapping the attention weights to the tags of the code snippet, the attention score of the code line can be obtained by normalizing the attention weights of each tag in the code line and combining the attention weights of each tag in the code line.
[0074] In this embodiment, the attention score of each code line in the code snippet is obtained, and the probability of each code line containing a vulnerability is determined. After the repair template is constructed in combination with the target code line, the large language model can refer to the repair template to enhance the large language model's ability to identify and locate vulnerabilities and optimize the vulnerability repair effect.
[0075] In one embodiment, a repair template of a pre-trained large language model is obtained according to a target code line, and the code to be repaired is input into the large language model to output the repaired code, including: obtaining a preset code containing a vulnerability to be repaired, a preset repair strategy, and a vulnerability example repaired based on the repair strategy; using the target code line, the preset code, the repair strategy, and the vulnerability example as a repair template, inputting the code to be repaired into the large language model, and obtaining the repaired code.
[0076] The preset code containing the vulnerability to be fixed is used to show the structure of a typical vulnerability to be fixed to the large language model. The repair strategy includes one or more strategies corresponding to the vulnerability to be fixed. The vulnerability example corresponding to the repair strategy includes: the preset code containing the vulnerability to be fixed is used to complete the vulnerability repair through one or more of the above repair strategies.
[0077] Preset codes, repair strategies, and vulnerability examples can help large language models identify vulnerabilities and provide repair ideas for large language models. Because the target code line can indicate the location of the vulnerability to be repaired, based on the preset code, repair strategy, and vulnerability examples, combined with the target code line, the domain knowledge of the vulnerability to be repaired and the code structure analysis are combined together to further enhance the model's ability to identify vulnerability types and improve the vulnerability repair effect of the large language model.
[0078] In one embodiment, after obtaining a repair template of a pre-trained large language model according to the target code line, inputting the code to be repaired into the large language model, and outputting the repaired code, the method further includes: detecting the repaired code to obtain a detection report of the repair of the large language model.
[0079] Optionally, the detection report includes the repair accuracy and / or success rate of the large language model for the vulnerabilities to be repaired. According to the vulnerabilities contained in the code to be repaired and the vulnerabilities contained in the repaired code, it is determined whether all the vulnerabilities contained in the code to be repaired are repaired. The more vulnerabilities are repaired, the higher the repair accuracy of the large language model; conversely, the lower the repair accuracy of the large language model.
[0080] Optionally, it is also detected whether the repaired vulnerability meets the preset vulnerability repair requirements. The more the repaired vulnerability meets the preset vulnerability repair requirements, the higher the large language model repair success rate; otherwise, the large language model repair success rate is lower.
[0081] In one embodiment, the code to be repaired is the source code of a smart contract. Among them, a smart contract is an automated program running in a blockchain system that can be executed spontaneously in a decentralized environment. Smart contracts can manage and maintain data and assets in a blockchain system, and therefore have achieved great success in fields such as finance, supply chain management, and insurance. With its wide application in many fields, the security issues of smart contracts have gradually attracted the attention of industry and academia. Smart contracts have the characteristics of being tamper-proof and automatically executable. Once there are logical loopholes or defects in the code, it may be attacked and cause serious economic losses. Taking the DAO attack as an example, hackers attacked the reentrant vulnerability in the smart contract and successfully stole digital assets up to $150 million.
[0082] At present, there are many types of vulnerabilities in smart contracts, which are complex and can occur at the blockchain system layer, contract layer, EVM (Ethereum Virtual Machine) layer, etc. The SWC registry summarizes 37 common vulnerability types, and the security company NCC Group summarizes the 10 most influential vulnerabilities. Therefore, using the source code of smart contracts as the code to be fixed will help maintain the blockchain system.
[0083] Furthermore, among the vulnerabilities of smart contracts, reentrancy vulnerabilities are the most representative. This vulnerability is not only the core reason for the DAO attack, but also a key topic in the research of smart contract security. Therefore, the vulnerability to be fixed can be set as a reentrancy vulnerability, thereby ensuring the security of smart contracts and blockchain systems.
[0084] In one embodiment, Figure 3A schematic diagram of a hybrid model-driven reentrancy vulnerability detection and repair method is provided. The code to be repaired is the source code of the smart contract, and the vulnerability to be repaired is a reentrancy vulnerability. The specific process is as follows Figure 3 As shown, it includes: obtaining the smart contract source code, performing root cause analysis, and determining the key functions that cause the vulnerability according to the cause of the reentrant vulnerability; designing a vulnerability feature extraction algorithm to extract a function slice set related to the reentrant vulnerability feature; designing a data enhancement algorithm to balance the function slice set; sending the balanced function slice set to the deep learning model for training, and performing vulnerability detection to obtain a function slice set that does not contain reentrant vulnerabilities and a function slice set that contains reentrant vulnerabilities. Obtain the detection results of the function slice set containing reentrant vulnerabilities, and extract the code line information with high attention scores in the function slices. According to the domain knowledge and code line information of the reentrant vulnerability, a template for repairing the reentrant vulnerability is constructed. Among them, the relevant domain knowledge of the reentrant vulnerability includes vulnerability samples, repair samples, repair methods, etc.; the code line information is the code line with high attention scores in the function slice. Send the repair template to the LLM (Large Language Model) large model repair to obtain the repair result of the large model for the reentrant vulnerability. And according to the repair evaluation report of the LLM model.
[0085] The following is a detailed description of the hybrid model-driven reentrancy vulnerability detection and repair method.
[0086] Among them, vulnerability detection includes four modules: root cause analysis module, feature extraction module, data enhancement module and vulnerability detection module based on deep learning. The root cause analysis module is used to analyze the root causes and key functions of reentrant vulnerabilities. The feature extraction module is used to design a feature extraction algorithm for reentrant vulnerabilities. The data enhancement module is used to design an algorithm to expand the unbalanced function slice set to a balanced function slice set. The deep learning detection module is used to design deep learning for function slice set detection of reentrant vulnerabilities.
[0087] Through root cause analysis, it can be seen that the occurrence of reentrancy vulnerabilities is closely related to the DAO attack. Among them, the DAO attack refers to an attack by external hackers against the logic of smart contracts. Reentrancy vulnerabilities occur when the internal call structure and fund transfer mechanism of smart contracts are improperly used. This reentrancy vulnerability can be exploited by meeting the following three conditions: 1. The attacked smart contract needs to contain the call.value keyword, which is a key function to implement reentrancy attacks because it allows funds to flow out of the contract; 2. The attacker must be able to directly or indirectly call the function containing call.value to achieve repeated calls to the smart contract, thereby achieving the purpose of stealing funds; 3. The attacker needs to build a fallback function that can repeatedly call the target function containing the call.value keyword. By identifying the code snippet containing the call.value keyword and its function call relationship, the characteristics of reentrancy vulnerabilities can be effectively extracted and located.
[0088] Based on the root cause analysis results, feature extraction can be performed on the source code of the smart contract. Figure 4 A flow chart of a feature extraction method is provided, such as Figure 4 As shown, it includes: obtaining the smart contract source code, performing a structured analysis on the functions of the smart contract source code, determining the functions containing the call.value keyword in the source code as core functions, and determining the functions not including call.value as non-core functions. Among them, the core function is the first function in the above embodiment. Extract the function slices of the core function as the core feature of the reentrancy vulnerability. Analyze the non-core function slice set to determine whether there is a functional dependency relationship between the non-core function and the core function. If there is a functional dependency relationship, extract the slices of the non-core function. Among them, the non-core function that has a dependency relationship with the core function is the second function in the above embodiment. Take the core function and the non-core function that has a functional dependency relationship as the output of the entire smart contract, traverse all smart contract data sets, and output the slice set of the target function related to the reentrancy vulnerability.
[0089] The feature extraction of reentrant vulnerabilities is one of the cores of the hybrid model-driven reentrant vulnerability detection in this embodiment. Traditional technologies focus on the vulnerability attack features in smart contracts, but ignore the functional dependency features of smart contract vulnerabilities. The new algorithm designed in the feature extraction module in this embodiment can not only extract the direct features of reentrant vulnerabilities, but also extract the functional dependency features of reentrant vulnerabilities.
[0090] After obtaining the slice set of the target function of the reentrancy vulnerability, a new algorithm is designed to perform data enhancement on the slice set of the target function, including: obtaining the function slice set, randomly selecting a line of the function, and inserting it to the end of the entire function slice; obtaining the function slice set, randomly selecting a word or symbol of the function, and inserting it to the end of the entire function slice. By enhancing the function slice set at the line level and the token level, the data for deep learning is balanced.
[0091] The data-enhanced reentrant vulnerability slice set is detected through deep learning. Figure 5 A schematic diagram of deep learning model training is provided, such as Figure 5 As shown in the figure, the deep learning model includes: pre-training layer, encoding layer, convolution layer and classification layer.
[0092] Pre-training layer: Considering that there is no specific pre-training model for the smart contract programming language Solidity, the Solidity code is vectorized with the help of a cross-language pre-trained code representation model. Among them, a pre-trained model refers to a model that has been trained in advance on a large-scale data set to capture the characteristics and patterns of the data and form a broad knowledge representation. The pre-trained model learns common features by training on general tasks such as language modeling or image recognition, and can be fine-tuned or directly applied in downstream specific tasks, thereby significantly improving the performance of the model on various tasks. Among them, code representation models such as CodeBERT, Roberta, and CodeBERTa can be selected, and the code representation model has been pre-trained on multiple programming languages to obtain a pre-training layer, thereby effectively extracting the semantic features of the code.
[0093] The encoder layer is a network layer in a deep learning model that converts input data into a high-dimensional representation suitable for subsequent processing. Its main function is to extract features of the input. The encoder layer can enable the model to capture the semantic and contextual information of the input data through a series of nonlinear transformations. Optionally, the global features of Solidity code snippets can be extracted by using a Transformer encoder based on the self-attention mechanism in the encoder layer. The self-attention mechanism converts the input vector into a query, key, and value matrix, calculates the correlation between the parts, and finally obtains the self-attention weight of the input vector. This attention mechanism enables the model to focus on the key parts of the code and improve the ability to identify vulnerability patterns.
[0094] Convolutional neural network layer, used to extract local features of input data. By applying convolution kernels on the input data, the convolutional neural network layer can recognize low-level features such as edges and textures in images or sequences, and gradually extract higher-level patterns or semantic information after stacking layers. Optionally, by using convolution kernels in the convolutional neural network layer, focus on extracting local features in the code snippets output by the encoder layer. Local features are an important part of analyzing vulnerability patterns because vulnerabilities are often confined to specific code structures. Optionally, in the convolutional neural network layer, 6 convolution kernels of size 2 and average pooling method are used to compress and optimize features, and finally retain the most valuable vulnerability feature information. Other convolution kernels can also be used to implement feature processing according to application requirements, and there is no limit on the number and size of convolution kernels.
[0095] The classification layer uses a multi-layer perceptron as a classifier to perform classification training on the feature set output by the convolutional layer. The multi-layer perceptron uses the SoftMax activation function and the cross entropy loss function to enable the deep learning model to effectively distinguish between code snippets containing reentrancy vulnerabilities and code snippets that do not contain reentrancy vulnerabilities.
[0096] The feature extraction algorithm, data enhancement algorithm, and deep learning model in this embodiment solve the problems of insufficient feature extraction and data imbalance in traditional deep learning methods in smart contract vulnerability detection, and achieve the effect of enhancing the detection of reentrant vulnerabilities.
[0097] After vulnerability detection, vulnerability repair can be performed. Figure 6 A schematic diagram of a vulnerability repair method is provided. A repair template is constructed by combining the deep learning model detection results and reentrant vulnerability knowledge. The LLMs model is input in combination with the LLM repair prompt. The LLMs model repairs the vulnerabilities of the smart contract and outputs a repair evaluation report.
[0098] Among them, combined with the detection results of the deep learning model, code line extraction based on attention scores can be achieved. Code line extraction based on attention scores specifically includes: extracting the score distribution of all tags in the code snippet containing reentrant vulnerabilities from the convolutional layer of the neural network to obtain the attention score of each code line; sending the code lines with high attention scores as target code lines into the large language model to guide the large language model to focus on these target code lines to achieve more accurate vulnerability positioning. Among them, during the training process of the deep learning network, the encoder layer in the deep learning model dynamically assigns attention weights to each token in the input code snippet. In order to improve the recognition accuracy of potential vulnerabilities, the attention weights of the tags are mapped to the code snippets containing reentrant vulnerabilities obtained by the classifier, and the attention scores of each code line of the code snippet containing reentrant vulnerabilities are obtained. Different attention scores reflect the different recognition capabilities of the deep learning network for vulnerability features. In the code snippet containing reentrant vulnerabilities, target code lines with higher attention scores are extracted as location hints to help the large language model locate vulnerabilities more accurately.
[0099] The repair template is constructed based on the deep learning model detection results and the knowledge of reentrant vulnerabilities. The repair template includes four core parts: vulnerability examples, vulnerability examples repaired based on the repair strategy, repair strategy, and code lines based on attention scores.
[0100] Among them, the vulnerability example is the preset code of the vulnerability to be fixed, specifically including the function code snippet of the reentrancy vulnerability. The vulnerability example is used to show the typical vulnerability structure to the model.
[0101] The vulnerability example after repair based on the repair strategy is the code without reentrancy vulnerability. Optionally, the function code snippet containing the vulnerability is repaired through three repair strategies to obtain the vulnerability example, thereby providing multiple repair ideas for repairing other vulnerabilities in the large model.
[0102] Repair strategies may include one or more, such as Figure 6 As shown, the repair strategy includes: Method 1, including locking the call.value code line; Method 2, replacing the call.value function with a transfer or send function; Method 3, using the checks-effects-interactions pattern to ensure that all smart contract parameters are checked before modifying the state of the smart contract.
[0103] The code behavior based on the attention score is the target code line mentioned above, which is used to enable the model to locate and fix reentrancy vulnerabilities more accurately. Based on the repair template, combined with the LLMs prompt: Please repair the vulnerable codegadget and output the fixed result, the large language model outputs the repaired code of the smart contract. Among them, the large language model refers to a large-scale deep neural network model trained on massive text data, which is specially used to process and generate natural language. Large language models usually have billions to hundreds of billions of parameters. They learn the grammar, semantics and context information of the language through a complex multi-layer neural network structure, and have natural language processing capabilities such as understanding, generation and reasoning.
[0104] In the above-mentioned vulnerability repair method, by extracting the target code line fragments in the deep learning method, constructing the vulnerability repair knowledge in the field of reentrant vulnerabilities, and combining the target code line fragments to construct the repaired vulnerabilities, a large model vulnerability repair method based on knowledge enhancement is realized, which significantly enhances the ability of the large model to locate and repair vulnerabilities, and achieves the effect of enhancing the model's ability to identify vulnerability types, optimizing vulnerability repair effects, and improving the accuracy of large language models in vulnerability repair tasks.
[0105] Furthermore, in order to comprehensively evaluate the effectiveness of the large language model in repairing reentrancy vulnerabilities in smart contracts, a structured evaluation process can be designed, including: providing a code snippet containing a reentrancy vulnerability to the large language model, and providing a designed prompt that combines vulnerability repair knowledge to guide the large language model to generate repaired code; then, strictly testing the generated repair code to ensure that it meets the requirements of vulnerability repair; finally, by statistically analyzing the repair success rate of the large language model and the accuracy of the repair code, the performance of the model in repairing reentrancy vulnerabilities is evaluated, providing a basis for further enhancing the application value of the model in the field of smart contract security.
[0106] In this embodiment, a hybrid model framework is proposed to detect and repair reentrancy vulnerabilities. The hybrid model framework combines the vulnerability detection of deep learning models and the vulnerability repair of knowledge-enhanced large models, improves the accuracy of reentrancy vulnerability detection and repair, and proposes a new idea for the security research of smart contracts.
[0107] Based on the same inventive concept, the embodiment of the present application also provides a hybrid model-driven reentrant vulnerability repair device for implementing the hybrid model-driven reentrant vulnerability repair method involved above. The implementation scheme for solving the problem provided by the device is similar to the implementation scheme recorded in the above method, so the specific limitations in one or more hybrid model-driven reentrant vulnerability repair device embodiments provided below can refer to the limitations of the hybrid model-driven reentrant vulnerability repair method above, and will not be repeated here.
[0108] In one embodiment, Figure 7 As shown, a hybrid model-driven reentrant vulnerability repair device is provided, comprising:
[0109] A feature extraction module, used for obtaining a slice set of a target function corresponding to a preset character in the code to be repaired, wherein the preset character is associated with the vulnerability to be repaired in the code;
[0110] A deep learning model is used to classify and train the features of the slice set to obtain the code snippets containing the vulnerabilities to be fixed;
[0111] A template building module is used to calculate the attention score of each code line in the code snippet, and select a specified number of target code lines with high attention scores from the code lines according to the order of the attention scores;
[0112] The large language model uses a repair template of a pre-trained large language model obtained according to the target code line, and inputs the code to be repaired into the large language model, and outputs the repaired code.
[0113] In some of the embodiments, obtaining a slice set of a target function corresponding to a preset character in the code to be repaired includes: obtaining a first function containing the preset character and a second function that has a dependency relationship with the first function in the code to be repaired; and obtaining a slice set of the target function based on the first function and the second function.
[0114] In some of the embodiments, the hybrid model-driven reentrant vulnerability repair device also includes a data enhancement module, which is used to insert code lines randomly selected from the slice set into the end of the slice set before classifying and training the features of the slice set to obtain the code fragment containing the vulnerability to be repaired; and / or insert characters randomly selected from the slice set into the end of the slice set.
[0115] In some of the embodiments, a deep learning model performs classification training on features of a slice set to obtain a code snippet containing a vulnerability to be fixed, including: vectorizing the slice set; calculating the correlation between vectors corresponding to the slice set based on a self-attention mechanism to obtain a matrix containing attention weights of each tag in the slice set; extracting features from the matrix containing the attention weights; and classifying and training based on the features to obtain a code snippet containing the vulnerability to be fixed.
[0116] In some of the embodiments, the attention score of each code line in the code snippet is calculated, including: obtaining the attention weight of each tag in the slice set; mapping the attention weight to the tag of the code snippet to obtain the attention score of each code line in the code snippet.
[0117] In some of the embodiments, the large language model obtains a repair template, a repair prompt repair template of a pre-trained large language model according to the target code line, and inputs the code to be repaired into the large language model, and outputs the repaired code, including: obtaining a preset code containing the vulnerability to be repaired, a preset repair strategy, and a vulnerability repair example after the repair based on the repair strategy; using the target code line, the preset code, the repair strategy, and the vulnerability repair example as a repair template, a repair prompt repair template, inputting the code to be repaired into the large language model, and obtaining the repaired code.
[0118] In some of the embodiments, the hybrid model-driven reentrant vulnerability repair device also includes an evaluation module, which is used to obtain a repair template of a pre-trained large language model according to a target code line, and input the code to be repaired into the large language model. After outputting the repaired code, the repaired code is detected to obtain an accuracy and success rate detection report of the large language model repair.
[0119] In some of these embodiments, the code to be repaired is the source code of the smart contract.
[0120] Each module in the hybrid model-driven reentrant vulnerability repair device can be implemented in whole or in part by software, hardware, or a combination thereof. Each module can be embedded in or independent of a processor in a computer device in the form of hardware, or can be stored in a memory in a computer device in the form of software, so that the processor can call and execute operations corresponding to each module.
[0121] In one embodiment, a computer device is provided. The computer device may be a server, and its internal structure diagram may be as follows: Figure 8As shown. The computer device includes a processor, a memory, an input / output interface (Input / Output, referred to as I / O) and a communication interface. Among them, the processor, the memory and the input / output interface are connected through a system bus, and the communication interface is connected to the system bus through the input / output interface. Among them, the processor of the computer device is used to provide computing and control capabilities. The memory of the computer device includes a non-volatile storage medium and an internal memory. The non-volatile storage medium stores an operating system, a computer program and a database. The internal memory provides an environment for the operation of the operating system and the computer program in the non-volatile storage medium. The database of the computer device is used to store data required for executing a hybrid model-driven reentrant vulnerability repair method, including code to be repaired, a large language model, etc. The input / output interface of the computer device is used to exchange information between the processor and an external device. The communication interface of the computer device is used to communicate with an external terminal through a network connection. When the computer program is executed by the processor, a hybrid model-driven reentrant vulnerability repair method is implemented.
[0122] Those skilled in the art will understand that Figure 8 The structure shown in the figure is only a block diagram of a part of the structure related to the solution of the present application, and does not constitute a limitation on the computer device to which the solution of the present application is applied. The specific computer device may include more or fewer components than those shown in the figure, or combine certain components, or have a different arrangement of components.
[0123] In one embodiment, a computer device is further provided, including a memory and a processor, wherein a computer program is stored in the memory, and the processor implements the steps in the above method embodiments when executing the computer program.
[0124] In one embodiment, a computer-readable storage medium is provided, on which a computer program is stored. When the computer program is executed by a processor, the steps in the above-mentioned method embodiments are implemented.
[0125] In one embodiment, a computer program product is provided, including a computer program, which implements the steps in the above method embodiments when executed by a processor.
[0126] Those skilled in the art can understand that all or part of the processes in the above-mentioned embodiment methods can be completed by instructing the relevant hardware through a computer program, and the computer program can be stored in a non-volatile computer-readable storage medium. When the computer program is executed, it can include the processes of the embodiments of the above-mentioned methods. Among them, any reference to the memory, database or other medium used in the embodiments provided in the present application can include at least one of non-volatile and volatile memory. Non-volatile memory can include read-only memory (ROM), magnetic tape, floppy disk, flash memory, optical memory, high-density embedded non-volatile memory, resistive random access memory (ReRAM), magnetoresistive random access memory (MRAM), ferroelectric random access memory (FRAM), phase change memory (PCM), graphene memory, etc. Volatile memory can include random access memory (RAM) or external cache memory, etc. As an illustration and not limitation, RAM can be in various forms, such as static random access memory (SRAM) or dynamic random access memory (DRAM). The database involved in each embodiment provided in this application may include at least one of a relational database and a non-relational database. Non-relational databases may include distributed databases based on blockchains, etc., but are not limited to this. The processor involved in each embodiment provided in this application may be a general-purpose processor, a central processing unit, a graphics processor, a digital signal processor, a programmable logic device, a data processing logic device based on quantum computing, etc., but are not limited to this.
[0127] The technical features of the above embodiments may be combined arbitrarily. To make the description concise, not all possible combinations of the technical features in the above embodiments are described. However, as long as there is no contradiction in the combination of these technical features, they should be considered to be within the scope of this specification.
[0128] The above-described embodiments only express several implementation methods of the present application, and the descriptions thereof are relatively specific and detailed, but they cannot be understood as limiting the scope of the present application. It should be pointed out that, for a person of ordinary skill in the art, several variations and improvements can be made without departing from the concept of the present application, and these all belong to the protection scope of the present application. Therefore, the protection scope of the present application shall be subject to the attached claims.
Claims
1. A hybrid model-driven reentrant vulnerability troubleshooting method, characterized in that: The method comprises: Obtaining a slice set of a target function corresponding to a preset character in the code to be repaired, wherein the preset character is associated with the vulnerability to be repaired in the code; Performing classification training on the features of the slice set to obtain a code snippet containing the vulnerability to be fixed; Calculating the attention score of each code line in the code snippet, and selecting a specified number of target code lines with high attention scores from the code lines according to the order of the attention scores; A repair template of a pre-trained large language model is obtained according to the target code line, and the code to be repaired is input into the large language model, and the repaired code is output.
2. The method according to claim 1, characterized in that The step of obtaining a slice set of a target function corresponding to a preset character in the code to be repaired includes: Obtaining, in the code to be repaired, a first function including a preset character and a second function having a dependency relationship with the first function; A slice set of the objective function is obtained according to the first function and the second function.
3. The method according to claim 1, characterized in that Before performing classification training on the features of the slice set to obtain the code fragment containing the vulnerability to be repaired, the method further includes: inserting a code line randomly selected from the slice set into the end of the slice set; and / or, A character randomly selected from the slice set is inserted at the end of the slice set.
4. The method according to claim 1, characterized in that: The classifying and training the features of the slice set to obtain the code snippet containing the vulnerability to be fixed includes: vectorizing the set of slices; Calculate the correlation between the vectors corresponding to the slice set based on the self-attention mechanism to obtain a matrix containing the attention weights of each tag in the slice set; Extracting features from the matrix containing the attention weights; According to the features, classification training is performed to obtain a code snippet containing the vulnerability to be repaired.
5. The method according to claim 1, characterized in that The calculating the attention score of each code line in the code snippet includes: Obtaining the attention weight of each marker in the slice set; The attention weight is mapped to the tag of the code snippet to obtain the attention score of each code line in the code snippet.
6. The method according to claim 1, characterized in that The method of obtaining a repair template of a pre-trained large language model according to the target code line, inputting the code to be repaired into the large language model, and outputting the repaired code includes: Obtaining a preset code including the vulnerability to be repaired, a preset repair strategy, and an example of the vulnerability repaired based on the repair strategy; The target code line, the preset code, the repair strategy and the vulnerability example are used as a repair template, and the code to be repaired is input into the large language model to obtain the repaired code.
7. The method according to claim 1, characterized in that After obtaining a repair template of a pre-trained large language model according to the target code line, inputting the code to be repaired into the large language model, and outputting the repaired code, the method further includes: The repaired code is detected to obtain the large language model repair detection report.
8. A hybrid model-driven reentrant vulnerability repair device, characterized in that: The device comprises: A feature extraction module, used to obtain a slice set of the target function corresponding to a preset character in the code to be repaired, wherein the preset character is associated with the vulnerability to be repaired in the code; A deep learning model is used to classify and train the features of the slice set to obtain a code snippet containing the vulnerability to be fixed; A template construction module, used to calculate the attention score of each code line in the code snippet, and select a specified number of target code lines with high attention scores from the code lines according to the order of the attention scores; A large language model is used to obtain a repair template of a pre-trained large language model according to the target code line, and the code to be repaired is input into the large language model, and the repaired code is output.
9. A computer device comprising a memory and a processor, wherein the memory stores a computer program, wherein: When the processor executes the computer program, the steps of the method according to any one of claims 1 to 7 are implemented.
10. A computer storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 7 are implemented.
Citation Information
Patent Citations
Intelligent contract vulnerability detection method based on graph neural network and electronic equipment
CN116467720A
Vulnerability fixing method based on large model and related device
CN118445812A
Automatic code defect repairing method based on static code analysis tool and artificial intelligence
CN118860864A
Source code vulnerability detection and repair through machine learning
WO2024019848A1
Cited By
Intelligent contract vulnerability detection method and system based on type perception and graph enhancement
CN120671153A
Attack method for vulnerability auditing tool based on big language model based on interference attention
CN120893051A