Vulnerability detection method and device based on graph neural network, and electronic equipment

Abstract

The invention relates to the technical field of information security, in particular to a vulnerability detection method and device based on a graph neural network and electronic device.The vulnerability detection method based on the graph neural network comprises the steps that program source codes are analyzed to obtain a program dependency graph, node embedding is conducted on the program dependency graph, and the program dependency graph is obtained; obtaining a node feature vector; inputting the node feature vectors into a trained gating graph neural network model, wherein the gating graph neural network model enhances node feature representation; and classifying the node feature vectors by using convolution pooling and a multi-layer perceptron, and predicting whether the program source code has vulnerabilities or not. The software vulnerability detection method and device can effectively improve the accuracy and efficiency of software vulnerability detection.

Description

Technical Field

[0001] The present application relates to the field of information security technology, and in particular to a vulnerability detection method, device and electronic device based on graph neural network. Background Art

[0002] In today's digital age, the complexity and application scenarios of software systems continue to expand. From personal privacy protection to financial transaction security, the consequences of software vulnerabilities can be extremely serious. Vulnerability detection technology is essential because it can help developers and security teams promptly identify potential security risks during software development and operation, thereby reducing the attack surface and mitigating the risk of malicious exploitation. Effective vulnerability detection allows for rapid remediation, shortening the window of opportunity for attackers to exploit vulnerabilities and ensuring the overall security of software systems.

[0003] Currently, vulnerability detection technologies are primarily categorized into two main categories: static analysis and dynamic analysis. Static analysis analyzes the syntax, semantics, and structure of code, identifying potential issues without executing the code, making it suitable for use in the early stages of development. Dynamic analysis focuses on runtime behavioral monitoring, such as fuzz testing, which triggers anomalies by inputting large amounts of random data to discover vulnerabilities. Furthermore, with the development of artificial intelligence (AI) technology, vulnerability detection methods based on machine learning and deep learning are gaining popularity. These methods can automatically learn vulnerability patterns and improve detection efficiency. The application of graph neural networks, in particular, further enhances detection accuracy and efficiency by capturing the structural and semantic information of code.

[0004] Despite the continuous advancement of vulnerability detection technology, challenges remain. Static analysis, while highly effective, is prone to false positives and false negatives and struggles with complex code logic. Dynamic analysis, on the other hand, requires a large number of test cases, is less efficient, and fails to cover all possible execution paths. While machine learning and deep learning-based methods have shown some improvement in accuracy, they rely on high-quality datasets and complex feature engineering, and in some cases still suffer from high false positive rates. Summary of the Invention

[0005] The present application provides a vulnerability detection method, device, electronic device and storage medium based on graph neural network to solve the problems in the above-mentioned background technology.

[0006] In a first aspect, the present application provides a vulnerability detection method based on a graph neural network, comprising:

[0007] Parsing the program source code to obtain a program dependency graph, and performing node embedding on the program dependency graph to obtain a node feature vector;

[0008] Inputting the node feature vector into a trained gated graph neural network model, wherein the gated graph neural network model enhances the node feature representation;

[0009] Convolution pooling and a multi-layer perceptron are used to classify the node feature vectors to predict whether there are vulnerabilities in the program source code.

[0010] Furthermore, before parsing the program source code to obtain the program dependency graph, the process further includes:

[0011] Preprocess the program source code to remove redundant information and normalize function names, parameter names, and variable names.

[0012] Furthermore, the parsing of the program source code to obtain a program dependency graph includes:

[0013] Use open source tools to parse program source code, combine control dependencies and data dependencies, and form a program dependency graph;

[0014] The nodes in the program dependency graph represent statements in the program source code, and the edges represent dependency relationships between nodes.

[0015] Furthermore, the performing node embedding on the program dependency graph to obtain a node feature vector includes:

[0016] Using a natural language processing method to convert each node in the program dependency graph into a low-dimensional vector of fixed length;

[0017] The degree centrality of the node is calculated, and the degree centrality is multiplied by the low-dimensional vector to obtain a node feature vector.

[0018] Furthermore, the step of inputting the node feature vector into a trained gated graph neural network model, wherein the gated graph neural network model enhances the node feature representation, includes:

[0019] The program dependency graph is represented and learned through a gated graph neural network model to handle long-distance dependencies between nodes, and the graph attention mechanism is used to enhance node feature representation.

[0020] Furthermore, after using convolution pooling and a multi-layer perceptron to classify the node feature vectors and predict whether the program source code has a vulnerability, the method further includes:

[0021] A graph neural network interpreter is used to generate key subgraphs, revealing the nodes and edges that have the greatest influence on the prediction results, and providing fine-grained guidance for vulnerability repair.

[0022] Furthermore, the graph neural network interpreter is implemented by optimizing an interpretable mask, which is a weight vector representing the importance of nodes and edges.

[0023] In a second aspect, the present application provides a vulnerability detection device based on a graph neural network, comprising:

[0024] An extraction module is used to parse the program source code to obtain a program dependency graph, and perform node embedding on the program dependency graph to obtain a node feature vector;

[0025] An input module, configured to input the node feature vector into a trained gated graph neural network model, wherein the gated graph neural network model enhances the node feature representation;

[0026] The prediction module is used to classify the node feature vectors using convolution pooling and a multi-layer perceptron to predict whether there is a vulnerability in the program source code.

[0027] In a third aspect, the present application provides an electronic device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor. When the processor executes the computer program, the vulnerability detection method based on graph neural network as described above is implemented.

[0028] In a fourth aspect, the present application provides a computer-readable storage medium, which stores a computer program. When the computer program is executed by a processor, it implements the vulnerability detection method based on graph neural network as described above.

[0029] The above technical solution of this application has the following advantages:

[0030] The first aspect of the present application provides a vulnerability detection method based on graph neural networks, which obtains a program dependency graph by parsing the program source code, and embeds nodes in the program dependency graph to obtain a node feature vector. The node feature vector is input into a trained gated graph neural network model, and the gated graph neural network model enhances the node feature representation, uses convolution pooling and a multi-layer perceptron to classify the node feature vector, and predicts whether the program source code has a vulnerability, thereby improving the accuracy and efficiency of source code vulnerability detection.

[0031] It can be understood that the beneficial effects of the second, third and fourth aspects mentioned above can be found in the relevant description of the first aspect mentioned above, and will not be repeated here. BRIEF DESCRIPTION OF THE DRAWINGS

[0032] In order to more clearly illustrate the specific implementation methods of the present application or the technical solutions in the prior art, the following is a brief introduction to the drawings required for use in the specific implementation methods or the description of the prior art. Obviously, the drawings described below are some implementation methods of the present application. For ordinary technicians in this field, other drawings can be obtained based on these drawings without any creative work.

[0033] Figure 1 Flowchart of the vulnerability detection method based on graph neural network provided in this application;

[0034] Figure 2 A system block diagram of the vulnerability detection method based on graph neural network provided in this application;

[0035] Figure 3 Schematic diagram of the vulnerability detection method based on graph neural network provided in this application. DETAILED DESCRIPTION

[0036] In the following description, specific details such as specific system structures and techniques are provided for purposes of illustration rather than limitation to facilitate a thorough understanding of the embodiments of the present application. However, it will be apparent to those skilled in the art that the present application may be implemented in other embodiments without these specific details. In other cases, detailed descriptions of well-known systems, devices, circuits, and methods are omitted to avoid obscuring the description of the present application with unnecessary detail.

[0037] It should be understood that when used in the present specification and the appended claims, the term "comprising" indicates the presence of described features, integers, steps, operations, elements and / or components, but does not preclude the presence or addition of one or more other features, integers, steps, operations, elements, components and / or collections thereof.

[0038] In addition, in the description of the present application specification and the appended claims, the terms "first", "second", "third", etc. are only used to distinguish the descriptions and cannot be understood as indicating or implying relative importance.

[0039] References to "one embodiment" or "some embodiments" in this specification mean that one or more embodiments of the present application include a particular feature, structure, or characteristic described in conjunction with that embodiment. Thus, phrases such as "in one embodiment," "in some embodiments," "in other embodiments," and "in other embodiments" appearing in various places in this specification do not necessarily refer to the same embodiment, but rather mean "one or more but not all embodiments," unless otherwise specifically emphasized. The terms "including," "comprising," "having," and variations thereof all mean "including but not limited to," unless otherwise specifically emphasized. "Multiple" means "two or more."

[0040] The main purpose of this application is to represent the source code as a graph structure, capture its rich structural and semantic information, and retain the global structural information in the original graph when embedding graph node features. A model based on gated graph neural network and graph attention mechanism is used to effectively capture long-distance dependency information between graph nodes, further mine vulnerability information of program fragments, and finally use graph neural network interpreter to provide fine-grained detection and improve the interpretability of vulnerability detection model.

[0041] The following embodiments are used to illustrate the present invention, but are not intended to limit the scope of the present invention.

[0042] like Figure 1 As shown, an embodiment of the present application provides a vulnerability detection method based on a graph neural network, which specifically includes the following steps: parsing the program source code to obtain a program dependency graph, and embedding nodes in the program dependency graph to obtain a node feature vector; inputting the node feature vector into a trained gated graph neural network model, and the gated graph neural network model enhances the node feature representation; using convolution pooling and a multi-layer perceptron to classify the node feature vector to predict whether there is a vulnerability in the program source code.

[0043] In some embodiments, before parsing the program source code to obtain the program dependency graph, the method further includes: preprocessing the program source code to remove redundant information and normalize function names, parameter names, and variable names.

[0044] Preprocess the program source code to normalize it so the neural network can complete the subsequent detection task. This step removes redundant information, including comments in the code, and reduces noise in the statement embeddings. Then, normalize the custom function names, parameter names, and variable names in the code.

[0045] In some embodiments, parsing the program source code to obtain a program dependency graph includes: using open source tools to parse the program source code, combining control dependencies and data dependencies to form a program dependency graph; the nodes in the program dependency graph represent statements in the program source code, and the edges represent the dependency relationships between nodes.

[0046] The open-source tool Joern is used to parse the source code and obtain a program dependency graph (PDG). Joern combines control dependencies and data dependencies to form a program dependency graph (PDG). The nodes of a PDG typically represent statements in a program, while the edges represent the dependencies between nodes.

[0047] In some embodiments, performing node embedding on the program dependency graph to obtain a node feature vector includes: using a natural language processing method to convert each node in the program dependency graph into a low-dimensional vector of a fixed length; calculating the degree centrality of the node, and multiplying the degree centrality with the low-dimensional vector to obtain a node feature vector.

[0048] Obtain embeddings for PDG statement nodes. Use the natural language processing method Sent2Vec to convert them into fixed-length, low-dimensional vectors for graph neural network learning. Through the embedding process, each line of code in the PDG graph is converted into a vector representation of length 128. Calculate the centrality of the PDG nodes and synthesize the node embeddings. Calculating node degree centrality ensures that the generated feature vectors are more representative of the graph structure, thereby enriching the node vector information and improving the efficiency and accuracy of vulnerability detection.

[0049] The node centrality calculation and embedding method is as follows: for a graph G = (V, E), where V is the node set, E is the edge set, and node v i The degree centrality of

[0050]

[0051] To calculate, where d(v i ) represents node v i The degree of a sentence (i.e., the number of edges connected to it), A is the adjacency matrix of the graph, and N is the total number of nodes in the graph. Centrality is embedded by multiplying it with the feature vector obtained after sentence embedding, so that the feature vector representation of a node can include the importance of the node in the graph structure.

[0052] In some embodiments, the node feature vector is input into a trained gated graph neural network model, and the gated graph neural network model enhances the node feature representation, including: performing representation learning on the program dependency graph through the gated graph neural network model, processing long-distance dependencies between nodes, and utilizing a graph attention mechanism to enhance the node feature representation.

[0053] Train the gated graph neural network model and save it. Use the node data and edge data as the input of GGNN. Given the graph data, for each node v i , the node state is initialized to Set T as the total number of time steps for information aggregation. Then, in each time step t, all nodes aggregate the information of their neighboring nodes according to the formula

[0054]

[0055] Calculate the new state of the node. After information aggregation in the gated graph recursive layer, the new state features of the graph node are obtained. Subsequently, the representation of the node features is enhanced through the graph attention layer. To ensure more stable model training, a multi-head attention mechanism is used, averaging the calculation results of K independent attention mechanisms.

[0056] Attention correlation coefficient α ij The calculation is as follows:

[0057]

[0058] α ij Represents node v i For node v j The importance of , W is the weight matrix, and LeakyReLU( ) is the activation function.

[0059] Vulnerability prediction is performed on the graph using all graph node features obtained from the graph attention layer. The graph prediction layer consists of two stages: convolution and pooling, and result generation. The convolution and pooling method uses the MaxPool method, and the result generation uses a multi-layer perceptron as a classifier to further learn features of the graph embedding vector and finally obtain the classification result.

[0060] In some embodiments, after using convolutional pooling and a multi-layer perceptron to classify the node feature vectors and predict whether the program source code has a vulnerability, the method further includes: using a graph neural network interpreter to generate key subgraphs, revealing the nodes and edges that have the greatest influence on the prediction results, and providing fine-grained guidance for vulnerability repair.

[0061] In some embodiments, the graph neural network interpreter is implemented by optimizing an interpretable mask, which is a weight vector representing the importance of nodes and edges.

[0062] Initialize the explainer by creating an Explainer object and find a subgraph G in the entire graph G of the function S , so that the entire graph G and the minimum graph G S The difference in prediction scores obtained after inputting the GNN model is the smallest. S It can directly serve as an explanation for the model's predictions, including the features that are most critical to the predictions. The explainer adjusts the threshold to visualize the key subgraphs that contribute most to the predictions, enabling fine-grained vulnerability detection.

[0063] The explainer visualizes the key subgraphs as follows: GNNExplainer is implemented by optimizing an interpretable mask, which is a weight vector representing the importance of nodes and edges, aiming to explain the parts of the graph structure that contribute most to the model prediction. By iteratively adjusting the mask weights, GNNExplainer strives to maximize the accuracy and reliability of the explanation. This process is given by the formula Indicates. Among them, W G =(W V , W E ) is the weight vector of node and edge importance, is the loss function, M represents the GNN model, G S It represents the subgraph structure of the current recognition that affects the prediction result, and y represents the predicted value.

[0064] This application provides a vulnerability detection method based on graph neural network, which aims to achieve the following goals: 1. Represent the source code as a graph structure, capture the rich structural information and semantic information of the source code, and effectively capture the long-distance dependency information between graph nodes through gated graph neural network and graph attention mechanism, and further mine the vulnerability information of program fragments. 2. Improve the transparency and interpretability of the model, and provide fine-grained detection results, giving the code location closely related to the vulnerability. Its system block diagram is as follows Figure 2 As shown, it mainly includes three parts, including feature extraction, model prediction and model interpretation.

[0065] Feature Extraction: The feature extraction component of vulnerability detection involves extracting a program dependency graph and embedding nodes. First, the source code is normalized to remove redundant information and unify function, parameter, and variable names. The Joern tool is then used to generate a program dependency graph, capturing the control and data dependencies within the code. Subsequently, the Sent2Vec model is used to transform each node in the program dependency graph into a fixed-length, low-dimensional vector. The node's degree centrality is calculated and multiplied by the feature vector from the sentence embedding. This enriches the node's global structural information, providing high-quality feature representation for subsequent vulnerability detection.

[0066] Model Prediction: In the model prediction phase, a gated graph neural network (GGNN) is used to learn representations of the program dependency graph based on the extracted node embedding features to handle long-distance dependencies between nodes. The GGNN updates node states through information propagation and aggregation mechanisms, and further enhances the representation of node features through a graph attention layer, highlighting the importance of key nodes. Subsequently, convolutional pooling and a multi-layer perceptron (MLP) are used to classify the graph embedding vectors to predict whether the code contains vulnerabilities. The MLP classifier classifies samples as vulnerable or non-vulnerable through feature mapping in the hidden layer and probability calculation in the output layer, thereby achieving efficient vulnerability detection.

[0067] Model Explanation: During the model explanation phase, GNNExplainer performs interpretable analysis on the prediction results of graph neural networks. By modifying the graph structure or node features and observing the impact of these changes on model predictions, GNNExplainer identifies the key nodes and edges that have the greatest impact on the predictions and generates subgraph explanations. This fine-grained explanation not only improves model transparency but also helps developers intuitively locate vulnerability-related code locations, enabling more efficient vulnerability remediation. By visualizing key subgraphs, GNNExplainer provides an understandable basis for vulnerability detection, ensuring the reliability and practicality of model predictions.

[0068] like Figure 3 As shown, the vulnerability detection method based on graph neural network provided by this application is specifically described as follows:

[0069] This method is a function-level vulnerability detection technology based on gated neural networks, treating vulnerability detection as a binary classification task for a neural network model. After normalizing the source code dataset according to custom rules, the open-source tool Joern is used to parse the source code and obtain a program dependency graph (PDG), abstracting the code's graph structure. Social network centrality analysis is then introduced during the node embedding phase to calculate the degree centrality of nodes, thereby enriching the global structural information of statement embedding. Each node in the PDG is treated as a statement, and the Sent2Vec model is used to obtain a vector representation of each node in the PDG graph. A gated graph neural network (GGNN) model is then trained using the processed dataset and saved. The GGNN model can learn vulnerability patterns and effectively capture useful information in complex graph structures. The code snippet to be tested is then processed in the same manner as the training set, treating vulnerability detection as a binary classification task to determine whether the code snippet is vulnerable. Finally, we use the graph neural network explanatory tool GNNExplainer to obtain the subgraph that contributes most to the results. This subgraph is part of the source code graph representation and reveals the key features and relationships that best represent the model's prediction results. This allows for more fine-grained vulnerability detection and infers relevant information about the vulnerable code lines, enabling developers to perform more efficient repairs.

[0070] GNNExplainer is implemented by optimizing an interpretable mask, which is a weight vector representing the importance of nodes and edges, aiming to explain the parts of the graph structure that contribute most to the model prediction. By iteratively adjusting the mask weights, GNNExplainer strives to maximize the accuracy and reliability of the explanation. This process is expressed as the following formula:

[0071]

[0072] Among them, WG =(W V , W E ) is the weight vector of node and edge importance, is the loss function, M represents the GNN model, G S The output of the algorithm is the classification result of the code snippet input by the GGNN model, that is, 1 or 0 (whether there is a vulnerability or not), and the model explanation result subgraph G S , which is the subgraph structure that contributes most to the classification results of the GGNN model (potential location of code vulnerabilities).

[0073] This application improves the accuracy and efficiency of source code vulnerability detection based on a graph neural network model. By introducing node degree centrality during the embedding process of program dependency graph nodes, this not only preserves the semantic information of statements but also enriches the global structural information of the graph. During model training, a gated neural network model is used to enhance information propagation between distant nodes, effectively capturing effective information in complex graph structures and improving the accuracy of vulnerability detection. The detection scheme proposed in this application also considers the interpretability and transparency of deep learning models. It uses GNNExplainer to obtain the subgraph that contributes most to the results. This subgraph is part of the source code graph representation and reveals the key features and relationships that best represent the model's prediction results. This fine-grained detection improves the efficiency of vulnerability remediation.

[0074] The following describes the invention through specific embodiments.

[0075] Example

[0076] 1. Program Dependency Graph Extraction Phase

[0077] The initialization phase includes steps 1 and 2, which complete the source code normalization and program dependency graph extraction. The specific steps are as follows:

[0078] Step 1: Preprocess the program source code to normalize it for the neural network to complete the subsequent detection task. This step removes redundant information, including comments in the code, and reduces noise in the statement embeddings. Then, the custom function names, parameter names, and variable names in the code are normalized.

[0079] Step 2: Use the open-source tool Joern to parse the source code and obtain a program dependency graph (PDG). The PDG generates control dependencies and data dependencies. Control dependencies represent the dependencies of control flow in a program, such as the impact of conditional statements (such as if and while) on subsequent statements. Data dependencies represent the dependencies between variable definitions and usages, that is, how the value of one variable affects other variables. Joern combines control dependencies and data dependencies to form a program dependency graph (PDG). The nodes of a PDG typically represent statements in a program, while the edges represent the dependencies between nodes.

[0080] 2. Node Embedding Stage

[0081] The node embedding stage includes steps 3-4, which includes using the Sent2Vec model to obtain the vector representation of each node in the PDG graph, and calculating the node centrality through social network centrality analysis to enrich the global structural information of sentence embedding.

[0082] Step 3: Sentence Embedding. Since each node in the program dependency graph (PDG) represents a sentence, Sent2Vec, a natural language processing method, is used to convert it into a fixed-length, low-dimensional vector for graph neural network learning. Sent2Vec is an unsupervised learning model for generating sentence embeddings. Its core function is to convert sentences or text paragraphs into fixed-length vectors. It is essentially the same structure as Word2Vec, but with words replaced by sentences. The sentence vector representation is obtained by summing and averaging the word vectors. Similar to Word2Vec, Sent2Vec models include skip-word models and continuous bag-of-words models. The skip-word model uses a word at a fixed position to infer the surrounding words. During training, a fixed-length continuous text sequence is often selected, with the middle word as input and the preceding and following words as output. The continuous bag-of-words model, on the other hand, uses the words surrounding a central word to predict itself. Through the embedding process, each line of code in the PDG is converted into a vector representation of length 128.

[0083] Step 4: Node centrality calculation and embedding. Calculating node degree centrality enables the feature vectors generated by the nodes to better represent the information of the graph structure, thereby enriching the node vector information and improving the efficiency and accuracy of vulnerability detection.

[0084] The degree centrality of a node can be calculated using the following mathematical formula:

[0085] For a graph G = (V, E), where V is a set of nodes and E is a set of edges, node v i The degree centrality C d (v i ) is defined as:

[0086]

[0087] Among them, d(v i ) represents node v i where A is the adjacency matrix of the graph and N is the total number of nodes in the graph.

[0088] The calculated centrality is embedded into the node. This step adopts a relatively simple and effective method: multiplying the centrality with the feature vector obtained after sentence embedding to achieve the embedding of the centrality, so that the feature vector representation of the node can include the importance of the node in the graph structure.

[0089] Graph Embedding Stage

[0090] The graph embedding phase includes steps 5-6. After obtaining the embedding vectors of the PDG nodes, a gated neural network (GGNN) is used to learn the PDG representation. GGNN is chosen as the gated graph recursive layer because it can effectively handle the long-distance dependencies between nodes in PDG. GGNN can handle long-distance dependencies between nodes in PDG through information propagation and aggregation mechanisms, ensuring that information about long-distance nodes is appropriately retained during feature extraction. The graph attention layer amplifies the weights of related nodes, allowing nodes to retain as much useful information as possible during feature extraction.

[0091] Step 5-6: Train the gated graph neural network model and save it. GGNN enhances the long-term memory capacity of the network through gated recurrent units, which is very effective for source code understanding tasks. The node information in PDG is converted into vector data through node embedding, and the relationship between nodes is used as edge data, and is stored by defining edge types and their corresponding adjacency matrices. The dependency relationship provided in the program dependency graph reflects both data dependency and control dependency. Then the node data and edge data are used as the input of GGNN. Given graph data, for each node v i , the node state is initialized to Set T to the total number of time steps for information aggregation. Thereafter, at each time step t, all nodes aggregate the information of their neighboring nodes and calculate the new state as follows:

[0092]

[0093] Among them, W E ∈R d×d is the learned weight matrix, b is the bias term, Represents the new state of the node after aggregating the information of adjacent nodes. Use the gated recurrent unit to update the hidden state of the current node As shown below:

[0094]

[0095] After the information aggregation of the gated graph recursive layer, the new state features of the graph nodes are obtained. Subsequently, the representation of the node features is enhanced by the graph attention layer. The core of the graph attention layer is to adaptively focus on the node features with high relevance based on the input graph data, while ignoring the node features with low relevance. In the process of learning graph data, the attention mechanism identifies adjacent nodes with dependencies (such as data dependency and control dependency). Through the graph attention layer, the weights of node features with high relevance are increased, while the weights of node features with low relevance are reduced, and the features of the node itself are updated through weighting. Attention is only added to the node v i The adjacent nodes, h j represents the node v obtained from the gated graph recursive layer j The hidden state, N i Represents node v j Adjacent nodes. Attention correlation coefficient α ij Calculated by the following formula:

[0096]

[0097] α ij Represents node v i For node v j The importance of , W is the weight matrix, and LeakyReLU( ) is the activation function. To ensure a more stable model training process, a multi-head attention mechanism is used, averaging the calculation results of K independent attention mechanisms.

[0098] Graph Prediction Layer: Vulnerabilities are predicted on the graph using all graph node features obtained from the graph attention layer. The graph prediction layer consists of two stages: convolutional pooling and result generation. Convolutional pooling is used to reduce the size of the input data, thereby reducing the number of parameters to avoid overfitting. The convolutional results are then used to predict vulnerabilities. The convolutional pooling method uses the MaxPool method to select key and important features as much as possible. Graph embeddings can be used as input to any prediction layer to train the classification model for the task. After the code attribute graph is subjected to feature extraction by the GGNN model, the graph embedding vectors of vulnerable and non-vulnerable code samples still overlap in the feature space, with no clear separation. To achieve better separability between vulnerable and non-vulnerable samples, a multilayer perceptron (MLP) is used as a classifier to further learn features from the graph embedding vectors.

[0099] The MLP classifier has a three-layer structure, including an input layer, a hidden layer, and an output layer, with full connections between layers. The input layer takes the sample's graph embedding vector as input. The hidden layer maps the graph embedding vector to an abstract feature vector through a weight matrix. The output layer uses the sigmoid function to calculate the probability that the sample contains a vulnerability using the hidden layer's feature vector. If the probability is close to 1, the sample is classified as a vulnerable sample. If the probability is close to 0, the sample is classified as a non-vulnerable sample. The classification model uses the cross entropy loss function as the loss function:

[0100]

[0101] in is the classification prediction of the sample.

[0102] 4. Model interpretation stage

[0103] The model interpretation phase includes steps 7 and 8. After the graph neural network model obtains the prediction results, the interpretability analysis of GNNExplainer is used to perform prediction interpretation on the input instance. By adjusting the weight threshold, the algorithm filters and returns a list of nodes and edges that have the greatest impact on the prediction results.

[0104] Traditional GNN models are black-box models, and the opacity and unexplainability of their decision-making processes greatly reduce developers' ability to locate and fix vulnerabilities. GNNExplainer, however, serves as a powerful explainer to reveal the decision-making process behind the model.

[0105] GNNExplainer operates directly on the graph, eliminating the need for additional proxy models. This reduces computational complexity and implementation difficulty, making it more efficient when processing large-scale graph data. Furthermore, the subgraph explanations generated by GNNExplainer intuitively highlight the nodes and edges that most influence the prediction results, providing easy-to-understand explanations. This fine-grained explanation provides developers with deep insights when processing complex graph structures, helping them to deeply understand the causes of vulnerabilities and more effectively remediate them. GNNExplainer also boasts excellent scalability, enabling efficient processing of large datasets to meet the needs of real-world applications.

[0106] Steps 7-8: The initialization of the explainer is achieved by creating an Explainer object. The core algorithm is to modify the structure of the graph G, including changing the edge connections or node features, and observe the impact of these changes on the prediction results of the GNN model, thereby generating an explanation for the model decision. The goal of the algorithm is to find a subgraph G in the entire graph G of the function. S , so that the entire graph G and the minimum graph G SThe difference in prediction scores obtained after inputting the GNN model between them is minimal. By iteratively constructing and analyzing different subgraphs and their impact on model predictions, the algorithm aims to determine the most critical part for model prediction. The resulting subgraph G S It can be directly used as an explanation for the model's prediction results, including the features that are most critical to the prediction results.

[0107] The interpreter visualizes the key subgraphs that contribute most to the prediction results by adjusting the threshold. This not only improves the interpretability of the model, but also, because the subgraphs are closely related to the source code, developers can find the code locations related to the vulnerabilities based on the key subgraphs, thereby improving the efficiency and accuracy of vulnerability repair, thereby achieving fine-grained vulnerability detection as a whole.

[0108] This application first preprocesses the program source code, removes redundant information and normalizes function names, parameter names and variable names; then uses the open source tool Joern to generate a program dependency graph (PDG) to capture the control dependencies and data dependencies in the code. Next, the Sent2Vec model is used to embed the statement nodes in the PDG into low-dimensional vectors, and its global structural information is enriched by combining the degree centrality of the nodes. Subsequently, the PDG is represented by learning through a gated graph neural network (GGNN) to handle long-distance dependencies between nodes, and the node feature representation is enhanced using the graph attention mechanism. Afterwards, convolution pooling and multi-layer perceptron (MLP) are used to classify the graph embedding vector to predict whether the code has vulnerabilities. Finally, with the help of the GNNExplainer interpreter, key subgraphs are generated to reveal the nodes and edges that have the greatest impact on the prediction results, providing fine-grained guidance for vulnerability repair. By combining the strong expressive power of graph neural networks with interpretability analysis, this application effectively improves the accuracy, efficiency and interpretability of vulnerability detection, providing an innovative solution for the field of software security.

[0109] Corresponding to the vulnerability detection method based on graph neural network described in the above embodiment, the embodiment of the present application further provides a vulnerability detection device based on graph neural network, and the vulnerability detection device based on graph neural network includes:

[0110] An extraction module is used to parse the program source code to obtain a program dependency graph, and perform node embedding on the program dependency graph to obtain a node feature vector;

[0111] An input module, configured to input the node feature vector into a trained gated graph neural network model, wherein the gated graph neural network model enhances the node feature representation;

[0112] The prediction module is used to classify the node feature vectors using convolution pooling and a multi-layer perceptron to predict whether there is a vulnerability in the program source code.

[0113] It should be noted that the information interaction, execution process, etc. between the above-mentioned modules / units are based on the same concept as the method embodiment of this application. Their specific functions and technical effects can be found in the method embodiment section and will not be repeated here.

[0114] Those skilled in the art can clearly understand that, for the convenience and brevity of description, only the division of the above-mentioned functional units and modules is used as an example for illustration. In actual applications, the above-mentioned functions can be distributed and completed by different functional units and modules as needed, that is, the internal structure of the device can be divided into different functional units or modules to complete all or part of the functions described above. The functional units and modules in the embodiment can be integrated into one processing unit, or each unit can exist physically alone, or two or more units can be integrated into one unit. The above-mentioned integrated unit can be implemented in the form of hardware or in the form of software functional units. In addition, the specific names of the functional units and modules are only for the convenience of distinguishing each other, and are not used to limit the scope of protection of this application. The specific working process of the units and modules in the above-mentioned system can refer to the corresponding process in the aforementioned method embodiment, and will not be repeated here.

[0115] An embodiment of the present application also provides an electronic device, including a memory, a processor, and a computer program stored in the memory and executable on the processor. When the processor executes the computer program, the steps of the vulnerability detection method based on graph neural network provided in the first aspect are implemented.

[0116] In applications, the processor may be a central processing unit (CPU), other general-purpose processors, digital signal processors (DSP), application-specific integrated circuits (ASIC), field-programmable gate arrays (FPGA), other programmable logic devices, discrete gate or transistor logic devices, discrete hardware components, etc. A general-purpose processor may be a microprocessor or any conventional processor, etc.

[0117] In applications, the memory may be an internal storage unit of an electronic device in some embodiments, such as a hard disk or memory of the electronic device. In other embodiments, the memory may also be an external storage device of the electronic device, such as a plug-in hard disk equipped on the electronic device, a smart memory card (Smart Media Card, SMC), a secure digital (SD) card, a flash card (Flash Card), etc. The memory may also include both an internal storage unit of the electronic device and an external storage device. The memory is used to store an operating system, application programs, a boot loader (BootLoader), data, and other programs, such as the program code of a computer program. The memory may also be used to temporarily store data that has been output or is to be output.

[0118] An embodiment of the present application further provides a computer-readable storage medium, which stores a computer program. When the computer program is executed by a processor, the steps in the above-mentioned method embodiments can be implemented.

[0119] The present application implements all or part of the processes in the above-mentioned embodiment method, which can be completed by instructing the relevant hardware through a computer program. The computer program can be stored in a computer-readable storage medium. When the computer program is executed by the processor, it can implement the steps of the above-mentioned various method embodiments. The computer program includes computer program code, which can be in source code form, object code form, executable file or some intermediate form. The computer-readable medium may include at least: any entity or device that can carry the computer program code to an electronic device, a recording medium, a computer memory, a read-only memory (ROM), a random access memory (RAM), an electric carrier signal, a telecommunication signal and a software distribution medium. For example, a USB flash drive, a mobile hard disk, a magnetic disk or an optical disk.

[0120] Those skilled in the art will appreciate that the devices and algorithm steps of each example described in conjunction with the embodiments disclosed herein can be implemented in electronic hardware, or a combination of computer software and electronic hardware. Whether these functions are performed in hardware or software depends on the specific application and design constraints of the technical solution. Professional and technical personnel can use different methods to implement the described functions for each specific application, but such implementation should not be considered beyond the scope of this application.

[0121] In the embodiments provided herein, it should be understood that the disclosed devices and methods may be implemented in other ways. For example, the device embodiments described above are merely illustrative. Furthermore, the coupling or direct coupling or communication connection shown or discussed may be through some interface, or the devices may be indirectly coupled or communicated in some manner, whether electrical, mechanical, or other.

[0122] The above-described embodiments are only used to illustrate the technical solutions of the present application, rather than to limit them. Although the present application has been described in detail with reference to the aforementioned embodiments, those skilled in the art should understand that they can still modify the technical solutions described in the aforementioned embodiments, or make equivalent replacements for some of the technical features therein. These modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the various embodiments of the present application, and should all be included in the scope of protection of the present application.

Claims

1. A vulnerability detection method based on graph neural network, characterized in that: include: Parsing the program source code to obtain a program dependency graph, and performing node embedding on the program dependency graph to obtain a node feature vector; Inputting the node feature vector into a trained gated graph neural network model, wherein the gated graph neural network model enhances the node feature representation; Convolution pooling and a multi-layer perceptron are used to classify the node feature vectors to predict whether there are vulnerabilities in the program source code.

2. The vulnerability detection method based on graph neural network according to claim 1, characterized in that: Before parsing the program source code to obtain the program dependency graph, the method further includes: Preprocess the program source code to remove redundant information and normalize function names, parameter names, and variable names.

3. The vulnerability detection method based on graph neural network according to claim 1, characterized in that: The step of parsing the program source code to obtain a program dependency graph includes: Use open source tools to parse program source code, combine control dependencies and data dependencies, and form a program dependency graph; The nodes in the program dependency graph represent statements in the program source code, and the edges represent dependency relationships between nodes.

4. The vulnerability detection method based on graph neural network according to claim 1, characterized in that: The step of embedding nodes in the program dependency graph to obtain a node feature vector includes: Using a natural language processing method to convert each node in the program dependency graph into a low-dimensional vector of fixed length; The degree centrality of the node is calculated, and the degree centrality is multiplied by the low-dimensional vector to obtain a node feature vector.

5. The vulnerability detection method based on graph neural network according to claim 1, characterized in that: Inputting the node feature vector into a trained gated graph neural network model, wherein the gated graph neural network model enhances the node feature representation, includes: The program dependency graph is represented and learned through a gated graph neural network model to handle long-distance dependencies between nodes, and the graph attention mechanism is used to enhance node feature representation.

6. The vulnerability detection method based on graph neural network according to claim 1, characterized in that: After using convolution pooling and a multi-layer perceptron to classify the node feature vectors and predict whether the program source code has a vulnerability, the method further includes: A graph neural network interpreter is used to generate key subgraphs, revealing the nodes and edges that have the greatest influence on the prediction results, and providing fine-grained guidance for vulnerability repair.

7. The vulnerability detection method based on graph neural network according to claim 6, characterized in that: The graph neural network interpreter is implemented by optimizing an interpretable mask, which is a weight vector representing the importance of nodes and edges.

8. A vulnerability detection device based on graph neural network, characterized in that: include: An extraction module is used to parse the program source code to obtain a program dependency graph, and perform node embedding on the program dependency graph to obtain a node feature vector; An input module, configured to input the node feature vector into a trained gated graph neural network model, wherein the gated graph neural network model enhances the node feature representation; The prediction module is used to classify the node feature vectors using convolution pooling and a multi-layer perceptron to predict whether there is a vulnerability in the program source code.

9. An electronic device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein: When the processor executes the computer program, it implements the vulnerability detection method based on graph neural network as described in any one of claims 1 to 7.

10. A computer-readable storage medium storing a computer program, characterized in that: When the computer program is executed by a processor, the vulnerability detection method based on graph neural network as described in any one of claims 1 to 7 is implemented.

Images