A vulnerability detection method based on graph embedding and bidirectional gated graph neural network
By using graph embedding and bidirectional gated graph neural networks, a program dependency graph is generated and its features are extracted, which solves the problem that existing technologies cannot effectively capture the semantic features of complex source code structures, and achieves more efficient vulnerability detection.
Patent Information
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- HANGZHOU DIANZI UNIV
- Filing Date
- 2022-11-23
- Publication Date
- 2026-06-02
AI Technical Summary
Existing deep learning-based program modeling methods cannot effectively capture the non-linear semantic features in the complex structure of source code, resulting in insufficient vulnerability detection performance.
We employ a graph embedding and bidirectional gated graph neural network approach. By generating a program dependency graph, we extract feature vectors of nodes and edges, use node2vec for feature training, and train a bidirectional gated graph neural network model for vulnerability detection.
It improves the accuracy of code vulnerability detection, can better learn the graph structure information of the code, capture more control and data dependency information, and improves detection performance.
Smart Images
Figure CN115935372B_ABST
Abstract
Description
Technical Field
[0001] This invention relates to the fields of source code preprocessing and vulnerability detection in software programs, and particularly to a code vulnerability detection method based on graph embedding and bidirectional gated graph neural networks. Background Technology
[0002] Software vulnerabilities are the cause of many system attacks and data breaches. Machine learning is a viable approach to identifying common software vulnerabilities by building tools and models. Because different vulnerabilities exhibit similar underlying patterns, machine learning can first learn the latent patterns expressed by vulnerable programs from training samples, and then apply these patterns to new software projects to identify potentially vulnerable code.
[0003] Recently, researchers have used deep learning to learn the program structure of source code in order to identify potential software vulnerabilities. Compared with classic machine learning techniques, deep learning has the advantage of automatically learning structural features from training samples, eliminating the need for expert intervention in manually tuning the program structure. Existing deep learning-based program modeling methods typically use recurrent neural networks (RNNs), such as Long Short-Term Memory (LSTM) or its variants. However, LSTMs are designed for sequential sequences and are not suitable for program structure control and data flow modeling. Therefore, previous LSTM-based methods could only capture shallow, surface-level structural or syntactic information from the source code text, failing to fully learn the semantic features that provide more effective information within the program structure.
[0004] To better learn features from complex code structures, this invention proposes a method that directly manipulates program structure graphs and learns semantic information from them. This allows the model to store a large amount of control and data dependency information, enabling it to capture the underlying code structures of many software vulnerabilities. Addressing the aforementioned issues and practical implications, this invention enhances data preprocessing capabilities, fully learns the graph structure information of the code, and trains a bidirectional gated graph neural network (BGGNN) by optimizing parameter settings to achieve better detection performance. Summary of the Invention
[0005] To address the problem that existing static vulnerability discovery methods cannot effectively represent nonlinear semantic information in code graph structures and to improve the performance of neural network models, this invention provides a vulnerability detection method based on graph embedding and bidirectional gated graph neural networks, which can effectively solve the above problems. The specific technical solution adopted by this invention is as follows:
[0006] A vulnerability detection method based on graph embedding and bidirectional gated graph neural networks includes the following steps:
[0007] S1. Dataset acquisition and annotation, which includes the following sub-steps:
[0008] S11. Obtain the source code dataset and extract the source code containing vulnerabilities and the source code without vulnerabilities at the function level, including a total of k functions;
[0009] S12. Mark each function to determine if it contains vulnerabilities, obtaining the mark Y for each function file. i ∈{0,1},i∈[1,k], where 0 indicates that there is no vulnerability and 1 indicates that there is a vulnerability.
[0010] S2. Generate a program dependency graph, obtaining the set G = {V, E} of program dependencies for all source code in the entire project, where V represents the set of nodes and E represents the set of edges. This includes the following sub-steps:
[0011] S21. After importing the source code into the source code analysis tool, use query statements as input based on the function names in the source code. This step generates a program dependency graph (PDG) corresponding to the function names and outputs it as a dot-type graph description file;
[0012] S22. Using a unified variable name mapping method, user-defined variable names and function names are mapped to symbol names in a one-to-one manner in the PDG graph description file to obtain the preprocessed PDG graph.
[0013] S3. For all program dependency graphs, extract the required edge information and node text information. This includes the following sub-steps:
[0014] S31. Use regular expression matching to extract the directed edge relationships E between nodes in a dot file. ij =V i →V j Obtain the set of all directed edges and save it as a text file;
[0015] S32. Use regular expression matching to extract the code text V corresponding to the node ID in the dot file. i =[Text1,Text2,...,Text n This function retrieves the text set of all nodes and saves it as a dictionary file.
[0016] S4. Use node2vec to train features and obtain a dictionary of feature vectors. This includes the following sub-steps:
[0017] S41. Using the preprocessed text file with directed edges from S31 as input, appropriately configure the sampling strategy parameters in the node2vec model, train the text features, and output the smallest text unit, Text. i The corresponding vector_t i,i∈[1,n];
[0018] S42. Store all the output text feature vectors in a dictionary, named Dict_t=∪ i∈[1,n] {key:Text i ,value:vector_t i};
[0019] S43. Using the preprocessed text file containing directed edges from S31 as input, appropriately set the sampling strategy parameters in the node2vec model, and identify each node with a unique node ID. i Instead of the aforementioned text attributes, the training focuses on node dependency features, outputting a vector of dependency features between graph nodes. i ,i∈[1,m];
[0020] S44. Store all output node-dependent feature vectors in a dictionary, named Dict_n = ∪ i∈[1,m] {key:ID i ,value:vector_n i}
[0021] S5. Based on the text feature vectors and edge feature vectors obtained from training in S4, convert all PDGs into a matrix representation of feature vectors at the function level. This includes the following sub-steps:
[0022] S51. Merge the text descriptions representing nodes into a single line, then split the string into several Text elements. Based on the text vector dictionary Dict_t obtained in S42, convert the node text attributes into corresponding embedding vectors nodeTextvec. i =[vector_t i1 vector_t i2 ,...,vector_t in This allows us to obtain the text vector for each node.
[0023] S52. A directed edge represents the existence of a head node and a tail node, determined by using the IDs of these two nodes. s ID e Using the node ID dictionary obtained from S44 as the key, we can retrieve the head node vector vector_n. s The tail node vector vector_n e ;
[0024] S53. Subtract the head node vector from the tail node vector to obtain the embedding vector v corresponding to a directed edge. s→e =vector_n e-vector_n s The above process is performed on each directed edge in the list of each program dependency graph to obtain the edge vectors in all PDGs.
[0025] S54. Encapsulate the node text vectors and edge vectors together into a JSON file corresponding to the program dependency graph, which will then serve as the input to the subsequent neural network model. This JSON file can be viewed as a combination of an N×16 two-dimensional vector matrix and an M×16 two-dimensional vector matrix, where N represents the number of nodes in the program dependency graph and M represents the number of edges.
[0026] S6. Using the multiple JSON files output from S5 as input, train the bidirectional gated graph neural network model, specifically including the following sub-steps:
[0027] S61. Split the training set and the test set: Select d% of the data samples in the JSON file dataset generated in S53 as the training set, and the rest as the test set;
[0028] S62. A bidirectional gated graph neural network (BGGNN) is applied to learn the features contained in the dataset. The BGGNN consists of two directional gated graph neural networks (GGNN): one is a forward L1-layer gated graph neural network GGNN1, which receives forward input; the other is a backward L2-layer gated graph neural network GGNN2, which learns the backward input. The formula is expressed as:
[0029]
[0030]
[0031]
[0032] In the above formula, y t It is the model output. It is a forward output. It is a reverse output;
[0033] S63. Perform l iterations of training based on the above network. After training, save the neural network model to facilitate fast model loading later.
[0034] S7. Perform code vulnerability detection on the target program, which includes the following sub-steps:
[0035] S71. First, preprocess the target program source code as in steps S2 and S3 to obtain the preprocessed PDG;
[0036] S72. Based on the pre-trained dictionaries Dict_t and Dict_n in S4, perform the transformation steps from PDG to feature vector matrix as in S5, and save it as a JSON file of the target program;
[0037] S73. Reuse the neural network model Model generated in S6, and use the function-level feature vector matrix generated by the target program as input to perform function-level code vulnerability detection;
[0038] S74. This invention will output a list of function names containing potential code vulnerabilities in the target program, for relevant personnel to check and improve the program.
[0039] Preferably, in step S62, the gated graph neural network GGNN1 with a forward L1 layer has an L1 value of 3.
[0040] Preferably, in step S62, the gated graph neural network GGNN2 with an inverse L2 layer has an L2 value of 3.
[0041] Preferably, in step S63, the training is performed for l iterations, where l is 150.
[0042] The beneficial effects of this invention are as follows:
[0043] This invention presents a code vulnerability detection method based on graph embedding learning. Using a real vulnerability dataset, it extracts control flow and dependencies from the source code, making the code representation more specific and comprehensive. After training the code with node2vec, it learns the graph embedding and text embedding information of the code. A bidirectional gated graph neural network (BGGNN) model architecture is used as the classifier. The recurrent structure effectively learns the neighborhood information of graph nodes, exhibiting good performance. The objective function is optimized using stochastic gradient ascent, effectively improving the feature discrimination ability. Appropriate training techniques, such as selecting ideal network parameters, optimization algorithms, and learning rate settings, result in a more stable network, more reliable results, and improved accuracy in code vulnerability detection. Attached Figure Description
[0044] Figure 1 This is a flowchart of the method of the present invention;
[0045] Figure 2 This is a schematic diagram of the bidirectional gated graph neural network used for vulnerability detection in this invention. Detailed Implementation
[0046] 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.
[0047] Conversely, this invention encompasses any substitutions, modifications, equivalent methods, and solutions made within the spirit and scope of the invention as defined in the claims. Furthermore, to provide a better understanding of the invention, certain specific details are described in detail below. However, those skilled in the art will fully understand the invention even without these detailed descriptions.
[0048] This embodiment provides a code vulnerability detection method based on graph embedding and bidirectional gated graph neural networks, such as Figure 1 As shown, it includes the following steps:
[0049] S1. Dataset acquisition and annotation, which includes the following sub-steps:
[0050] S11. Obtain the source code dataset and extract the source code containing vulnerabilities and the source code without vulnerabilities at the function level, including a total of k functions;
[0051] S12. Mark each function to determine if it contains vulnerabilities, obtaining the mark Y for each function file. i ∈{0,1},i∈[1,k], where 0 indicates that there is no vulnerability and 1 indicates that there is a vulnerability.
[0052] S2. Generate a program dependency graph, obtaining the set G = {V, E} of program dependencies for all source code in the entire project, where V represents the set of nodes and E represents the set of edges. This includes the following sub-steps:
[0053] S21. After importing the source code into the source code analysis tool, use query statements as input based on the function names in the source code. This step generates a program dependency graph (PDG) corresponding to the function names and outputs it as a dot-type graph description file;
[0054] S22. Using a unified variable name mapping method, user-defined variable names and function names are mapped to symbol names in a one-to-one manner in the PDG graph description file to obtain the preprocessed PDG graph.
[0055] S3. For all program dependency graphs, extract the required edge information and node text information. This includes the following sub-steps:
[0056] S31. Use regular expression matching to extract the directed edge relationships E between nodes in a dot file. ij =V i →V j Obtain the set of all directed edges and save it as a text file;
[0057] S32. Use regular expression matching to extract the code text V corresponding to the node ID in the dot file. i=[Text1,Text2,...,Text n This function retrieves the text set of all nodes and saves it as a dictionary file.
[0058] S4. Use node2vec to train features and obtain a dictionary of feature vectors. This includes the following sub-steps:
[0059] S41. Using the preprocessed text file with directed edges from S31 as input, appropriately configure the sampling strategy parameters in the node2vec model, train the text features, and output the smallest text unit, Text. i The corresponding vector_t i ,i∈[1,n];
[0060] S42. Store all the output text feature vectors in a dictionary, named Dict_t=∪ i∈[1,n] {key:Text i ,value:vector_t i};
[0061] S43. Using the preprocessed text file containing directed edges from S31 as input, appropriately set the sampling strategy parameters in the node2vec model, and identify each node with a unique node ID. i Instead of the aforementioned text attributes, the training focuses on node dependency features, outputting a vector of dependency features between graph nodes. i ,i∈[1,m];
[0062] S44. Store all output node-dependent feature vectors in a dictionary, named Dict_n = ∪ i∈[1,m] {key:ID i ,value:vector_n i}
[0063] In this embodiment, in S41 and S43, the parameters in the node2vec model are set appropriately. The specific parameters include walk_length, num_walks, p, q, and window_size, which represent the walk length, the number of walks, the probability of visiting the previous node, whether the walk direction is depth-first or breadth-first (depth-first when q<1, breadth-first when q>1), and the window size, respectively. Among them, walk_length is 10, num_walks is 10, p is 0.1, q is 0.8, and window_size is 5.
[0064] S5. Based on the text feature vectors and edge feature vectors obtained from training in S4, convert all PDGs into a matrix representation of feature vectors at the function level. This includes the following sub-steps:
[0065] S51. Merge the text descriptions representing nodes into a single line, then split the string into several Text elements. Based on the text vector dictionary Dict_t obtained in S42, convert the node text attributes into corresponding embedding vectors nodeTextvec. i =[vector_t i1 vector_t i2 ,...,vector_t in This allows us to obtain the text vector for each node.
[0066] S52. A directed edge represents the existence of a head node and a tail node, determined by using the IDs of these two nodes. s ID e Using the node ID dictionary obtained from S44 as the key, we can retrieve the head node vector vector_n. s and head node vector_n e ;
[0067] S53. Subtract the head node vector from the tail node vector to obtain the embedding vector v corresponding to a directed edge. s→e =vector_n e -vector_n s The above process is performed on each directed edge in the list of each program dependency graph to obtain the edge vectors in all PDGs.
[0068] S54. Encapsulate the node text vectors and edge vectors together into a JSON file corresponding to the program dependency graph, which will then serve as the input to the subsequent neural network model. This JSON file can be viewed as a combination of an N×16 two-dimensional vector matrix and an M×16 two-dimensional vector matrix, where N represents the number of nodes in the program dependency graph and M represents the number of edges.
[0069] S6. Using the multiple JSON files output from S5 as input, train the bidirectional gated graph neural network model, specifically including the following sub-steps:
[0070] S61. Split the training set and the test set: Select d% of the data samples in the JSON file dataset generated in S53 as the training set, and the rest as the test set; where d is 70.
[0071] S62. A bidirectional gated graph neural network (BGGNN) is applied to learn the features contained in the dataset. The BGGNN consists of two directional gated graph neural networks (GGNN): one is a forward L1-layer gated graph neural network GGNN1, which receives forward input; the other is a backward L2-layer gated graph neural network GGNN2, which learns the backward input. The formula is expressed as:
[0072]
[0073]
[0074]
[0075] In the above formula, y t It is the model output. It is a forward output. It is a reverse output, where L1 is 3 and L2 is 3;
[0076] S63. Based on the above network, perform l iterations of training. After training, save the neural network model Model to facilitate fast model loading later. l is 150.
[0077] S7. Perform code vulnerability detection on the target program, which includes the following sub-steps:
[0078] S71. First, preprocess the target program source code as in steps S2 and S3 to obtain the preprocessed PDG;
[0079] S72. Based on the pre-trained dictionaries Dict_t and Dict_n in S4, perform the transformation steps from PDG to feature vector matrix as in S5, and save it as a JSON file of the target program;
[0080] S73. Reuse the neural network model Model generated in S6, and use the function-level feature vector matrix generated by the target program as input to perform function-level code vulnerability detection;
[0081] S74. This invention will output a list of function names containing potential code vulnerabilities in the target program, for relevant personnel to check and improve the program.
[0082] The embodiments of the present invention have been described in detail above with reference to the accompanying drawings, but the present invention is not limited to the described embodiments. For those skilled in the art, various changes, modifications, substitutions, and variations can be made to these embodiments, including components, without departing from the principles and spirit of the present invention, and these variations still fall within the protection scope of the present invention.
Claims
1. A vulnerability detection method based on graph embedding and bidirectional gated graph neural networks, characterized in that: S1. Obtain the dataset of source code and label it for vulnerabilities; S2. Generate a program dependency graph, obtaining a set of program dependency graphs corresponding to all source code in the entire project. , The set representing nodes, and The set representing edges; S3. For all program dependency graphs, extract the required set of directed edges and set of nodes, and save them as a text file and a dictionary file, respectively; S4. Use node2vec for feature training to obtain a dictionary of feature vectors; The feature vector dictionary in S4 includes a text vector dictionary and a node ID dictionary. The text vector dictionary is obtained by taking a text file as input to obtain text feature vectors and is stored in a dictionary. The node ID dictionary is obtained by taking a dictionary file as input to obtain node dependency feature vectors and is stored in a dictionary. The method for obtaining the text vector dictionary is as follows: using the preprocessed text file containing directed edges from S31 as input, setting the sampling strategy parameters in the node2vec model, training text features, and outputting the smallest text unit. corresponding vector ; All output text feature vectors are stored in a dictionary, which is: ; The method for obtaining the node ID dictionary is as follows: Using the preprocessed text file containing directed edges from S31 as input, the sampling strategy parameters in the node2vec model are appropriately set to identify each node as unique. Instead of the aforementioned text attributes, the training focuses on node dependency features, outputting a dependency feature vector between graph nodes. ; All output node-dependent feature vectors are stored in a dictionary, which is: ; S5. Based on the text feature vectors and edge feature vectors obtained from training in S4, convert all PDGs into a matrix representation of feature vectors at the function level; S6. Use the multiple JSON files output from S5 as input to train a bidirectional gated graph neural network model; S7. Perform code vulnerability detection on the target program. Complete the vulnerability detection on the target program code through steps S1-S6.
2. The vulnerability detection method based on graph embedding and bidirectional gated graph neural network according to claim 1, characterized in that, The method for annotating the source code in the dataset S1 is as follows: extract function-level source code containing vulnerabilities and source code without vulnerabilities from the source code dataset, including a total of Each function is labeled to determine if it contains vulnerabilities, resulting in a labeled file for each function. , where 0 indicates that there is no vulnerability and 1 indicates that there is a vulnerability.
3. The vulnerability detection method based on graph embedding and bidirectional gated graph neural network according to claim 1, characterized in that, In S2, the method for generating the program dependency graph is as follows: S21. After importing the source code into the source code analysis tool, based on the function names in the source code, use query statements as input to generate a program dependency graph corresponding to the function names, and output it as a dot type graph description file; S22. Using a unified variable name mapping method, user-defined variable names and function names are mapped to symbol names in a one-to-one manner in the description file of the program dependency graph to obtain the preprocessed program dependency graph.
4. The vulnerability detection method based on graph embedding and bidirectional gated graph neural network according to claim 1, characterized in that, In S3, the required set of directed edges and nodes are extracted using regular expression matching.
5. The vulnerability detection method based on graph embedding and bidirectional gated graph neural network according to claim 4, characterized in that, S5 specifically includes the following sub-steps: S51. Combine the text descriptions representing the nodes into one line, then split the string into several parts. Based on text vector dictionary Convert node text attributes into corresponding embedding vectors. This yields the text vector for each node; S52. A directed edge represents the existence of a head node and a tail node, which are used to define the head and tail nodes. Using the node ID dictionary as the key, we obtain the head node vector. Sum of tail node vectors ; S53. Subtract the head node vector from the tail node vector to obtain the embedding vector corresponding to a directed edge. Perform the above processing on each directed edge in the list of each program dependency graph to obtain the edge vectors in all PDGs. ; S54. Encapsulate the node text vectors and edge vectors together into a JSON file corresponding to the program dependency graph, and use this JSON file as input to the subsequent neural network model. This JSON file is considered as a... A two-dimensional vector matrix and a A combination of two-dimensional vector matrices, where This represents the number of nodes in a program's dependency graph. Represents the number of edges.
6. The vulnerability detection method based on graph embedding and bidirectional gated graph neural network according to claim 5, characterized in that, S6 specifically includes the following sub-steps: S61. Split the training and test sets: Select the dataset from the JSON file generated in S53. % of the data samples are used as the training set, and the remainder as the test set; S62. A bidirectional gated graph neural network (BGGNN) is used to learn the features contained in the dataset. BGGNN consists of two directional gated graph neural networks: one forward and one backward. The gated graphical neural network GGNN1 has two layers, one accepting positive input and the other accepting negative input. The gated graphical neural network GGNN2, which learns the inverse input, is expressed by the following formula: ； ； ； In the above formula, It is the model output. It is a forward output. It is a reverse output; S63. Based on the above network, After one iteration of training, the neural network model is saved. This facilitates faster model loading later.
7. The vulnerability detection method based on graph embedding and bidirectional gated graph neural network according to claim 6, characterized in that, S7 specifically includes the following sub-steps: S71. First, preprocess the target program source code as in steps S2 and S3 to obtain the preprocessed program dependency graph; S72. Using the dictionary pre-trained in S4 and Based on this, perform the steps of converting PDG to feature vector matrix as in S5, and save it as a JSON file of the target program; S73. Reusing the neural network model generated in S6 The function-level feature vector matrix generated by the target program is used as input to perform function-level code vulnerability detection. S74. Output a list of function names that contain potential code vulnerabilities in the target program, so that relevant personnel can check and improve the program.