Function level vulnerability detection method based on source code annotation and program dependency graph

Abstract

The invention relates to the field of software security technology and the like, discloses a function-level vulnerability detection method based on source code annotation and a program dependency graph, and aims to realize function-level vulnerability detection by fusing code annotation and code statement features and using the program dependency graph as a code representation mode. A graph convolutional neural network is used for learning structure information and grammar semantic information of codes, the vulnerability detection rate is effectively improved, meanwhile, the mode has good expandability, and the method comprises the following steps that nonstandard annotation information existing in the codes is removed, and the codes are subjected to standardization processing; generating a program dependency graph (PDG) of the normalized code; generating annotation information for each node by using the large model to obtain a program dependency graph containing the code statements and the annotation information of the code statements; the vector representation of the code annotation information and the vector representation of the code statement are generated, and then the two vector representations are fused into a fused feature vector which is used for subsequently inputting the overall architecture of the graph neural network to carry out model training.

Description

TECHNICAL FIELD

[0001] The application relates to the field of software security technology, and particularly relates to a function-level vulnerability detection method based on source code annotation and program dependency graph. BACKGROUND

[0002] With the rapid development of the Internet and software, software is widely used in various fields, and the security and quality of source code have become an important issue that cannot be ignored in modern software development process. According to the report of IBM, the average cost of data leakage in 2023 reached 4.45 million US dollars. In the past ten years, the number of common vulnerabilities and exposures (CVE) has increased significantly, reaching 25,227 in 2022, an increase of 25.1% over 2021. Vulnerabilities in source code can cause serious security problems, such as system crashes, data leaks and malicious attacks. Therefore, source code vulnerability detection has become an indispensable part of the software development process.

[0003] At present, the source code vulnerability detection methods mainly include two types: traditional defect detection methods and deep learning-based defect detection methods. The traditional defect detection methods are divided into static analysis and dynamic analysis, and machine learning methods. Static analysis discovers defects by scanning source code, while dynamic analysis discovers defects by analyzing program state and execution path during program execution. However, these methods rely on expert-defined rules and are difficult to deal with complex software systems and new defects. Although early machine learning methods can reduce the dependence on expert knowledge, they still need to manually construct features, and the detection effect is limited, with high false positive rate and false negative rate.

[0004] In recent years, deep learning has achieved great success in image recognition, natural language processing and other fields, and researchers have begun to apply it to source code defect detection. Deep learning can automatically mine deep features in code and reduce the dependence on manual feature engineering. Deep learning (DL) based methods have achieved remarkable success, reducing the dependence on expert knowledge and being able to detect various types of software vulnerabilities.

[0005] Inspired by deep learning-based image classification technology, VulCNN converts the source code of a function into an image while preserving the program semantics, and combines deep learning technology to achieve efficient and accurate vulnerability detection. The experimental results show that VulCNN outperforms existing tools in vulnerability detection tasks and can discover new vulnerabilities in large-scale codebases.

[0006] SySeVR generates vulnerable code slices according to certain rules through program dependency graph (PDG), and finally encodes the code slices (set of program statements) into vectors to input BLSTM (bidirectional long short-term memory network) for training, and finally obtains the classification result.

[0007] LineVul uses CodeBERT as its base model, which is pre-trained on CodeSearchNet (a 20GB code corpus). This model can capture rich semantic and syntactic features of the code and utilizes a self-attention mechanism to capture long-distance dependencies in the code, overcoming the shortcomings of RNN-like models in processing long sequences. By analyzing the attention weights of each layer of the Transformer, it can directly locate potentially vulnerable lines of code.

[0008] However, the above code vulnerability detection methods have the following shortcomings: they treat code only as a text sequence and use natural language processing techniques to detect vulnerabilities. While sequence-based representation is simple and intuitive, it loses structural information, leading to poor performance. Furthermore, they use single code features, while actual code vulnerabilities may require assessment from multiple aspects. Additionally, they do not incorporate code comments as a code feature, making them ineffective in practical tasks. Existing vulnerability detection methods focus only on code structure or syntax, ignoring the semantic information in comments. Integrating comment information into the PDG (Programmable Logic Generation Group) could help the model better understand the code's intent, thereby improving detection performance. Code comments play a crucial role in helping developers understand code information; therefore, code comments can also be used as a feature for neural networks to learn. Before the widespread application of large models, adding code comments to every line of code was extremely difficult, but now large models can add comments to each statement accurately, which can be used as a way to represent code, allowing neural networks to better understand code information. Summary of the Invention

[0009] The purpose of this invention is to address the shortcomings of existing technologies and provide a function-level vulnerability detection method based on source code comments and program dependency graphs. This method aims to improve the vulnerability detection rate by integrating code comments and code statement features, using program dependency graphs as a code representation method, and employing graph convolutional neural networks to learn the structural and syntactic semantic information of the code. At the same time, this method has good scalability.

[0010] This invention is achieved through the following technical solution: a function-level vulnerability detection method based on source code comments and program dependency graphs, comprising the following specific steps: 1) Remove non-standard comments from the code and standardize the code in the following ways: map user-defined variables to symbol names (e.g., VAR1) in a one-to-one manner, and map user-defined functions to symbol names (e.g., FUN1) in a one-to-one manner. 2) After step 1), a program dependence graph PDG of the normalized code is generated; each node of the program dependence graph PDG is a complete code statement, and control flow information and data flow information existing between each node exist in the form of edges; each program dependence graph PDG has a file suffix of.dot, and the dot file is a graph representation file; 3) After step 2), after generating the program dependence graph PDG of the code (the code in each file is divided at the function level, because only the code at the function level can generate the program dependence graph, and hereinafter the code is uniformly referred to as code), a large model is used to generate annotation information for each node to obtain a program dependence graph (code-explian PDG) containing code statements and annotation information of the code statements; each node of the program dependence graph (code-explian PDG) containing code statements and annotation information of the code statements contains two attributes: one is a code statement, and the other is annotation information of the code statement. The original program dependence graph has each node being a complete code statement, and control flow information and data flow information existing between each node exist in the form of edges; after using the large model to generate the annotation information of the code statement, the content of each node becomes code statement + code statement annotation information, and the edge information remains unchanged; 4) After step 3), the code annotation information part of each node is segmented using the tokenizer of the BERT model, and then the code annotation information of each node is converted into a vector representation with rich natural language semantic information using the BERT model. Since the segmentation operation converts the original text into the input of the BERT model, the BERT model uses natural language in the pre-training data, and performs better on natural language tasks, so the BERT is selected for the embedding of the code annotation; 5) After step 3), the code statement part of each node is segmented, and then the code statement part of each node is converted into a vector representation with rich code language semantic information using the CodeBERT pre-training model. Since the segmentation operation converts the original text into the input of the CodeBERT model, the CodeBERT model uses code and natural language in the pre-training data, and performs better on tasks such as code generation and vulnerability judgment, so the CodeBERT is selected for the embedding of the code; 6) The vector representation of the generated code statement part and the vector representation of the code statement annotation information are fused into a fusion feature vector, which is used for subsequent input into the graph neural network overall architecture for model training. The vector representation of the code statement and the code annotation information of each node are spliced to obtain a feature fusion vector, and the node information of each program dependence graph (code-explian PDG) containing code statements and annotation information of the code statements becomes a feature vector, and the edge information does not change; 7) Repeat steps 1)~6) for each source code file to obtain a vectorized program dependence graph containing code and code comments, and train a vulnerability detection model based on a graph attention neural network using all the vectorized program dependence graphs containing code and code comments as training data sets. The vulnerability detection model based on the graph attention neural network is a specific implementation of the overall architecture of the graph neural network.

[0011] Each source code file contains a code function. For example, I have an add.c source code file that contains a function add(int a, int b) {......} for calculating the sum of a and b. Then I generate a program dependence graph for the add() function and generate comments.

[0012] To better implement the function-level vulnerability detection method based on source code comments and program dependence graphs, the following settings are particularly used: the step 1) includes the following specific steps: 1.1) Remove all comment information from the code; 1.2) Renaming user-defined variables and functions according to a unified symbol naming rule: uniformly standardizing user-defined function names to "FUN_i", where i is the order of the corresponding function name, and i∈(1,...,n); uniformly standardizing user-defined variable names to "VAR_j", where j is the order of the corresponding variable name, j∈(1,...,n).

[0013] To better implement the function-level vulnerability detection method based on source code comments and program dependence graphs, the following settings are particularly used: the step 2) is specifically: using the code analysis tool Joern to generate a program dependence graph PDG for each function in each standardized code file, the content of each PDG node is the code statement of the source function, and each PDG edge represents the relationship between nodes, which is saved as a dot graph file.

[0014] To better implement the function-level vulnerability detection method based on source code comments and program dependence graphs, the following settings are particularly used: the step 3) includes the following specific steps: 3.1) Transfer the code statements of each node of the program dependence graph PDG to the large model through the API of the large model, and let the large model generate comment information for each node by setting the prompt word; 3.2) Accept the comment information of the large model generated code, find the corresponding source code statement node, save the comment information to the corresponding code statement node, after saving, each node contains two contents, one is the source code statement, and one is the code comment information, the information of the original PDG edge remains unchanged, thereby obtaining a program dependence graph containing code statements and code statement annotation information, and saving in the form of a dot file.

[0015] Further to better implement the function level vulnerability detection method based on source code annotation and program dependence graph, the following setting mode is particularly adopted: the step 4) comprises the following specific steps: 4.1) The tokenizer of the BERT model first checks whether the complete word is in the vocabulary table, if the complete word exists, it is kept as a single token; if the complete word does not exist, it is split into subwords, and then each Token is converted into a corresponding vocabulary table ID to obtain input_ids, then the tokenizer automatically generates a token_type_ids vector and an attention_mask vector; 4.2) input_ids, token_type_ids vector, and attention_mask vector are input into the pre-trained BERT model, and finally a sentence embedding vector is obtained; wherein the pre-trained BERT model is obtained by pre-training a BERT model with twelve Transformer encoder layers. Assuming that the input sentence is "I love NLP" (after tokenization, ["[CLS]", "i", "love", "nlp", "[SEP]"]), the shape of the output of the last layer of the model is a [1, 5, 768] tensor, 1 represents the batch size, that is, the number of sentences processed at a time, 5 represents the number of tokens, and the length of each token is a 768-dimensional vector.

[0016] 4.3) The [CLS] vector of the last layer of the sentence embedding vector is selected as the vector representation of the code statement annotation information, and the dimension is 768.

[0017] The output vector form of the BERT model is [CLS] + the vector representation of each word. The [CLS] vector is the information fused from all the words in the entire sentence, covering the semantics of the entire sentence, so the [CLS] vector is selected as the vector representation of the sentence.

[0018] Further to better implement the function level vulnerability detection method based on source code annotation and program dependence graph, the following setting mode is particularly adopted: the step 5) comprises the following specific steps: 5.1) Tokenization: Tokenization is performed on each node's code statement using the tokenizer of CodeBERT: The tokenizer first checks whether the entire word (such as IdentifierDeclStatement) is in the pre-trained vocabulary, and if it exists, it is directly reserved as a single token (such as "for" which is reserved in its entirety), and if it does not exist (such as long variable name IdentifierDeclStatement), the word is decomposed into subwords or characters that exist in the vocabulary, for example "IdentifierDeclStatement" -> ["Identifier", "Decl","Statement"], and then each Token is converted into the corresponding vocabulary ID to obtain input_ids, and then the tokenizer automatically generates a token_type_ids vector and an attention_mask vector; wherein the input_ids word corresponds to the vocabulary ID of the CodeBERT model input. token_type_ids marks whether it belongs to the first sentence or the second sentence. attention_mask indicates which positions are real tokens and which are padding (filled with 0).

[0019] 5.2) input_ids, token_type_ids vector, and attention_mask vector are input into the pre-trained CodeBERT model, and finally the sentence embedding vector is obtained; wherein the pre-trained CodeBERT model is obtained by pre-training the CodeBERT model with twelve Transformer encoder layers; 5.3) The [CLS] vector of the last layer of the sentence embedding vector is selected as the vector representation of the code statement annotation information, which has a dimension of 768.

[0020] The output vector of CodeBERT is in the form of [CLS] + the vector representation of each word. The [CLS] vector is the information of all words in the entire sentence, which covers the semantics of the entire sentence, and it is a vector representation with rich natural language semantic information, so the [CLS] vector is selected as the vector representation of the code statement.

[0021] Further, in order to better implement the function-level vulnerability detection method based on source code annotation and program dependency graph according to the present application, the following setting mode is particularly adopted: the step 6) comprises the following specific steps: 6.1) The vector representation of the code statement of each node in the generated program dependence graph containing code and code comment and the vector representation of the comment information of the code statement are connected by the library function cat of Pytorch to become a fusion feature vector that fuses the code statement feature and the code statement comment information feature, and the dimension is 1536 dimensions; assuming that the a vector is [1, 2, 3], the b vector is [4, 5, 6], and the new vector c connected by cat is [1, 2, 3, 4, 5, 6].

[0022] 6.2) After obtaining the fusion feature vector of each node (because each graph has many nodes, and each node contains code + code comment, therefore the code and code comment of each node of the graph are spliced after being generated into vectors by BERT and CodeBERT. Therefore, the vector after splicing contains the code and code explanation, and each node has a fusion vector), the feature vector is used as the attribute of each node to replace the node of the original program dependence graph (code-explian PDG) containing the code statement and the comment information of the code statement, and then saved as a dot file.

[0023] The generated graph file can be directly input into a neural network for training, because each node is a vector representation, and the graph neural network can automatically read the node information and edge information of the graph and automatically process.

[0024] Further, in order to better realize the function level vulnerability detection method based on source code annotation and program dependence graph, the following setting mode is particularly adopted: 7.1) Repeat steps 1) to 6) for each source code file to obtain a training data set; 7.2) Use the dataset class of the deep learning framework pytorch to pack the training data set into a file and save it as a pth file; the pth file contains all the vectorized program dependence graphs, and the model can be directly obtained from the file during model calling, without reading single files one by one, which accelerates the training speed and is a function provided by the pytorch framework for users to use; 7.3) Input the packed training data set (pth file) into the designed graph attention neural network for training to obtain a vulnerability detection model based on the graph attention neural network.

[0025] Compared with the prior art, the present application has the following advantages and beneficial effects: The present application uses a program dependence graph that fuses code information and code comment information as a code intermediate representation, the nodes of which fuse code statements and code comments with natural language attributes, the edges of which represent control flow and data flow relationships between code statements, and simultaneously uses a graph attention neural network with an efficient neighborhood information aggregation mechanism to automatically learn the semantics of code elements and their topological relationships. Therefore, the present application effectively improves the accuracy of source code vulnerability detection.

[0026] The present application uses a code-explian PDG that contains code statements and comment information of code statements as an intermediate representation of code.

[0027] The present application improves the vulnerability detection rate while fusing source code semantics, syntax and comment information of source code, making the complement of natural language and code semantics enhanced, and achieving more efficient prediction on the basis of multi-feature fusion.

[0028] The present application realizes a hierarchical graph neural network based on an attention mechanism, automatically identifies vulnerability code patterns through the attention mechanism, and simultaneously adopts a cascaded architecture of "graph attention convolution-dynamic pooling" to enable the model to more accurately classify vulnerabilities. BRIEF DESCRIPTION OF DRAWINGS

[0029] Figure 1 The graph attention neural network designed for the present application.

[0030] Figure 2 The model training flowchart of the present application for function-level vulnerability detection based on source code comments and program dependence graphs.

[0031] Figure 3 The function-level vulnerability detection flowchart of the present application based on source code comments and program dependence graphs. DETAILED DESCRIPTION

[0032] The present application will be further described in detail below in conjunction with embodiments, but the embodiments of the present application are not limited thereto.

[0033] In order to make the purposes, technical solutions and advantages of the embodiments of the present application clearer, the technical solutions in the embodiments of the present application will be described clearly and completely below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are only some of the embodiments of the present application but not all the embodiments. Based on the embodiments in the present application, all other embodiments obtained by a person of ordinary skill in the art without creative work fall within the protection scope of the present application. Therefore, the following detailed description of the embodiments of the present application provided in the drawings is not intended to limit the scope of the claimed present application, but only represents selected embodiments of the present application. Based on the embodiments in the present application, all other embodiments obtained by a person of ordinary skill in the art without creative work fall within the protection scope of the present application.

[0034] In addition, the terms "first", "second" are only for descriptive purposes and cannot be understood as indicating or implying relative importance or implicitly indicating the number of the indicated technical features. Therefore, the features defined with "first", "second" can explicitly or implicitly include one or more of the features. In the description of the present application, the meaning of "multiple" is two or more, unless otherwise explicitly and specifically limited.

[0035] Noun explanation: Joern: a specific open source source code analysis tool.

[0036] PDG: Program Dependence Graph, a graph structure representing data and control dependence relationships between statements in a program, each node is a statement, and the edge represents data dependence and control dependence relationship.

[0037] BERT: a pre-training model for natural language.

[0038] CodeBERT: a pre-training model for mixed programming language and natural language scenarios, which can be regarded as an extended version of BERT in the field of software engineering.

[0039] Pytorch: an open source deep learning framework.

[0040] API: Application Programming Interface, a set of pre-defined functions, the purpose is to provide the ability of application program to access a set of routines based on certain software or hardware without accessing the source code or understanding the details of the internal working mechanism.

[0041] dot file: a text file format used to represent graph structure, which belongs to the input format of Graphviz software, and is widely used to draw and analyze graph structure.

[0042] pth file: save any PyTorch-supported data structure.

[0043] ELU (Exponential Linear Unit): an activation function used in deep learning models.

[0044] TopK pooling: a pooling method for graph neural networks that reduces the size of a graph while preserving key information by retaining the most important K nodes (or features) in the graph.

[0045] Example 1: The present application designs a function-level vulnerability detection method based on source code annotations and program dependency graphs, which aims to fuse code annotations and code statement features, use program dependency graphs as code representation methods, use graph convolutional neural networks to learn the structural information and syntax semantic information of the code, effectively improve the vulnerability detection rate, and at the same time, this method has good scalability, including the following specific steps: 1) Code normalization: remove non-standard annotation information in the code, and normalize the code in the following way: map user-defined variables to symbolic names (e.g. VAR1) in a one-to-one manner, and map user-defined functions to symbolic names (e.g. FUN1) in a one-to-one manner; 2) After step 1), generate the program dependency graph PDG of the normalized code; each node of the program dependency graph PDG is a complete code statement, and the control flow information and data flow information between each node exist in the form of edges, each program dependency graph PDG has a file suffix of.dot, and the dot file is a graph representation file; 3) After generating the program dependency graph PDG of the code (the code in each file is divided by function level, because only function-level code can generate a program dependency graph, hereinafter referred to as code), use a large model to generate annotation information for each node to obtain a program dependency graph containing code statements and code statement annotation information (code-explian PDG); Each node of the program dependency graph containing code statements and code statement annotation information (code-explian PDG) contains two attributes: one is the code statement, and the other is the code statement annotation information. The original program dependency graph has complete code statements at each node, and the control flow information and data flow information between each node exist in the form of edges. After using a large model to generate code statement annotation information, the content of each node becomes code statement + code statement annotation information, and the edge information remains unchanged; 4) After step 3), the code annotation information part of each node is segmented using the tokenizer of the BERT model, and then the code annotation information of each node is converted into a vector representation with rich natural language semantic information using the BERT model. Since the segmentation operation will convert the original text into the input of the BERT model, the natural language used by the BERT model pre-training data performs better on natural language tasks, so BERT is selected for the embedding of code annotations; 5) After step 3), the code statement part of each node is segmented, and then the code statement part of each node is converted into a vector representation with rich code language semantic information using the CodeBERT pre-training model. Since the segmentation operation will convert the original text into the input of the CodeBERT model, the code and natural language used by the CodeBERT model pre-training data perform better on code generation, vulnerability judgment and other tasks, so CodeBERT is selected for the embedding of code; 6) The vector representation of the generated code statement part and the vector representation of the code statement annotation information are fused into a fusion feature vector, which is used for subsequent input into the graph neural network overall architecture for model training. The vector representation of the code annotation information and the code statement generated by each node are spliced to obtain a feature fusion vector. At this time, the node information of each program dependence graph (code-explian PDG) containing code statements and code statement annotation information becomes a feature vector, and the information of the edge does not change; 7) Repeat steps 1) to 6) for each source code file to obtain a vectorized program dependence graph containing code and code annotations. All vectorized program dependence graphs containing code and code annotations are used as a training data set to train and obtain a vulnerability detection model based on a graph attention neural network. The vulnerability detection model based on the graph attention neural network is a specific implementation scheme of the graph neural network overall architecture.

[0046] Each source code file contains a code function. For example, I have an add.c source code file that contains a function add(int a, int b) {......} for calculating the sum of a+b. Then I generate a program dependence graph for the add() function and generate annotations.

[0047] Embodiment 2 This embodiment is further optimized on the basis of the above embodiment, and the same as the foregoing technical solutions will not be repeated here. To better implement the function-level vulnerability detection method based on source code annotations and program dependence graphs, the following setting method is particularly adopted: 1.1) Remove all annotation information from the code; 1.2) Renaming of user-defined variables and functions according to a unified symbol naming rule: the user-defined function name is uniformly standardized as "FUN_i", where i is the order of the corresponding function name, and i e (1,...,n); the user-defined variable name is uniformly standardized as "VAR_j", where j is the order of the corresponding variable name, j e (1,...,n).

[0048] Embodiment 3 This embodiment is further optimized on the basis of any of the above embodiments, and the same as the foregoing technical solutions will not be repeated here. Further, in order to better implement the function-level vulnerability detection method based on source code annotation and program dependence graph, the following setting method is adopted: the step 2) is specifically: using the code analysis tool Joern to generate a program dependence graph PDG for each function in each code file that has been standardized, the content of each PDG node is the code statement of the source function, and each PDG edge represents the relationship between the nodes. Save it as a dot graph file.

[0049] Embodiment 4 This embodiment is further optimized on the basis of any of the above embodiments, and the same as the foregoing technical solutions will not be repeated here. Further, in order to better implement the function-level vulnerability detection method based on source code annotation and program dependence graph, the following setting method is adopted: the step 3) includes the following specific steps: 3.1) The code statement of each node of the program dependence graph PDG is transmitted to the large model through the API of the large model, and the large model is prompted to generate annotation information for each node by setting the prompt word; 3.2) Accept the annotation information generated by the large model, find the corresponding source code statement node, save the annotation information to the corresponding code statement node, and after saving, each node contains two contents, one is the source code statement, and the other is the code annotation information. The information of the original PDG edge remains unchanged, thereby obtaining a program dependence graph containing code statements and code statement annotation information, and saving it in the form of a dot file.

[0050] Embodiment 5 This embodiment is further optimized on the basis of any of the above embodiments, and the same as the foregoing technical solutions will not be repeated here. Further, in order to better implement the function-level vulnerability detection method based on source code annotation and program dependence graph, the following setting method is adopted: the step 4) includes the following specific steps: 4.1) The tokenizer of the BERT model first checks whether the complete word is in the vocabulary, and if the complete word exists, it is kept as a single token; if the complete word does not exist, it is split into subwords, and then each Token is converted into the corresponding vocabulary ID to obtain input_ids, and then the tokenizer automatically generates a token_type_ids vector and an attention_mask vector; 4.2) input_ids, token_type_ids vector, and attention_mask vector are input into the pre-trained BERT model, and finally a sentence embedding vector is obtained; wherein the pre-trained BERT model is pre-trained by using a BERT model with twelve Transformer encoder layers. Assuming that the input sentence is "I love NLP" (after tokenization, it is ["[CLS]", "i", "love", "nlp", "[SEP]"]), the shape of the output of the last layer of the model is a tensor of [1, 5, 768], 1 represents the batch size, that is, the number of sentences processed at a time, 5 represents the number of tokens, and the length of each token is a 768-dimensional vector.

[0051] 4.3) The [CLS] vector of the last layer of the sentence embedding vector is selected as the vector representation of the code statement annotation information, and the dimension is 768.

[0052] The output vector of the BERT model is [CLS] + the vector representation of each word. The [CLS] vector is the information of all the words in the entire sentence, which covers the semantics of the entire sentence, so the [CLS] vector is selected as the vector representation of the sentence.

[0053] Embodiment 6: This embodiment is further optimized on the basis of any of the above embodiments, and the same as the foregoing technical solutions will not be repeated here. In order to better realize the function of the function level vulnerability detection method based on source code annotation and program dependency graph, the following setting method is particularly adopted: the step 5) comprises the following specific steps: 5.1) Tokenization operation is performed on the code statements of each node using the tokenizer of CodeBERT: the tokenizer first checks whether the entire word (such as IdentifierDeclStatement) is in the pre-trained vocabulary, if it exists, it is directly reserved as a single token (such as "for" will be reserved in its entirety), if it does not exist (such as long variable name IdentifierDeclStatement), the word will be decomposed into subwords or characters that exist in the vocabulary, for example "IdentifierDeclStatement" -> ["Identifier", "Decl","Statement"], then each Token is converted into the corresponding vocabulary ID to obtain input_ids, then the tokenizer will automatically generate a token_type_ids vector and an attention_mask vector; wherein the input_ids word corresponds to the vocabulary ID of the CodeBERT model input. token_type_ids marks whether it belongs to the first sentence or the second sentence. attention_mask indicates which positions are real tokens and which are padding (filled with 0).

[0054] 5.2) input_ids, token_type_ids vector, and attention_mask vector are input into the pre-trained CodeBERT model, and finally the sentence embedding vector is obtained; wherein the pre-trained CodeBERT model is obtained by pre-training the CodeBERT model with twelve Transformer encoder layers; 5.3) The [CLS] vector of the last layer of the sentence embedding vector is selected as the vector representation of the code statement annotation information, which has a dimension of 768.

[0055] The output vector of CodeBERT is in the form of [CLS] + the vector representation of each word. The [CLS] vector is the information of all words in the entire sentence, which covers the semantics of the entire sentence, and it is a vector representation with rich natural language semantic information, so the [CLS] vector is selected as the vector representation of the code statement.

[0056] Embodiment 7: This embodiment is further optimized on the basis of any of the above embodiments, and the same as the foregoing technical solutions will not be repeated here. In order to better realize the function of the function level vulnerability detection method based on source code annotation and program dependency graph, the following setting method is particularly adopted: 6.1) The vector representation of the code statement of each node in the generated program dependency graph containing code and code comments and the vector representation of the annotation information of the code statement are connected by the library function cat of Pytorch to become a fusion feature vector that fuses the code statement features and the code statement annotation information features, and the dimension is 1536 dimensions; assuming that the a vector is [1, 2, 3], the b vector is [4, 5, 6], and the new vector c connected by cat is [1, 2, 3, 4, 5, 6].

[0057] 6.2) After obtaining the fusion feature vector of each node (because each graph has many nodes, and each node contains code + code annotation, therefore the code and code annotation of each node of the graph are spliced after being generated into vectors by BERT and CodeBERT. Therefore, the vector after splicing contains the code and code explanation, and each node has a fusion vector), the feature vector is used as the attribute of each node to replace the original code-explian PDG node containing the code statement and the annotation information of the code statement, and then saved as a dot file.

[0058] The generated graph file can be directly input into the neural network for training, because each node is a vector representation, and the graph neural network can automatically read the node information and edge information of the graph and automatically process it.

[0059] Embodiment 8: The embodiment is further optimized on the basis of any of the above embodiments, and the same as the foregoing technical solutions will not be repeated here. In order to better realize the function level vulnerability detection method based on source code annotation and program dependency graph, the following setting method is particularly adopted: 7.1) Repeat steps 1) to 6) for each source code file to obtain a training data set; 7.2) Use the dataset class of the deep learning framework pytorch to pack the training data set into a file and save it as a pth file; the pth file contains all the vectorized program dependency graphs, which can be directly obtained from the file when the model is called, without reading individual files one by one, which speeds up the training speed and is a function provided by the pytorch framework for users to use; 7.3) Input the packed training data set (pth file) into the designed graph attention neural network for training to obtain a vulnerability detection model based on the graph attention neural network.

[0060] Embodiment 9: A function-level vulnerability detection method based on source code annotations and program dependency graphs aims to effectively improve the vulnerability detection rate by fusing code annotations and code statement features, using program dependency graphs as the code representation method, and using graph convolutional neural networks to learn the structural information and syntax semantic information of the code. At the same time, this method has good scalability. Combined with Figure 1 、 Figure 2 、 Figure 3 As shown in the drawings, the method comprises the following steps: S1, code normalization processing: S1.1, remove all annotation information in the code; S1.2, user-defined variables and functions are renamed according to a unified symbol naming rule: user-defined function names are uniformly standardized to "FUN_i", where i is the order of the corresponding function name, and i∈(1,..., n); user-defined variable names are uniformly standardized to "VAR_j", where j is the order of the corresponding variable name, j∈(1,..., n).

[0061] S2, generating program dependency graph PDG: after step S1, use the code analysis tool Joern to generate a program dependency graph PDG for each function in each code file that has been normalized, the content of each PDG node is the code statement of the source function, and each PDG edge represents the relationship between nodes (such as whether there is a control flow dependency or data dependency between two nodes, if there is a node between two nodes, there is an edge), which is saved as a dot graph file, and the generated.dot file is output to a specified folder, which is used to store the original program dependency graph PDG. Each node of the program dependency graph PDG is a complete code statement, and the control flow information and data flow information between each node exist in the form of edges, and each program dependency graph PDG has a.dot file suffix, and the dot file is a graph representation file.

[0062] S3, generating node annotations for large models and constructing code-explian PDG: using a large model to generate annotation information for each node, obtaining a program dependency graph (code-explian PDG) containing code statements and annotation information of code statements, comprising the following specific steps: S3.1, after generating the program dependency graph PDG for each function code, traverse all nodes of the PDG, extract the code statement corresponding to the node, and transmit the code statement of each node of the program dependency graph PDG to the large model through the API of the large model, and through the set prompt word, let the large model generate annotation information for each node (i.e. each line of code statement); S3.2, accept the annotation information of the large model generated code, map the annotation returned by the large model with the PDG node, each node corresponds to an annotation information, then extend each node attribute, that is, find the source code statement node corresponding to each annotation, save the annotation information to the corresponding code statement node, and save it in the form of dot file, each node has two attributes, one is the source code statement, and the other is the code annotation information (the original node without generating annotation has only one attribute-source code statement), finally output the program dependence graph (code-explian PDG) containing code statement and code statement annotation information to the specified folder, this folder is used to save the program dependence graph code-explian PDG after generating annotation information.

[0063] Each node of the program dependence graph (code-explian PDG) containing code statement and code statement annotation information contains two attributes: one is the code statement, and the other is the annotation information of the code statement. The original program dependence graph PDG is a complete code statement, and the control flow information and data flow information existing between each node exist in the form of edge. After generating the annotation information of the code statement by using the large model, the content of each node becomes code statement+code statement annotation information, and the information of the edge remains unchanged.

[0064] Since the segmentation operation will convert the original text into the input of the BERT model, the BERT model uses natural language in the pre-training data, and performs better on natural language tasks, therefore, BERT is selected for code annotation embedding.

[0065] S4, code annotation branch, segmentation, BERT vectorization, natural language semantic vector: S4.1, the tokenizer of the BERT model first checks whether the complete word is in the vocabulary, if there is a complete word (such as "function", "return"), it is kept as a single token; if there is no complete word (such as long compound words or technical terms), it is split into subwords (for example, "initialization" → ["initial", "##ization"]), then each Token is converted into the corresponding vocabulary ID to obtain input_ids, then the tokenizer will automatically generate a token_type_ids vector and an attention_mask vector; wherein, the vocabulary ID corresponding to the input_ids word is the input of the BERT model. token_type_ids marks whether it belongs to the first sentence or the second sentence; attention_mask indicates which positions are real tokens and which are padding (padding 0 part).

[0066] S4.2, input input_ids, token_type_ids vector, attention_mask vector into the pre-trained BERT model, and finally get the sentence embedding vector; wherein the pre-trained BERT model is pre-trained by using the BERT model with twelve Transformer encoder layer version. Assuming that the input sentence is "I love NLP" (after tokenization, ["[CLS]", "i", "love", "nlp", "[SEP]"]), then the shape of the output of the last layer model is [1, 5, 768] tensor, 1 represents batch size, that is, the number of sentences sent into the model for processing at a time, 5 represents the number of tokens, and the length of each token is a 768-dimensional vector.

[0067] S4.3, select the [CLS] vector of the last layer of the sentence embedding vector as the vector representation of the code statement annotation information, which has a dimension of 768.

[0068] The output vector form of the BERT model is [CLS] + the vector representation of each word. Wherein the [CLS] vector is the information of all words in the whole sentence, which covers the semantics of the whole sentence, so the [CLS] vector is selected as the vector representation of the sentence.

[0069] Because the tokenization operation will convert the original text into the input of the CodeBERT model, the code and natural language used in the pre-training data of the CodeBERT model perform better in code generation, vulnerability judgment and other tasks, so CodeBERT is selected for code embedding.

[0070] S5, code statement branch, tokenization, CodeBERT vectorization, code semantic vector: S5.1, the tokenizer using CodeBERT first checks whether the entire word (such as IdentifierDeclStatement) is in the pre-trained vocabulary. If it exists, it is directly reserved as a single token (such as "for" will be reserved in its entirety). If it does not exist (such as long variable name IdentifierDeclStatement), the word will be decomposed into subwords or characters that exist in the vocabulary, for example "IdentifierDeclStatement" -> ["Identifier", "Decl", "Statement"], and then each Token will be converted into the corresponding vocabulary ID to obtain input_ids. Then the tokenizer will automatically generate a token_type_ids vector and an attention_mask vector; wherein the input_ids word corresponds to the input ID of the CodeBERT model. token_type_ids marks whether it belongs to the first sentence or the second sentence. attention_mask indicates which positions are real tokens and which are padding (filled with 0).

[0071] S5.2, input input_ids, token_type_ids vector, and attention_mask vector into the pre-trained CodeBERT model, and finally obtain a sentence embedding vector; wherein the pre-trained CodeBERT model is obtained by pre-training a CodeBERT model with twelve Transformer encoder layers. S5.3, select the [CLS] output vector of the last layer of the sentence embedding vector as the vector representation of the code statement annotation information, which has a dimension of 768. Wherein, [CLS] aggregates the semantic information of the entire sentence.

[0072] The output vector of CodeBERT is in the form of [CLS] + the vector representation of each word. The [CLS] vector is the information of all the words in the entire sentence, which covers the semantic information of the entire sentence. It is a vector representation with rich natural language semantic information, so the [CLS] vector is selected as the vector representation of the code statement.

[0073] S6, feature vector fusion: The vector of each node's code annotation information and code statement is concatenated to obtain a feature fusion vector. At this time, the node information of each program dependence graph (code-explian PDG) containing code statements and code statement annotation information becomes a feature vector, and the edge information does not change. Specifically: S6.1, The vector representation of the code statement of each node in the generated program dependency graph containing code and code comment and the vector representation of the annotation information of the code statement are connected by the library function cat of Pytorch to become a fusion feature vector that fuses the code statement features and the code statement annotation information features, and the dimension is 1536; assuming that the a vector is [1, 2, 3] and the b vector is [4, 5, 6], the new vector c connected by cat is [1, 2, 3, 4, 5, 6].

[0074] S6.2, After obtaining the fusion feature vector of each node (because each graph has many nodes, and each node contains code + code annotation, therefore the code and code annotation of each node of the graph are spliced after being generated into vectors by BERT and CodeBERT. Therefore, the vector after splicing contains the two parts of code and code explanation, and each node has a fusion vector), the feature vector is used as the attribute of each node to replace the original node of the program dependency graph (code-explian PDG) containing the code statement and the annotation information of the code statement, and then saved as a dot file, thereby obtaining the vectorized code-explian PDG.

[0075] The generated graph file can be directly input into the neural network for training or reasoning, because each node is represented by a vector, the graph neural network can automatically read the node information and edge information of the graph and automatically process them.

[0076] S7, Vulnerability monitoring model: S7.1, Repeat steps S1-S6 for each source code file to obtain the vectorized code-explian PDG of each source code file to be detected; S7.2, input each vectorized code-explian PDG into the trained hierarchical graph attention neural network for reasoning to obtain the detection result.

[0077] S8, output the detection result: after the code file to be detected is subjected to steps S1-S7, the final detection result is obtained, and it is output whether the code has a vulnerability.

[0078] As Figure 1As shown, the hierarchical graph attention network structure includes a hierarchical module 1, a hierarchical module 2, and a classifier module. First, the code-explain PDG containing code and code annotation information which has been vectorized is input into the graph attention layer (GAT) in the hierarchical module 1. The layer dynamically captures the semantic dependency relationship between nodes through the multi-head attention mechanism to weight and aggregate the adjacent information of the nodes, thereby enhancing the recognition ability of the key structure. Subsequently, the ELU activation function (exponential linear unit) introduces a nonlinear transformation, which not only improves the expression ability of the model, but also preserves the gradient through the negative input to prevent the information of the low-frequency or weakly connected nodes in the code-explain PDG from being ignored in the propagation. Next, the TopK pooling operation retains the nodes ranked in the top 80% according to the importance score (such as the attention weight) of the node features, effectively eliminates redundant information such as repeated code or supplementary annotation nodes, thereby compressing the graph structure and highlighting the core semantic area. After the first layer of abstraction, the remaining simplified graph continues to be transmitted to the hierarchical module 2, and a higher level of feature aggregation is performed in the second GAT layer to capture more complex semantic combination relationships in the program structure, and then the key nodes are further filtered through the ELU activation and TopK pooling to build a compact graph structure with stronger semantic representation ability. Finally, the graph representation after two layers of semantic extraction and structure compression is input into the classifier module. First, the global average pooling integrates the features of all remaining nodes in the graph to form a unified graph-level vector representation, and then the feature dimension reduction and regularization processing are realized through the fully connected layer, ReLU activation function and Dropout layer in sequence, and finally the classification result is generated by the output layer, and the final prediction probability is output by the LogSoftmax function. The structure gradually extracts the most critical semantic information in the graph through layer-by-layer abstraction, which is suitable for structural understanding and efficient classification of the graph.

[0079] The above is only a preferred embodiment of the present application, and does not limit the present application in any form. Any simple modification or equivalent change based on the technical essence of the present application to the above embodiment falls within the protection scope of the present application.

Claims

1. A function-level vulnerability detection method based on source code annotations and program dependency graphs, characterized in that: Comprise the following specific steps: 1) remove the code in the presence of non-standard comment information, and the code is standardized in the following way: in a one-to-one manner, the user-defined variables are mapped to the symbol name, and the user-defined functions are mapped to the symbol name in a one-to-one manner; 2) after step 1), generate the program dependence graph PDG of the code after standardization; 3) after step 2), use a large model to generate annotation information for each node, and obtain a program dependence graph containing code statements and code statement annotation information; 4) after step 3), use the tokenizer of the BERT model to tokenize the code annotation information part of each node, and then use the BERT model to convert the code annotation information of each node into a vector representation with rich natural language semantic information; 5) after step 3), tokenize the code statement part of each node, and then use the CodeBERT pre-training model to convert the code statement part of each node into a vector representation with rich code language semantic information; 6) fuse the generated code statement part vector representation and code statement annotation information vector representation into a fusion feature vector, which is used as input for subsequent graph neural network architecture model training; 7) repeat steps 1)~6) for each source code file to obtain a vectorized program dependence graph containing code and code annotations, and use all the vectorized program dependence graphs containing code and code annotations as training data sets to train and obtain a vulnerability detection model based on graph attention neural network. 2.The function-level vulnerability detection method based on source code annotation and program dependency graph according to claim 1, characterized in that: The step 1) comprises the following specific steps: 1.1) remove all the comment information in the code; 1.2) user-defined variables and functions are renamed according to the unified symbol naming rules: user-defined function names are uniformly standardized to "FUN_i", where i is the order of the corresponding function name, and i∈(1,...,n); user-defined variable names are uniformly standardized to "VAR_j", where j is the order of the corresponding variable name, j∈(1,...,n).

3. The function-level vulnerability detection method based on source code annotations and program dependency graphs according to claim 1, characterized in that: The step 2) is specifically: using the code analysis tool Joern to generate a program dependence graph PDG for each function in each standardized code file, and the content of each PDG node is the code statement of the source function, and each PDG edge represents the relationship between nodes, which is saved as a dot graph file.

4. The function-level vulnerability detection method based on source code annotations and program dependence graphs according to claim 3, characterized in that: The step 3) comprises the following specific steps: 3.1) transmit the code statement of each node of the program dependence graph PDG to the large model through the API of the large model, and let the large model generate annotation information for each node by setting the prompt word; 3.2) accept the code annotation information generated by the large model, and find the corresponding source code statement node, save the annotation information to the corresponding code statement node, after saving, each node contains two contents, one is the source code statement, and the other is the code annotation information, the information of the original PDG edge remains unchanged, thereby obtaining a program dependence graph containing code statements and code statement annotation information, and saving it in the form of a dot file.

5. The function-level vulnerability detection method based on source code annotations and program dependence graphs according to claim 1 or 2 or 3 or 4, characterized in that: The step 4) comprises the following specific steps: 4.1) The tokenizer of the BERT model first checks whether the complete word is in the vocabulary, and if the complete word exists, it is kept as a single token; if the complete word does not exist, it is split into subwords, and then each Token is converted into a corresponding vocabulary ID to obtain input_ids, and then the tokenizer automatically generates a token_type_ids vector and an attention_mask vector; 4.2) input_ids, token_type_ids vector, and attention_mask vector are input into the pre-trained BERT model, and finally a sentence embedding vector is obtained; wherein the pre-trained BERT model is pre-trained by using a BERT model with twelve Transformer encoder layers; 4.3) the [CLS] vector of the last layer of the sentence embedding vector is selected as the vector representation of the annotation information of the code statement, and the dimension is 768.

6. The function-level vulnerability detection method based on source code annotations and program dependence graphs according to claim 1, characterized in that: The step 5) comprises the following specific steps: 5.1) The tokenizer of the CodeBERT model first checks whether the entire word is in the pre-trained vocabulary, and if the word exists, it is directly kept as a single token, and if the word does not exist, the word is decomposed into subwords or characters that exist in the vocabulary, and then each Token is converted into a corresponding vocabulary ID to obtain input_ids, and then the tokenizer automatically generates a token_type_ids vector and an attention_mask vector; 5.2) input_ids, token_type_ids vector, and attention_mask vector are input into the pre-trained CodeBERT model, and finally a sentence embedding vector is obtained; wherein the pre-trained CodeBERT model is pre-trained by using a CodeBERT model with twelve Transformer encoder layers; 5.3) the [CLS] vector of the last layer of the sentence embedding vector is selected as the vector representation of the annotation information of the code statement, and the dimension is 768.

7. The function-level vulnerability detection method based on source code annotations and program dependence graphs according to claim 1, characterized in that: The step 6) comprises the following specific steps: 6.1) The vector representation of the code statement and the vector representation of the annotation information of the code statement of each node in the generated program dependence graph containing code and code annotation are connected by the library function cat of Pytorch to become a fusion feature vector that fuses the code statement feature and the code statement annotation information feature, and the dimension is 1536; 6.2) After obtaining the fusion feature vector of each node, the feature vector is used as the attribute of each node to replace the original node of the program dependence graph containing the code statement and the annotation information of the code statement, and then saved as a dot file.

8. The function-level vulnerability detection method based on source code annotations and program dependence graphs according to claim 1, characterized in that: The step 7) comprises the following specific steps: 7.1) Steps 1) to 6) are repeated for each source code file to obtain a training data set; 7.2) Pack the training dataset into a file using the dataset class of the deep learning framework pytorch, and save it as a pth file; 7.3) Input the packed training dataset into the designed graph attention neural network for training, and obtain the vulnerability detection model based on the graph attention neural network.

Images