Enhanced graph representation learning and transformer combined vulnerability detection method, system and device

Abstract

The invention belongs to the field of software testing, and particularly relates to a vulnerability detection method, system and device combining enhanced graph representation learning and transformer. According to the method, a vulnerability detection model comprising a preprocessing module, a feature extraction module and a classification module is constructed. A large amount of source codes containing vulnerability label information are obtained to form a training set and a test set, and a ternary loss set containing focusing loss is utilized to train and test a vulnerability detection module; and finally, performing vulnerability detection on the source code by using the tested vulnerability detection model. The preprocessing module generates code slices and slice sub-graphs according to source codes. The feature extraction module carries out feature extraction by adopting CodeBERT and EA-GGNN models and obtains a fusion feature vector. And the classification module adopts the MLP and generates a vulnerability detection result of the source code based on the fused feature vector. According to the method, the problem that the vulnerability detection efficiency, precision and generalization are insufficient due to the fact that an existing method cannot fully mine the dependency relationship between codes and global information is solved.

Description

Technical Field

[0001] The present invention belongs to the field of software testing, and in particular relates to a vulnerability detection method, system and device combining enhanced graph representation learning with transformer. Background Art

[0002] Source code vulnerability detection is an important research direction in software engineering and network security. With the increase in the scale and complexity of software systems, potential vulnerabilities in the code pose a serious threat to the security and stability of the system. At present, vulnerability detection methods are mainly divided into static analysis methods, dynamic analysis methods and deep learning-based methods.

[0003] Among them, static analysis methods try to locate possible vulnerabilities by checking the structure and syntax characteristics of the source code. Such methods include rule-based analysis and symbolic execution. However, such methods have a strong dependence on code complexity and diversity, and are prone to high false positive and false negative rates when faced with vulnerabilities with complex code logic or cross-function calls. Dynamic analysis methods discover potential vulnerabilities by actually running the program and monitoring its behavior, such as fuzz testing and sandbox testing. This method can effectively capture abnormal behavior in dynamic execution, but it often requires high computing resources and it is difficult to cover all execution paths.

[0004] At present, there are two main types of vulnerability detection models based on deep learning. One is a sequence-based model, which converts source code into a token sequence and inputs it into a deep learning model to learn vulnerability features and detect vulnerabilities. The serialized input makes the code representation have serialized information. The other is a graph-based model, which converts source code into a graph and inputs it into a vulnerability detection model. The graph includes AST abstract syntax tree, PDG program dependency graph, CPG code property graph, etc. The code representation has structured information by converting it into a graph. The above two models have their own advantages and disadvantages. For example, the sequence-based model has advantages in learning the global information of the code, especially when using the large pre-trained model CodeBERT to learn code features, but the sequence-based model often ignores the structured information of the code, such as the dependencies between codes, such as data dependency and control dependency. The graph-based model can represent the code in a structured manner, but it can only be used to represent the dependencies between codes, and cannot represent the global information of the code.

[0005] In addition, the performance of vulnerability detection models based on deep learning will also be affected by the quality of the dataset. This is because real-world programs are often very long, with hundreds or even thousands of lines, while vulnerability-related codes only account for dozens or even a few lines of them, and the vulnerable functions in the datasets we can use to train models only account for a small part of them, which leads to an extremely unbalanced ratio of vulnerability classes and non-vulnerability classes in the data used to train the model, which makes it easier for our model to learn features that are not related to vulnerabilities, resulting in a decrease in the accuracy and generalization of the model. Summary of the invention

[0006] In order to solve the problem that the existing vulnerability detection methods based on deep learning cannot fully mine the dependencies and global information between codes, which leads to low vulnerability detection efficiency and precision, and insufficient accuracy and generalization of the model, the present invention provides a vulnerability detection method, system and device combining enhanced graph representation learning with transformer.

[0007] The technical solution provided by the present invention is:

[0008] A vulnerability detection method combining enhanced graph representation learning and transformer includes the following process:

[0009] A vulnerability detection model including a preprocessing module, a feature extraction module and a classification module is constructed. The preprocessing module is used to generate code slices and their corresponding slice subgraphs containing all vulnerability-related risk nodes based on the source code. The feature extraction module includes a sequence branch and a graph branch. The sequence branch uses the CodeBERT model to extract features from the code slices to obtain a sequence feature vector. The graph branch uses the GGNN model containing the edge attention (EA) mechanism to extract features from the slice subgraph to obtain a graph feature vector. In the graph branch, the EA mechanism is used to dynamically assign weights to the input and output edges of the data flow and control flow and the self-loop edges, and fuse them into the embedded representation of the node, thereby obtaining a node representation that fuses heterogeneous edge information. Finally, the feature extraction module fuses the sequence feature vectors obtained in the two branches with the graph feature vector according to the preset weights to obtain the corresponding fused feature vector. The classification module uses MLP and is used to generate vulnerability detection results for the source code based on the input fused feature vector.

[0010] A large amount of source code containing vulnerability label information is obtained as sample data to form the original dataset, which is divided into a training set and a test set. focal , cosine loss L p and the regularization loss L reg The ternary loss L trpAs the loss function of MLP in the training phase, the training set and test set are used to train and test the vulnerability detection module.

[0011] The model parameters of the tested vulnerability detection model with the best performance are retained and used to perform vulnerability detection on the source code.

[0012] As a further improvement of the present invention, the process of the preprocessing module generating code slices and slice subgraphs according to the source code includes:

[0013] First, the code parsing tool joern is used to generate the AST and PDG of the source code. Then, the vulnerability risk code lines in the AST are matched according to the vulnerability syntax features provided by Checkmarx, and used as the slice center. Next, Joern is used to traverse the PDG forward and backward along the edges of the data flow and the control flow starting from the slice center, and then the corresponding slice subgraph is obtained. Finally, the same method is used to generate the slice subgraphs corresponding to all the matched slice centers; the code line numbers of the slice subgraphs are extracted, and the required code slices are formed by all the code lines containing the slice center.

[0014] As a further improvement of the present invention, the vulnerability detection model includes the following two data processing strategies:

[0015] (1) First, the preprocessing module identifies all slice centers in the source code and generates a set of corresponding code slices and slice subgraphs for each slice center. Then, the feature extraction module extracts the sequence feature vectors and graph feature vectors corresponding to each set of code slices and slice subgraphs, and fuses them to obtain multiple fused feature vectors. Finally, the MLP performs vulnerability detection on each fused feature vector and locates the vulnerability based on the detection results.

[0016] (2) First, the preprocessing module identifies all slice centers in the source code, and uses the shortest code behavior containing all slice centers as the code slice, and generates the slice subgraph corresponding to the code slice. Then, the feature extraction module extracts the sequence feature vector and graph feature vector corresponding to the code slice and slice subgraph, and fuses them to obtain a fused feature vector. Finally, the MLP performs vulnerability detection on the fused feature vector and generates the corresponding vulnerability detection result.

[0017] As a further improvement of the present invention, the CodeBERT model is a large-scale pre-trained model based on transformer, which is used to extract the serialized code global information contained in the code slices, and then obtain a 256-dimensional sequence feature vector.

[0018] and / or

[0019] The CodeBERT model consists of an Embedding module for vocabulary embedding and position encoding, and a 12-layer encoder consisting of a feedforward neural network and a bidirectional multi-head self-attention mechanism.

[0020] As a further improvement of the present invention, the slice subgraph contains multiple nodes, as well as data inflow edges, data outflow edges, control inflow edges, control outflow edges and self-loop edges between nodes. The EA-GGNN adopted by the graph branch dynamically assigns weights to different types of edges through the EA mechanism, effectively integrating heterogeneous edge information into the node representation; and combining the iterative mechanism of the traditional GGNN to capture local dependencies and remote node dependencies; and then obtain the 256-dimensional graph feature vector corresponding to the slice subgraph.

[0021] As a further improvement of the present invention, the EA mechanism dynamically learns the contribution of each edge type to the node feature update. The edge e between any two neighboring nodes i and j ij The attention weight The expression is as follows:

[0022]

[0023] In the above formula, and Respectively represent the feature vectors of nodes i and j in the initial state; W k represents the trainable weight matrix corresponding to k types of edges; a represents the attention vector used to determine the importance of the edge; σ represents the activation function.

[0024] As a further improvement of the present invention, GGNN updates the attention weights of all edges according to The expression for updating the feature vector of any node is as follows:

[0025]

[0026] In the above formula, represents the feature vector of node i after the t+1th iteration, N k (i) represents the set of neighbor nodes connected to node i through edges of type k; M is the set of all edge types; represents the feature vector of node j after the tth iteration.

[0027] As a further improvement of the present invention, the ternary loss L reg The expression is as follows:

[0028]

[0029] In the above formula, β and δ represent the cosine loss L p and the regularization loss Lreg In the ternary loss L trp The weight in t represents the predicted probability of the true category; represents a balance factor for adjusting the importance of positive and negative samples; γ represents a focusing parameter for adjusting the contribution of correctly classified samples; x g represents the input sample; h(x g ) represents x g The potential representation of h(x same ) represents the same as h(x g ) belongs to the potential representation of samples of the same category; h(x diff ) represents the same as h(x g ) is the potential representation of samples belonging to different categories; ε is a hyperparameter used to define the minimum separation boundary; D(h(x g ),h(x same )) represents h(x g ) and h(x same ) between the cosine distance; D(h(x g ),h(x diff )) represents h(x g ) and h(x diff ) is the cosine distance between .

[0030] The present invention also includes a vulnerability detection system, which is used to generate corresponding vulnerability detection results according to the output software source code. The vulnerability detection system includes a code acquisition unit and a vulnerability detection model trained in the vulnerability detection method combining enhanced graph representation learning and transformer as described above. The code acquisition unit is used to acquire the source code of the software to be detected; the vulnerability detection model is used to generate a detection result of whether the input source code includes a vulnerability according to the input source code.

[0031] The vulnerability detection module includes a preprocessing module, a feature extraction module, and a classification module. The preprocessing module is used to generate code slices and their corresponding slice subgraphs containing all vulnerability-related risk nodes based on the source code. The feature extraction module includes a sequence branch and a graph branch; the sequence branch uses the CodeBERT model to extract features from the code slices and obtain a sequence feature vector. The graph branch uses the GGNN model containing the EA mechanism to extract features from the slice subgraph and obtain a graph feature vector. In the graph branch, the EA mechanism is used to dynamically assign weights to the input and output edges of the data flow and control flow and the self-loop edges, and fuse them into the embedded representation of the node. The feature extraction module fuses the sequence feature vector with the graph feature vector according to the preset weights to obtain the corresponding fused feature vector; the classification module uses MLP and is used to generate the vulnerability detection results of the source code based on the input fused feature vector.

[0032] The present invention also includes a vulnerability detection device, which includes a memory, a processor, and a computer program stored in the memory and running in the processor. When the processor executes the computer program, a vulnerability detection system as described above is created, thereby realizing the detection of whether there are vulnerabilities in the input software source code.

[0033] The technical solution provided by the present invention has the following beneficial effects:

[0034] The present invention designs a new EA-GGNN network in the feature extraction part of the vulnerability detection model. Compared with the traditional GGNN module, the EA-GGNN designed by the present invention uses slicing to reduce the input size of the GGNN on the one hand; on the other hand, it uses the edge attention mechanism to fuse heterogeneous edge information into the node representation, thereby being able to obtain more accurate and reliable vulnerability-related feature information.

[0035] In contrast, many current solutions often use SMOTE to oversample the minority classes in the dataset when faced with an unbalanced dataset. However, SMOTE has great limitations, especially when dealing with complex vulnerability data in the real world. For example, if there is noise in the minority class samples, SMOTE may use the noise samples to generate new synthetic samples, thereby reducing model performance. In addition, the samples generated by SMOTE are likely to be repetitive, reducing model performance.

[0036] The provided vulnerability detection model introduces an MLP module with Focal Loss in the classification part. This improvement can overcome the defect that the traditional cross entropy loss function gives too high weight to negative samples (code without vulnerabilities) in unbalanced data, making it easier for the model to learn features unrelated to vulnerabilities. The Focal Loss used in the present invention dynamically allocates weights of positive and negative samples, so that the model pays more attention to vulnerability codes when facing unbalanced data sets, which greatly increases the effectiveness of the model. BRIEF DESCRIPTION OF THE DRAWINGS

[0037] Figure 1 This is a flowchart of the steps of a vulnerability detection method combining enhanced graph representation learning and transformer provided in Example 1 of the present invention.

[0038] Figure 2 This is a model architecture diagram of the vulnerability detection model including sequence branches and graph branches designed in Example 1 of the present invention.

[0039] Figure 3 is an example of a simple PDG graph containing all different types of edges.

[0040] Figure 4 This is a module diagram of a vulnerability detection system provided in Example 2 of the present invention. DETAILED DESCRIPTION

[0041] In order to make the purpose, technical solution and advantages of the present invention more clearly understood, the present invention is further described in detail below in conjunction with the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain the present invention and are not used to limit the present invention.

[0042] Example 1

[0043] This embodiment provides a vulnerability detection method that combines enhanced graph representation learning with transformer, such as Figure 1 As shown, the method includes the following three steps:

[0044] S1: Build a vulnerability detection model with an innovative architecture.

[0045] The vulnerability detection model created in this embodiment first generates code slices and slice subgraphs related to each risk node with potential vulnerabilities contained in the source code of the input software program. Then, the code slices and slice subgraphs are extracted using a transformer-based serialization model (such as CodeBERT) and a gated graph neural network (GGNN) model (i.e., EA-GGNN) containing an edge attention (EA) mechanism; and a fused feature vector containing two types of feature information is generated. Finally, the improved multi-layer perceptron (MLP) uses the fused feature vector to detect whether the source code contains vulnerabilities.

[0046] S2: Train and test the constructed vulnerability detection model.

[0047] In this embodiment, a large amount of source code containing vulnerability label information is obtained as sample data to form an original data set, and the original data set is divided into a training set and a test set. focal , cosine loss L p and the regularization loss L reg The ternary loss L trp As the loss function of MLP in the training phase, the training set and test set are used to train and test the vulnerability detection module.

[0048] S3: The model parameters of the tested vulnerability detection model with the best performance are retained and used to perform vulnerability detection on the source code.

[0049] In the vulnerability detection method provided in this embodiment, if Figure 2 As shown in FIG. 1 , the designed vulnerability detection module includes a vulnerability detection model of a preprocessing module, a feature extraction module and a classification module. The preprocessing module is used to generate code slices and their corresponding slice subgraphs containing all vulnerability-related risk nodes according to the source code.

[0050] It is well known to those skilled in the art that for the source code of most software programs, vulnerabilities are mainly caused by some key lines of code. In order to more accurately extract features and identify and classify codes with potential vulnerability risks, and reduce the interference of redundant information corresponding to a large number of vulnerability-irrelevant codes on the solution, this embodiment chooses to preprocess the source code of the software. Based on this purpose, the preprocessing process of this embodiment is mainly used to identify potential vulnerability risk nodes in the software program, and then retain the code fragments containing each vulnerability risk node and having complete functions. And delete the code lines other than the code of this part of the function, so as to obtain a more "concentrated" "code slice (Program Slice)". In addition, in order to enable the vulnerability detection model to further extract the local dependencies and long-distance dependencies contained in the code slice, the preprocessing module of this embodiment also uses the source code to generate the program dependency graph (PDG) corresponding to the code slice, and uses it as a slice subgraph.

[0051] In actual application, code slices and slice subgraphs can be generated in the following ways:

[0052] First, the code parsing tool Joern is used to generate the abstract syntax tree (AST) and PDG of the source code. Then, according to the vulnerability syntax feature matching function provided by Checkmarx, the vulnerability risk code lines of the source code are obtained from the AST. The vulnerability features provided by Checkmarx include the use of abbreviation arrays, the use of pointers, expression declarations, etc. This embodiment uses these features to match the risk nodes related to vulnerabilities contained in the AST, and uses these nodes as slice centers. Next, Joern is used to traverse forward and backward on the PDG according to the edges of the data flow and the edges of the control flow with the slice center as the starting point, and then the corresponding slice subgraph is obtained. The same method is used to generate the slice subgraphs corresponding to all the matched slice centers.

[0053] After completing the slicing operation of the source code and obtaining the corresponding slice subgraph, the code line number of the slice subgraph can be further extracted. All the code lines corresponding to each slice subgraph constitute the required code slice.

[0054] It should be noted that in actual application, the source code of each software to be tested may be long or short, and the source code may include one function or multiple functions. Therefore, after matching and identifying by Checkmarx, there may be more than one key node (or slice center) with potential vulnerability risks. In view of this situation, the preprocessing module of this embodiment may also include the following multiple situations when generating code slices and slice subgraphs using source code:

[0055] (1) When the source code has only one slice center, a set of code slices and slice subgraphs are generated.

[0056] (2) When the source code contains more than one slice center, multiple groups of code slices and slice subgraphs corresponding to the number of slice centers can be generated. Alternatively, multiple code slices and multiple slice subgraphs can be merged separately to obtain a group of code slices containing multiple slice centers and multiple slice subgraphs.

[0057] by Figure 2 Taking the software source code containing 19 lines of code as an example, when generating code slices and slice subgraphs, the 5th element of the array arr in the example() function in the Checkmarx analysis code is directly assigned to 10, and the pointer ptr is initialized to point to the first address of the array arr. In the pointer dereference operation on line 12, the pointer ptr is used to modify the value of the array, and the array operations on lines 5 and 8 are the key to the vulnerability. The real cause of the vulnerability is the potential insecurity of the pointer dereference operation and the illegal access to the array, and the rest of the code is irrelevant to the vulnerability. Therefore, the array arr and the pointer ptr are selected as the slice center, and the code slices and slice subgraphs related to them are extracted. The final code slice only contains lines 1 to 4, 9 to 12, and 17 and 18 of the source code, and the total number of code lines of the code slice is 10. In addition, the PDG (i.e., slice subgraph) corresponding to the code slice contains 6 nodes, 5 control dependency edges and 1 data dependency edge.

[0058] In the provided vulnerability detection model, after obtaining the initial representation of the source code consisting of code slices and slice subgraphs, this embodiment introduces two models to extract feature information. One is a transformer-based serialization model, and the other is a graph-based structured model. Specifically, Figure 2 As shown in the figure, the feature extraction module includes a sequence branch and a graph structure branch (referred to as the graph branch). The sequence branch uses the CodeBERT model to extract features from code slices, and then obtains the corresponding sequence feature vector. The graph branch uses the GGNN model containing the EA mechanism to extract features from the slice subgraph, and then obtains the corresponding graph feature vector. According to the introduction of the above process, since code slices and slice subgraphs are data with more "concentrated" feature information generated based on the source code, the sequence branch and the graph branch can obtain feature information that is more highly correlated with the vulnerability features in the code, and reduce the interference of a large amount of irrelevant information contained in the input sample data.

[0059] Specifically, in the feature extraction module of the vulnerability detection model of this embodiment, code slices, like source code, are still data in sequence format. This embodiment inputs them into the large-scale pre-trained model Codebert based on transformer to obtain serialized code global information. After learning, Codebert generates the overall representation of the code slice in the form of a 256-dimensional vector.

[0060] The Codebert model used in this example performs well in various natural language processing tasks. Figure 2 As shown, the CodeBERT model used in this embodiment includes, in sequence, an Embedding module for vocabulary embedding and position encoding, and a 12-layer encoder including a feedforward neural network and a bidirectional multi-head self-attention mechanism. The pre-training model used in this embodiment generally follows the paradigm of training on a large-scale general corpus. For a specialized source code corpus, a customized method is required. CodeBERT is the first model to adopt a bimodal intervention training strategy that can capture the semantic relationship between natural language (NL) and programming language (PL). The model has been pre-trained on six programming languages. This embodiment uses CodeBERT to encode program slices, thereby providing more accurate initialization input for the detection model. Specifically, in practical applications, CodeBERT includes two pre-training tasks: masked language modeling (MLM) and replacement tag detection (RTD). In the pre-training stage of the CodeBERT model, the loss function it adopts is defined as:

[0061] min θ L MIM (θ)+L RTD (θ)

[0062] Among them, L RTD (θ) denotes the replacement label detection loss; L MIM (θ) represents the masked language modeling loss and satisfies:

[0063]

[0064] In the above formula, x i represents the input sample; p D1 represents a discriminator that predicts tokens from a large vocabulary; m w represents the set of masked natural language tokens, m c represents the set of masked code tokens; w masked represents the masked natural language context, c masked Represents the masked code context;.

[0065] It should be emphasized that: although CodeBERT performs well in natural language processing tasks and can learn the global features of the code very well. However, due to its serialized input characteristics, it is destined to be unable to directly capture the dependencies between nodes. In order to overcome the defects of a single sequence branch in extracting source code feature information, this embodiment inputs the generated slice subgraph corresponding to the code slice into the gated graph neural network (Edge Aware Gated Graph Neural Network: EA-GGNN) with edge attention mechanism newly proposed in this embodiment. EA-GGNN is used to extract the dependencies between nodes contained in the slice subgraph related to each code slice.

[0066] In many existing methods, technicians usually use GGNN as a model. GGNN performs well in source code vulnerability detection tasks and can better capture the dependencies between nodes, but it also has limitations. For example, GGNN learns features by iteratively propagating information from neighboring nodes, but it is difficult to capture global structural information. In particular, most software code lines with vulnerabilities in actual applications are often very long, which is a challenge for GGNN iterative learning. In addition, GGNN iteration ignores the impact of heterogeneous edge information on vulnerability detection. The EA-GGNN model with edge attention mechanism provided in this embodiment overcomes the above-mentioned defects of GGNN.

[0067] Specifically, Figure 3 As shown, the slice subgraph contains multiple nodes, as well as data inflow edges, data outflow edges, control inflow edges, control outflow edges, and self-loop edges between nodes. In the EA-GGNN provided in this embodiment, through the edge attention mechanism, weights are dynamically assigned to different types of edges (for example, it can be distinguished whether the edge between nodes is a control edge or a data edge, and whether it belongs to an in-edge or an out-edge, and variable weights are set accordingly), and heterogeneous edge information is effectively integrated into the node representation, and the iterative mechanism of the traditional GGNN is combined to capture local dependencies and remote node dependencies.

[0068] In addition, for the slice subgraph mentioned above, it belongs to the preprocessed data. Compared with the program dependency graph PDG that directly generates the source code, the input slice subgraph greatly reduces the size of the graph structure data input into the EA-GGNN model, making the data processing efficiency of the enhanced graph representation learning (EA-GGNN) module provided by this embodiment higher. Finally, in the graph branch, this embodiment uses graph embedding to fuse the vector representation of each node of the slice subgraph after EA-GGNN learning into an overall representation of the slice subgraph.

[0069] Specifically, in the graph branch, EA-GGNN first uses word2vec to perform vector embedding on the extracted code slices, so that each node is embedded as a 128-dimensional feature vector. For each embedded node, only the information of the node itself is contained. On this basis, the feature extraction module of this embodiment further uses the edge attention mechanism to fuse the heterogeneous edge information related to the node into the node representation. For a node, this embodiment calculates the attention weights of all its neighboring nodes and the edges between the node, and ensures that the sum of the attention weights of all calculated edges is equal to 1. Specifically, the EA mechanism dynamically learns the contribution of each edge type to the node feature update. The edge e between any two neighboring nodes i and j ij The attention weight The expression is as follows:

[0070]

[0071] In the above formula, and Respectively represent the feature vectors of nodes i and j in the initial state; W k represents the trainable weight matrix corresponding to k types of edges; a represents the attention vector used to determine the importance of the edge; σ represents the activation function.

[0072] After calculating the attention weights of all edges, the feature representation of the node is updated by aggregating the features of its neighboring nodes weighted by the attention scores. This step is to assign different weights to different edges according to the edge type, and the result is a 128-dimensional vector for each node after the update. Specifically, the representation of each node depends not only on the features of its neighboring nodes, but also combines the edge weight and edge type information. This design enables the message passing mechanism to distinguish different types of dependencies and dynamically adjust the influence of neighboring nodes on the target node. After multiple iterations, the node features are gradually updated and finally capture the global dependencies in the program graph. Specifically, GGNN updates the feature representation of each node based on the attention weights of all edges after the update. The expression for updating the feature vector of any node is as follows:

[0073]

[0074] In the above formula, represents the feature vector of node i after the t+1th iteration, N k (i) represents the set of neighbor nodes connected to node i through edges of type k; M is the set of all edge types; represents the feature vector of node j after the tth iteration.

[0075] Next, this embodiment inputs the 128-dimensional vector of each node into the GGNN, and the GGNN assigns a GRU to each node to iteratively update the feature representation of the node, and finally outputs the 256-dimensional hidden state of each node. In summary, the EA-GGNN model designed in this embodiment uses graph embedding to fuse the 256-dimensional feature vector of each node in the graph into a feature vector as the feature vector of the final slice subgraph.

[0076] In the feature extraction module, two sets of 256-dimensional feature vectors can be obtained through branches composed of two different network models. In order to construct a comprehensive representation of the input code, the feature extraction module of this embodiment fuses the outputs of the sequence-based model (CodeBERT) and the graph-based model (EA-GGNN). The sequence-based model captures the global context semantics, while the graph-based model focuses on local dependencies and structural relationships. The fusion process is as follows:

[0077] E fusion =λ1E seq +λ2E graph

[0078] Among them, E fusion represents the fused feature vector; E seq represents the sequence feature vector; E graph represents the graph feature vector; λ1 and λ2 are weights controlling the contribution of the two models. In this embodiment, λ1=λ2=0.5 is set to ensure that the contribution of the feature representation extracted by the two models is equal.

[0079] In the vulnerability detection model constructed in this embodiment, a multi-layer perceptron (MLP) is used in the final detection stage to classify the final representation of the code obtained in the previous stage. In the training stage of MLP, traditional cross entropy loss functions often perform poorly on the problem of class imbalance because they are overly affected by easy-to-classify samples, thereby reducing the training weights assigned to difficult-to-classify samples. To solve this problem, this embodiment uses a ternary loss function L containing a focal loss (FocalLoss) in the training stage of the network model. reg The cross entropy loss is replaced by , which can pay more attention to samples that are difficult to classify.

[0080] Specifically, the ternary loss function L used in this embodiment is reg The expression is as follows:

[0081]

[0082] In the above formula, β and δ represent the cosine loss L p and the regularization loss L reg In the ternary loss L trp The weight int represents the predicted probability of the true category; represents a balance factor for adjusting the importance of positive and negative samples; γ represents a focusing parameter for adjusting the contribution of correctly classified samples; x g represents the input sample; h(x g ) represents x g The potential representation of h(x same ) represents the same as h(x g ) belongs to the potential representation of samples of the same category; h(x diff ) represents the same as h(x g ) is the potential representation of samples belonging to different categories; ε is a hyperparameter used to define the minimum separation boundary; D(h(x g ),h(x same )) represents h(x g ) and h(x same ) between the cosine distance; D(h(x g ),h(x diff )) represents h(x g ) and h(x diff ) is the cosine distance between .

[0083] In addition, it is necessary to further explain that: considering that the source code of each software can be preprocessed into one or more sets of code slices and slice subgraphs according to the number of high-risk nodes with potential vulnerability risks contained therein, when the source code can be converted into multiple sets of code slices and slice subgraphs, the vulnerability detection model can include the following two data processing strategies:

[0084] Strategy 1: Partition Identification

[0085] First, the preprocessing module identifies all slice centers in the source code and generates a set of corresponding code slices and slice subgraphs for each slice center. Then, the feature extraction module extracts the sequence feature vectors and graph feature vectors corresponding to each set of code slices and slice subgraphs, and fuses them to obtain multiple fused feature vectors. Finally, MLP performs vulnerability detection on each fused feature vector. Among them, since each fused feature vector corresponds to one part of the source code, when any fused feature vector is identified as having a vulnerability, the vulnerability contained in the detection result can be located according to the position of the code slice corresponding to the fused feature vector in the source code. Correspondingly, when all fused feature sequences are identified as not containing vulnerabilities, it means that the software corresponding to the source code does not have logical vulnerabilities.

[0086] Strategy 2: Overall Identification

[0087] First, the preprocessing module identifies all slice centers in the source code, and uses the shortest code behavior containing all slice centers as the code slice, and generates the slice subgraph corresponding to the code slice. Then, the feature extraction module extracts the sequence feature vector and graph feature vector corresponding to the code slice and slice subgraph, and fuses them to obtain a fused feature vector. Finally, the MLP performs vulnerability detection on the fused feature vector and generates the corresponding vulnerability detection result.

[0088] In summary, in view of the problem that the existing deep learning-based methods cannot obtain effective code representation and cannot effectively deal with the problem that the model cannot effectively learn the characteristics of vulnerability code due to data imbalance, the solution proposed by the present invention integrates the enhanced graph representation learning with the code representation obtained by the transformer-based serialization model to obtain a code representation that integrates serialization and structured information, and performs vulnerability detection through an improved multi-layer perceptron, thereby improving the effectiveness and generalization of the model in dealing with real-world source code vulnerabilities.

[0089] Example 2

[0090] Based on the solution of Example 1, this embodiment further provides a vulnerability detection system, which is used to generate corresponding vulnerability detection results according to the output software source code. Figure 4 As shown, the vulnerability detection system includes a code acquisition unit and a vulnerability detection model trained in the vulnerability detection method combining enhanced graph representation learning and transformer as described above. The code acquisition unit is used to acquire the source code of the software to be detected; the vulnerability detection model is used to generate a detection result of whether the input source code includes a vulnerability according to the input source code.

[0091] Specifically, the vulnerability detection module includes a preprocessing module, a feature extraction module, and a classification module. The preprocessing module is used to generate code slices containing all vulnerability-related risk nodes and their corresponding slice subgraphs based on the source code. The preprocessing module implements this process by built-in or calling the code parsing tool Joern and the commercial Checkmarx tool.

[0092] The feature extraction module includes a sequence branch and a graph branch. The sequence branch uses the CodeBERT model to extract features from code slices and obtain sequence feature vectors. The sequence branch implements this process by building in or calling an existing mature pre-trained CodeBERT model.

[0093] The graph branch uses the newly designed GGNN model including the EA mechanism in this embodiment to extract features from the slice subgraph and obtain the graph feature vector. In the graph branch, the EA mechanism is used to dynamically assign weights to the input and output edges of the data flow and control flow and the self-loop edges during the iterative update of the node vector in the GGNN model, and fuse them into the embedded representation of the node. The feature extraction module also includes a feature fusion unit, which fuses the sequence feature vector with the graph feature vector according to the preset weight (1:1) to obtain the corresponding fused feature vector.

[0094] The classification module adopts MLP and is used to generate vulnerability detection results of source code based on the input fused feature vector.

[0095] Example 3

[0096] On the basis of the scheme in Example 1, this embodiment further provides a vulnerability detection device, which includes a memory, a processor, and a computer program stored in the memory and running in the processor. When the processor executes the computer program, a vulnerability detection system as described above is created, thereby realizing the detection of whether there are vulnerabilities in the input software source code. The vulnerability detection device provided in this embodiment is essentially a computer device for implementing the scheme in Example 1. In actual application, the computer device can adopt an intelligent terminal, a tablet computer, a laptop computer, a desktop computer, a rack server, a blade server, a tower server or a cabinet server (including an independent server, or a server cluster composed of multiple servers) that can execute a program.

[0097] The computer device indicated in this embodiment includes at least but is not limited to: a memory and a processor that can be connected to each other through a system bus. Among them, the memory (i.e., a readable storage medium) includes a flash memory, a hard disk, a multimedia card, a card-type memory (e.g., an SD or DX memory, etc.), a random access memory (RAM), a static random access memory (SRAM), a read-only memory (ROM), an electrically erasable programmable read-only memory (EEPROM), a programmable read-only memory (PROM), a magnetic memory, a disk, an optical disk, etc. In some embodiments, the memory can be an internal storage unit of a computer device, such as a hard disk or a memory of the computer device. In other embodiments, the memory can also be an external storage device of a computer device, such as a plug-in hard disk equipped on the computer device, a smart memory card (Smart Media Card, SMC), a secure digital (SecureDigital, SD) card, a flash card (Flash Card), etc. Of course, the memory can also include both the internal storage unit of the computer device and its external storage device. In this embodiment, the memory is generally used to store an operating system and various application software installed on the computer device. In addition, the memory can also be used to temporarily store various types of data that have been output or are to be output.

[0098] The processor may be a central processing unit (CPU), a graphics processing unit (GPU), a controller, a microcontroller, a microprocessor, or other data processing chips in some embodiments. The processor is generally used to control the overall operation of a computer device. In this embodiment, the processor is used to run program codes stored in a memory or process data.

[0099] Simulation test

[0100] In order to verify the performance and advantages of the vulnerability detection method combining enhanced graph representation learning and transformer provided by the present invention, technicians formulated an experimental plan and simulated and tested relevant solutions.

[0101] 1. Performance comparison

[0102] First, this experiment selected four existing mature solutions, VulDeePecker, SySeVR, Devign, and Reveal, as the control group, and trained and tested the vulnerability detection performance of the control group and the solution of the present invention on three datasets, FFmpeg+Qemu, REVEAL, and Fan et al. During the test, Accuracy, Precision, Recall, and F1 were selected as performance evaluation indicators. During the comparative test, the vulnerability detection performance of the five models is shown in the following table:

[0103] Table 1: Comparison of vulnerability detection performance between the present invention and the control group solution By analyzing the data in Table 1, it can be found that the solution of the present invention has significantly improved performance on the three data sets compared with multiple representative control group models. The solution of the present invention has reached the highest level in almost all evaluation indicators, and is significantly higher than the existing solutions in accuracy and F1 score.

[0104] 2. Ablation Experiment

[0105] The main improvements made by the solution of the present invention include: (1) An enhanced gated graph neural network (EA-GGNN) combined with an edge attention mechanism is used in the graph branch to integrate heterogeneous edge information into the node representation. (2) The present invention combines the feature representation extracted from the graph branch sequence branch to achieve heterogeneous coding fusion. (3) The present invention uses a ternary loss function including focal loss to train the MLP in the network model training stage to alleviate the data imbalance in large-scale unbalanced data sets. In order to evaluate the contribution of the above improvements to the outstanding performance of the solution of the present invention, this experiment also conducted improved ablation experiments on the above three. The experimental content and results are as follows;

[0106] 2.1 Effectiveness of EA-GGNN

[0107] This experiment compares the scheme of the present invention with three schemes in which EA-GGNN is replaced with GGNN, graph convolutional network (GCN) and relational graph convolutional network (R-GCN) in three data sets. The final experimental data are shown in the following table:

[0108] Table 2: Ablation experiment results related to EA-GGNN By analyzing the data in the above table, we can find that the accuracy of the EA-GGNN module is 5.05% higher than that of other models, and the F1 score is 2.72% higher. This shows that EA-GGNN effectively integrates edge information into node representation and significantly enhances the model's graph representation learning ability.

[0109] 2.2 Effectiveness of Heterogeneous Coding Fusion

[0110] In order to verify whether the model of the solution of the present invention can improve performance by fusing graph-based and sequence-based models, this experiment was conducted on the FFmpeg+Qemu dataset, and the performance of the fusion solution of the present invention was compared with three sequence-based models (BGNN4VD, VulDeePecker, SySeVR) and three graph-based models (Devign, Reveal, EA-GGNN only). The experimental results are shown in the following table:

[0111] Table 2: Ablation experiment results related to heterogeneous coding fusion

[0112]

[0113] By analyzing the above experimental data, it can be found that compared with a single model, the accuracy of the fusion method of the present invention is improved by 5.46% to 18.69%, and the F1 score is improved by 9.09% to 34.57%. This proves that the heterogeneous encoding fusion method adopted by the present invention effectively integrates structured and serialized information and enhances the code representation ability of the model.

[0114] 2.3 Effectiveness of Ternary Loss + MLP

[0115] In order to verify whether the MLP trained by the ternary loss including focal loss adopted by the present invention can improve the performance of the model on the minority class. This experiment compares and analyzes the unbalanced data set proposed by predecessors with other three detection methods using cross entropy loss (SVM, RF, MLP). Among them, in the data set used in this experiment, the code containing vulnerabilities only accounts for 5.88% of the total data. The experimental results are shown in the following table:

[0116] Table 2: Ablation experiments on classification models and losses

[0117]

[0118] By analyzing the data in the above table, we can find that compared with the existing solutions, the accuracy and F1 score of the solution of the present invention are the best on the data set with fewer vulnerabilities. This shows that the ternary loss + MLP adopted by the present invention can significantly improve the performance of the model on the unbalanced data set and prove its effectiveness in the vulnerability detection task.

[0119] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions and improvements made within the spirit and principles of the present invention should be included in the protection scope of the present invention.

Claims

1. A vulnerability detection method combining enhanced graph representation learning and transformer, characterized in that: It includes: Build a vulnerability detection model including preprocessing module, feature extraction module and classification module; The preprocessing module is used to generate code slices and their corresponding slice subgraphs containing all vulnerability-related risk nodes according to the source code; the feature extraction module includes a sequence branch and a graph branch; the sequence branch uses the CodeBERT model to extract features from the code slice to obtain a sequence feature vector; the graph branch uses the GGNN model including the edge attention mechanism to extract features from the slice subgraph to obtain a graph feature vector; in the graph branch, the edge attention mechanism is used to dynamically assign weights to the input and output edges of the data flow and control flow and the self-loop edges, and fuse them into the embedded representation of the node; the feature extraction module fuses the sequence feature vector with the graph feature vector according to the preset weight to obtain the corresponding fused feature vector; the classification module uses MLP and is used to generate the vulnerability detection result of the source code according to the input fused feature vector; A large amount of source code containing vulnerability label information is obtained as sample data to form the original data set, and the original data set is divided into a training set and a test set; the focus loss L focal , cosine loss L p and the regularization loss L reg The ternary loss L trp As the loss function of MLP in the training phase, the vulnerability detection module is trained and tested using the training set and the test set; The model parameters of the tested vulnerability detection model with the best performance are retained and used to perform vulnerability detection on the source code.

2. The vulnerability detection method combining enhanced graph representation learning and transformer as claimed in claim 1, characterized in that: The process of the preprocessing module generating code slices and slice subgraphs according to the source code includes: First, the code parsing tool joern is used to generate the AST and PDG of the source code. Then, the vulnerability risk code lines in the AST are matched as the slice center according to the vulnerability syntax features provided by Checkmarx. Next, Joern is used to traverse the PDG forward and backward along the edges of the data flow and the control flow starting from the slice center to obtain the corresponding slice subgraph. Finally, the same method is used to generate the slice subgraphs corresponding to all the matched slice centers. The code line numbers of the slice subgraphs are extracted, and the required code slices are formed by all the code lines containing the slice center.

3. The vulnerability detection method combining enhanced graph representation learning and transformer as claimed in claim 2, characterized in that: The vulnerability detection model includes two data processing strategies: (1) First, the preprocessing module identifies all slice centers in the source code and generates a set of corresponding code slices and slice subgraphs for each slice center. Then, the feature extraction module extracts the sequence feature vectors and graph feature vectors corresponding to each set of code slices and slice subgraphs, and fuses them to obtain multiple fused feature vectors. Finally, the MLP performs vulnerability detection on each fused feature vector and locates the vulnerability based on the detection results. (2) First, the preprocessing module identifies all slice centers in the source code, takes the shortest code behavior containing all slice centers as the code slice, and generates the slice subgraph corresponding to the code slice; then, the feature extraction module extracts the sequence feature vector and graph feature vector corresponding to the code slice and slice subgraph, and fuses them to obtain a fused feature vector; finally, the MLP performs vulnerability detection on the fused feature vector to generate the corresponding vulnerability detection result.

4. The vulnerability detection method combining enhanced graph representation learning and transformer as claimed in claim 1, characterized in that: The CodeBERT model is a large-scale pre-trained model based on transformer, which is used to extract the serialized code global information contained in the code slice, and then obtain a 256-dimensional sequence feature vector; and / or The CodeBERT model includes an Embedding module for vocabulary embedding and position encoding, and a 12-layer encoder including a feedforward neural network and a bidirectional multi-head self-attention mechanism.

5. The vulnerability detection method combining enhanced graph representation learning and transformer as claimed in claim 1, characterized in that: The slice subgraph contains multiple nodes, as well as data inflow edges, data outflow edges, control inflow edges, control outflow edges and self-loop edges between nodes; the EA-GGNN adopted in the graph branch dynamically assigns weights to different types of edges through the EA mechanism, effectively integrates heterogeneous edge information into the node representation, and combines the iterative mechanism of traditional GGNN to capture local dependencies and remote node dependencies; and then obtains a 256-dimensional graph feature vector corresponding to the slice subgraph.

6. The vulnerability detection method combining enhanced graph representation learning and transformer as claimed in claim 5, characterized in that: The EA mechanism dynamically learns the contribution of each edge type to the node feature update. The edge e between any two neighbor nodes i and j ij The attention weight The expression is as follows: In the above formula, and Respectively represent the feature vectors of nodes i and j in the initial state; W k represents the trainable weight matrix corresponding to k types of edges; a represents the attention vector used to determine the importance of the edge; σ represents the activation function.

7. The vulnerability detection method combining enhanced graph representation learning and transformer as claimed in claim 6, characterized in that: The GGNN is based on the updated attention weights of all edges The expression for updating the feature vector of any node is as follows: In the above formula, represents the feature vector of node i after the t+1th iteration, N k (i) represents the set of neighbor nodes connected to node i through edges of type k; M is the set of all edge types; represents the feature vector of node j after the tth iteration.

8. The vulnerability detection method combining enhanced graph representation learning and transformer as claimed in claim 7, characterized in that: The ternary loss L reg The expression is as follows: In the above formula, β and δ represent the cosine loss L p and the regularization loss L reg In the ternary loss L trp The weight in t represents the predicted probability of the true category; represents a balance factor for adjusting the importance of positive and negative samples; γ represents a focusing parameter for adjusting the contribution of correctly classified samples; x g represents the input sample; h(x g ) represents x g The potential representation of h(x same ) represents the same as h(x g ) belongs to the potential representation of samples of the same category; h(x diff ) represents the same as h(x g ) is the potential representation of samples belonging to different categories; ε is a hyperparameter used to define the minimum separation boundary; D(h(x g ),h(x same )) represents h(x g ) and h(x same ) between the cosine distance; D(h(x g ),h(x diff )) represents h(x g ) and h(x diff ) is the cosine distance between .

9. A vulnerability detection system, characterized in that: It comprises a code acquisition unit and a vulnerability detection model trained in the vulnerability detection method combining enhanced graph representation learning and transformer as described in any one of claims 1 to 8; the code acquisition unit is used to acquire the source code of the software to be detected; the vulnerability detection model is used to generate a detection result of whether the input source code includes a vulnerability according to the input source code; The vulnerability detection model includes a preprocessing module, a feature extraction module and a classification module; the preprocessing module is used to generate a code slice containing all vulnerability-related risk nodes and its corresponding slice subgraph according to the source code; the feature extraction module includes a sequence branch and a graph branch; the sequence branch uses the CodeBERT model to extract features from the code slice to obtain a sequence feature vector; the graph branch uses the GGNN model including the edge attention mechanism to extract features from the slice subgraph to obtain a graph feature vector; in the graph branch, the edge attention mechanism is used to dynamically assign weights to the input and output edges of the data flow and control flow and the self-loop edges, and fuse them into the embedded representation of the node; the feature extraction module fuses the sequence feature vector with the graph feature vector according to the preset weight to obtain the corresponding fused feature vector; the classification module uses MLP and is used to generate a vulnerability detection result of the source code according to the input fused feature vector.

10. A vulnerability detection device, comprising a memory, a processor, and a computer program stored in the memory and executed in the processor, characterized in that: When the processor executes the computer program, it creates the vulnerability detection system as claimed in claim 9, thereby realizing the detection of whether there are vulnerabilities in the input software source code.

Images