Vulnerability patch generation method and system based on retrieval enhancement and storage medium

Abstract

The invention discloses a vulnerability patch generation method and system based on retrieval enhancement and a storage medium. The method comprises the following steps: S1, constructing a repair template retriever: extracting a repair template from a vulnerability training set, storing mapping of the repair template and historical repair submission, and constructing a template database; matching the defect code snippets with templates, and retrieving a group of related historical repair behaviors through the successfully matched templates; s2, a patch generator training stage: training the similar historical repair fragments retrieved in the step S1 to obtain a patch generation model; and S3, generating a vulnerability patch by using the patch generation model in the step S2. The system and the storage medium are both used for executing the method. The method has the advantages that the principle is simple, the application range is wide, and the repairing template can be automatically used for repairing.

Description

Technical Field

[0001] The present invention mainly relates to the technical field of vulnerability patch generation, and in particular to a vulnerability patch generation method, system and storage medium based on retrieval enhancement. Background Art

[0002] Automatic vulnerability repair technology (AVR) includes three main stages: vulnerability location, patch generation, and patch verification. Among them, the quality of the patch generated by the repair program directly determines the repair result and is crucial to the entire repair process. Therefore, many existing automatic vulnerability repair technologies are committed to generating more accurate patches.

[0003] According to the technology and technical route adopted in the patch generation stage, the existing automatic vulnerability repair technical methods can be divided into the following four types: (1) Search-based methods, also known as heuristic solutions, search for the correct patch in a predefined patch space; (2) Constraint-based methods guide the repair process by formulating a set of constraint specifications. These methods target vulnerabilities with clear specifications, such as buffer overflows and out-of-bounds reads; (3) Template-based methods use predefined program repair templates to generate repair patches; (4) Learning-based methods, also known as data-driven inpainting solutions.

[0004] Due to the low efficiency of search-based methods and the limited application scenarios of constraint-based methods, practitioners' research directions are mainly focused on template-based and learning-based AVR technologies.

[0005] The basic idea of ​​template-based automatic repair technology is to abstract a series of code changes into a template that can be applied to the code with defects. Therefore, the template-based automatic repair process includes utilizing the context information of the defective code node to match the context constraints defined in the given repair template.

[0006] With the booming development of deep learning technology, some practitioners have used the power of language models to fix vulnerabilities. Specifically, the learning-based AVR technology treats the vulnerability repair task as a code-to-code translation task and learns repair actions from historical data.

[0007] As can be seen from the above, the above technologies are heavily dependent on the quantity and quality of historical vulnerability repair data, and it is difficult for existing vulnerability datasets to meet this demand. The current mainstream public vulnerability datasets Big-Vul and CVEFixes contain 3,754 and 5,465 pairs of vulnerabilities, respectively, covering 91 and 180 CWE types. However, there is redundancy in these datasets, which makes it difficult to achieve the data diversity required for learning-based vulnerability repair, thus bringing challenges to the study of feasible learning-based vulnerability repair solutions with limited datasets.

[0008] Therefore, in learning-based automatic vulnerability repair technology, the problem of optimizing the use of existing vulnerability datasets becomes a crucial technical issue. By making full use of existing vulnerability datasets, retrieving relevant code snippets and establishing connections between similar data, it is helpful to generate the required code, which can reduce the search space and provide the necessary code in the code generation process. This is consistent with human thinking, that is, when faced with a vulnerability that needs to be repaired, experienced developers usually first search for relevant code snippets in other large-scale code repositories (such as Github) as references. Developers have accumulated rich knowledge and experience in the development and maintenance stages, and usually produce various documents to classify problems encountered in team collaboration. This enables them to easily access similar code snippets or information to help their repair tasks. However, it is a challenge for automatic vulnerability repair tools to establish a mapping between vulnerabilities and their corresponding code snippets for the patch generation stage.

[0009] To this end, the above traditional technologies still have some shortcomings: 1. Existing template-based automatic vulnerability repair methods are all targeted at specific types of vulnerabilities; 2. Existing template-based automatic repair methods have not achieved true automation. Although template mining technology has greatly reduced the workload of manually summarizing repair templates, the mined templates cannot be directly used in the patch generation process. Developers are still required to apply templates to patch generation. Summary of the invention

[0010] The technical problem to be solved by the present invention is: in response to the technical problems existing in the prior art, the present invention provides a retrieval-enhanced vulnerability patch generation method, system and storage medium which have a simple principle, a wide range of applications and can automatically use repair templates for repair.

[0011] In order to solve the above technical problems, the present invention adopts the following technical solutions: A vulnerability patch generation method based on retrieval enhancement, comprising: Step S1: Construction of repair template retriever: Extract repair templates from the vulnerability training set, store the mapping between repair templates and historical repair submissions, and build a template database; match defect code snippets with templates, and retrieve a set of related historical repair behaviors through successfully matched templates; Step S2: Patch generator training phase: train similar historical repair fragments retrieved from step S1 to obtain a patch generation model; Step S3: Generate a vulnerability patch using the patch generation model of step S2.

[0012] As a further improvement of the method of the present invention: the step S1 comprises: Step S11: converting code changes into a rich editing script; Step S12: clustering and extracting repair templates; Step S13: Repair template database construction; Step S14: Matching the defect code with the repair template; Step S15: Combine the defect code with the repair pair retrieved from the template to form training data.

[0013] As a further improvement of the method of the present invention: in the step S11, all code modifications are converted into a rich editing script RES, and the rich editing script RES uses four operations and node types in AST to represent the code modifications as editing scripts; the rich editing script RES retains the context details of the code changes, thereby providing richer information for template matching; in the step S12, a set of trees with the same structure is identified based on the collected rich editing script RES; when the set contains at least two trees, a subtree structure of a set of identical modification operations is extracted to form a template; each template is an abstraction of several code changes in the same cluster, and encapsulates the same repair modifications applied to the same type of nodes.

[0014] As a further improvement of the method of the present invention: in step S13, the repair template is an abstraction of a group of similar repair behaviors, and a dictionary is established to store the mapping relationship between the template and the relevant vulnerability repair pair group, wherein the template is used as the key and the vulnerability repair pair group is used as the value.

[0015] As a further improvement of the method of the present invention: in the step S14, for the defective function Vb to be repaired, it is first matched level by level with the templates in the mined template library at the level of the abstract syntax tree; after the match is successful, a set of related historical code changes are accessed through the dictionary constructed in the template mining process; the modified statements and the surrounding code lines are extracted from the defective code and the repaired code to form a retrieved historical repair pair for the subsequent patch generation stage.

[0016] As a further improvement of the method of the present invention: the step S2 comprises: Step S21: model initialization; Step S22: using the defect function and the retrieved vulnerability repair pair as model input; Step S23: Code segmentation of input; training input<Vt , Ft> , the code is tokenized at the subword level through byte pair encoding, and then the embedding layer generates an embedding vector for each subword to capture the semantic features of the code token; the word embedding result is used as the input of the encoder module to derive the hidden state of the encoder, which in turn notifies the decoder module to generate and output the predicted patch; Step S24: Calculate the cross entropy loss function; the cross entropy loss function is used in the training stage; Step S25: Back propagation, update parameters.

[0017] As a further improvement of the method of the present invention: in the step S21, CodeT5 is used as the model of the patch generation module, and CodeT5 adopts an encoder-decoder architecture, including 12 Transformer layers of the encoder and decoder respectively; each layer uses 12 attention heads for multi-head attention calculation.

[0018] As a further improvement of the method of the present invention: in step S22, given the vulnerability pair from the training data set<Vt,Ft> , namely the defect function Vt and its correct fix Ft. The template index retriever outputs several related vulnerability fix code changes when the template is successfully matched.<Vi,Fi> ,<Vj,Fj> …; During the training phase of the patch generator, all retrieved vulnerability-fix pairs, Vt, are provided as input to the patch generator to generate predicted patches.

[0019] The present invention further provides a method and system for generating vulnerability patches based on retrieval enhancement, comprising a memory, a processor and a computer program stored in the memory, wherein the processor executes the computer program to implement the steps of any one of the above methods.

[0020] The present invention further provides a storage medium, which can be read by a computer or a processor, and stores a computer program for executing any one of the above methods.

[0021] Compared with the prior art, the advantages of the present invention are: 1. Since the existing template-based automatic vulnerability repair methods are all targeted at specific types of vulnerabilities, the vulnerability patch generation method, system and storage medium based on retrieval enhancement of the present invention are a general type of vulnerability patch generation method, which is used to generate correct patches for multiple CWE type vulnerabilities. In the present invention, the template is used to retrieve the historical repair behavior related to the defective code fragment, and is directly used for patch generation, so it is applicable to all types of vulnerabilities, and is a patch generation method for general types of vulnerabilities.

[0022] 2. The vulnerability patch generation method, system and storage medium based on retrieval enhancement of the present invention realize an automated vulnerability patch generation method, which realizes the automatic application of repair templates in the patch generation stage without manual coding, thus reducing the manual burden. However, the existing template-based automatic repair method has not achieved true automation. Although the template mining technology has greatly reduced the workload of manually summarizing the repair template, the mined template cannot be directly used in the patch generation process, and developers are still required to apply the template to patch generation.

[0023] 3. The retrieval-enhanced vulnerability patch generation method, system and storage medium of the present invention, by constructing a retriever based on a repair template, alleviates the limitation of the current learning-based automatic repair method that is overly dependent on the quantity and quality of historical vulnerability repair data, and makes full use of the historical repair data in the existing data set. By retrieving related code snippets and establishing connections between similar data, the search space can be reduced and necessary code can be provided in the code generation process, thereby generating higher quality vulnerability patches. BRIEF DESCRIPTION OF THE DRAWINGS

[0024] Figure 1 It is a schematic diagram of the process of the method of the present invention in a specific application example.

[0025] Figure 2 It is a schematic diagram of the process of constructing a repair template retriever in a specific application example of the present invention.

[0026] Figure 3 It is a flowchart of the patch generator training phase in a specific application example of the present invention.

[0027] Figure 4 It is a schematic diagram of a repair template and its corresponding code change in a specific application example of the present invention. DETAILED DESCRIPTION

[0028] The present invention will be further described in detail below with reference to the accompanying drawings and specific embodiments.

[0029] like Figure 1-Figure 3 As shown, the present invention provides a vulnerability patch generation method based on retrieval enhancement, and the process includes: Step S1: Repair template retriever construction; Extract repair templates from the vulnerability training set, store the mapping between repair templates and historical repair submissions, and build a template database; Match the defective code snippet with the template, and retrieve a set of related historical repair behaviors through the successfully matched template; Step S2: patch generator training phase; Step S3: Generate a vulnerability patch using the patch generation model of step S2.

[0030] The patch generation model is obtained by training on similar historical repair fragments retrieved from step S1.

[0031] See also Figure 2 In a specific application example, step S1 may include the following process according to actual needs: Step S11: Convert all code changes into rich editing scripts; Convert all code modifications into Rich Edit Scripts. RES uses the four operations (i.e., update, insert, delete, and move) and node types in AST to represent code modifications as edit scripts.

[0032] Among them, RES preserves the contextual details of code changes, such as AST structure and changed node types, thus providing richer information for template matching.

[0033] Step S12: clustering and extracting repair templates; For trees with the same structure, based on the collected RES, the present invention identifies a set of trees with the same structure.

[0034] When the set includes at least two trees, the present invention extracts a group of subtree structures with the same modification operations therein to form a template.

[0035] Each template is an abstraction of several code changes in the same cluster and encapsulates the same fix modifications applied to nodes of the same type.

[0036] Step S13: Repair template database construction; A repair template is an abstraction of a group of similar repair behaviors. The present invention establishes a dictionary for storing the mapping relationship between the template and the related vulnerability repair pair group, wherein the template is used as the key and the vulnerability repair pair group is used as the value.

[0037] Step S14: matching the defect code with the repair template; For the defective function Vb to be repaired, the present invention first matches it with the templates in the template library mined by the present invention at the level of the abstract syntax tree step by step; after the match is successful, the present invention can access a set of related historical code changes through the dictionary constructed in the template mining process.

[0038] Given that the differences between the defective code and the repaired code usually appear in consecutive code lines, which contain key information of the patch, the present invention extracts the modified statements and the surrounding code lines from the defective code and the repaired code to form retrieved historical repair pairs for the subsequent patch generation stage.

[0039] Step S15: Combine the defect code with the repair pair retrieved from the template to form training data.

[0040] See also Figure 4 , is a schematic diagram of a repair template and its corresponding code changes, showing an example of a template and two related vulnerability repair pairs expressed in GNU diff format. It can be observed that the present invention mines a template that modifies the expression in the if statement condition, and Figure 4 The two bug-fixing diff patches in (b) of FIG. 1 both involve deleting the operator and updating the calling function in the condition. The template mining tool of the present invention captures this feature and extracts the common edits on the same type of nodes to form Figure 4 The template shown in (a) in FIG. 1 is a template shown in (a) in FIG. 1 . The dictionary constructed by the present invention can be directly Figure 4 The template in (a) is retrieved Figure 4 (b) shows two vulnerability repair pairs. The present invention establishes a dictionary for storing the mapping relationship between templates and related vulnerability repair pairs, where the template is used as a key and the vulnerability repair pair is used as a value.

[0041] See also Figure 3 In a specific application example, step S2 may include the following process according to actual needs: Step S21: model initialization; As a preferred embodiment, the present invention uses CodeT5 as the model of the patch generation module. CodeT5 adopts an encoder-decoder architecture, including 12 Transformer layers for the encoder and decoder respectively. Each layer uses 12 attention heads for multi-head attention calculation, and the total parameter size is 220M. This architecture is widely used in the most advanced pre-trained models and various application scenarios.

[0042] As a preferred embodiment, the present invention initializes the parameters of each generator using the weights of CodeT5. The purpose of this initialization is to provide the generator with specific knowledge in the field of programming languages, which can help the present invention generate more reliable patches. Subsequently, these generators were fine-tuned for the patch generation task.

[0043] Step S22: using the defect function and the retrieved vulnerability repair pair as model input; Given a vulnerability pair from the training dataset<Vt, Ft> , that is, the defect function Vt and its correct fix Ft. The template index retriever will output several related vulnerability fix code changes if the template is successfully matched.<Vi,Fi> ,<Vj,Fj> ….

[0044] During the training phase of the patch generator, the present invention provides all retrieved vulnerability repair pairs, as well as Vt, as input to the patch generator to generate predicted patches.

[0045] The present invention further uses a search field R to represent the retrieved code change pair.<Vi,Fi> ,<Vj,Fj> …, using the special marker [SEP] to separate different pairs. That is: Rt =<Vi,Fi> [SEP]<Vj , Fj> …. For a given pair of training data<Vt,Ft> , the input of the patch generator of the present invention is enhanced to: It = VtRt.

[0046] In addition, in the specific application process, there are some vulnerabilities that cannot match any mined templates, which indicates that similar historical fixes cannot be retrieved from the code base for reference during the patch generation stage. In this case, the search field R will be NULL, and the input of the patch generator of the present invention will only include vulnerable code, that is: It=Vt.

[0047] After concatenating and enhancing the search field with the defect function input, the present invention uses the enhanced training data for the patch generator. This is expressed as: It = Vt[RETR] Vi[SEP]Fi…, where [RETR] is a special marker used to separate the retrieved vulnerability repair pairs from the original vulnerable code, and [SEP] is a special marker used to separate the vulnerability repair pairs.

[0048] Step S23: code segmentation of the input; Training input pairs<Vt,Ft> ; First, the code is tokenized at the subword level through byte pair encoding (BPE), and then the embedding layer generates an embedding vector for each subword to capture the semantic features of the code token. The word embedding result is used as the input of the encoder module of the present invention to derive the hidden state of the encoder, which in turn notifies the decoder module to generate and output the predicted patch.

[0049] Step S24: Calculate the cross entropy loss function; The main goal of model training is to fine-tune the model to obtain specific weights for the patch generation task on the vulnerability dataset.

[0050] Therefore, the present invention adopts the cross entropy loss function in the training stage, and its loss function is:

[0051] Where N represents the tokens in the patch. are the parameters of the model, I is the baseline value, and Ft is the tokens generated so far; Step S25: Back propagation, update parameters.

[0052] The present invention further includes a derivation phase for patch generation, during which the present invention adopts a beam search strategy to generate a list of candidate repair patches for a given defective function. Specifically, for a given vulnerable function Vd, Tree-sitter, a widely used AST parser generator, is used to build an AST and extract its defective node list. Then, each defective code is matched with the mined template. The present invention combines the defective function Vd with the retrieved code changes mapped from the matching template. Then, the present invention splits Vd using a tokenization component and inputs the subword tokenized sequence into the training model. The model extracts and decodes the features of the sequence and then generates a probability distribution of the tokens that constitute the candidate patches. The present invention implements beam search to select multiple candidate patches based on the probability of each decoding step. That is, the beam size parameter k is determined to select the patches predicted by the Top-k probability in the encoding step. These candidate patches will then be verified for their correctness.

[0053] In summary, the present invention combines the template-based automatic vulnerability repair technology with the deep learning-based automatic repair technology to achieve complementary advantages of the two types of repair methods: (1) Among them, the repair template plays the role of retrieving similar historical repairs. It clusters a group of repair code changes that perform similar repair operations on the same abstract syntax tree (AST) node, extracts their same repair operations to form a template, which is an abstraction of a group of historical code changes with similar repair behaviors. By matching the defect function to be repaired with the template at the AST level, a group of historical repair behaviors corresponding to the template are retrieved through the successfully matched template for subsequent patch generation; (2) The existing deep learning-based automatic repair technology shows that pre-trained models (such as CodeT5) have good application prospects in patch generation. However, they only treat the automatic repair task as a translation task and do not target the patch generation task itself. The present invention uses a repair template that represents a set of historical repair code changes to guide model training, making it more targeted at repair tasks.

[0054] The present invention is actually a template-guided retrieval-enhanced patch generation method, which uses repair templates to retrieve relevant vulnerability repair pairs. These templates are high-level abstractions of a set of similar repair operations, and then use the retrieved repair pairs to guide patch generation, thereby realizing a general type of vulnerability patch generation method, and realizing the generation of correct patches for multiple CWE type vulnerabilities. At the same time, by using templates for retrieval, the repair templates are automatically applied to the repair process, and the patch generation process is automated.

[0055] The present invention further provides a vulnerability patch generation system based on retrieval enhancement, comprising a memory, a processor and a computer program stored in the memory, wherein the processor executes the computer program to implement the steps of the above method.

[0056] The present invention further provides a storage medium, which can be read by a computer or a processor and stores a computer program for executing the above method.

[0057] Those skilled in the art should understand that the above-mentioned embodiments of the present application can be provided as methods, systems, or computer program products. Therefore, the present application can take the form of a complete hardware embodiment, a complete software embodiment, or an embodiment combining software and hardware. Moreover, the present application can take the form of a computer program product implemented on one or more computer-readable storage media (including but not limited to disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code. The present application is described with reference to the flowcharts and / or block diagrams of the methods, devices (systems), and computer program products according to the embodiments of the present application. It should be understood that each process and / or box in the flowchart and / or block diagram, as well as the combination of the processes and / or boxes in the flowchart and / or block diagram, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, a special-purpose computer, an embedded processor, or other programmable data processing device to produce a machine, so that the instructions executed by the processor of the computer or other programmable data processing device generate instructions for implementing the processes in the process Figure 1 A process or multiple processes and / or boxes Figure 1These computer program instructions can also be stored in a computer-readable memory that can guide a computer or other programmable data processing device to work in a specific way, so that the instructions stored in the computer-readable memory produce a product including an instruction device, which implements the functions specified in the process. Figure 1 A process or multiple processes and / or boxes Figure 1 These computer program instructions can also be loaded onto a computer or other programmable data processing device, so that a series of operation steps are executed on the computer or other programmable device to produce a computer-implemented process, so that the instructions executed on the computer or other programmable device provide for implementing the process in the process. Figure 1 A process or multiple processes and / or boxes Figure 1 The steps for the functions specified in one or more boxes.

[0058] The above are only preferred embodiments of the present invention, and the protection scope of the present invention is not limited to the above embodiments. All technical solutions under the concept of the present invention belong to the protection scope of the present invention. It should be pointed out that for ordinary technicians in this technical field, some improvements and modifications without departing from the principle of the present invention should be regarded as the protection scope of the present invention.

Claims

1. A vulnerability patch generation method based on retrieval enhancement, It is characterized in that include: Step S1: Repair template retriever construction: extract repair templates from the vulnerability training set, store the mapping between repair templates and historical repair submissions, and build a template database; Match the defective code snippet with the template, and retrieve a set of related historical repair behaviors through the successfully matched template; Step S2: Patch generator training phase: train similar historical repair fragments retrieved from step S1 to obtain a patch generation model; Step S3: Generate a vulnerability patch using the patch generation model of step S2.

2. The method for generating vulnerability patches based on retrieval enhancement according to claim 1, It is characterized in that The step S1 comprises: Step S11: converting code changes into a rich editing script; Step S12: clustering and extracting repair templates; Step S13: Repair template database construction; Step S14: Matching the defect code with the repair template; Step S15: Combine the defect code with the repair pair retrieved from the template to form training data.

3. The method for generating vulnerability patches based on retrieval enhancement according to claim 2, It is characterized in that In the step S11, all code modifications are converted into a rich editing script RES, and the rich editing script RES uses four operations and node types in AST to represent the code modifications as editing scripts; the rich editing script RES retains the context details of the code changes, thereby providing richer information for template matching; in the step S12, a set of trees with the same structure is identified based on the collected rich editing script RES; when the set contains at least two trees, a subtree structure of a set of identical modification operations is extracted to form a template; each template is an abstraction of several code changes in the same cluster, and encapsulates the same repair modifications applied to the same type of nodes.

4. The method for generating vulnerability patches based on retrieval enhancement according to claim 2, It is characterized in that In step S13, the repair template is an abstraction of a group of similar repair behaviors, and a dictionary is established to store the mapping relationship between the template and the relevant vulnerability repair pair group, wherein the template is used as the key and the vulnerability repair pair group is used as the value.

5. The method for generating vulnerability patches based on retrieval enhancement according to claim 2, It is characterized in that In step S14, for the defective function Vb to be repaired, firstly, it is matched with the templates in the mined template library at the level of the abstract syntax tree step by step; after the match is successful, a set of related historical code changes is accessed through the dictionary constructed in the template mining process; The modified statements and surrounding code lines are extracted from the defective code and the fixed code to form retrieved history-fix pairs for the subsequent patch generation stage.

6. The method for generating vulnerability patches based on retrieval enhancement according to any one of claims 1 to 5, It is characterized in that The step S2 comprises: Step S21: model initialization; Step S22: using the defect function and the retrieved vulnerability repair pair as model input; Step S23: Code segmentation of input; training input<Vt , Ft> , the code is tokenized at the subword level through byte pair encoding, and then the embedding layer generates an embedding vector for each subword to capture the semantic features of the code token; the word embedding result is used as the input of the encoder module to derive the hidden state of the encoder, which in turn notifies the decoder module to generate and output the predicted patch; Step S24: Calculate the cross entropy loss function; the cross entropy loss function is used in the training stage; Step S25: Back propagation, update parameters.

7. The method for generating vulnerability patches based on retrieval enhancement according to claim 6, It is characterized in that In the step S21, CodeT5 is used as the model of the patch generation module. CodeT5 adopts an encoder-decoder architecture, including 12 Transformer layers of the encoder and decoder respectively; each layer uses 12 attention heads for multi-head attention calculation.

8. The method for generating vulnerability patches based on retrieval enhancement according to claim 6, It is characterized in that In step S22, given a vulnerability pair from a training data set<Vt,Ft> , namely the defect function Vt and its correct fix Ft. The template index retriever outputs several related vulnerability fix code changes when the template is successfully matched.<Vi,Fi> ,<Vj,Fj> …; During the training phase of the patch generator, all retrieved vulnerability-fix pairs, Vt, are provided as input to the patch generator to generate predicted patches.

9. A vulnerability patch generation system based on retrieval enhancement, comprising a memory, a processor, and a computer program stored in the memory, It is characterized in that The processor executes the computer program to implement the steps of the method described in any one of claims 1 to 8.

10. A storage medium capable of being read by a computer or a processor, It is characterized in that The storage medium stores a computer program for executing the method described in any one of claims 1 to 8.

Images