Vulnerability repairing method and device based on self-attention mask and storage medium
Through the vulnerability repair method based on self-attention masks, using neural networks and vulnerability masks, the problems of complex context vulnerability detection and high false negative rates in the existing technology are solved, and efficient and accurate vulnerability repair is achieved.
Patent Information
- Application Number
- CN202411940863.7
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-12-26
- Publication Date
- 2025-04-22
AI Technical Summary
Existing static application security testing tools are difficult to detect complex context vulnerabilities, resulting in high false negative rates and frequent false positives.
The vulnerability repair method based on self-attention mask is adopted, and the embedding sequence is encoded by generating code token embedding and utilizing the self-attention layer and feedforward neural network of the neural network, and applied to the model with the vulnerability mask, to achieve functional-level vulnerability repair.
It effectively improves the efficiency of vulnerability repair, can accurately identify and repair vulnerable code snippets, reduces false alarms and missed alarm rates, and realizes a fully automated vulnerability repair process.
Smart Images
Figure CN119862583A_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the field of security technology, and in particular to a vulnerability repair method, device and storage medium based on self-attention mask. Background Art
[0002] Software vulnerabilities are usually introduced into the system due to flawed security control design or developer errors during the implementation of software specifications. Such defects and errors are inevitable during the design and implementation phase of the software life cycle. Static software analysis techniques or vulnerability pattern libraries, matching rules, etc., to scan source code and identify potential errors or vulnerabilities, often have false positives. This is due to the limitations or incompleteness of rules and pattern matching algorithms, as well as the presence of vulnerabilities in inaccessible code. Static application security testing tools have difficulty detecting complex contextual vulnerabilities, resulting in a high false negative rate. Summary of the invention
[0003] The present invention proposes a vulnerability repair method, device and storage medium based on self-attention mask, which can at least solve one of the technical problems in the background technology.
[0004] To achieve the above object, the present invention adopts the following technical solutions:
[0005] A vulnerability repair method based on self-attention mask, comprising the following steps:
[0006] Step 101: Vulnerability repair coding generates a code token embedding for a given source code;
[0007] Step 102: Vulnerability repair decodes the input vulnerability query Q, where the purpose of each vulnerability query Q is to match and capture information about the vulnerability code token in a given source code;
[0008] Step 103: Initialize the first Q embedding E 0 This is the initial vulnerability query embedding vector, each element is a vector representing the initial embedding of each token in the vulnerability query; represents the initial embedding vector of the i-th token; these embedding vectors are used for subsequent encoding and decoding;
[0009] Step 104: For the data set data = {(c i ,y i ),i=1,...,m}, each fragile function c i =[t1,...,t m ] is a code token sequence, and the data set data contains multiple sample pairs (c i ,y i ), ci is the code token sequence of the i-th sample, y i is the label of the i-th sample, i.e., vulnerable or not vulnerable, c i is a code token sequence, expressed as [t1,……,tm], t j is the jth token in the sequence;
[0010] Code token vc j Belongs to a vulnerable statement and obtains the vulnerable scope or information of the code token; token-level vulnerable label label i =[l1,...,l m ] Among them l j =1 indicates code token vc j belong to vulnerable sentences; use the information to learn vulnerability masks;
[0011] Step 105: Apply the vulnerability mask to the model, i.e. replace the vulnerable parts with the repair tokens.
[0012] Furthermore, each token in step 101 is embedded into a vector through an embedding layer and input into the first encoder block, and a stack of encoder blocks are used to encode the representation of the embedded sequence through its self-attention layer and feedforward neural network. Each encoder block is described as follows:
[0013]
[0014] The hidden state from the previous encoder block is sequentially passed through a multi-head self-attention mechanism Two layers of feedforward neural network FNN and one layer of normalized GN are iterated until the final encoder hidden state is obtained to represent the vulnerable function, L is the number of encoder blocks applied, and Contains the embedding of the code token.
[0015] Further, step 102 specifically includes first performing an initial vulnerability query embedding And it is input into the subsequent multiple decoder blocks, in each block, the multi-head self-attention mechanism and the cross-attention mechanism are applied as follows:
[0016]
[0017] in For the encoder output, cross-attention CA can help cross-match Q embeddings and code token embedding; Q embedding Focus on and emphasize the vulnerable code token embedding in the vulnerable function, which contains enough information to generate repair tokens; get the output Q embedding Where U is the number of decoder blocks applied; based on the Q embedding, the predicted repair labels rep1,..., rep k , a linear layer is embedded on each Q The purpose is to predict rep1,...,rep k and rep k+1 , the mini-batch c is maximized by the following formula i The likelihood value of is:
[0018]
[0019] Where t1,...,t m is the source code, rep1,...,rep k For the corresponding repair patch.
[0020] Furthermore, in step 103, initialization By setting To start fixing the specific embedding of tokens, Through information initialization, the vulnerability repair problem is defined as the task of generating a repair patch in the source code. Given the source code c i =[t1,...,t m ], pass it through the vulnerability repair encoder, and calculate the encoder output Embed from the first Q Start, and provide it to the vulnerability repair decoder to generate the first repair token rep1; set Q embedding q1 = rep1, and provide it to the vulnerability repair decoder to generate the second repair token rep2, repeat until all reps are generated j .
[0021] Furthermore, the token-level vulnerable label label in step 104 i =[l1,...,l m ] Among them l j =1 indicates code token vc j Belong to vulnerable sentences; use information to learn vulnerability masks, use the pre-trained BERT model to learn vulnerability masks, c i Each token in is embedded into a vector and forwarded through the BERT architecture layers; a global max pooling layer and sigmoid activation are used to obtain the probability mask (c i ), and minimize relative to the mini-batch c i The cross entropy loss is:
[0022]
[0023] Apply the following transformation with a threshold of T
[0024]
[0025] When the predicted probability exceeds T, the model classifies the token as vulnerable; otherwise it is classified as benign.
[0026] Further, in step 105, in the scenario involving replacement, the vulnerability mask emphasizes replacing the vulnerable part with the repair token, and in the context of deletion, the mask emphasizes the vulnerable code to be deleted. The vulnerability mask is merged into the self-attention output of the encoder and the cross-attention output of the decoder. For the encoder, it is applied as follows:
[0027]
[0028] For the cross attention in the decoder, the following method is adopted:
[0029]
[0030] Encapsulates the vulnerability mask vulnerability repair encoder and vulnerability repair decoder, and finally outputs the vulnerability repair prediction result.
[0031] On the other hand, the present invention further discloses a computer-readable storage medium storing a computer program, wherein when the computer program is executed by a processor, the processor executes the steps of the above method.
[0032] On the other hand, the present invention further discloses a computer device, including a memory and a processor, wherein the memory stores a computer program, and when the computer program is executed by the processor, the processor executes the steps of the above method.
[0033] From the above technical solution, it can be seen that the present invention proposes a vulnerability repair method based on self-attention mask, generates code token embedding for a given source code, uses a neural network to encode the representation of the embedded sequence through its self-attention layer and feedforward neural network, matches the relevant vulnerability code through query, and implements function-level vulnerability repair by applying the vulnerability mask to the model. This method can locate vulnerable code tokens, implement software vulnerability matching and output vulnerability repair prediction results, effectively improving the efficiency of vulnerability repair.
[0034] The present invention can be used for software vulnerability code matching and vulnerability repair, and the present invention is fully automatic and does not rely on manual participation. The present invention realizes function-level vulnerability repair and effectively improves the efficiency of vulnerability repair.
[0035] Specifically, the advantages of the present invention are as follows:
[0036] 1. Fully automated vulnerability repair process
[0037] No need for human intervention: The present invention generates code token embedding, constructs vulnerability query Q in an automated way, and uses a self-attention mechanism to perform code matching and repair. The entire process does not require human intervention, greatly reducing the reliance on security experts and reducing labor costs.
[0038] Real-time response: Since it is a fully automated process, the system can start the repair process immediately after detecting a vulnerability, shortening the time from vulnerability discovery to repair completion and improving the security of the system.
[0039] 2. High-precision vulnerability matching and repair
[0040] Precise matching based on self-attention mechanism: The present invention uses a self-attention mechanism to process code token sequences, which can capture contextual information and semantic relationships in the code. This mechanism enables the system to more accurately identify vulnerable code snippets and avoid false positives and negatives.
[0041] Functional-level vulnerability repair: The present invention can not only identify vulnerabilities, but also repair them at the code level. By generating repair suggestions for specific vulnerabilities or directly applying repair patches, it ensures that the repaired code not only solves the current vulnerability but also maintains the original functional logic, which helps to improve the quality and reliability of the repair.
[0042] 3. Efficient vulnerability repair efficiency
[0043] Parallel processing capability: The present invention can process multiple vulnerability queries Q at the same time, taking advantage of parallel computing to quickly scan and repair large amounts of code. Compared to the traditional line-by-line checking method, the processing speed of the present invention is significantly improved, especially when processing large code bases.
[0044] Reduce duplication of effort: The automated repair process avoids duplication of effort and human errors that may occur in manual repair. The system can learn best practices based on historical repair records, further optimize repair strategies, and improve overall efficiency.
[0045] 4. Highly adaptable, suitable for multiple programming languages and frameworks
[0046] Cross-language support: The design of the present invention has good versatility and is applicable to a variety of programming languages and development frameworks. By tokenizing codes in different languages, the system can flexibly respond to different types of holes and meet diverse needs.
[0047] Scalability: As new vulnerabilities continue to emerge, the present invention can quickly adapt to new security threats by updating model parameters and training data sets. This scalability enables the system to maintain efficiency and accuracy over the long term.
[0048] 5. Enhanced security and compliance
[0049] Reduce security risks: Through automated vulnerability repair, the system can repair known vulnerabilities in the first place, reducing the time window in which the system is exposed to attackers, thereby effectively reducing security risks. Comply with industry standards: The present invention follows the best practices of secure development to ensure that the repaired code meets industry standards and compliance requirements. This is particularly important for enterprises, especially in industries with high security requirements such as finance and healthcare.
[0050] 6. Continuous learning and improvement
[0051] Feedback-based learning mechanism: The present invention can continuously optimize the model by collecting feedback on the effect of the repaired code. The system will adjust the parameters according to the actual repair results, and gradually improve the accuracy and success rate of the repair; Self-evolution: As time goes by, the system will accumulate more vulnerability repair experience and form a self-evolving knowledge base. This enables the system to better cope with future security challenges and maintain technological leadership. BRIEF DESCRIPTION OF THE DRAWINGS
[0052] Figure 1 It is a flowchart of the present invention. DETAILED DESCRIPTION
[0053] In order to make the purpose, technical solutions and advantages of the embodiments of the present invention clearer, the technical solutions in the embodiments of the present invention will be clearly and completely described below in conjunction with the drawings in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments.
[0054] like Figure 1 As shown, the vulnerability repair method based on self-attention mask described in this embodiment includes the following steps:
[0055] Step 101: Vulnerability Fix Encoding Generate code token embeddings for the given source code. Each token is embedded into a vector through the embedding layer and input to the first encoder block. A stack of encoder blocks are used to encode the representation of the embedded sequence through its self-attention layer and feed-forward neural network. Each encoder block can be described as follows:
[0056]
[0057] The hidden state from the previous encoder block is sequentially passed through a multi-head self-attention mechanism Two layers of feedforward neural network FNN and one layer of normalized GN. Iterate until the final encoder hidden state is obtained to represent the vulnerable function. L is the number of encoder blocks applied, and Contains the embedding of the code token.
[0058] Step 102: Vulnerability repair decodes the input vulnerability query Q, and the purpose of each Q is to match and capture information about the vulnerability code token in the given source code. First, perform the initial vulnerability query embedding And it is input into the subsequent multiple decoder blocks. In each block, the multi-head self-attention mechanism and cross-attention mechanism are applied as follows:
[0059]
[0060] in is the encoder output. Cross-attention CA can help cross-match Q embedding and code token embedding. Q embedding Focusing on and emphasizing the vulnerable code token embeddings in the vulnerable function, it contains enough information to generate repair tokens. The output Q embedding is obtained Where U is the number of decoder blocks applied. Based on the Q embedding, the predicted repair labels rep1,..., rep k A linear layer is embedded on each Q The purpose is to predict rep1,...,rep k and rep k+1 , the mini-batch c is maximized by the following formula i The likelihood value of is:
[0061]
[0062] Where t1,...,t m is the source code, rep1,...,rep k For the corresponding repair patch.
[0063] Step 103: Initialize the first Q embedding Initialize the first Q embedding E 0 This is the initial vulnerability query embedding vector, each element is a vector representing the initial embedding of each token in the vulnerability query; represents the initial embedding vector of the i-th token; these embedding vectors are used for subsequent encoding and decoding;
[0064] Unlike some object detection methods, the first Q embedding is not randomly initialized. In fact, the initialization By setting To start fixing the specific embedding of tokens, Through information initialization, the vulnerability repair problem is defined as the task of generating repair patches in the source code. Given the source code c i =[t1,...,t m ], pass it through the vulnerability repair encoder, and calculate the encoder output Embed from the first Q Start, and provide it to the vulnerability repair decoder to generate the first repair token rep1. Set Q embedding q1 = rep1, and provide it to the vulnerability repair decoder to generate the second repair token rep2, repeat until all reps are generated j .
[0065] Step 104: For data data={(c i ,y i ),i=1,...,m}, each fragile function c i =[t1,...,t m ] is a code token sequence, and the data set data contains multiple sample pairs (c i ,y i ), c i is the code token sequence of the i-th sample, y i is the label of the i-th sample, i.e., vulnerable or not vulnerable, c i is a code token sequence, expressed as [t1,……,tm], t j is the jth token in the sequence; if a code token vc j For a vulnerable statement, and get the vulnerable scope or information of the code token. Token-level vulnerable label label i =[l1,...,l m ] Among them l j =1 indicates code token vc j The vulnerable sentences are used to learn the vulnerability mask. The vulnerability mask is learned by using the pre-trained BERT model. i Each token in is embedded into a vector and forwarded through the BERT architecture layers. A global max pooling layer and sigmoid activation are used to obtain the probability mask (c i ), and minimize relative to the mini-batch c i The cross entropy loss is:
[0066]
[0067] Apply the following transformation with a threshold of T
[0068]
[0069] When the predicted probability exceeds T, the model classifies the token as vulnerable; otherwise it is classified as benign.
[0070] Step 105: Apply the vulnerability mask to the model. In scenarios involving replacement, the vulnerability mask emphasizes replacing the vulnerable parts with repair tokens. In the context of deletion, the mask emphasizes the vulnerable code to be deleted. The vulnerability mask is merged into the self-attention output of the encoder and the cross-attention output of the decoder. For the encoder, it is applied as follows:
[0071]
[0072] For the cross attention in the decoder, the following method is adopted:
[0073]
[0074] Encapsulates the vulnerability mask vulnerability repair encoder and vulnerability repair decoder, and finally outputs the vulnerability repair prediction result.
[0075] In summary, the embodiment of the present invention proposes a vulnerability repair method based on self-attention mask, generates code token embedding for a given source code, uses a neural network to encode the representation of the embedded sequence through its self-attention layer and feedforward neural network, matches the relevant vulnerability code through query, and implements function-level vulnerability repair by applying the vulnerability mask to the model. This method can locate vulnerable code tokens, implement software vulnerability matching and output vulnerability repair prediction results, effectively improving the efficiency of vulnerability repair.
[0076] On the other hand, the present invention further discloses a computer-readable storage medium storing a computer program, wherein when the computer program is executed by a processor, the processor executes the steps of the above method.
[0077] On the other hand, the present invention further discloses a computer device, including a memory and a processor, wherein the memory stores a computer program, and when the computer program is executed by the processor, the processor executes the steps of the above method.
[0078] In another embodiment provided by the present application, a computer program product including instructions is also provided. When the computer program product is run on a computer, the computer executes any vulnerability repair method based on self-attention mask in the above embodiments.
[0079] It is understandable that the system, device and storage medium provided in the embodiments of the present invention correspond to the method provided in the embodiments of the present invention, and the explanation, examples and beneficial effects of the relevant contents can refer to the corresponding parts in the above methods.
[0080] In the above embodiments, it can be implemented in whole or in part by software, hardware, firmware or any combination thereof. When implemented using software, it can be implemented in whole or in part in the form of a computer program product. The computer program product includes one or more computer instructions. When the computer program instructions are loaded and executed on a computer, the process or function described in the embodiment of the present application is generated in whole or in part. The computer may be a general-purpose computer, a special-purpose computer, a computer network, or other programmable device. The computer instructions may be stored in a computer-readable storage medium, or transmitted from one computer-readable storage medium to another computer-readable storage medium, for example, the computer instructions may be transmitted from a website site, computer, server or data center by wired (e.g., coaxial cable, optical fiber, digital subscriber line (DSL)) or wireless (e.g., infrared, wireless, microwave, etc.) mode to another website site, computer, server or data center. The computer-readable storage medium may be any available medium that a computer can access or a data storage device such as a server or data center that includes one or more available media integrated. The available medium may be a magnetic medium, (e.g., a floppy disk, a hard disk, a tape), an optical medium (e.g., a DVD), or a semiconductor medium (e.g., a solid-state drive Solid State Disk (SSD)), etc.
[0081] It should be noted that, in this article, relational terms such as first and second, etc. are only used to distinguish one entity or operation from another entity or operation, and do not necessarily require or imply any such actual relationship or order between these entities or operations. Moreover, the terms "include", "comprise" or any other variants thereof are intended to cover non-exclusive inclusion, so that a process, method, article or device including a series of elements includes not only those elements, but also other elements not explicitly listed, or also includes elements inherent to such process, method, article or device. In the absence of further restrictions, the elements defined by the sentence "comprise a ..." do not exclude the existence of other identical elements in the process, method, article or device including the elements.
[0082] Each embodiment in this specification is described in a related manner, and the same or similar parts between the embodiments can be referred to each other, and each embodiment focuses on the differences from other embodiments. In particular, for the system embodiment, since it is basically similar to the method embodiment, the description is relatively simple, and the relevant parts can be referred to the partial description of the method embodiment.
[0083] The above embodiments are only used to illustrate the technical solutions of the present invention, rather than to limit the same. Although the present invention has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that the technical solutions described in the aforementioned embodiments may still be modified, or some of the technical features may be replaced by equivalents. However, these modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the embodiments of the present invention.
Claims
1. A vulnerability repair method based on self-attention mask, characterized in that: The following steps are included: Step 101: Vulnerability fix coding generates code token embedding for a given source code; Step 102: Vulnerability repair decodes the input vulnerability query Q, where the purpose of each vulnerability query Q is to match and capture information about the vulnerability code token in a given source code; Step 103: Initialize the first Q embedding E 0 This is the initial vulnerability query embedding vector, each element is a vector representing the initial embedding of each token in the vulnerability query; Represents the initial embedding vector of the i-th token; These embedding vectors are used for subsequent encoding and decoding; Step 104: For the data set data = {(c i ,y i ),i=1,...,m}, each fragile function c i =[t1,...,t m ] is a code token sequence, and the data set data contains multiple sample pairs (c i ,y i ), c i is the code token sequence of the i-th sample, y i is the label of the i-th sample, i.e., vulnerable or not vulnerable, c i is a code token sequence, expressed as [t1,……,tm], t j is the jth token in the sequence; code tokenvc j Belongs to a vulnerable statement and obtains the vulnerable scope or information of the code token; token-level vulnerable label label i =[l1,...,l m ] Among them l j =1 indicates code tokenvc j belong to vulnerable sentences; use the information to learn vulnerability masks; Step 105: Apply the vulnerability mask to the model, i.e. replace the vulnerable parts with the repair tokens.
2. The vulnerability repair method based on self-attention mask according to claim 1, characterized in that: In step 101, each token is embedded into a vector through an embedding layer and input into the first encoder block. A stack of encoder blocks are used to encode the representation of the embedded sequence through its self-attention layer and feed-forward neural network. Each encoder block is described as follows: The hidden state from the previous encoder block is sequentially passed through a multi-head self-attention mechanism Two layers of feedforward neural network FNN and one layer of normalized GN are iterated until the final encoder hidden state is obtained to represent the vulnerable function, L is the number of encoder blocks applied, and Contains the embedding of the code token.
3. The vulnerability repair method based on self-attention mask according to claim 2, characterized in that: Step 102 specifically includes first performing an initial vulnerability query embedding And it is input into the subsequent multiple decoder blocks, in each block, the multi-head self-attention mechanism and the cross-attention mechanism are applied as follows: in For the encoder output, cross-attention CA helps cross-match Q embeddings and code token embedding; Q embedding Focus on and emphasize the vulnerable code token embedding in the vulnerable function, which contains enough information to generate repair tokens; get the output Q embedding Where U is the number of decoder blocks applied; based on the Q embedding, the predicted repair labels rep1,..., rep k , a linear layer is embedded on each Q The purpose is to predict rep1,...,rep k and rep k+1 , the mini-batch c is maximized by the following formula i The likelihood value of is: Where t1,...,t m is the source code, rep1,...,rep k For the corresponding repair patch.
4. The vulnerability repair method based on self-attention mask according to claim 3, characterized in that: Initialization in step 103 By setting To start fixing the specific embedding of tokens, Through information initialization, the vulnerability repair problem is defined as the task of generating a repair patch in the source code. Given the source code c i =[t1,...,t m ], pass it through the vulnerability repair encoder, and calculate the encoder output Embed from the first Q Start, and provide it to the vulnerability repair decoder to generate the first repair token rep1; set Q embedding q1 = rep1, and provide it to the vulnerability repair decoder to generate the second repair token rep2, repeat until all reps are generated j .
5. The vulnerability repair method based on self-attention mask according to claim 4, characterized in that: The label of the token-level vulnerable tag in step 104 i =[l1,...,l m ] Among them l j =1 indicates code token vc j Belong to vulnerable sentences; use information to learn vulnerability masks, use the pre-trained BERT model to learn vulnerability masks, c i Each token in is embedded into a vector and forwarded through the BERT architecture layers; Use the global maximum pooling layer and sigmoid activation to obtain the probability mask (c i ), and minimize relative to the mini-batch c i The cross entropy loss is: Apply the following transformation with a threshold of T When the predicted probability exceeds T, the model classifies the token as vulnerable; otherwise it is classified as benign.
6. The vulnerability repair method based on self-attention mask according to claim 5, characterized in that: In step 105, in the scenario involving replacement, the vulnerability mask emphasizes replacing the vulnerable part with the repair token. In the context of deletion, the mask emphasizes the vulnerable code to be deleted. The vulnerability mask is merged into the self-attention output of the encoder and the cross-attention output of the decoder. For the encoder, it is applied as follows: For the cross attention in the decoder, the following method is adopted: Encapsulates the vulnerability mask vulnerability repair encoder and vulnerability repair decoder, and finally outputs the vulnerability repair prediction result.
7. A computer-readable storage medium storing a computer program, characterized in that: When the computer program is executed by a processor, the processor is caused to perform the steps of the method according to any one of claims 1 to 6.
8. A computer device comprising a memory and a processor, wherein the memory stores a computer program, wherein: When the computer program is executed by the processor, the processor is caused to perform the steps of the method according to any one of claims 1 to 6.