A source code vulnerability detection method using a heterogeneous graph attention network

Abstract

The application provides a source code vulnerability detection method based on a heterogeneous graph attention network, and belongs to the field of intelligent software engineering. The technical problem of insufficient semantic information in the feature extraction process of a deep learning model is solved. The technical scheme comprises the following steps: S1: data preprocessing is performed on an open source data set; S2: a code parser is used to parse the source code into a graph structure; S3: a plurality of code structure graphs are fused, and a pruning operation is performed on the graph structure; S4: Word2Vec is used to initialize the features of the heterogeneous graph nodes; S5: the heterogeneous graph is divided into 12 meta graphs based on a meta path; S6: the heterogeneous graph information is input into a two-layer attention network for model training; and S7: the best model is selected based on a verification set. The application has the beneficial effect that the source code is parsed into a heterogeneous graph structure, and a two-layer attention network is used for model training, thereby effectively improving the performance of vulnerability detection.

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of intelligent software engineering, and particularly relates to a source code vulnerability detection method using a heterogeneous graph attention network. BACKGROUND

[0002] Software vulnerabilities, also known as software defects or weaknesses, pose a significant risk to confidentiality, integrity, and data availability. Therefore, it is crucial to effectively and efficiently detect and fix software vulnerabilities. In order to prevent vulnerabilities, many techniques have been proposed. Among them, vulnerability detection is one of the most important defense measures. Many vulnerability detection techniques have been proposed and have achieved success in evaluations. Traditional vulnerability detection methods are mainly based on static analysis and dynamic analysis. Static analysis relies on fixed rules to detect vulnerabilities, effectively identifying vulnerabilities that conform to these rules. However, when dealing with highly complex code, especially in the case of obfuscation or polymorphism, accuracy will be significantly reduced. Dynamic analysis methods allow real-time detection of vulnerabilities, but they usually have strict real-time performance requirements and may miss vulnerabilities that are not displayed during execution or depend on specific execution paths. Recent research has mostly been based on deep learning models, achieving good results. However, these methods have the problems of limited code representation, redundant information in data, and insufficient model learning.

[0003] How to solve the above technical problems becomes a problem faced by the present application. SUMMARY

[0004] The purpose of the present application is to provide a source code vulnerability detection method using a heterogeneous graph attention network. This method can automatically detect whether there are vulnerabilities in the source code.

[0005] The idea of the present application is that the present application proposes a source code vulnerability detection method using a heterogeneous graph attention network, which extracts semantic feature information of vulnerabilities by establishing a heterogeneous graph information network; uses a code parser Joern to extract graph structure information of the source code; performs pruning operations on the graph through HGS1 and HGS2 algorithms; divides the heterogeneous graph into multiple meta-graphs based on meta-paths; and finally constructs a two-layer attention network for testing, thereby more accurately detecting whether there are vulnerabilities in the function. The method of the present application achieves better performance than other methods in the same data set.

[0006] The present application is realized by the following measures, and the technical scheme is as follows: a source code vulnerability detection method using a heterogeneous graph attention network, comprising the following steps:

[0007] (1) Collect the required C and C++ source code of the project through an open source website, preprocess it, and construct it into a vulnerability data set D, with a format of <function id, function, label>;

[0008] (1-1) Select the open source Devign, Reveal, and Fan et al. datasets;

[0009] (1-2) Preprocess the source code, including removing redundant content such as comments, blank lines, macro definitions, unifying variables into the form of var1, var2, etc., and unifying functions into the form of func1, func2;

[0010] (1-3) Delete functions with truncation exceptions, delete functions that cannot be parsed by the Joern tool, and delete functions with more than 500 nodes after parsing;

[0011] (2) Use the Joern tool to parse the code of each function in dataset D and extract the function's abstract syntax tree, control flow graph, and data flow graph;

[0012] (3) Construct each graph in the previous step into a heterogeneous graph G, and then use the code pruning algorithms HGS1 and HGS2 to perform graph pruning operations on the heterogeneous graph G;

[0013] (3-1) Extract the node set N based on the function's abstract syntax tree, control flow graph, and data flow graph, whose format is <node id, node information, node type>;

[0014] (3-2) The function's abstract syntax tree, control flow graph, and data flow graph are concatenated based on the node ID. Each node is traversed, and then an edge set is constructed. Edges from different graphs are added to the edge set and recorded as different variable types, such as syntax edges, control flow edges, data flow edges, etc. The format of the edge set E is <start node, end node, edge type>;

[0015] (3-3) For the HGS1 algorithm, first perform a breadth-first traversal of the input graph. When two nodes v and j have the same value and share a common parent node, delete the edge between v and j and node j, and change the type of v to the parent node type;

[0016] (3-4) For the HGS2 algorithm, when the values ​​of nodes v and j are the same and v is the only parent node of j, the two nodes are merged.

[0017] (4) Input the processed heterogeneous graph G into the Word2Vec pre-training model to generate the initial feature vector of each node, and obtain the final data set structure D1, whose format is <function id, node set, edge set, feature vector set, label>;

[0018] (5) Construct a meta-path-based heterogeneous graph attention network and integrate a multi-head self-attention mechanism to mine deeper semantic and syntactic information in the source code and obtain predicted labels;

[0019] (5-1) Based on 12 fixed semantic meta-paths, the heterogeneous graph is constructed into a heterogeneous information network, whose format is <meta Figure 1 ,Yuan Figure 2 ,…,metagraph12>, where each metagraph corresponds to a metapath;

[0020] (5-2) Construct a node-level attention layer. Calculate the correlation coefficient e between each node and its neighboring nodes in each meta-graph, and then calculate the attention coefficient between each neighboring node and the node. After obtaining the attention coefficient, perform a linear transformation on the initial node representation, and then update the node representation based on the attention coefficient.

[0021] (5-3) Construct a semantic-level attention layer and calculate the weight of each meta-path. Perform a nonlinear transformation on the node-level embeddings and use Softmax normalization as the semantic-level attention. After obtaining the weight coefficients, all node-level embeddings are fused together to obtain the final embedding.

[0022] (5-4) After obtaining the final embedding vector, it is used as the graph-level feature representation of the function file, fed back into a multi-layer perceptron, and classified using the sigmoid function.

[0023] (6) Divide the dataset D1 into training, validation, and test sets at a random ratio of 80%:10%:10%, and then use the training set for model training in step 1.5. Based on the specified performance metric, retain the model M with the highest performance based on the validation set.

[0024] (7) Treat each function in the test set as the source code to be tested, and use the best model M to predict whether these functions contain vulnerabilities.

[0025] Compared with existing technologies, the present invention offers the following advantages: A source code vulnerability detection method using a heterogeneous graph attention network employs a novel heterogeneous graph construction method, employing Joern to generate a code structure diagram at the function level. Two heterogeneous graph simplification algorithms are proposed to remove redundant information. Word2Vec is then used to generate an initial vectorized representation of the nodes. Finally, a two-layer attention network is employed to detect software vulnerabilities. Multiple experiments demonstrate the superiority of this method over existing methods. BRIEF DESCRIPTION OF THE DRAWINGS

[0026] The accompanying drawings are used to provide further understanding of the present invention and constitute a part of the specification. They are used to explain the present invention together with the embodiments of the present invention and do not constitute a limitation of the present invention.

[0027] Figure 1This is a system framework diagram of a source code vulnerability detection method using a heterogeneous graph attention network provided by the present invention.

[0028] Figure 2 This is an example diagram of a source code heterogeneous graph of the present invention. DETAILED DESCRIPTION

[0029] In order to make the purpose, technical solutions and advantages of the present invention more clearly understood, the present invention will be further described in detail below in conjunction with the accompanying drawings and embodiments. Of course, the specific embodiments described herein are only used to explain the present invention and are not intended to limit the present invention.

[0030] Example 1

[0031] See also Figure 1 As shown, this embodiment provides a source code vulnerability detection method using a heterogeneous graph attention network, which includes the following contents:

[0032] (1) Collect the C and C++ source codes required for the project from open source websites, select the open source Devign dataset, preprocess it, and construct a vulnerability dataset D. The format of D is <function id, function, label>;

[0033] (1-1) The Devign dataset is manually labeled and comes from two open-source C language projects. It contains approximately 10,000 vulnerability instances and 12,000 vulnerability-free instances. This example preprocesses the source code, including removing redundant content such as comments, blank lines, and macro definitions, unifying variables into the form of var1, var2, and unifying functions into the form of func1, func2.

[0034] (1-2) Delete functions with truncation anomalies, delete functions that cannot be parsed by the Joern tool, and delete functions with more than 500 nodes after parsing. The statistical results of the dataset after preprocessing are shown in Table 1;

[0035] Table 1 Dataset information

[0036]

[0037] (2) Use the Joern tool to parse the code of each function in dataset D and extract the function's abstract syntax tree, control flow graph, and data flow graph;

[0038] (3) Construct the graphs from the previous step into a heterogeneous graph G. Extract the node set N based on the abstract syntax tree, control flow graph, and data flow graph of the function. Its format is <node id, node information, node type>. Then construct an edge set. Add the edges from different graphs to the edge set and record them as different variable types, such as syntax edge, control flow edge, data flow edge, etc. The format of the edge set E is <start node, end node, edge type>. Then use the code pruning algorithms HGS1 and HGS2 to perform graph pruning operations on the heterogeneous graph G. The processing results are shown in the following example. Figure 2 As shown;

[0039] (4) Input the processed heterogeneous graph G into the Word2Vec pre-training model to generate the initial feature vector of each node, and obtain the final data set structure D1, whose format is <function id, node set, edge set, feature vector set, label>;

[0040] (5) Construct a meta-path-based heterogeneous graph attention network and integrate a multi-head self-attention mechanism to mine deeper semantic and syntactic information in the source code and obtain predicted labels;

[0041] (5-1) Based on 12 fixed semantic meta-paths, the heterogeneous graph is constructed into a heterogeneous information network, whose format is <meta Figure 1 ,Yuan Figure 2 ,…,metagraph12>, where each metagraph corresponds to a metapath;

[0042] (5-2) Construct a node-level attention layer. Calculate the correlation coefficient e between each node and its neighboring nodes in each meta-graph, and then calculate the attention coefficient between each neighboring node and the node. After obtaining the attention coefficient, perform a linear transformation on the initial node representation, and then update the node representation based on the attention coefficient.

[0043] (5-3) Construct a semantic-level attention layer and calculate the weight of each meta-path. Perform a nonlinear transformation on the node-level embeddings and use Softmax normalization as the semantic-level attention. After obtaining the weight coefficients, all node-level embeddings are fused together to obtain the final embedding.

[0044] (5-4) After obtaining the final embedding vector, it is used as the graph-level feature representation of the function file, fed back into a multi-layer perceptron, and classified using the sigmoid function.

[0045] (6) Divide the dataset D1 into training, validation, and test sets at a random ratio of 80%:10%:10%, and then use the training set for model training in step 1.5. Based on the specified performance metric, retain the model M with the highest performance based on the validation set.

[0046] (7) Treat each function in the test set as the source code to be tested, use the best model M to predict whether these functions contain vulnerabilities, evaluate the method of this embodiment and the existing vulnerability detection method on the same data set, and use four performance indicators from the field of vulnerability detection research (i.e., Accuracy, Precision, Recall, and F1-score) to automatically evaluate the quality of the model. Accuracy, which means "the number of correctly predicted samples divided by the total number of samples"; Precision, which means the accuracy or precision, refers to the number of samples predicted as True by the model that are actually True; Recall, which means the recall rate, refers to the number of samples that are actually True that are selected; F1-score comprehensively considers the factors of precision and recall, and achieves a balance between the two.

[0047] Table 2 Comparison of results between the method of this embodiment and other methods

[0048]

[0049] Experiments demonstrate that the source code vulnerability detection method proposed in this embodiment, using a heterogeneous graph attention network, can provide more reliable vulnerability detection compared to other baseline methods. This embodiment's heterogeneous graph-based approach effectively captures syntactic and semantic information from source code, outperforming all baseline models. Specifically, compared to AMPLE, this embodiment's method achieves a 7.7% improvement in accuracy and a 3.8% improvement in F1 score on the Devign dataset.

[0050] Example 2

[0051] Based on Example 1, this example selects another data set for experimentation, specifically including the following:

[0052] (1) We select the open-source Reveal dataset and construct it into vulnerability dataset D. The Reveal dataset contains two open-source projects: Linux Debian Kernel and Chromium, and consists of approximately 2k vulnerability instances and 20k non-vulnerable instances. We then perform preprocessing on the source code. The results are shown in Table 3.

[0053] Table 3 Dataset information

[0054]

[0055] (2) Similar to the steps in Example 1, the Joern tool is used to parse the code of each function in dataset D and construct a heterogeneous graph G. A meta-path-based heterogeneous graph attention network is constructed, incorporating a multi-head self-attention mechanism to obtain predicted labels. Finally, the dataset is divided for model training. The results are shown in Table 4.

[0056] Table 4 Comparison of results between the method of this embodiment and other methods

[0057]

[0058] Experiments show that the source code vulnerability detection method proposed in this embodiment, using a heterogeneous graph attention network, can provide more reliable vulnerability detection compared to other baseline methods. Specifically, compared with AMPLE, this embodiment's method improves accuracy by 1.5% and F1 score by 12.2% on the Reveal dataset.

[0059] Example 3

[0060] Based on Example 1, this example selects another data set for experimentation, specifically including the following:

[0061] (1) We select the open-source Fan et al. dataset and construct it into vulnerability dataset D. The Fan et al. dataset comes from over 300 open-source C / C++ GitHub projects and consists of approximately 10,000 vulnerability instances and 177,000 vulnerability-free instances. We then preprocess the source code, and the results are shown in Table 5.

[0062] Table 5 Dataset information

[0063]

[0064] (2) We parse the code of each function in dataset D and construct a heterogeneous graph G. We then build a meta-path-based heterogeneous graph attention network, incorporating a multi-head self-attention mechanism to obtain predicted labels. Finally, we divide the dataset for model training. The results are shown in Table 6.

[0065] Table 6 Comparison of results between the method of this embodiment and other methods

[0066]

[0067]

[0068] Experiments show that the source code vulnerability detection method proposed in this embodiment, using a heterogeneous graph attention network, can provide more reliable vulnerability detection than other baseline methods, outperforming all baseline models. Specifically, compared with AMPLE, this embodiment's method improves accuracy by 2% and F1 score by 2.01% on the Fan et al. dataset. Furthermore, the vulnerability detection performance of this embodiment's method is even more significantly improved than that of the devign method.

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

Claims

1. A source code vulnerability detection method using a heterogeneous graph attention network, characterized by comprising the following steps: 1.1: Collect the C and C++ source codes required for the project from open source websites, preprocess them, and construct a vulnerability dataset D. The format of D is <function id, function, label>; 1.2: Use the Joern tool to parse the code of each function in dataset D and extract the function's abstract syntax tree, control flow graph, and data flow graph; 1.3: Construct each graph in 1.2 into a heterogeneous graph G, and then use the code pruning algorithms HGS1 and HGS2 to perform graph pruning operations on the heterogeneous graph G; 1.4: Input the processed heterogeneous graph G into the Word2Vec pre-training model to generate the initial feature vector of each node. The final dataset structure is D1, which has the format of <function id, node set, edge set, feature vector set, label>; 1.5: Build a meta-path-based heterogeneous graph attention network and integrate a multi-head self-attention mechanism to mine deeper semantic and syntactic information in the source code and obtain predicted labels; 1.6: Divide the dataset D1 obtained in step 1.4 into training, validation, and test sets at a random ratio of 80%:10%:10%. Then use the training set to train the model in step 1.

5. Based on the specified performance metric, retain the model M with the best performance based on the validation set. 1.7: Treat each function in the test set as the source code to be tested, and use the best model M to predict whether these functions contain vulnerabilities.

2. A source code vulnerability detection method using a heterogeneous graph attention network according to claim 1, characterized in that: The collection and preprocessing of open source C and C++ datasets in step 1.1 include the following steps: 2.1: Select the open source Devign, Reveal, and Fan et al. datasets; 2.2: Preprocess the source code, including removing comments, blank lines, and redundant macro definitions, unifying variables into var1 and var2, and unifying functions into func1 and func2; 2.3: Delete functions with truncation exceptions, delete functions that cannot be parsed by the Joern tool, and delete functions with more than 500 nodes after parsing.

3. A source code vulnerability detection method using a heterogeneous graph attention network according to claim 1, characterized in that: In step 1.3, each graph is constructed into a heterogeneous graph and two graph pruning algorithms, HGS1 and HGS2, are used, including the following steps: 3.1: Extract the node set N based on the function's abstract syntax tree, control flow graph, and data flow graph. Its format is <node id, node information, node type>; 3.2: Concatenate the function's abstract syntax tree, control flow graph, and data flow graph based on the node ID, traverse each node, and then build an edge set. Add edges from different graphs to the edge set and record them as different variable types: syntax edge, control flow edge, and data flow edge. The format of the edge set E is <start node, end node, edge type>; 3.3: For the HGS1 algorithm, first perform a breadth-first traversal of the input graph. When two nodes v and j have the same value and share a common parent node, delete the edge between v and j and the node j, and change the type of v to the parent node type. 3.4: For the HGS2 algorithm, when the values ​​of nodes v and j are the same and v is the only parent node of j, the two nodes are merged.

4. A source code vulnerability detection method using a heterogeneous graph attention network according to claim 1, characterized in that: In step 1.5, constructing a meta-path-based heterogeneous graph attention network includes the following steps: 4.1: Based on 12 fixed semantic meta-paths, heterogeneous graphs are constructed into a heterogeneous information network with the format of <meta-graph 1, meta-graph 2, …, meta-graph 12>, where each meta-graph corresponds to a meta-path; 4.2: Construct a node-level attention layer. Calculate the correlation coefficient e between each node and its neighboring nodes in each meta-graph. Then calculate the attention coefficient between each neighboring node and the node. After obtaining the attention coefficient, perform a linear transformation on the initial node representation. Then, combine the attention coefficient to update the node representation. 4.3: Construct a semantic-level attention layer, calculate the weight of each meta-path, perform a nonlinear transformation on the node-level embedding, use Softmax normalization as the semantic-level attention, obtain the weight coefficient, and then fuse all the node-level embeddings together to obtain the final embedding; 4.4: After obtaining the final embedding vector, it is used as the graph-level feature representation of the function file, fed back into a multi-layer perceptron, and classified using the sigmoid function.

Images