Depending attention-based statement-level vulnerability positioning method and system and storage medium

Abstract

The invention discloses a statement-level vulnerability positioning method and system based on dependency attention and a storage medium, and the method comprises the steps: S1, training data preprocessing: carrying out the preprocessing of a large amount of vulnerability data, and obtaining the training data meeting the requirements of a model; s2, model training: training the model by using the training data obtained in the step S1 to obtain a neural network model; and S3, vulnerability statement positioning: for a given defect function, using the trained model to predict a specific vulnerability statement line number. The system and the storage medium are both used for executing the method. The method has the advantages that the principle is simple, the defect statement in the vulnerability program can be accurately positioned, and a finer-grained positioning result can be provided.

Description

Technical Field

[0001] The present invention mainly relates to the technical field of software program vulnerability location, and in particular to a statement-level vulnerability location method, system and storage medium based on attention reliance. Background Art

[0002] As the scale of software continues to grow, software vulnerabilities become inevitable. Software vulnerability location technology is a very important task and research topic in software debugging activities. Accurate vulnerability location can reduce the workload of developers and help them further understand the root causes and hazards of vulnerabilities.

[0003] At present, practitioners have proposed various technical methods or technical routes to perform software vulnerability detection, which can be mainly divided into two categories: (1) program analysis-based methods; (2) deep learning (DL)-based methods. Among them, the program analysis-based method uses manually predefined vulnerability templates to detect specific types of vulnerabilities, such as out-of-bounds write, cross-site scripting, and use-after-free. However, due to the diversity and complexity of software vulnerabilities, it is not feasible to define accurate and comprehensive templates for each vulnerability. Therefore, deep learning-based methods are increasingly proposed to identify a wide range of vulnerabilities without manual labor. In the prior art, a typical deep learning-based method is: first, code representation learning is used to implicitly capture vulnerability patterns in the source code, and then the vulnerability is detected through a classifier model.

[0004] Recently, in order to provide developers with more accurate detection results, other practitioners have proposed some statement-level deep learning vulnerability detection methods, that is, for identifying vulnerable statements within a function. Specifically, some methods represent the source code as multiple views, such as data flow graphs (DFGs), code property graphs (CPGs), and program dependency graphs (PDGs), and then use various graph neural networks (GNNs) to predict suspicious nodes corresponding to vulnerable statements. Other methods directly convert the source code into labeled sequences and use sequence-based deep learning models, such as recurrent neural networks (RNNs) and Transformers, to capture the grammatical features and semantic information of statements.

[0005] Although Transformer has been shown to be superior in capturing the syntactic and semantic features of source code, existing Transformer-based methods have a limitation: they simply treat the source code as a set of sequential code statements and directly input them into the Transformer. As a result, the self-attention mechanism in the Transformer treats the code statements as fully connected relations, ignoring the intrinsic structural information of the source code (such as data or control dependencies). This omission may lead to a gap between the vulnerability features learned by the model and the underlying semantics of the code, potentially affecting the performance and accuracy of vulnerability detection. Summary of the invention

[0006] 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 sentence-level vulnerability localization method, system and storage medium based on attention reliance, which have simple principle, wide application range, high detection accuracy and good detection performance.

[0007] In order to solve the above technical problems, the present invention adopts the following technical solutions: A sentence-level vulnerability localization method based on attention reliance includes: Step S1: training data preprocessing: preprocess a large amount of vulnerability data to obtain training data that meets the model requirements; for each function in the training data set, use the Joern program analysis tool to extract the program dependency graph of the function; Step S2: Model training: using the function and the corresponding program dependency graph as training data, that is, using the training data obtained in step S1 to train the model, and obtaining a neural network model; Step S3: Vulnerability statement location: For a given defective function, use the model trained in step S2 to predict the specific vulnerability statement line number.

[0008] As a further improvement of the method of the present invention: the step S2 comprises: Step S21: Use the pre-trained model to obtain the embedding vector [x_1^v,x_2^v,…,x_n^v] of each statement in the function, where x_i^v represents the embedding vector of each statement; Step S22: After obtaining the embedding vector [x_1^v,x_2^v,…,x_n^v] of each sentence, use the attention masking matrix and the Transformer encoder to extract the data / control dependency between the sentences, and calculate the dependency-based sentence representation vector [s_1^v,s_2^v,…,s_n^v], where s_i^v is the dependency-based sentence representation vector; Step S23: After obtaining the sentence representation vector [s_1^v,s_2^v,…,s_n^v] based on the dependency information, use the Transformer encoder to obtain the function representation vector; Step S24: Model training: The model is trained using a cross entropy loss function.

[0009] As a further improvement of the method of the present invention: in step S21, the process of obtaining the embedding vector of each statement in the function includes: Step S211: Split the function source code into lines, perform token segmentation on each line of the source code, and insert an additional token at the first position of the token sequence after token segmentation; Step S212: Use the pre-trained model to obtain the embedding vector of the token sequence, and take the first embedding vector of each sequence as the embedding vector of each line of the statement. .

[0010] As a further improvement of the method of the present invention: in step S22, the process of calculating the dependency-based statement representation vector includes: Step S221: construct an attention masking matrix based on the program dependency graph; assuming that the size of the attention masking matrix M is n*n, if there is data / control dependency between two lines of statements, the size of the element corresponding to the attention masking matrix is ​​negative infinity, and if there is no data / control dependency between two lines of statements, the size of the element corresponding to the attention masking matrix is ​​0; Step S222: Use the attention masking matrix and the Transformer encoder to extract the data / control dependency between sentences and calculate the dependency-based sentence representation vector .

[0011] As a further improvement of the method of the present invention: the step S222 includes: adding the attention masking matrix to the attention matrix in the original Transformer encoder to obtain a dependency-based attention matrix, and then using SoftMax to obtain the attention weight of each element. The specific process is:

[0012]

[0013] Among them, a ij is the element at position (i, j) in the attention matrix, , is the embedding vector of the i-th sentence and the j-th sentence, is the attention weight of the 𝑗th element, dk is the embedding dimension, mij is the element at position (i, j) in the attention mask matrix, and c iis the attention score of the i-th element, , , is the weight matrix.

[0014] As a further improvement of the method of the present invention: in step S23, the function representation vector is obtained using the Transformer encoder The specific process is: after obtaining the sentence representation vector based on dependency information , the first bit is inserted into an additional vector and input into the Transformer encoder to obtain the function representation vector .

[0015] As a further improvement of the method of the present invention: in step S24, the cross entropy loss function is:

[0016] During the model training process, the adaptive moment estimation method is used to reduce the Loss value in the above formula. Where N is the number of samples in the data set, L func , L statement is the loss function of the function and the statement, X i is the i-th sample, Y i is the function-level label of the i-th sample, Z i is the sentence-level label of the i-th sample.

[0017] As a further improvement of the method of the present invention: the step S3 comprises: Step S31: for a given function, extract its program dependency graph, and input the program dependency graph and function source code into the model being trained in S2; Step S32: predict whether the input function has a vulnerability. If the function is predicted to have a vulnerability, further predict the defective statement in the function to obtain the predicted defective statement line number.

[0018] The present invention further provides a sentence-level vulnerability location system based on attention reliance, 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.

[0019] 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.

[0020] Compared with the prior art, the advantages of the present invention are: The present invention's statement-level vulnerability location method, system and storage medium based on dependent attention have simple principles, wide application range, high detection accuracy and good detection performance. The existing Transformer-based statement-level vulnerability location technology only regards the code as a sequence of statements and does not utilize the data / control dependencies in the function. Therefore, the present invention innovatively proposes a statement-level vulnerability location technology based on dependent attention, which can utilize the data / control dependencies in the function, so that the model can better learn the dependencies in the function and produce more accurate statement location results. The present invention can accurately locate the line number of the defective statement in the defective program, and can provide more fine-grained location results compared to other software defect location methods. BRIEF DESCRIPTION OF THE DRAWINGS

[0021] Figure 1 It is a schematic diagram of the process of the present invention.

[0022] Figure 2 It is a function source code with a vulnerability in a specific application example of the present invention.

[0023] Figure 3 It is a program dependency graph of functions extracted in a specific application example of the present invention. DETAILED DESCRIPTION

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

[0025] like Figure 1 As shown, the sentence-level vulnerability localization method based on attention dependence of the present invention comprises the following steps: Step S1: training data preprocessing: preprocess a large amount of vulnerability data to obtain training data that meets the model requirements; Step S2: Model training: Use the training data obtained in step S1 to train the model and obtain a neural network model; Step S3: Vulnerability statement location: For a given defective function, use the model trained in step S2 to predict the specific vulnerability statement line number.

[0026] In a specific application example, in step S1, a program dependency graph is first extracted during data preprocessing. For each function in the training data set, the Joern program analysis tool is used to extract the program dependency graph of the function. Further, the function and the corresponding program dependency graph are used as training data for step S2.

[0027] In a specific application example, the detailed process of step S2 may include: Step S21: Use the pre-trained model to get the embedding vector of each statement in the function ; Step S22: Obtain the embedding vector of each sentence , use the attention masking matrix and Transformer encoder to extract the data / control dependencies between sentences and calculate the dependency-based sentence representation vector ,in is a dependency-based statement representation vector; Step S23: After obtaining the sentence representation vector based on the dependency information After that, use the Transformer encoder to get the function representation vector ; Step S24: Model training: The model is trained using a cross entropy loss function.

[0028] Further, as a preferred embodiment, the specific steps of obtaining the embedding vector of each statement in the function in step S21 include: Step S211: Split the function source code into lines, perform token segmentation on each line of source code, and insert an additional token at the first position of the token sequence after token segmentation; Step S212: Use the pre-trained model to obtain the embedding vector of the token sequence, and take the first embedding vector of each sequence as the embedding vector of each line of the statement. .

[0029] Further, as a preferred embodiment, the specific process of calculating the dependency-based statement representation vector in step S22 is: Step S221: Construct an attention masking matrix based on the program dependency graph.

[0030] Assume the attention mask matrix , where n is the number of statements in the function. If there is a data / control dependency between two rows of statements, the element size corresponding to the attention mask matrix is ​​negative infinity. If there is no data / control dependency between two rows of statements, the element size corresponding to the attention mask matrix is ​​0.

[0031] Step S222: Use the attention masking matrix and the Transformer encoder to extract the data / control dependency between sentences and calculate the dependency-based sentence representation vector .

[0032] The specific process is: The attention mask matrix is ​​added to the attention matrix in the original Transformer encoder to obtain the dependency-based attention matrix, and then SoftMax is used to obtain the attention weight of each element. The specific process is:

[0033]

[0034] Among them, a ij is the element at position (i, j) in the attention matrix, , is the embedding vector of the i-th sentence and the j-th sentence, is the attention weight of the 𝑗th element, dk is the embedding dimension, mij is the element at position (i, j) in the attention mask matrix, and c i is the attention score of the i-th element, , , is the weight matrix.

[0035] Further, as a preferred embodiment, in step S23, the function representation vector is obtained by using a Transformer encoder The specific process is: after obtaining the sentence representation vector based on dependency information , the first bit is inserted into an additional vector and input into the Transformer encoder to obtain the function representation vector .

[0036] Further, as a preferred embodiment, the cross entropy loss function in step S24 is:

[0037] During the model training process, the adaptive moment estimation method is used to reduce the Loss value in the above formula. Where N is the number of samples in the data set, L func , L statement is the loss function of the function and the statement, X i is the i-th sample, Y i is the function-level label of the i-th sample, Z i is the sentence-level label of the i-th sample.

[0038] In a specific application example, step S3 may include the following process according to actual needs: Step S31: for a given function, extract its program dependency graph, and input the program dependency graph and function source code into the model being trained in S2; Step S32: predict whether the input function has a vulnerability. If the function is predicted to have a vulnerability, further predict the defective statement in the function to obtain the predicted defective statement line number.

[0039] In a specific application example, Figure 2 As shown, it is a function source code with a vulnerability. The specific application steps of the statement-level vulnerability location method based on attention dependence of the present invention are as follows: Step S1000: Use the Joern program analysis tool to extract the program dependency graph of the function, such as Figure 3 As shown; each node represents a statement, the red line represents the control dependency edge, and the green line represents the data dependency edge.

[0040] Step S2000: Segment each sentence and generate sentence embedding vectors Construct an attention mask matrix M with a size of 8*8. If there is data / control dependency between two lines of statements, the element size corresponding to the attention mask matrix is ​​negative infinity. If there is no data / control dependency between two lines of statements, the element size corresponding to the attention mask matrix is ​​0.

[0041] Step S3000: After encoding the dependent attention, it is input into the trained model, and the encoder network outputs the vulnerability sentence line numbers as 1, 3, 5, 6, which can correctly predict the vulnerability sentence.

[0042] The present invention further provides a sentence-level vulnerability location system based on attention reliance, 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.

[0043] 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.

[0044] 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 A step that specifies a function in one or more boxes.

[0045] 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 sentence-level vulnerability localization method based on attention dependence, characterized in that: include: Step S1: training data preprocessing: preprocess a large amount of vulnerability data to obtain training data that meets the model requirements; For each function in the training dataset, the Joern program analysis tool is used to extract the program dependency graph of the function; Step S2: Model training: using the function and the corresponding program dependency graph as training data, that is, using the training data obtained in step S1 to train the model, and obtaining a neural network model; Step S3: Vulnerability statement location: For a given defective function, use the model trained in step S2 to predict the specific vulnerability statement line number.

2. The method for locating sentence-level vulnerabilities based on attention reliance according to claim 1 is characterized in that: The step S2 comprises: Step S21: Use the pre-trained model to get the embedding vector of each statement in the function ,in Embedding vector representing each sentence; Step S22: Obtain the embedding vector of each sentence , use the attention masking matrix and Transformer encoder to extract the data / control dependencies between sentences and calculate the dependency-based sentence representation vector ,in Dependency-based statement representation vector; Step S23: After obtaining the sentence representation vector based on the dependency information , use the Transformer encoder to get the function representation vector; Step S24: Model training: The model is trained using a cross entropy loss function.

3. The method for locating sentence-level vulnerabilities based on attention reliance according to claim 2 is characterized in that: In step S21, the process of obtaining the embedding vector of each statement in the function includes: Step S211: Split the function source code into lines, and perform token segmentation on each line of source code, and insert an additional token, namely [CLS], at the first position of the token sequence after token segmentation; Step S212: Use the pre-trained model to obtain the embedding vector of each token in the token sequence, and take the first embedding vector of each sequence as the embedding vector of each line of the statement .

4. The method for locating sentence-level vulnerabilities based on attention reliance according to claim 2 is characterized in that: In step S22, the process of calculating the dependency-based statement representation vector includes: Step S221: construct an attention masking matrix based on the program dependency graph; assuming that the size of the attention masking matrix is ​​n*n, if there is data / control dependency between two lines of statements, the size of the element corresponding to the attention masking matrix is ​​negative infinity, and if there is no data / control dependency between two lines of statements, the size of the element corresponding to the attention masking matrix is ​​0; Step S222: Use the attention masking matrix and the Transformer encoder to extract the data / control dependency between sentences and calculate the dependency-based sentence representation vector.

5. The method for locating sentence-level vulnerabilities based on attention dependence according to claim 4 is characterized in that: The step S222 includes: adding the attention masking matrix to the attention matrix in the original Transformer encoder to obtain a dependency-based attention matrix, and then using SoftMax to obtain the attention weight of each element. The specific process is: Among them, a ij is the element at position (i, j) in the attention matrix, , is the embedding vector of the i-th sentence and the j-th sentence, is the attention weight of the 𝑗th element, dk is the embedding dimension, mij is the element at position (i, j) in the attention mask matrix, and c i is the attention score of the i-th element, , , is the weight matrix.

6. The method for locating sentence-level vulnerabilities based on attention reliance according to claim 2, characterized in that: In step S23, the function representation vector is obtained using the Transformer encoder. The specific process is: after obtaining the sentence representation vector based on dependency information , the first bit is inserted into an additional vector and input into the Transformer encoder to obtain the function representation vector.

7. The method for locating sentence-level vulnerabilities based on attention reliance according to claim 2 is characterized in that: In step S24, the cross entropy loss function is: During the model training process, the adaptive moment estimation method is used to reduce the Loss value in the above formula; where N is the number of samples in the data set, L func , L statement is the loss function of the function and the statement, X i is the i-th sample, Y i is the function-level label of the i-th sample, Z i is the sentence-level label of the i-th sample.

8. The method for locating a sentence-level vulnerability based on attention according to any one of claims 1 to 7, characterized in that: The step S3 comprises: Step S31: for a given function, extract its program dependency graph, and input the program dependency graph and function source code into the model being trained in S2; Step S32: predict whether the input function has a vulnerability. If the function is predicted to have a vulnerability, further predict the defective statement in the function to obtain the predicted defective statement line number.

9. A sentence-level vulnerability localization system based on attention dependence, comprising a memory, a processor, and a computer program stored in the memory, 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, which can be read by a computer or a processor, characterized in that: The storage medium stores a computer program for executing the method described in any one of claims 1 to 8.

Images