Method, system and device for vulnerability detection combining enhanced graph representation learning and transformers

Abstract

The present application 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. The method constructs a vulnerability detection model comprising a preprocessing module, a feature extraction module and a classification module. A large amount of source code containing vulnerability label information is obtained to form a training set and a test set, and the vulnerability detection module is trained and tested by using a ternary loss set containing a focus loss; finally, the tested vulnerability detection model is used to detect the vulnerabilities of the source code. The preprocessing module generates code slices and slice subgraphs according to the source code. The feature extraction module uses the CodeBERT and EA-GGNN models to extract features and obtain fusion feature vectors. The classification module uses MLP and generates the vulnerability detection results of the source code based on the fusion feature vectors. The present application solves the problem that the existing method cannot fully mine the dependency relationship and global information between codes, resulting in insufficient vulnerability detection efficiency, precision and generalization.

Description

Technical Field

[0001] This invention belongs to the field of software testing, specifically relating to a vulnerability detection method, system, and apparatus that combines enhanced graph representation learning with transformers. Background Technology

[0002] Source code vulnerability detection is an important research area in software engineering and cybersecurity. As software systems become larger and more complex, potential vulnerabilities in the code pose a serious threat to system security and stability. Currently, vulnerability detection methods are mainly divided into static analysis methods, dynamic analysis methods, and deep learning-based methods.

[0003] Static analysis methods examine the structure and syntax of source code to attempt to locate potential vulnerabilities. These methods include rule-based analysis and symbolic execution. However, these methods are highly dependent on code complexity and diversity, and are prone to high false positive and false negative rates when dealing with vulnerabilities involving complex code logic or cross-function calls. Dynamic analysis methods discover potential vulnerabilities by actually running the program and monitoring its behavior, such as fuzzing and sandboxing. This method can effectively capture anomalous behavior during dynamic execution, but often requires significant computational resources and struggles to cover all execution paths.

[0004] Currently, deep learning-based vulnerability detection models mainly fall into two categories: sequence-based models and graph-based models. Sequence-based models convert source code into a sequence of tokens, which is then input into the deep learning model to learn vulnerability features and detect vulnerabilities. This serialized input gives the code representation sequential information. Graph-based models convert source code into a graph format, such as Abstract Syntax Trees (ASTs), Program Dependency Graphs (PDGs), and Code Attribute Graphs (CPGs). This graph format gives the code representation structured information. Each model has its own advantages and disadvantages. For example, sequence-based models excel at learning global code information, especially when using large pre-trained models like CodeBERT to learn code features. However, they often overlook structural information, such as dependencies between code segments, including data and control dependencies. Graph-based models can represent code in a structured way, but they can only represent dependencies between code segments and cannot represent global code information.

[0005] In addition, the performance of the vulnerability detection model based on deep learning is also affected by the quality of the data set. Because the real-world program is often long, with hundreds or even thousands of lines, and the vulnerability-related code only accounts for tens or even a few lines, and the vulnerability functions in the data set available for training the model only account for a small part, which leads to the imbalance of the proportion of vulnerability and non-vulnerability classes in the data used for training the model, which makes the model more likely to learn features irrelevant to vulnerabilities, resulting in a decrease in the accuracy and generalization of the model. SUMMARY

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

[0007] The technical scheme provided by the present application is:

[0008] A vulnerability detection method combining enhanced graph representation learning and transformer, comprising the following processes:

[0009] A vulnerability detection model comprising a preprocessing module, a feature extraction module and a classification module is constructed. The preprocessing module is used to generate code slices containing all vulnerability-related risk nodes and their corresponding slice subgraphs from source code. The feature extraction module includes a sequence branch and a graph branch. The sequence branch uses the CodeBERT model to extract features of the code slices to obtain a sequence feature vector. The graph branch uses a GGNN model containing an Edge Aware (EA) mechanism to extract features of the slice subgraphs to obtain a graph feature vector. In the graph branch, the EA mechanism is used to dynamically allocate weights for the incoming and outgoing edges of the data flow and control flow and the self-loop edge, and fuse them into the node embedding representation, thereby obtaining the node representation fused with the heterogeneous edge information. Finally, the feature extraction module fuses the sequence feature vector and the graph feature vector obtained in the two branches according to the preset weight, thereby obtaining 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.

[0010] A large number of source codes containing vulnerability label information are 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. A three-element loss L focal containing a focal loss L p , a cosine loss L reg and a regularization loss L trpAs the loss function of the MLP in the training stage, the training set and the test set are used to train and test the vulnerability detection module.

[0011] The model parameters of the tested vulnerability detection model with the optimal performance are reserved and used for vulnerability detection on the source code.

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

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

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

[0015] (1) Firstly, all the slice centers in the source code are identified by the preprocessing module, and a set of corresponding code slices and slice subgraphs are generated for each slice center. Then, the sequence feature vectors and graph feature vectors corresponding to each set of code slices and slice subgraphs are extracted by the feature extraction module, and a plurality of fusion feature vectors are obtained by fusion. Finally, the MLP performs vulnerability detection on each fusion feature vector, and realizes vulnerability positioning according to the detection result.

[0016] (2) Firstly, all the slice centers in the source code are identified by the preprocessing module, and the shortest code containing all the slice centers is used as the code slice, and the slice subgraph corresponding to the code slice is generated. Then, the sequence feature vectors and graph feature vectors corresponding to the code slice and the slice subgraph are extracted by the feature extraction module, and a fusion feature vector is obtained by fusion. Finally, the MLP performs vulnerability detection on the fusion feature vector to generate the corresponding vulnerability detection result.

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

[0018] And / or

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

[0020] As a further improvement of this invention, the sliced ​​subgraph contains multiple nodes, as well as data inflow edges, data outflow edges, control inflow edges, control outflow edges, and self-looping edges between nodes. The graph branching uses EA-GGNN to dynamically assign weights to different types of edges through the EA mechanism, effectively integrating heterogeneous edge information into the node representation; and combines the iterative mechanism of traditional GGNN to capture local dependencies and long-distance node dependencies; thus obtaining a 256-dimensional graph feature vector corresponding to the sliced ​​subgraph.

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

[0022]

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

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

[0025]

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

[0027] As a further improvement to 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, respectively. p and regularization loss Lreg In the ternary loss L trp , the weight; p t represents the predicted probability of the true class; represents a balance factor for adjusting the importance of positive and negative samples; represents a focusing parameter for adjusting the contribution degree of correctly classified samples; x g represents an input sample; h(x g ) represents the latent representation of x g ; h(x same ) represents the latent representation of the sample belonging to the same class as h(x g ); h(x diff ) represents the latent representation of the sample belonging to a different class from h(x g ); epsilon is a hyperparameter for defining the minimum separation boundary; D(h(x g ), h(x same )) represents the cosine distance between h(x g ) and h(x same ); D(h(x g ), h(x diff )) represents the cosine distance between h(x g ) and h(x diff ).

[0030] The application also includes a vulnerability detection system for generating corresponding vulnerability detection results according to the source code of the output software. The vulnerability detection system includes a code acquisition unit and a vulnerability detection model trained by 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 the detection result of whether the source code includes vulnerabilities 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 containing all vulnerability-related risk nodes and their corresponding slice subgraphs from 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 of the code slices to obtain sequence feature vectors. The graph branch uses the GGNN model containing the EA mechanism to extract features of the slice subgraphs to obtain graph feature vectors. In the graph branch, the EA mechanism is used to dynamically assign weights to the incoming and outgoing edges of the data flow and the control flow and the self-loop edge, and fuse them into the node embedding representation. The feature extraction module fuses the sequence feature vectors and the graph feature vectors according to the preset weights to obtain the corresponding fusion feature vectors; the classification module uses MLP and is used to generate the vulnerability detection result of the source code according to the input fusion feature vectors.

[0032] The application further comprises a vulnerability detection device, which comprises a memory, a processor, and a computer program stored in the memory and running in the processor.

[0033] The technical scheme provided by the application has the following beneficial effects:

[0034] The application designs a novel 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 application reduces the input size of the GGNN by using slicing, and fuses heterogeneous edge information into node representation by using edge attention mechanism, so as to obtain more accurate and reliable vulnerability-related feature information.

[0035] On the contrary, many current schemes often use SMOTE to oversample the minority class in the data set when facing an unbalanced data set. 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 sample, SMOTE may generate new synthetic samples using noise samples, thereby reducing the performance of the model, and the samples generated by SMOTE are likely to have duplication, thereby reducing the performance of the model.

[0036] The vulnerability detection model provided in the application 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 the negative samples (non-vulnerability code) in the unbalanced data, so that the model is more likely to learn vulnerability-irrelevant features. The Focal Loss used in the application dynamically allocates the weights of positive and negative samples, so that the model pays more attention to vulnerability code when facing an unbalanced data set, which greatly increases the effectiveness of the model. BRIEF DESCRIPTION OF DRAWINGS

[0037] Figure 1 A step flowchart of a vulnerability detection method combining enhanced graph representation learning and a transformer provided in embodiment 1 of the application.

[0038] Figure 2 A model architecture diagram of a vulnerability detection model designed in embodiment 1 of the application, which comprises a sequence branch and a graph branch.

[0039] Figure 3 An example of a simple PDG graph comprising all different types of edges.

[0040] Figure 4 A module schematic diagram of a vulnerability detection system provided in embodiment 2 of the application. Detailed Implementation

[0041] To make the objectives, technical solutions, and advantages of this invention clearer, the invention will be further described in detail below with reference to the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are merely illustrative and not intended to limit the invention.

[0042] Example 1

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

[0044] S1: Construct an innovative vulnerability detection model.

[0045] The vulnerability detection model created in this embodiment first generates code slices and slice subgraphs associated with each potentially vulnerable node in the input software program's source code. Then, it extracts features from the code slices and slice subgraphs using a transformer-based serialization model (such as CodeBERT) and a gated graph neural network (GGNN) model incorporating edge attention (EA) (i.e., EA-GGNN), respectively, generating a fused feature vector containing both types of feature information. Finally, an improved multilayer 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 the original dataset, which is then divided into a training set and a test set. The dataset includes focusing loss L... focal Cosine loss L p and regularization loss L reg The three-dimensional loss L trp As the loss function of MLP during the training phase, the vulnerability detection module is trained and tested using the training set and the test set.

[0048] S3: Retain the model parameters of the best-performing vulnerability detection model that has been tested, and use it to perform vulnerability detection on the source code.

[0049] In the vulnerability detection method provided in this embodiment, such as Figure 2 As shown, the designed vulnerability detection module comprises a preprocessing module, a feature extraction module, and a classification module. The preprocessing module generates code slices containing all vulnerability-related risk nodes and their corresponding slice subgraphs based on the source code.

[0050] It is well known to those skilled in the art that for most software programs, vulnerabilities are mainly caused by some key lines of code. In order to more accurately extract and identify the code with potential vulnerability risks and reduce the interference of a large amount of vulnerability-irrelevant code on the scheme, the embodiment selects the source code of the software for preprocessing. Based on this purpose, the preprocessing process of the embodiment is mainly used to identify potential vulnerability risk nodes in the software program, and then retain code segments containing each vulnerability risk node and having complete functions. In addition, the lines of code other than the part of the function are deleted, and a more "concentrated" "code slice" is obtained. In addition, in order to enable the vulnerability detection model to further extract the local dependency and long-distance dependency contained in the code slice, the preprocessing module of the embodiment also generates a program dependency graph (PDG) corresponding to the code slice by using the source code, and uses the PDG as a slice subgraph.

[0051] In actual application, the code slice and the slice subgraph can be generated in the following manner:

[0052] First, the abstract syntax tree (AST) and the PDG of the source code are generated by using the code parsing tool joern. Then, the vulnerability risk lines of code in the source code are matched from the AST according to the vulnerability syntax feature matching function provided by Checkmarx. The vulnerability features provided by Checkmarx include array use, pointer use, expression declaration and the like, and the embodiment uses these features to match the vulnerability-related risk nodes contained in the AST, and uses these nodes as the slice center. Then, the joern is used to perform forward and backward traversal on the PDG according to the data flow edges and the control flow edges starting from the slice center, and then the corresponding slice subgraph is obtained. The slice subgraphs corresponding to all the matched slice centers are generated by using the same method.

[0053] After the slicing operation of the source code is completed and the corresponding slice subgraph is obtained, the code line numbers 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 additionally noted that in actual application, the source code of each software to be detected has different lengths, and the source code can include one function or multiple functions. Therefore, after being matched and identified by Checkmarx, the key nodes (or slice centers) with potential vulnerability risks can also be more than one. In view of this situation, the preprocessing module of the embodiment can also include the following multiple situations when generating the code slice and the slice subgraph by using the 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 sets 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 set of code slices and multiple slice subgraphs that simultaneously contain multiple slice centers.

[0057] by Figure 2 Taking a software source code containing 19 lines of code as an example, during the generation of code slices and slice subgraphs, Checkmarx analysis revealed that in the `example()` function, the 5th element of the array `arr` was directly assigned the value 10, and the pointer `ptr` was initialized to point to the starting address of the array `arr`. In the pointer dereference operation on line 12, the pointer `ptr` was used to modify the array value, 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; the rest of the code is unrelated to the vulnerability. Therefore, the array `arr` and the pointer `ptr` were chosen as the slice center, and the related code slices and slice subgraphs were extracted. The final code slice only contains lines 1-4, 9-12, and 17 and 18 of the source code, with a total of 10 lines of code. Furthermore, the PDG (slice subgraph) corresponding to this code slice contains 6 nodes, 5 control dependency edges, and 1 data dependency edge.

[0058] In the provided vulnerability detection model, after obtaining an initial representation of the source code consisting of code slices and slice subgraphs, this embodiment introduces two models for feature information extraction: a transformer-based serialization model and a graph-based structured model. Specifically, as follows... Figure 2 As shown, the feature extraction module includes a sequence branch and a graph branch (referred to as the graph branch). The sequence branch uses the CodeBERT model to extract features from code slices, thereby obtaining the corresponding sequence feature vectors. The graph branch uses a GGNN model with an EA mechanism to extract features from slice subgraphs, thereby obtaining the corresponding graph feature vectors. As described above, since both code slices and slice subgraphs are generated from source code and contain more "condensed" feature information, the sequence branch and 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 the embodiment, the code slice and the source code are still sequence format data, and the embodiment inputs the code slice into the large-scale pre-training model Codebert based on the transformer to obtain the global information of the serialized code. The overall representation of the code slice is generated after learning by Codebert, which is in the form of a 256-dimensional vector.

[0060] The Codebert model used in the embodiment performs well in various natural language processing tasks. In actual applications, as shown in Figure 2 The Codebert model used in the embodiment sequentially includes an Embedding module for performing word embedding and position encoding, and 12 layers of encoders containing feedforward neural networks and bidirectional multi-head self-attention mechanisms. The pre-training model used in the embodiment generally follows the training paradigm on a large general corpus. For a specialized source code corpus, a customized method is needed. CodeBERT is the first model to use a dual-modal intervention training strategy, which can capture the semantic relationship between natural language (NL) and programming language (PL). The model has been pre-trained on six programming languages. The CodeBERT is used in the embodiment to encode the program slice, thereby providing more accurate initialization input for the detection model. Specifically, in actual applications, CodeBERT includes two pre-training tasks: masked language modeling (MLM) and replacement token detection (RTD). In the pre-training phase of the CodeBERT model, the loss function used is defined as:

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

[0062] where L RTD (θ) represents the replacement token detection loss; L MIM (θ) represents the masked language modeling loss, and satisfies:

[0063]

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

[0065] It is emphasized that: CodeBERT, although it performs well in natural language processing tasks, can learn the global features of the code well. But due to its sequential input characteristics, it is doomed that it cannot directly capture the dependency between nodes. In order to overcome the defects of single sequence branch in source code feature information extraction, the embodiment will input the generated slice subgraph corresponding to the code slice into the newly proposed edge-aware gated graph neural network (Edge Aware Gated Graph Neural Network: EA-GGNN). The EA-GGNN is used to extract the dependency between nodes contained in the slice subgraph related to each code slice.

[0066] In some existing methods, technicians usually use GGNN as a model, which performs well in source code vulnerability detection tasks and can better capture the dependency between nodes, but it also has limitations, such as GGNN learns features by iteratively propagating neighbor node information, but it is difficult to capture global structure information, especially in practical application process, most of the software code lines with vulnerabilities are often long, which is a challenge for the iterative learning of GGNN, and the iteration of GGNN ignores the influence of heterogeneous edge information on vulnerability detection. The EA-GGNN model provided in the embodiment overcomes the above defects of GGNN.

[0067] Specifically, as shown in Figure 3 The slice subgraph contains multiple nodes, as well as data flow-in edges, data flow-out edges, control flow-in edges, control flow-out edges, and self-loop edges between nodes. In the EA-GGNN provided in the embodiment, the edge attention mechanism is used to dynamically allocate weights to different types of edges (such as distinguishing whether the edge between nodes is a control edge or a data edge, and whether it belongs to an incoming edge or an outgoing edge, and setting the variable weight accordingly), effectively integrating heterogeneous edge information into node representation, and combining the iterative mechanism of traditional GGNN to capture local dependency and remote node dependency.

[0068] In addition, for the slice subgraph mentioned earlier, it belongs to preprocessed data. Compared with directly generating a program dependence graph (PDG) of 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 in the embodiment higher. Finally, in the graph branch, the vector representation of each node of the slice subgraph learned by the EA-GGNN is fused into a global representation of the slice subgraph through graph embedding.

[0069] Specifically, in the graph branch, the EA-GGNN first uses word2vec for 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 included. On this basis, the feature extraction module of the embodiment further uses an edge attention mechanism to fuse the heterogeneous edge information related to the node into the node representation. For a node, the embodiment scheme calculates the attention weights of all neighbor nodes and the edges between the nodes, and ensures that the sum of all calculated edge attention weights is equal to 1. Specifically, the EA mechanism dynamically learns the contribution of each edge type to the update of the node features. The expression of the attention weight of the edge e ij between any two neighbor nodes i and j 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 a trainable weight matrix corresponding to the k-type edge; a represents an attention vector for determining the importance of the edge; and σ represents an 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 adjacent nodes weighted by the attention scores. This step is to assign different weights to different edges according to their types, and the result is a 128-dimensional vector of each node after updating. Specifically, the representation of each node depends not only on the features of its adjacent nodes, but also on the edge weight and edge type information. This design enables the message passing mechanism to distinguish between different types of dependency relationships and dynamically adjust the influence of adjacent nodes on the target node. After multiple iterations, the node features are gradually updated, and finally the global dependency relationships in the program graph are captured. Specifically, the GGNN updates the feature vector of any node according to the updated attention weights of all edges 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+1 iteration, N k (i) represents the set of neighbor nodes connected to node i through edges of type k; and M is the set of all edge types. represents the feature vector of node j after the t iteration.

[0075] ​Next, the 128-dimensional vector of each node is input into the GGNN, which assigns a GRU to each node to iteratively update the node's feature representation, and finally outputs a 256-dimensional hidden state for 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 final feature vector of the slice subgraph.

[0076] In the feature extraction module, two groups of 256-dimensional feature vectors can be obtained through the 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 global contextual 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] where 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 that control the contribution of the two models. In this embodiment, λ1 = λ2 = 0.5 is set to ensure that the contributions of the features extracted by the two models are equal.

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

[0080] Specifically, the expression of the tri- loss function L reg used in this embodiment is as follows:

[0081]

[0082] In the above formula, β and δ represent the cosine loss L p and the regularization loss L reg respectively, and p is the weight in the tri- loss L trp .t a predicted probability representing a true class; represents a balance factor for adjusting the importance of positive and negative samples; γ represents a focusing parameter for adjusting the contribution degree of correctly classified samples; x g represents an input sample; h(x g represents a latent representation of x g ; h(x same ) represents a latent representation of a sample belonging to the same class as h(x g ); h(x diff ) represents a latent representation of a sample belonging to a different class as h(x g ); ε is a hyperparameter for defining a minimum separation boundary; D(h(x g ), h(x same )) represents a cosine distance between h(x g ) and h(x same ); D(h(x g ), h(x diff )) represents a cosine distance between h(x g ) and h(x diff ).

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

[0084] Strategy One: Partition Identification

[0085] First, all slice centers in the source code are identified by the preprocessing module, and a corresponding group of code slices and slice subgraphs is generated for each slice center. Then, the sequence feature vector and the graph feature vector corresponding to each group of code slices and slice subgraphs are extracted by the feature extraction module, and multiple fusion feature vectors are obtained by fusion. Finally, the MLP performs vulnerability detection on each fusion feature vector. Since each fusion feature vector corresponds to a part of the source code, when any fusion feature vector is identified as containing a vulnerability, the location of the vulnerability contained in the detection result can be determined according to the location of the code slice corresponding to the fusion feature vector in the source code. Accordingly, when all fusion feature sequences are identified as not containing vulnerabilities, it means that the software corresponding to the source code does not contain logical vulnerabilities.

[0086] Strategy Two: Overall Identification

[0087] Firstly, all slice centers in the source code are identified by the preprocessing module to include the shortest code behavior code slice of all slice centers, and the slice subgraph corresponding to the code slice is generated. Then, the sequence feature vector and the graph feature vector corresponding to the code slice and the slice subgraph are extracted by the feature extraction module, and a fusion feature vector is obtained by fusion. Finally, the MLP performs vulnerability detection on the fusion feature vector to generate the corresponding vulnerability detection result.

[0088] In summary, in view of the problems that the existing deep learning-based method cannot obtain effective code representation and cannot effectively process the data imbalance to make the model unable to effectively learn the vulnerability code features, the scheme provided by the present application fuses the enhanced graph representation learning and the code representation obtained based on the transformer-based serialization model to obtain a code representation fused with serialization and structural information, and performs vulnerability detection through an improved multilayer perceptron to improve the effectiveness and generalization of the model in dealing with source code vulnerabilities in the real world.

[0089] Embodiment 2

[0090] Based on the scheme of embodiment 1, the present embodiment further provides a vulnerability detection system for generating a corresponding vulnerability detection result according to the output source code of the software. As shown in the figure, the vulnerability detection system 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 above. The code acquisition unit is used to acquire the source code of the software to be detected; and the vulnerability detection model is used to generate the detection result of whether the input source code includes vulnerabilities. Figure 4

[0091] Specifically, the vulnerability detection module comprises 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 according to the source code. The preprocessing module realizes this process by built-in or calling the code analysis 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 of the code slice to obtain a sequence feature vector. The sequence branch realizes this process by built-in or calling the existing mature pre-trained CodeBERT model.

[0093] ​The graph branch adopts the GGNN model containing the EA mechanism newly designed in the embodiment to extract features of the slice subgraph and further obtain a graph feature vector. In the graph branch, the EA mechanism is used to dynamically assign weights to the in-edges and out-edges of the data flow and the control flow and the self-loop edges in the process of iterative updating of the node vector by the GGNN model, and fuse the weights into the embedding representation of the node. The feature extraction module further includes a feature fusion unit, which fuses the sequence feature vector and the graph feature vector according to a preset weight (1:1) to obtain a corresponding fusion feature vector.

[0094] The classification module adopts the MLP and is used to generate a vulnerability detection result of the source code according to the input fusion feature vector.

[0095] Embodiment 3

[0096] On the basis of the scheme of the embodiment 1, the 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, and further the detection of whether there is a vulnerability in the input software source code is implemented. The vulnerability detection device provided in the embodiment is essentially a computer device for implementing the scheme of the embodiment 1. In actual application, the computer device can adopt an intelligent terminal, a tablet computer, a notebook computer, a desktop computer, a rack server, a blade server, a tower server or a cabinet server (including a stand-alone server, or a server cluster composed of multiple servers).

[0097] The computer device indicated in the embodiment includes at least but is not limited to: a memory and a processor which can be connected to each other in communication through a system bus. 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 magnetic disk, an optical disk, etc. In some embodiments, the memory can be an internal storage unit of the 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 the computer device, such as a plug-in hard disk, a smart media card (SMC), a secure digital (SD) card, a flash card, etc. equipped on the computer device. Of course, the memory can include both the internal storage unit and the external storage device of the computer device. In the embodiment, the memory is usually used to store the operating system and various application software installed on the computer device, etc. In addition, the memory can also be used to temporarily store various data that have been output or will be output.

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

[0099] Simulation test

[0100] In order to verify the performance and advantages of the vulnerability detection method provided by the application, which combines enhanced graph representation learning with transformer, the technical personnel formulate an experimental plan and simulate and test the related scheme.

[0101] I. Performance comparison

[0102] Firstly, the present experiment selects VulDeePecker, SySeVR, Devign, Reveal four existing mature schemes as the control group, and compares the control group with the present application scheme in the performance training and testing of vulnerability detection on FFmpeg+Qemu, REVEAL and Fan et al. three data sets. In the test process, Accuracy, Precision, Recall and F1 are selected as performance evaluation indexes. In the comparison test process, the vulnerability detection performance of the five kinds of models is as shown in the following table:

[0103] Table 1: Comparison of vulnerability detection performance of the present application and the control group scheme Analysis of table 1 data can find that: the scheme of the present application has significant performance improvement compared with multiple representative control group models on three data sets. The scheme of the present application reaches the highest level in almost all evaluation indexes, and is significantly higher than the existing scheme in accuracy and F1 score.

[0104] II. Ablation experiment

[0105] The main improvements made by the present application scheme include: (1) an enhanced gated graph neural network (EA-GGNN) combined with edge attention mechanism is adopted in the graph branch to further integrate heterogeneous edge information into node representation. (2) The present application combines the feature representation extracted by the graph branch sequence branch to realize heterogeneous coding fusion. (3) The present application adopts a ternary loss function containing focal loss to train the MLP in the network model training stage to alleviate the data imbalance in large-scale imbalanced data set. In order to evaluate the contribution of the above improvements to the outstanding performance of the present application scheme, the present experiment also carries out ablation experiment on the above three improvements. The experimental content and results are as follows;

[0106] 2.1, effectiveness of EA-GGNN

[0107] The present application scheme is compared with three schemes in which EA-GGNN is replaced by GGNN, graph convolutional network (GCN) and relation graph convolutional network (R-GCN) in three data sets, and the final experimental data are as shown in the following table:

[0108] Table 2: Ablation experiment results related to EA-GGNN Analysis of the above table data can find that: the accuracy of EA-GGNN module is improved by 5.05%, and the F1 score is improved by 2.72%, which shows that EA-GGNN effectively integrates edge information into node representation, and significantly enhances the graph representation learning ability of the model.

[0109] 2.2, effectiveness of heterogeneous code fusion

[0110] To verify whether the model of the present application scheme can improve performance by fusing graph-based and sequence-based models, experiments were performed on the FFmpeg+Qemu dataset, and the performance of the fusion scheme of the present application and three sequence-based models (BGNN4VD, VulDeePecker, and SySeVR) and three graph-based models (Devign, Reveal, and EA-GGNN only) were compared. The experimental results are shown in the following table:

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

[0112]

[0113] Analyzing the above experimental data, it can be found that compared with a single model, the accuracy of the fusion method of the present application 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 code fusion method adopted by the present application effectively integrates structured and sequential information and enhances the code representation ability of the model.

[0114] 2.3, effectiveness of triple loss + MLP

[0115] To verify whether the MLP trained by the triple loss containing focal loss adopted by the present application can improve the performance of the model on the minority class, the present experiment compared and analyzed other three detection methods (SVM, RF, and MLP) using cross-entropy loss on the imbalanced dataset proposed by the previous person. Among them, the code containing vulnerabilities in the dataset used in the present experiment only accounts for 5.88% of the total data. The experimental results are shown in the following table:

[0116] Table 2: Ablation experiment related to classification model and loss

[0117]

[0118] Analyzing the above table data, it can be found that compared with existing schemes, the accuracy and F1 score of the present application scheme are optimal on the dataset with fewer vulnerabilities. This shows that the triple loss + MLP adopted by the present application can significantly improve the performance of the model on the imbalanced dataset, and proves its effectiveness in the vulnerability detection task.

[0119] The above only describes the preferred embodiments of the present application and does not limit the present application. Any modification, equivalent replacement, and improvement made within the spirit and principles of the present application shall be included in the protection scope of the present application.

Claims

1. A vulnerability detection method combining enhanced graph representation learning and transformer, characterized in that, It comprises: The vulnerability detection model comprising 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 from source code; the feature extraction module comprises a sequence branch and a graph branch; the sequence branch uses the CodeBERT model to extract features of the code slices to obtain sequence feature vectors; the graph branch uses the GGNN model containing edge attention mechanism to extract features of the slice subgraphs to obtain graph feature vectors; in the graph branch, the edge attention mechanism is used to dynamically allocate weights for the incoming and outgoing edges of the data flow and the control flow and the self-loop edges, and fuse them into the embedding representation of the nodes; the feature extraction module fuses the sequence feature vectors and the graph feature vectors according to a preset weight to obtain corresponding fusion feature vectors; the classification module uses MLP and is used to generate vulnerability detection results of the source code according to the input fusion feature vectors. A large amount of source code containing vulnerability label information was obtained as sample data to form the original dataset, which was then divided into a training set and a test set; a set containing focusing loss L... focal Cosine loss L p and regularization loss L reg The three-dimensional loss L trp As the loss function of MLP during the training phase, the vulnerability detection model is trained and tested using the training set and the test set; The model parameters of the vulnerability detection model with the best performance after testing are reserved and used for vulnerability detection of the source code.

2. The vulnerability detection method of claim 1, wherein: The process of generating code slices and slice subgraphs by the preprocessing module comprises: Firstly, the AST and PDG of the source code are generated using the code parsing tool joern; then, the vulnerability risk code lines in the AST are matched as slice centers according to the vulnerability syntax features provided by Checkmarx; next, the slice subgraphs corresponding to the slice centers are obtained by using Joern to perform forward and backward traversal on the PDG according to the edges of the data flow and the edges of the control flow; finally, all the slice subgraphs corresponding to the matched slice centers are generated by using the same method; 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 centers.

3. The vulnerability detection method of claim 2, wherein the method is enhanced by combining graph representation learning and transformers. The vulnerability detection model comprises two data processing strategies: (1) Firstly, all slice centers in the source code are identified by the preprocessing module, and a set of corresponding code slices and slice subgraphs are generated for each slice center; then, the sequence feature vectors and the graph feature vectors corresponding to each set of code slices and slice subgraphs are extracted by the feature extraction module, and a plurality of fusion feature vectors are obtained by fusion; finally, the vulnerability detection is performed on each fusion feature vector by the MLP, and the vulnerability positioning is realized according to the detection results; (2) Firstly, all slice centers in the source code are identified by the preprocessing module, and the shortest code containing all slice centers is taken as a code slice, and the slice subgraph corresponding to the code slice is generated; then, the sequence feature vectors and the graph feature vectors corresponding to the code slice and the slice subgraph are extracted by the feature extraction module, and a fusion feature vector is obtained by fusion; finally, the vulnerability detection is performed on the fusion feature vector by the MLP to generate the corresponding vulnerability detection result.

4. The vulnerability detection method of claim 1, wherein the method is enhanced by combining graph representation learning and transformers. The CodeBERT model is a large-scale pre-training model based on transformer, which is used to extract the global information of the serialized code contained in the code slices to obtain a 256-dimensional sequence feature vector; And / or The CodeBERT model sequentially comprises an Embedding module for performing vocabulary embedding and position encoding, and 12 layers of an encoder comprising a feedforward neural network and a bidirectional multi-head self-attention mechanism.

5. The vulnerability detection method of claim 1, wherein: The slice subgraph comprises a plurality of nodes, and data inflow edges, data outflow edges, control inflow edges, control outflow edges and self-loop edges between the nodes; the EA-GGNN used by the graph branch dynamically allocates weights to different types of edges through an EA mechanism, effectively fuses heterogeneous edge information into node representation, and captures local dependency and long-range node dependency in combination with an iteration mechanism of a traditional GGNN; and a 256-dimensional graph feature vector corresponding to the slice subgraph is obtained.

6. The vulnerability detection method of claim 5, wherein the graph representation learning is enhanced by combining the graph representation learning with a transformer. The EA mechanism dynamically learns the contribution of each edge type to the node feature update, and the edge e between any two neighbor nodes i and j ij The attention weight of 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 the edge of type k; a represents the attention vector for determining the importance of the edge; and σ represents the activation function.

7. The vulnerability detection method of claim 6, wherein the graph representation learning is enhanced by combining the graph representation learning with a transformer. The GGNN updates the attention weights of all edges according to the updated attention weights of all edges The expression for updating the feature vector of any node is as follows: In the above equation, denotes the feature vector of node i after the (t+1)th iteration, N k (i) denotes the set of neighbor nodes connected to node i through edges of type k; M is the set of all edge types; denotes the feature vector of node j after the tth iteration.

8. The vulnerability detection method of claim 7, wherein the method is enhanced by combining graph representation learning and transformers. The ternary loss L reg The expression of the ternary loss L is as follows: In the above formula, β and δ respectively represent the cosine loss L p and the regularization loss L reg The weight in the triplet loss L trp ; p t represents the predicted probability of the real category; represents a balance factor for adjusting the importance of positive and negative samples; γ represents a focusing parameter for adjusting the contribution degree of correct classification samples; x g represents an input sample; h(x g ) represents the latent representation of x g , h(x same ) represents the latent representation of the sample belonging to the same category as h(x g ), h(x diff ) represents the latent representation of the sample belonging to different categories as h(x g ), ε is a hyperparameter for defining the minimum separation boundary; D(h(x g ), h(x same )) represents the cosine distance between h(x g ) and h(x same ); D(h(x g ), h(x diff )) represents the cosine distance between h(x g ) and h(x diff ).

9. A vulnerability detection system characterized by: The vulnerability detection system comprises a code acquisition unit and a vulnerability detection model trained by the vulnerability detection method combining enhanced graph representation learning and a transformer; the code acquisition unit is configured to acquire source code of software to be detected; and the vulnerability detection model is configured to generate a detection result of whether the source code includes a vulnerability according to the input source code. The vulnerability detection model comprises a preprocessing module, a feature extraction module and a classification module; the preprocessing module is configured to generate code slices comprising all vulnerability-related risk nodes and corresponding slice subgraphs according to the source code; the feature extraction module comprises a sequence branch and a graph branch; the sequence branch uses a CodeBERT model to extract features of the code slices and obtain sequence feature vectors; the graph branch uses a GGNN model comprising an edge attention mechanism to extract features of the slice subgraphs and obtain graph feature vectors; in the graph branch, the edge attention mechanism is configured to dynamically allocate weights to data inflow edges, data outflow edges and self-loop edges of control flows, and fuse the weights into embedded representations of the nodes; the feature extraction module fuses the sequence feature vectors and the graph feature vectors according to preset weights to obtain corresponding fused feature vectors; and the classification module uses an MLP and is configured to generate a vulnerability detection result of the source code according to the input fused feature vectors.

10. A vulnerability detection apparatus comprising a memory, a processor, and a computer program stored in the memory and running in the processor, characterized by: When the processor executes the computer program, the vulnerability detection system of claim 9 is created, and detection of whether the input software source code includes a vulnerability is realized.

Images