Source code vulnerability detection method using heterogeneous graph attention network

Abstract

The invention 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 that semantic information of a deep learning model is insufficient in the feature extraction process is solved. According to the technical scheme, the method comprises the following steps of S1, performing data preprocessing on an open source data set; s2, parsing the source code into a graph structure through a code parser; s3, fusing various code structure diagrams, and pruning the diagram structures; s4, carrying out feature initialization on the nodes of the heterogeneous graph by utilizing Word2Vec; s5, dividing the heterogeneous graph into 12 meta-graphs based on the meta-path; s6, inputting heterogeneous graph information into a two-layer attention network for model training; and S7, screening out an optimal model based on the verification set. The method has the advantages that the source code is analyzed into the heterogeneous graph structure, model training is carried out in combination with the two layers of attention networks, and the performance of vulnerability detection is effectively improved.

Description

Technical Field

[0001] The present invention relates to the technical field of intelligent software engineering, and in particular to a source code vulnerability detection method using a heterogeneous graph attention network. Background Art

[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 detect and fix software vulnerabilities effectively and efficiently. Many techniques have been proposed to prevent vulnerabilities. Among them, vulnerability detection is one of the most important defense measures. Many vulnerability detection techniques have been proposed and have been successful in evaluation. Traditional vulnerability detection methods are mainly based on static analysis and dynamic analysis. Static analysis relies on fixed rules to detect vulnerabilities and effectively identifies vulnerabilities that meet these rules. However, when dealing with highly complex code, especially in the case of obfuscation or polymorphism, the accuracy drops significantly. 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. Most of the research in recent years has achieved good results based on deep learning models. However, these methods have problems such as limited code representation, redundant information in data, and insufficient model learning.

[0003] How to solve the above technical problems becomes the subject faced by the present invention. Summary of the invention

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

[0005] The idea of ​​the present invention is: the present invention proposes a source code vulnerability detection method using a heterogeneous graph attention network, that is, by establishing a heterogeneous graph information network to extract the semantic feature information of the vulnerability; using the code parser Joern to extract the graph structure information of the source code; pruning the graph through the HGS1 and HGS2 algorithms; dividing the heterogeneous graph into multiple meta-graphs based on meta-paths; and finally constructing a two-layer attention network for testing, so as to more accurately detect whether there are vulnerabilities in the function. The method of the present invention achieves better performance than other methods in the same data set.

[0006] The present invention is implemented by the following measures, and the technical solution is specific: a source code vulnerability detection method using a heterogeneous graph attention network, which includes the following steps:

[0007] (1) Collect the C and C++ source codes required by the project from open source websites, preprocess them, and construct a vulnerability dataset D. The format of D is <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 contents 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 data set D and extract the abstract syntax tree, control flow graph, and data flow graph of the function;

[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 according to the abstract syntax tree, control flow graph and data flow graph of the function, whose format is <node id, node information, node type>;

[0014] (3-2) The abstract syntax tree, control flow graph, and data flow graph of the function are spliced ​​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 <starting node, ending 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 the 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 ,…,meta graph 12>, where each meta graph corresponds to a meta path;

[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 in combination with 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 embedding, use Softmax normalization as the semantic-level attention, obtain the weight coefficient, and then fuse all node-level embeddings 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 set, validation set, and test set at a ratio of 80%:10%:10% randomly, 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 best 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 the prior art, the beneficial effects of the present invention are as follows: the present invention proposes a source code vulnerability detection method using a heterogeneous graph attention network, uses a new heterogeneous graph construction method, uses Joern to generate a code structure diagram at the function level, constructs a heterogeneous graph for the code, proposes two heterogeneous graph simplification algorithms to remove redundant information, then uses Word2Vec to generate an initial vectorized representation of the node, and finally uses a two-layer attention network to implement software vulnerability detection. This method has been proven to be superior to other existing methods through multiple sets of experiments. 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 1A 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 solution and advantages of the present invention more clearly understood, the present invention is further described in detail below in conjunction with the accompanying drawings and embodiments. Of course, the specific embodiments described here are only used to explain the present invention and are not used 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, containing about 10k vulnerability instances and 12k vulnerability-free instances. This embodiment preprocesses the source code, including removing redundant content such as comments, blank lines, macro definitions, etc., unifying variables into the form of var1, var2, etc., 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 data set 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 data set D and extract the abstract syntax tree, control flow graph, and data flow graph of the function;

[0038] (3) Construct the graphs in the previous step into a heterogeneous graph G, extract the node set N according to the abstract syntax tree, control flow graph and data flow graph of the function, and 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 <starting node, ending 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 ,…,meta graph 12>, where each meta graph corresponds to a meta path;

[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 in combination with 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 embedding, use Softmax normalization as the semantic-level attention, obtain the weight coefficient, and then fuse all node-level embeddings 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 set, validation set, and test set at a ratio of 80%:10%:10% randomly, 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 best 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 samples predicted correctly ÷ the total number of samples"; Precision, which means the accuracy or precision, refers to how many samples are actually True among the samples predicted as True by the model; Recall, which means the recall rate, refers to how many samples that are actually True are selected; F1-score comprehensively considers the factors of precision and recall, and achieves a compromise between the two.

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

[0048]

[0049] Experiments show that the source code vulnerability detection method using a heterogeneous graph attention network proposed in this embodiment can perform more reliable vulnerability detection compared to other baseline methods. The heterogeneous graph-based method of this embodiment can effectively capture grammatical and semantic information from source code, which is better than all baseline models. Specifically, compared with AMPLE, the accuracy of this embodiment method on the Devign dataset is improved by 7.7%, and the F1 score is improved by 3.8%.

[0050] Example 2

[0051] Based on Example 1, this example selects another data set for experiment, which specifically includes the following contents:

[0052] (1) 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, which consists of about 2k vulnerability instances and 20k non-vulnerable instances. Then preprocess the source code, and the results after processing are shown in Table 3;

[0053] Table 3 Dataset information

[0054]

[0055] (2) The steps are the same as those in Example 1. The Joern tool is used to parse the code of each function in the dataset D and construct a heterogeneous graph G. A meta-path-based heterogeneous graph attention network is constructed and a multi-head self-attention mechanism is integrated to obtain the prediction label. 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 using a heterogeneous graph attention network proposed in this embodiment can perform more reliable vulnerability detection compared to other baseline methods. Specifically, compared with AMPLE, the accuracy of the method in this embodiment on the Reveal dataset is improved by 1.5%, and the F1 score is improved by 12.2%.

[0059] Example 3

[0060] Based on Example 1, this example selects another data set for experiment, which specifically includes the following contents:

[0061] (1) Select the open source Fan et al. dataset and construct it into vulnerability dataset D. The Fan et al. dataset comes from more than 300 open source C / C++ GitHub projects and consists of about 10k vulnerability instances and 177k non-vulnerable instances. Then preprocess the source code. The results after processing are shown in Table 5.

[0062] Table 5 Dataset information

[0063]

[0064] (2) The code of each function in the dataset D is parsed and constructed into a heterogeneous graph G. A meta-path-based heterogeneous graph attention network is constructed, and a multi-head self-attention mechanism is integrated to obtain the prediction label. Finally, the dataset is divided 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 using a heterogeneous graph attention network proposed in this embodiment can perform more reliable vulnerability detection than other baseline methods and is superior to all baseline models. Specifically, compared with AMPLE, the accuracy of the method in this embodiment on the Fan et al. dataset is improved by 2%, and the F1 score is improved by 2.01%. Moreover, the vulnerability detection performance of the method in this embodiment is 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 principle of the present invention should be included in the protection scope of the present invention.

Claims

1. A source code vulnerability detection method using a heterogeneous graph attention network, characterized in that it includes the following steps: 1.1: Collect the C and C++ source codes required by 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 data set D and extract the abstract syntax tree, control flow graph and data flow graph of the function; 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, and obtain the final data set structure D1, whose format is <function id, node set, edge set, feature vector set, label>; 1.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; 1.6: Divide the data set D1 obtained in step 1.4 into training set, validation set and test set at a ratio of 80%:10%:10% randomly, and then use the training set for model training in step 1.

5. According to 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 step 1.1 includes collecting open source C and C++ data sets and preprocessing, including 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 according to the abstract syntax tree, control flow graph and data flow graph of the function, whose 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 edges, control flow edges, and data flow edges. 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 on 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, the heterogeneous graphs are constructed into a heterogeneous information network, whose format is <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 adjacent nodes in each meta-graph, and then calculate the attention coefficient between each neighbor 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. 4.3: Construct a semantic-level attention layer, calculate the weight of each meta-path, perform nonlinear transformation on the node-level embedding, use Softmax normalization as the semantic-level attention, obtain the weight coefficient, and fuse all 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