Attention-dependent sentence-level vulnerability locating method and system, and storage medium
By introducing an attention masking matrix of the program dependency graph into the Transformer model, the dependencies between statements are extracted, solving the problem of existing methods ignoring code structure, achieving higher accuracy in vulnerability localization, and providing more granular vulnerability detection results.
Patent Information
- Application Number
- CN202410984507.9
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-07-22
- Publication Date
- 2025-11-07
- Estimated Expiration
- 2044-07-22
AI Technical Summary
Existing Transformer-based statement-level vulnerability localization methods ignore the inherent structural information of the source code, resulting in a gap between the vulnerability features learned by the model and the underlying semantics of the code, which affects the performance and accuracy of vulnerability detection.
We employ a dependency attention-based approach, constructing an attention masking matrix through a program dependency graph to extract data/control dependencies between statements. We then combine this with a Transformer encoder to calculate dependency-based statement representation vectors and use a cross-entropy loss function for model training, thereby improving the accuracy of vulnerability localization.
It achieves more accurate statement-level vulnerability location, provides more granular defect location results, and improves detection accuracy and performance.
Smart Images

Figure CN119127662B_ABST
Abstract
Description
TECHNICAL FIELD
[0001] The present application mainly relates to the technical field of software program vulnerability positioning, and particularly relates to a sentence-level vulnerability positioning method and system based on a dependency attention mechanism and a storage medium. BACKGROUND
[0002] With the continuous increase in the size of software, software vulnerabilities have become inevitable. Software vulnerability positioning technology is a very important task and research topic in software debugging activities. Accurate vulnerability positioning can reduce the workload of developers and help developers further understand the root cause and harm of vulnerabilities.
[0003] Currently, various technical methods or technical routes have been proposed by practitioners to detect software vulnerabilities, which can be mainly divided into two categories: (1) program analysis-based methods; and (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-bound 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 type of vulnerability. Therefore, deep learning-based methods are increasingly being proposed to identify a wide range of vulnerabilities without manual labor. In the prior art, a typical deep learning-based method is as follows: first, code representation learning is used to implicitly capture vulnerability patterns in source code, and then a classifier model is used to detect vulnerabilities.
[0004] Recently, in order to provide more accurate detection results for developers, some practitioners have proposed some deep learning vulnerability detection methods based on the sentence level, i.e., for identifying vulnerable statements within a function. Specifically, some methods represent source code as multiple views, such as data flow graphs (DFG), code property graphs (CPG), and program dependency graphs (PDG), and then use various graph neural networks (GNNs) to predict suspicious nodes corresponding to vulnerable statements. In addition, some methods directly convert source code into token sequences and use sequence-based deep learning models, such as recurrent neural networks (RNNs) and Transformers, to capture the syntactic features and semantic information of statements.
[0005] Although the Transformer has been proven to be superior in capturing the syntactic and semantic features of source code, there is a limitation in existing Transformer-based methods: they simply treat source code as a sequence of code statements and directly input them into the Transformer. Therefore, the self-attention mechanism in the Transformer treats code statements as fully connected relationships, ignoring the intrinsic structural information of source code (such as data or control dependencies). This oversight can result in 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
[0006] The technical problem to be solved by the present application is to provide a dependency attention-based sentence-level vulnerability positioning method, system and storage medium with simple principle, wide application range, high detection accuracy and good detection performance.
[0007] To solve the above technical problems, the present application adopts the following technical solutions:
[0008] A dependency attention-based sentence-level vulnerability positioning method comprises:
[0009] Step S1: training data preprocessing: a large amount of vulnerability data is preprocessed to obtain training data meeting the requirements of the model; for each function in the training data set, a Joern program analysis tool is used to extract the program dependency graph of the function;
[0010] Step S2: model training: the function and the corresponding program dependency graph are used as training data, i.e., the model is trained using the training data obtained in step S1 to obtain a neural network model;
[0011] Step S3: vulnerability statement positioning: for a given defective function, the trained model in step S2 is used to predict the specific vulnerability statement line number.
[0012] As a further improvement of the method of the present application, step S2 comprises:
[0013] Step S21: use a 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;
[0014] Step S22: After obtaining the embedding vector [x_1^v, x_2^v, …, x_n^v] of each statement, the data / control dependency relationship between the statements is extracted using the attention mask matrix and the Transformer encoder, and the dependency-based statement representation vector [s_1^v, s_2^v, …, s_n^v] is calculated, where s_i^v is the dependency-based statement representation vector.
[0015] Step S23: After obtaining the dependency-based statement representation vector [s_1^v, s_2^v, …, s_n^v], the function representation vector is obtained using the Transformer encoder.
[0016] Step S24: Model training; the model is trained using the cross-entropy loss function.
[0017] As a further improvement of the method of the application: in step S21, the process of obtaining the embedding vector of each statement in the function includes:
[0018] Step S211: The function source code is divided by line, and each line of source code is tokenized, and an additional token is inserted at the first position of the token sequence after tokenization.
[0019] Step S212: The embedding vector of the token sequence is obtained using the pre-trained model, and the first embedding vector of each sequence is taken as the embedding vector of each line of statement .
[0020] As a further improvement of the method of the application: in step S22, the process of calculating the dependency-based statement representation vector includes:
[0021] Step S221: Construct an attention mask matrix based on the program dependency graph; assume that the size of the attention mask matrix M is n*n, if there is data / control dependency between the two lines of statements, the size of the corresponding element of the attention mask matrix is negative infinity, and if there is no data / control dependency between the two lines of statements, the size of the corresponding element of the attention mask matrix is 0;
[0022] Step S222: The data / control dependency relationship between the statements is extracted using the attention mask matrix and the Transformer encoder, and the dependency-based statement representation vector is calculated .
[0023] As a further improvement of the method of the application: in step S222, it includes: adding the attention mask matrix to the original attention matrix in the 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:
[0024]
[0025]
[0026] wherein, a ij is the element of the (i, j) position in the attention matrix, , is the embedding vector of the i-th sentence and the j-th sentence, is the attention weight of the j-th element, dk is the embedding dimension, mij is the element of the (i, j) position in the attention mask matrix, c i is the attention score of the i-th element, , , is the weight matrix.
[0027] As a further improvement of the method of the present application: in the step S23, the function representation vector is obtained using the Transformer encoder. The specific process is: after obtaining the sentence representation vector based on the dependency information , an additional vector is inserted in the first position and input into the Transformer encoder to obtain the function representation vector
[0028] As a further improvement of the method of the present application: in the step S24, the cross-entropy loss function is:
[0029]
[0030] In the model training process, an adaptive matrix estimation method is used to reduce the Loss value in the above formula. Wherein, N is the number of samples in the data set, L func , L statement is the loss function of the function and the sentence, 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.
[0031] As a further improvement of the method of the present application: the step S3 comprises:
[0032] Step S31: for a given function, extract its program dependency graph, and input the program dependency graph and the function source code into the model in the step S2 training;
[0033] Step S32: predict whether the input function is vulnerable, if the function is predicted to be vulnerable, further predict the defect sentence in the function, and obtain the predicted defect sentence line number.
[0034] The application further provides a sentence-level vulnerability positioning system based on dependency attention, comprising a memory, a processor and a computer program stored in the memory, wherein the processor executes the computer program to realize the steps of the above method.
[0035] The application further provides a storage medium which can be read by a computer or a processor, and the storage medium stores a computer program for executing any one of the above methods.
[0036] Compared with the prior art, the application has the following advantages:
[0037] The sentence-level vulnerability positioning method, system and storage medium based on dependency attention of the application have simple principles, wide application range, high detection precision and good detection performance. BRIEF DESCRIPTION OF DRAWINGS
[0038] Figure 1 is a flowchart of the method of the application.
[0039] Figure 2 is a function source code with a vulnerability in a specific application example of the application.
[0040] Figure 3 is a program dependency graph of a function extracted in a specific application example of the application. DETAILED DESCRIPTION
[0041] The application will be further described in detail below in combination with the drawings and specific embodiments.
[0042] As shown in Figure 1 the steps of the sentence-level vulnerability positioning method based on dependency attention of the application include:
[0043] Step S1: training data preprocessing: preprocessing a large amount of vulnerability data to obtain training data meeting the requirements of the model;
[0044] Step S2: model training: training the model using the training data obtained in step S1 to obtain a neural network model;
[0045] Step S3: Vulnerability statement positioning: for a given defect function, use the model trained in step S2 to predict the specific vulnerability statement line number.
[0046] In a specific application example, in step S1, the program dependency graph is extracted when pre-processing the data. 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 in step S2.
[0047] In a specific application example, the detailed process of step S2 can include:
[0048] Step S21: Obtain the embedding vector of each statement in the function using the pre-trained model ;
[0049] Step S22: After obtaining the embedding vector of each statement , use the attention mask matrix and the Transformer encoder to extract the data / control dependency relationship between the statements, and calculate the dependency-based statement representation vector , where is the dependency-based statement representation vector.
[0050] Step S23: After obtaining the statement representation vector based on the dependency information , use the Transformer encoder to obtain the function representation vector ;
[0051] Step S24: Model training. Train the model using the cross-entropy loss function.
[0052] Further, as a preferred embodiment, the specific steps of obtaining the embedding vector of each statement in the function in step S21 include:
[0053] Step S211: Cut the function source code by line, tokenize each line of source code, and insert an additional token at the first position of the token sequence after tokenization;
[0054] 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 statement .
[0055] Further, as a preferred embodiment, the specific process of calculating the dependency-based statement representation vector in step S22 is:
[0056] Step S221: Construct an attention mask matrix based on the program dependency graph.
[0057] Assumed attention mask matrix where n is the number of statements in the function. If there is data / control dependency between two rows of statements, the size of the element corresponding to the attention mask matrix is negative infinity, and if there is no data / control dependency between two rows of statements, the size of the element corresponding to the attention mask matrix is 0.
[0058] Step S222: Extract data / control dependency relationships between statements using the attention mask matrix and the Transformer encoder, and calculate the dependency-based statement representation vector .
[0059] The specific process is:
[0060] The attention mask matrix is added to the original attention matrix in the Transformer encoder to obtain a dependency-based attention matrix, and then the SoftMax is used to obtain the attention weight of each element. The specific process is:
[0061]
[0062]
[0063] where a ij is the element in the (i, j) position of the attention matrix, , is the embedding vector of the i-th statement and the j-th statement, is the attention weight of the j-th element, dk is the embedding dimension, mij is the element in the (i, j) position of the attention mask matrix, c i is the attention score of the i-th element, , , is the weight matrix.
[0064] Further, as a preferred embodiment, the specific process of obtaining the function representation vector in step S23 using the Transformer encoder is: after obtaining the statement representation vector based on the dependency information, an additional vector is inserted at the first position, and the Transformer encoder is input to obtain the function representation vector .
[0065] Further, as a preferred embodiment, the cross-entropy loss function in step S24 is:
[0066]
[0067] In the model training process, an adaptive matrix estimation method is used to reduce the loss value in the above formula. Wherein, 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 statement level label of the i th sample.
[0068] In a specific application example, the step S3 can include the following processes according to actual needs:
[0069] Step S31: for a given function, extract its program dependency graph, and input the program dependency graph and the function source code into the model in S2 training;
[0070] Step S32: predict whether the input function is vulnerable, if the function is predicted to be vulnerable, further predict the defect statement in the function, and obtain the predicted defect statement line number.
[0071] In a specific application example, as shown in Figure 2 , it is a function source code with a vulnerability, and the specific application steps of the statement level vulnerability positioning method based on dependency attention of the application are as follows:
[0072] Step S1000: use the Joern program analysis tool to extract the program dependency graph of the function, as shown in Figure 3 ; wherein each node represents a statement, the red line represents the control dependency edge, and the green line represents the data dependency edge.
[0073] Step S2000: perform word segmentation on each statement, and generate statement embedding vectors . The attention mask matrix M is constructed, which has a size of 8*8. If there is data / control dependency between two rows of statements, the corresponding element size of the attention mask matrix is negative infinity, and if there is no data / control dependency between two rows of statements, the corresponding element size of the attention mask matrix is 0.
[0074] Step S3000: input the dependency attention encoding into the trained model, and the encoder network outputs the vulnerability statement line number as 1, 3, 5, and 6, which can correctly predict the vulnerability statement.
[0075] The application further provides a statement level vulnerability positioning system based on dependency attention, which comprises a memory, a processor and a computer program stored in the memory, and the processor executes the computer program to realize the steps of the above method.
[0076] The application further provides a storage medium, which can be read by a computer or a processor, and has a computer program stored therein for executing the above method.
[0077] Those skilled in the art will understand that the above-mentioned embodiments of the present application can be provided as a method, a system or a computer program product. Therefore, the present application can take the form of an entirely hardware embodiment, an entirely software embodiment or an embodiment combining software and hardware aspects. 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-ROMs, optical storage devices, etc.) containing computer usable program code. The present application is described with reference to flowcharts and / or block diagrams of methods, apparatus (systems) and computer program products according to embodiments of the present application. It should be understood that each flow and / or block in the flowcharts and / or block diagrams, as well as combinations of flows and / or blocks in the flowcharts and / or block diagrams, 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 apparatus to produce a machine, so that the instructions executed by the processor of the computer or other programmable data processing apparatus produce a device that implements the functions specified in the flowcharts and / or block diagrams. Figure 1 one or more flows and / or blocks Figure 1 one or more flows and / or blocks Figure 1 one or more flows and / or blocks Figure 1 one or more flows and / or blocks Figure 1 one or more flows and / or blocks Figure 1 one or more flows and / or blocks
[0078] The above is only the preferred embodiment of the present application, and the protection scope of the present application is not limited to the above-mentioned embodiment. Any technical solution falling within the concept of the present application belongs to the protection scope of the present application. It should be noted that, for ordinary skilled in the art, some improvements and refinements without departing from the principle of the present application should be considered as the protection scope of the present application.
Claims
1. A method for attention-dependent sentence-level vulnerability localization, characterized in that, Comprise: Step S1: training data preprocessing: a large number of vulnerability data are preprocessed to obtain training data meeting the requirements of the model; For each function in the training data set, the Joern program analysis tool is used to extract the program dependence graph of the function; Step S2: model training: the function and the corresponding program dependence graph are taken as training data, i.e., the model is trained using the training data obtained in step S1 to obtain a neural network model; specifically: step S21: using a pre-trained model to obtain an embedding vector of each sentence in the function wherein represents the embedding vector of each sentence; Step S22: obtaining an embedding vector of each sentence , using an attention mask matrix and a Transformer encoder to extract data / control dependency between sentences, and calculating a dependency-based sentence representation vector , wherein the dependency-based sentence representation vector; The flow of calculating the dependence-based sentence representation vector includes: Step S221: constructing an attention mask matrix based on the program dependence graph; assuming that the size of the attention mask matrix is If there is data / control dependence between the two statements, the size of the corresponding element of the attention mask matrix is negative infinity, and if there is no data / control dependence between the two statements, the size of the corresponding element of the attention mask matrix is 0. Step S222: use the attention mask matrix and the Transformer encoder to extract the data / control dependence relationship between sentences, and calculate the dependence-based sentence representation vector; Step S23: obtaining a sentence representation vector based on the dependency information using a Transformer encoder to obtain a function representation vector; Step S24: model training; the model is trained using the cross-entropy loss function; Step S3: vulnerability sentence positioning: for a given defective function, the model trained in step S2 is used to predict the specific vulnerability sentence line number.
2. The attention-dependent based sentence-level vulnerability locating method according to claim 1, wherein, In step S21, the flow of obtaining the embedding vector of each statement in the function includes: Step S211: split the function source code by line, and perform tokenization on each line of source code, and insert an additional token, namely [CLS], at the first position of the token sequence after tokenization; Step S212: Obtain the embedding vector of each token in the token sequence using the pre-trained model, and take the first embedding vector of each sequence as the embedding vector of each row of sentences . 3.The attention-dependent based sentence-level vulnerability locating method according to claim 1, wherein, In step S222, the attention mask matrix is added to the original attention matrix in the Transformer encoder to obtain a dependence-based attention matrix, and then the SoftMax is used to obtain the attention weight of each element, and the specific process is as follows: wherein, is an element in the attention matrix at position (i, j), , is an embedding vector for the i-th sentence and the j-th sentence, is an attention weight for the j-th element, is an embedding dimension, is an element in the attention mask at position (i, j), is an attention score for the i-th element, , , is a weight matrix.
4. The attention-dependent based sentence-level vulnerability locating method according to claim 1, wherein, In the step S23, the function representation vector is obtained using the Transformer encoder The specific process is as follows: obtaining the sentence representation vector based on the dependency information The first additional vector is inserted, and the function representation vector is obtained by inputting the Transformer encoder.
5. The attention-dependent based sentence-level vulnerability locating method according to claim 1, wherein, In step S24, the cross-entropy loss function is: In the model training process, the adaptive matrix estimation method is used to reduce the loss value in the above formula; wherein N is the number of samples in the data set, L func , statement L is the loss function of function and sentence, X i is the i-th sample, Y i is the function level label of the i-th sample, and Z i is the sentence level label of the i-th sample.
6. The attention-dependent based sentence-level vulnerability locating method according to any one of claims 1-5, characterized in that, The step S3 includes: Step S31: for a given function, extract its program dependence graph, and input the program dependence graph and function source code into the model in step S2 training; Step S32: predict whether the input function is vulnerable, if the function is predicted to be vulnerable, further predict the defective statement in the function, and obtain the predicted defective statement line number.
7. An attention-dependent-based sentence-level vulnerability localization system comprising a memory, a processor, and a computer program stored on the memory, wherein, The processor executes the computer program to realize the steps of the method of any one of claims 1-6.
8. A storage medium capable of being read by a computer or processor, characterized in that, The storage medium stores a computer program for executing the method of any one of the above claims 1-6.
Citation Information
Patent Citations
SAT-based function level source code vulnerability detection method
CN117150495A