Multimodal source code vulnerability static detection method based on Transform and GAT

Abstract

The invention relates to the field of software security technology and the like, and discloses a multi-modal source code vulnerability static detection method based on Transform and GAT, comprising the following steps: processing an original code into a sequence input matrix and a graph input matrix which can be received by a deep learning model; respectively inputting the sequence input matrix and the graph input matrix into two parallel Transform and GAT models to obtain a sequence hidden state representation vector and a graph hidden state representation vector; splicing the sequence hidden state representation vector and the graph hidden state representation vector by using concat to obtain a fusion representation vector of the code function; performing dichotomy training on the fusion representation vector of the code function, and updating model parameters for a classification result in a cross validation mode; and after the training is completed, the test set is used for testing, and the deep learning model passing the test is stored as a final vulnerability detection model, so that the beneficial effects of effectively reducing the labor cost and having better generalization are achieved.

Description

Technical Field

[0001] The present invention relates to the fields of software security technology, and in particular to a multi-modal source code vulnerability static detection method based on Transformer and GAT. Background Art

[0002] Existing source code vulnerability static detection technologies can be divided into two categories: rule-based methods and deep learning technology-based methods.

[0003] Among them, the rule-based method relies on expert experience to manually define vulnerability rules and relies on the following technologies to detect vulnerabilities: 1) Code similarity: By converting the source code into code representation form and grading it, and then using the comparison code segment representation to represent the similarity with the known vulnerability code base, common methods include character matching, string search, hash matching, vector distance calculation, etc. The main idea of ​​this method is that similar source code segments are likely to contain the same vulnerabilities. 2) Symbolic execution: By converting the variables in the source code into a composite expression composed of special symbol values ​​and constants, the source code vulnerability is regarded as the value of some variables cannot meet the corresponding constraints, and then judging whether the expression representing the variable value meets the corresponding constraints to detect whether there are corresponding vulnerabilities in the source code. 3) Taint analysis: By analyzing the lexical and grammatical aspects of the code, combined with data dependency and control dependency, by detecting whether the flow of tainted data is propagated from the taint source to the taint sink, it is determined whether there are vulnerabilities in the source code.

[0004] Methods based on deep learning technology can be divided into the following two categories: 1) Methods based on sequence representation: treat the code as a continuous text sequence, and use recurrent neural network RNN ​​to obtain the context-dependent syntax and semantics of the code to extract vulnerability features. 2) Methods based on graph representation: Convert the code into a graph to represent the lexical and semantic information of the code. Commonly used graphs include data flow graphs DFG, control flow graphs CFG, and program dependency graphs PDG, and then use graph neural networks GNN to build a vulnerability detection model.

[0005] The shortcomings of existing rule-based methods are: 1. High false positive and false negative rates of vulnerabilities; 2. Labor costs are high.

[0006] The defects of existing methods based on deep learning technology are: 1. The extraction of code vulnerability feature information is not comprehensive enough. Existing methods are only based on a single code representation of a sequence or graph. Among them, the method based on sequence representation is limited by the RNN model and cannot pay attention to the long-term context dependency features of long sequences. On the one hand, the method based on graph representation disrupts the order relationship of the code, making it difficult for the model to understand the basic semantics of the code; on the other hand, the existing methods simply use GNN and do not carry out targeted processing on the graph representation of the code.

[0007] 2. Poor generalization ability. Specifically, the false positive rate and missed negative rate of standard vulnerability library detection are better than those of rule-based methods, but the accuracy rate drops significantly when detecting real code. Summary of the invention

[0008] The purpose of the present invention is to solve the shortcomings of the prior art and provide a multimodal source code vulnerability static detection method based on Transformer and GAT, which can effectively reduce labor costs while having better generalization.

[0009] The present invention is implemented by the following technical solution: a multimodal source code vulnerability static detection method based on Transformer and GAT, comprising the following steps: A. Code preprocessing: Process the original code into sequence input matrices and graph input matrices that can be accepted by the deep learning model; B. Feature extraction: The sequence input matrix and the graph input matrix obtained in step A are respectively input into two parallel Transformer and GAT models to obtain the sequence hidden state representation vector and the graph hidden state representation vector; C. Multimodal feature fusion: Use concat to concatenate the sequence hidden state representation vector and the graph hidden state representation vector to obtain the fused representation vector of the code function; D. Model training and testing: Perform binary classification training on the fused representation vector of the code function, and update the model parameters by cross-validation based on the classification results. After the training is completed, use the test set to test the deep learning model with updated model parameters, and save the deep learning model that passes the test as the final vulnerability detection model.

[0010] In order to better implement the multimodal source code vulnerability static detection method based on Transformer and GAT described in the present invention, the following setting is particularly adopted: the step A includes the following specific steps: A.1. Data cleaning: Split the code in the original code file into function level, delete comments and other contents irrelevant to the code semantics, and normalize the function names and variable names defined by users in the code; A.2. Serialization and graph extraction: Use the word segmentation technology nltk to split the function-level code obtained in step A.1 into the smallest semantic unit tokens, and splice these tokens into a token sequence; at the same time, use joern to generate the program dependency graph PDG of the function-level code, and extract three attributes from the node statements of the program dependency graph PDG: type Type, value Value, and line number Number; A.3 Word embedding: Use all the minimum semantic unit tokens in step A.2 and the node attributes of the program dependency graph as the corpus to pre-train a word embedding model using word2vec; then obtain the token sequence and the vector of each token and node in the program dependency graph PDG from the word embedding model to obtain the sequence input matrix and graph input matrix of the function-level code.

[0011] In order to better implement the multimodal source code vulnerability static detection method based on Transformer and GAT described in the present invention, the following setting method is particularly adopted: the function names and variable names defined by users in the code are normalized, specifically: the function names are uniformly named FUNCTION, and the corresponding numbers are added at the end according to the order in the code file; the variable names are uniformly named VAR, and the corresponding numbers are added at the end according to the order in the code file.

[0012] In order to better implement the multimodal source code vulnerability static detection method based on Transformer and GAT described in the present invention, the following setting is particularly adopted: the step B includes the following specific steps: B.1. Sequence representation learning: In the Pytorch framework, the sequence input matrix is ​​positionally encoded. The encoding formula is: (1); (2); in, POS is the position of the token vector in the entire sequence; i Pos divided by 2 and rounded down to indicate even and odd positions; d model is the dimension of the token vector; The sequence input matrix and the sequence input matrix position code are added and input into a 2-layer Transformer network; the first layer is a TransformerEncoder layer with a multi-head attention number of 5, which is used to extract sequence features; the second layer is a linear transformation layer that maps the sequence input matrix to a hidden state representation vector with a length of 256; B.2. Graph representation learning: In the Pytorch framework, the graph input matrix is ​​input into a 2-layer graph attention network GAT network; the first layer is the GATconv convolution layer, which is used to transfer the adjacent node information in the graph input matrix; the second layer is the linear transformation layer, which maps the graph input matrix to a 256*3 matrix and reduces it to one dimension through average pooling to obtain a hidden state representation vector with a length of 256.

[0013] In order to better implement the multimodal source code vulnerability static detection method based on Transformer and GAT described in the present invention, the following setting is particularly adopted: the step D includes the following specific steps: D.1. Classification: Use a linear transformation layer to perform binary classification on the fused representation vector obtained in step C, and use the ReLU activation function to map the classification result to the range of 0 to 1; D.2. Cross-validation: Divide the dataset into training set, validation set, and test set in a ratio of 7:2:1; use the training set to train the classification process to update the model parameters, and test the deep learning model once every 10 rounds of training through the validation set to monitor its classification accuracy: if qualified, test it again on the test set; if not qualified, continue training; save the deep learning model that has qualified detection accuracy on both the validation set and the test set as the final vulnerability detection model.

[0014] Compared with the prior art, the present invention has the following advantages and beneficial effects: (1) The present invention can effectively reduce labor costs: the deep learning technology can automatically learn vulnerability features from the vulnerability library, without the need to manually define vulnerability detection rules, and has a lower false alarm rate and missed alarm rate.

[0015] (2) The present invention has better generalization and better performance in detecting real-world codes. Because it has two semantic sources of code representation, and introduces the Transformer model to replace the traditional recurrent neural network to extract sequence information, it can effectively solve the problem that the recurrent neural network is insufficient in extracting long-term contextual dependencies. This technical solution can comprehensively extract sequential contextual dependencies and structural features. BRIEF DESCRIPTION OF THE DRAWINGS

[0016] Figure 1 It is a flow chart of the present invention. DETAILED DESCRIPTION

[0017] The present invention is further described in detail below in conjunction with examples, but the embodiments of the present invention are not limited thereto.

[0018] 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 part of the embodiments of the present invention, rather than all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative work are within the scope of protection of the present invention. Therefore, the following detailed description of the embodiments of the present invention provided in the drawings is not intended to limit the scope of the invention claimed for protection, but merely represents the selected embodiments of the present invention. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without making creative work are within the scope of protection of the present invention.

[0019] In addition, the terms "first" and "second" are used for descriptive purposes only and should not be understood as indicating or implying relative importance or implicitly indicating the number of the indicated technical features. Therefore, the features defined as "first" and "second" may explicitly or implicitly include one or more of the features. In the description of the present invention, the meaning of "plurality" is two or more, unless otherwise clearly and specifically defined.

[0020] Glossary: Transformer: A neural network model architecture commonly used in natural language processing. GAT: Graph Attention Network; concat: feature concatenation technology; nltk: word segmentation technology; joern: a tool for generating program dependency graphs. PDG: Program Dependency Graph; Token: the smallest semantic unit token; word2vec: word embedding generation technology; Pytorch - deep learning framework. TransformerEncoder: Transformer encoder; GATconv: Graph Attention Convolutional Layer.

[0021] Embodiment 1: The present invention designs a multi-modal source code vulnerability static detection method based on Transformer and GAT, which includes the following steps: A. Code preprocessing: Process the original code into sequence input matrices and graph input matrices that can be accepted by the deep learning model; B. Feature extraction: The sequence input matrix and the graph input matrix obtained in step A are respectively input into two parallel Transformer and GAT models to obtain the sequence hidden state representation vector and the graph hidden state representation vector; C. Multimodal feature fusion: Use concat to concatenate the sequence hidden state representation vector and the graph hidden state representation vector to obtain the fused representation vector of the code function; D. Model training and testing: Perform binary classification training on the fused representation vector of the code function, and update the model parameters by cross-validation based on the classification results. After the training is completed, use the test set to test the deep learning model with updated model parameters, and save the deep learning model that passes the test as the final vulnerability detection model.

[0022] Embodiment 2: This embodiment is further optimized on the basis of the above embodiment, and the same points as the above technical solution are not repeated here. In order to better implement the multimodal source code vulnerability static detection method based on Transformer and GAT described in the present invention, the following setting is particularly adopted: Step A includes the following specific steps: A.1. Data cleaning: Split the code in the original code file into function level, delete comments and other contents irrelevant to the code semantics, and normalize the function names and variable names defined by users in the code; A.2. Serialization and graph extraction: Use the word segmentation technology nltk to split the function-level code obtained in step A.1 into the smallest semantic unit tokens, and splice these tokens into a token sequence; at the same time, use joern to generate the program dependency graph PDG of the function-level code, and extract three attributes from the node statements of the program dependency graph PDG: type Type, value Value, and line number Number; A.3 Word embedding: Use all the minimum semantic unit tokens in step A.2 and the node attributes of the program dependency graph as the corpus to pre-train a word embedding model using word2vec; then obtain the token sequence and the vector of each token and node in the program dependency graph PDG from the word embedding model to obtain the sequence input matrix and graph input matrix of the function-level code.

[0023] Embodiment 3: This embodiment is further optimized on the basis of any of the above embodiments, and the similarities with the above technical solutions are not repeated here. In order to better realize the multimodal source code vulnerability static detection method based on Transformer and GAT described in the present invention, the following setting method is particularly adopted: the function names and variable names defined by the user in the code are normalized, specifically: the function names are uniformly named FUNCTION, and the corresponding numbers are added at the end according to the order in the code file; the variable names are uniformly named VAR, and the corresponding numbers are added at the end according to the order in the code file.

[0024] Embodiment 4: This embodiment is further optimized on the basis of any of the above embodiments, and the same points as the above technical solutions are not repeated here. In order to better implement the multimodal source code vulnerability static detection method based on Transformer and GAT described in the present invention, the following setting is particularly adopted: the step B includes the following specific steps: B.1. Sequence representation learning: In the Pytorch framework, the sequence input matrix is ​​positionally encoded. The encoding formula is: (1); (2); in, POS is the position of the token vector in the entire sequence; i Pos divided by 2 and rounded down to indicate even and odd positions; d model is the dimension of the token vector; The sequence input matrix and the sequence input matrix position code are added and input into a 2-layer Transformer network; the first layer is a TransformerEncoder layer with a multi-head attention number of 5, which is used to extract sequence features; the second layer is a linear transformation layer that maps the sequence input matrix to a hidden state representation vector with a length of 256; B.2. Graph representation learning: In the Pytorch framework, the graph input matrix is ​​input into a 2-layer graph attention network GAT network; the first layer is the GATconv convolution layer, which is used to transfer the adjacent node information in the graph input matrix; the second layer is the linear transformation layer, which maps the graph input matrix to a 256*3 matrix and reduces it to one dimension through average pooling to obtain a hidden state representation vector with a length of 256.

[0025] Embodiment 5: This embodiment is further optimized on the basis of any of the above embodiments, and the same points as the above technical solutions are not repeated here. In order to better implement the multimodal source code vulnerability static detection method based on Transformer and GAT described in the present invention, the following setting is particularly adopted: the step D includes the following specific steps: D.1. Classification: Use a linear transformation layer to perform binary classification on the fused representation vector obtained in step C, and use the ReLU activation function to map the classification result to the range of 0 to 1; D.2. Cross-validation: Divide the dataset into training set, validation set, and test set in a ratio of 7:2:1; use the training set to train the classification process to update the model parameters, and test the deep learning model once every 10 rounds of training through the validation set to monitor its classification accuracy: if qualified, test it again on the test set; if not qualified, continue training; save the deep learning model that has qualified detection accuracy on both the validation set and the test set as the final vulnerability detection model.

[0026] Embodiment 6: A multimodal source code vulnerability static detection method based on Transformer and GAT, such as Figure 1 As shown, the following steps are included: 1. Data cleaning: Split the code in the original code file into function level, and delete comments and other contents that are irrelevant to the code semantics; normalize the function names and variable names defined by users in the code, that is, uniformly name the function names as FUNCTION1, FUNCTION2, etc. according to the order in the code file, and uniformly name the variable names as VAR1, VAR2, etc.; 2. Serialization and graph extraction: Use the word segmentation technology nltk to split the function-level code obtained in step 1 into the smallest semantic unit token, and splice these tokens into a token sequence; at the same time, use joern to generate the program dependency graph PDG of the function-level code, and extract three attributes from the node statements of the program dependency graph PDG: type Type, value Value, and line number Number; 3. Word embedding: Use word2vec to pre-train a word embedding model using all the minimum semantic unit tokens in step 2 and the node attributes of the program dependency graph as a corpus; then obtain the token sequence and the vector of each token and node in the program dependency graph PDG from the word embedding model to obtain the sequence input matrix and graph input matrix of the function-level code.

[0027] 4. Sequence representation learning: In the Pytorch framework, the sequence input matrix is ​​positionally encoded. The encoding formula is: (1); (2); in, POS is the position of the token vector in the entire sequence; i Pos divided by 2 and rounded down to indicate even and odd positions; d model is the dimension of the token vector; The sequence input matrix and the sequence input matrix position code are added and input into a 2-layer Transformer network; the first layer is a TransformerEncoder layer with a multi-head attention number of 5, which is used to extract sequence features; the second layer is a linear transformation layer that maps the sequence input matrix to a hidden state representation vector with a length of 256; 5. Graph representation learning: In the Pytorch framework, the graph input matrix is ​​input into a two-layer graph attention network GAT network; the first layer is the GATconv convolution layer, which is used to transfer the adjacent node information in the graph input matrix; the second layer is the linear transformation layer, which maps the graph input matrix to a 256*3 matrix and reduces it to one dimension through average pooling to obtain a hidden state representation vector with a length of 256.

[0028] 6. Multimodal feature fusion: Use concat to concatenate the sequence hidden state representation vector and the graph hidden state representation vector to obtain the fusion representation vector of the code function; 7. Classification: Use a linear transformation layer to classify the fused representation vector obtained in step 6, and use the Relu activation function to map the classification result to the range of 0 to 1. 8. Cross-validation: Divide the data set into training set, validation set, and test set in a ratio of 7:2:1; use the training set to train the classification process to update the model parameters. 9. Classification accuracy: After every 10 rounds of training, the deep learning model is tested on the validation set to monitor its classification accuracy. If it passes, it is tested again on the test set; if it fails, training continues.

[0029] 10. Save model: The deep learning model with qualified detection accuracy in both the validation set and the test set is saved as the final vulnerability detection model.

[0030] The above description is only a preferred embodiment of the present invention and does not limit the present invention in any form. Any simple modification or equivalent change made to the above embodiment based on the technical essence of the present invention shall fall within the protection scope of the present invention.

Claims

1. A multimodal source code vulnerability static detection method based on Transformer and GAT, characterized by: The steps include: A. Code preprocessing: Process the original code into sequence input matrices and graph input matrices that can be accepted by the deep learning model; B. Feature extraction: The sequence input matrix and the graph input matrix obtained in step A are respectively input into two parallel Transformer and GAT models to obtain the sequence hidden state representation vector and the graph hidden state representation vector; C. Multimodal feature fusion: Use concat to concatenate the sequence hidden state representation vector and the graph hidden state representation vector to obtain the fused representation vector of the code function; D. Model training and testing: Perform binary classification training on the fusion representation vector of the code function, and update the model parameters in a cross-validation manner based on the classification results; After the training is completed, the test set is used to test the deep learning model after updating the model parameters, and the deep learning model that passes the test is saved as the final vulnerability detection model.

2. According to claim 1, a multimodal source code vulnerability static detection method based on Transformer and GAT is characterized by: The step A comprises the following specific steps: A.

1. Data cleaning: Split the code in the original code file into function level, delete comments and other contents irrelevant to the code semantics, and normalize the function names and variable names defined by users in the code; A.

2. Serialization and graph extraction: Use the word segmentation technology nltk to split the function-level code obtained in step A.1 into the smallest semantic unit tokens, and splice these tokens into a token sequence; at the same time, use joern to generate the program dependency graph PDG of the function-level code, and extract three attributes from the node statements of the program dependency graph PDG: type Type, value Value, and line number Number; A.3 Word embedding: Use all the minimum semantic unit tokens in step A.2 and the node attributes of the program dependency graph as the corpus to pre-train a word embedding model using word2vec; then obtain the token sequence and the vector of each token and node in the program dependency graph PDG from the word embedding model to obtain the sequence input matrix and graph input matrix of the function-level code.

3. According to claim 2, a multi-modal source code vulnerability static detection method based on Transformer and GAT is characterized in that: Normalize the user-defined function names and variable names in the code as follows: uniformly name the function as FUNCTION and add corresponding numbers at the end according to the order in the code file; uniformly name the variable as VAR and add corresponding numbers at the end according to the order in the code file.

4. According to claim 1, a multimodal source code vulnerability static detection method based on Transformer and GAT is characterized by: The step B comprises the following specific steps: B.

1. Sequence representation learning: In the Pytorch framework, the sequence input matrix is ​​positionally encoded. The encoding formula is: （1）； （2）； in, POS is the position of the token vector in the entire sequence; i Pos divided by 2 and rounded down to indicate even and odd positions; d model is the dimension of the token vector; The sequence input matrix and the sequence input matrix position code are added and input into a 2-layer Transformer network; the first layer is a TransformerEncoder layer with a multi-head attention number of 5, which is used to extract sequence features; the second layer is a linear transformation layer that maps the sequence input matrix to a hidden state representation vector with a length of 256; B.

2. Graph representation learning: In the Pytorch framework, the graph input matrix is ​​input into a 2-layer graph attention network GAT network; the first layer is the GATconv convolution layer, which is used to transfer the adjacent node information in the graph input matrix; the second layer is the linear transformation layer, which maps the graph input matrix to a 256*3 matrix and reduces it to one dimension through average pooling to obtain a hidden state representation vector with a length of 256.

5. According to claim 1, a multimodal source code vulnerability static detection method based on Transformer and GAT is characterized by: The step D comprises the following specific steps: D.

1. Classification: Use a linear transformation layer to perform binary classification on the fused representation vector obtained in step C, and use the ReLU activation function to map the classification result to the range of 0 to 1; D.

2. Cross-validation: Divide the dataset into training set, validation set, and test set in a ratio of 7:2:1; use the training set to train the classification process to update the model parameters, and test the deep learning model once every 10 rounds of training through the validation set to monitor its classification accuracy: if qualified, test it again on the test set; if not qualified, continue training; save the deep learning model that has qualified detection accuracy on both the validation set and the test set as the final vulnerability detection model.

Images