A source code vulnerability detection method and system based on contrastive learning

Abstract

The present invention discloses a source code vulnerability detection method based on contrastive learning. The method generates required sample pairs by customizing semantic-preserving heuristics, represents vulnerability codes from the perspective of code tokens to statements to graphs, and uses code attribute graphs to better represent the deep semantic information of vulnerabilities, fully explore the relationship between vulnerability codes and contexts, adopts contrastive learning technology, performs training on small samples, and realizes the high-precision requirements of small sample data. The present invention correspondingly provides a source code vulnerability detection system based on contrastive learning.

Description

Technical Field

[0001] The present invention relates to the field of software security, and in particular to a source code vulnerability detection method and system based on contrastive learning. Background Art

[0002] Vulnerability detection is an important part of the software maintenance process. Existing deep learning models rely on a large number of labeled vulnerability samples, which are difficult to obtain in the actual development process. Manual labeling will cause a lot of resource waste and the correctness needs to be studied. The sample size of the existing labeled dataset is not large enough to train a detection model with a high detection rate. Therefore, the detection model trained based on the existing labeled dataset method has low detection accuracy and is difficult to use for vulnerability detection tasks in actual environments.

[0003] At present, some works have used contrastive learning methods to detect software vulnerabilities and their types. For example, the paper "CONTRASTIVE LEARNING FOR SOURCE CODE WITH STRUCTURAL AND FUNCTIONAL PROPERTIES" constructs vulnerability samples by manually injecting vulnerability codes. This method still relies on prior expert experience and still requires manual work. Some works also use deep learning methods to detect software vulnerabilities, such as the paper "Deep Learning based Vulnerability Detection: Are We There Yet?", which does not fully consider the rich semantic information in the source code and is difficult to cover different types of vulnerability semantic features. Therefore, it has certain limitations in capturing deep semantic information of the code. Summary of the invention

[0004] Purpose of the invention: The purpose of the present invention is to provide a source code vulnerability detection method and system based on contrastive learning, so as to achieve the high-precision requirements of small sample data, and at the same time automatically generate the required sample pairs through customized semantic preservation heuristics, utilize the deep semantic information of the vulnerability, and fully explore the relationship between the vulnerability code and the context.

[0005] Technical solution: The present invention provides a source code vulnerability detection method based on contrastive learning, comprising the following steps:

[0006] 1) Collect vulnerability data from the vulnerability database NVD and the open source code library GitHub to build a vulnerability dataset, and perform data enhancement on the vulnerability dataset to form positive and negative sample pairs;

[0007] 2) Perform code feature representation learning on positive and negative sample pairs, and convert the positive and negative sample pairs into code attribute graphs. The control flow graph CFG and program dependency graph PDG contained in the code attribute graph together form a joint graph;

[0008] 3) Embed the joint graph and apply contrastive learning to obtain the feature joint graph vector of the positive and negative sample pairs. At the same time, add a classifier and use the feature joint graph vector of the positive and negative sample pairs as the input of the classifier to obtain the predicted label, so as to build a vulnerability detection model;

[0009] 4) Extract vulnerable code from the vulnerability database NVD and the open source code library GitHub as the code to be tested, and use the vulnerability detection model to analyze whether the code to be tested contains vulnerabilities.

[0010] Furthermore, in step 1), a vulnerability dataset is constructed and data enhancement is performed, which includes the following steps:

[0011] 1.1) Collect vulnerability data from the vulnerability database NVD and the open source code library GitHub, extract vulnerability files and corresponding vulnerability patch files;

[0012] 1.2) Preprocess the extracted vulnerability files and corresponding vulnerability patch files, extract the vulnerability code in the vulnerability files and the patch code in the corresponding vulnerability patch files from the function level, remove redundant information, including header files, comments, and declared global parameters, and obtain the vulnerability dataset;

[0013] 1.3) Select vulnerability codes from the vulnerability dataset as hard negative samples, and the corresponding patch codes as original samples. Apply semantic preservation heuristics to the original samples to change the grammatical structure of the original samples while keeping the semantics unchanged, thereby generating positive samples. The corresponding positive samples, original samples, and hard negative samples are combined as positive and negative sample pairs.

[0014] Furthermore, step 2) specifically includes the following steps:

[0015] 2.1) Perform code representation on all samples in each positive and negative sample pair, and generate a code property graph CPG for the positive and negative sample pairs using the tool Joern, where the code property graph CPG includes an abstract syntax tree AST, a control flow graph CFG, and a program dependency graph PDG;

[0016] 2.2) Use byte pair encoding (BPE) to segment the leaf nodes of the abstract syntax tree AST, use Word2vec to embed the obtained code tokens, and use the tree-based convolutional neural network (TBCNN) to generate sentence embedding vectors according to the hierarchical structure in the abstract syntax tree AST;

[0017] 2.3) After the tree aggregation layer, perform maximum pooling to aggregate the vectors of all nodes of the tree into a single vector containing the semantics of the corresponding statement, and use this single vector as the initial feature vector of the corresponding statement node in the control flow graph CFG and the program dependency graph PDG;

[0018] 2.4) The control flow graph CFG and the program dependency graph PDG are combined into a joint graph, and the control flow graph CFG, the program dependency graph PDG and the vectors of the statements in the joint graph are input into the graph-level code representation layer.

[0019] Furthermore, step 3) specifically includes the following steps:

[0020] 3.1) Calculate the corresponding attention weight coefficient of each node in the joint graph according to the control flow, control dependency and data dependency edges in the control flow graph CFG and the program dependency graph PDG;

[0021] 3.2) The graph-level code representation layer consists of multiple graph attention layers. One graph attention layer corresponds to one edge type. According to the attention weight coefficient, the graph attention layer updates the node vector representation in the joint graph.

[0022] 3.3) Averaging multiple graph attention layers to form a multi-head attention layer to obtain the final vector representation of the nodes in the graph;

[0023] 3.4) According to the different importance of nodes in the joint graph, the feature joint graph vector of the positive and negative sample pairs is obtained;

[0024] 3.5) The original sample and the positive sample are regarded as similar pairs, and the original sample and the hard negative sample are regarded as dissimilar pairs;

[0025] 3.6) Use cosine similarity to minimize the loss;

[0026] 3.7) Add a classifier, use the feature joint graph vector of the positive and negative sample pairs as the input of the classifier, and use the activation function softmax to train the classifier to obtain the predicted label. If the predicted label is 1, it is a vulnerable code, and if the label is 0, it is a non-vulnerable code.

[0027] Furthermore, step 4) specifically includes the following steps:

[0028] 4.1) Extract the vulnerability code from the vulnerability database NVD and the open source code library GitHub as the code to be tested, and execute step 2) on the code file to be tested to generate the corresponding code attribute graph;

[0029] 4.2) The code attribute graph obtained through code feature representation learning is used as the input of the vulnerability detection model to obtain the predicted label of the code to be tested;

[0030] 4.3) Output the predicted label. If the predicted label is 1, it is a vulnerable code, and if the predicted label is 0, it is a non-vulnerable code.

[0031] The present invention correspondingly provides a source code vulnerability detection system based on contrastive learning, comprising a data set construction module, a feature learning module, a model construction module, and a vulnerability detection module;

[0032] The dataset construction module is used to collect vulnerability data from the vulnerability database NVD and the open source code library GitHub to build a vulnerability dataset, and perform data enhancement on the vulnerability dataset to form positive and negative sample pairs;

[0033] The feature learning module is used to learn the code feature representation of positive and negative sample pairs, and convert the positive and negative sample pairs into a code attribute graph. The control flow graph CFG and program dependency graph PDG contained in the code attribute graph together form a joint graph;

[0034] The model building module is used to embed the joint graph and apply contrastive learning to obtain the feature joint graph vector of the positive and negative sample pairs. At the same time, a classifier is added and the feature joint graph vector of the positive and negative sample pairs is used as the input of the classifier to obtain the predicted label, thereby building a vulnerability detection model.

[0035] The vulnerability detection module is used to extract vulnerability code from the vulnerability database NVD and the open source code library GitHub as the code to be tested, and use the vulnerability detection model to analyze whether the code to be tested contains vulnerabilities.

[0036] Furthermore, the dataset construction module also includes a collection unit, an extraction unit, and a generation unit;

[0037] The collection unit is used to collect vulnerability data from the vulnerability database NVD and the open source code library GitHub, and extract vulnerability files and corresponding vulnerability patch files;

[0038] The extraction unit is used to pre-process the extracted vulnerability files and corresponding vulnerability patch files, extract the vulnerability code in the vulnerability files and the patch code in the corresponding vulnerability patch files from the function level, remove redundant information, including header files, comments, and declared global parameters, and obtain a vulnerability data set;

[0039] The generation unit is used to select vulnerability codes from the vulnerability dataset as hard negative samples and the corresponding patch codes as original samples. The semantic preservation heuristic applied to the original samples changes the grammatical structure of the original samples while keeping the semantics unchanged, thereby generating positive samples. The corresponding positive samples, original samples and hard negative samples are combined as positive and negative sample pairs.

[0040] Furthermore, the feature learning module also includes a representation unit, a sentence embedding unit, a node aggregation unit, and a graph embedding preparation unit;

[0041] The representation unit is used to perform code representation on all samples in each positive and negative sample pair, and generate a code property graph CPG for the positive and negative sample pairs through the tool Joern, where the code property graph CPG includes an abstract syntax tree AST, a control flow graph CFG, and a program dependency graph PDG;

[0042] The sentence embedding unit is used to segment the leaf nodes of the abstract syntax tree AST using byte pair encoding BPE, embed the obtained code tokens using Word2vec, and use the tree-based convolutional neural network TBCNN to generate sentence embedding vectors according to the hierarchical structure in the abstract syntax tree AST;

[0043] The node aggregation unit is used to perform maximum pooling through the tree aggregation layer to aggregate the vectors of all nodes of the tree into a single vector containing the semantics of the corresponding statement, and use the single vector as the initial feature vector of the corresponding statement node in the control flow graph CFG and the program dependency graph PDG;

[0044] The graph embedding preparation unit is used to combine the control flow graph CFG and the program dependency graph PDG into a joint graph, and input the control flow graph CFG, the program dependency graph PDG and the vectors of the statements in the joint graph into the graph-level code representation layer.

[0045] Furthermore, the model building module also includes an attention weight calculation unit, a graph attention layer unit, a multi-head attention layer unit, a graph node aggregation unit, a comparison rule unit, a loss minimization unit, and a training classifier unit;

[0046] The attention weight calculation unit is used to calculate the corresponding attention weight coefficient of each node in the joint graph according to the three edges of control flow, control dependency and data dependency in the control flow graph CFG and the program dependency graph PDG;

[0047] The graph attention layer unit is used to represent the layer according to the graph-level code, which consists of multiple graph attention layers. One graph attention layer corresponds to one type of edge. According to the attention weight coefficient, the graph attention layer updates the node vector representation in the joint graph.

[0048] The multi-head attention layer unit is used to average multiple graph attention layers to form a multi-head attention layer to obtain the final vector representation of the nodes in the graph;

[0049] The graph node aggregation unit is used to obtain the feature joint graph vector of the positive and negative sample pairs according to the different importance of the nodes in the joint graph;

[0050] The contrast rule unit is used to treat the original sample and the positive sample as similar pairs, and the original sample and the hard negative sample as dissimilar pairs;

[0051] The loss minimization unit is used to minimize the loss using cosine similarity;

[0052] The training classifier unit is used to add a classifier. The feature joint graph vector of the positive and negative sample pairs is used as the input of the classifier. The classifier is trained with the activation function softmax to obtain the predicted label. If the predicted label is 1, it is a vulnerable code, and if the label is 0, it is a non-vulnerable code.

[0053] Furthermore, the vulnerability detection module also includes a graph generation unit, a detection unit, and an output unit;

[0054] The graph generation unit is used to extract vulnerability codes from the vulnerability database NVD and the open source code library GitHub as the code to be tested, and execute the feature learning module on the code file to be tested to generate the corresponding code attribute graph;

[0055] The detection unit is used to use the code attribute graph obtained through code feature representation learning as the input of the vulnerability detection model to obtain the predicted label of the code to be tested;

[0056] The output unit is used to output a predicted label. If the predicted label is 1, it is a vulnerable code, and if the predicted label is 0, it is a non-vulnerable code.

[0057] Beneficial effect: Compared with the prior art, the present invention has the remarkable feature of generating the required sample pairs through customized semantic-preserving heuristics, representing the vulnerability code from the perspective of code tokens to statements to graphs, and using code attribute graphs to better represent the deep semantic information of the vulnerability, fully mining the relationship between the vulnerability code and the context, and using contrastive learning technology to train on small samples to achieve high-precision requirements for small sample data. BRIEF DESCRIPTION OF THE DRAWINGS

[0058] Figure 1 It is a schematic diagram of the process of the present invention;

[0059] Figure 2 It is a schematic diagram of the architecture of the tree-based convolutional neural network in the present invention;

[0060] Figure 3 It is the multi-head graph attention layer in this invention. DETAILED DESCRIPTION

[0061] The present invention is further described in detail below in conjunction with the accompanying drawings and specific implementation methods.

[0062] Example 1

[0063] The present invention provides a source code vulnerability detection method based on contrastive learning, please refer to Figure 1 As shown, the following steps are included:

[0064] 1) Collect vulnerability data from the vulnerability database NVD and the open source code library GitHub to build a vulnerability dataset, and perform data enhancement on the vulnerability dataset to form positive and negative sample pairs;

[0065] 1.1) Collect vulnerability data from the vulnerability database NVD and the open source code library GitHub, extract vulnerability files and corresponding vulnerability patch files;

[0066] 1.2) Preprocess the extracted vulnerability files and corresponding vulnerability patch files, extract the vulnerability code in the vulnerability files and the patch code in the corresponding vulnerability patch files from the function level, remove redundant information, including header files, comments, and declared global parameters, and obtain the vulnerability dataset;

[0067] 1.3) Select the vulnerability code from the vulnerability dataset as the hard negative sample, and the corresponding patch code as the original sample. Apply the semantic preservation heuristic to the original sample to change the grammatical structure of the original sample while keeping the semantics unchanged, thereby generating a positive sample. The corresponding positive sample, original sample, and hard negative sample are combined as a positive-negative sample pair;

[0068] Among them, the semantic preservation heuristic rule is:

[0069] Rule 1: Introduce a new else branch, which introduces a new else branch statement into an existing if conditional statement, or reverses the condition in an if conditional statement;

[0070] Rule 2: Conditional expressions with AND, for if statements with AND containing conditions C1 and C2, are rewritten as nested if structures;

[0071] Rule 3: Introduce a guard outside the for loop and change the program containing the for loop statement to a new program structure by adding an if(True) statement outside the loop.

[0072] Rule 4: Introduce a guard in the outer layer of the while loop, and change the program containing the for loop statement into a new program structure by adding a while judgment statement in the outer layer of the loop;

[0073] Rule 5: Add a conditional break statement in the while loop. A while loop with a condition C1 can be changed into another infinite loop statement. A C1 conditional break statement needs to be added to the loop statement body.

[0074] 2) Perform code feature representation learning on positive and negative sample pairs, and convert the positive and negative sample pairs into code attribute graphs. The control flow graph CFG and program dependency graph PDG contained in the code attribute graph together form a joint graph;

[0075] 2.1) Perform code representation on all samples in each positive and negative sample pair, and generate a code property graph CPG for the positive and negative sample pairs using the tool Joern, where the code property graph CPG includes an abstract syntax tree AST, a control flow graph CFG, and a program dependency graph PDG;

[0076] 2.2) Please refer to Figure 2As shown, the leaf nodes of the abstract syntax tree AST are segmented using byte pair encoding BPE, the obtained code tokens are embedded using Word2vec, and the tree-based convolutional neural network TBCNN is used to generate sentence embedding vectors according to the hierarchical structure in the abstract syntax tree AST;

[0077] Among them, the execution process of the tree-based convolutional neural network TBCNN is as follows:

[0078] By sliding a triangular convolution kernel on the tree, the parent node representation is updated according to the child node representation. The update calculation formula is:

[0079]

[0080] In the formula, is the updated vector representation of the father node, H is the node set containing the father node itself in the triangular convolution kernel, is the vector representation of the nodes in the H set, σ represents the activation function, and W q and b represent the weight matrix and bias term respectively; the main components of TBCNN include: vector representation and encoding, tree-based convolution, dynamic pooling, a fully connected hidden layer and an output layer;

[0081] 2.3) After the tree aggregation layer, perform maximum pooling to aggregate the vectors of all nodes of the tree into a single vector containing the corresponding statement semantics, and use this single vector as the initial feature vector of the corresponding statement node in the control flow graph CFG and the program dependency graph PDG:

[0082]

[0083] In the formula, h G is the vector representation set of nodes in CFG and PDG after the tree aggregation layer, TE() represents the tree embedding process, V G It is a node set of CFG, PDG, AST n is the AST whose root node is n;

[0084] 2.4) The control flow graph CFG and the program dependency graph PDG are combined into a joint graph, and the control flow graph CFG, the program dependency graph PDG and the vectors of the statements in the joint graph are input into the graph-level code representation layer.

[0085] 3) Embed the joint graph and apply contrastive learning to obtain the feature joint graph vector of the positive and negative sample pairs. At the same time, add a classifier and use the feature joint graph vector of the positive and negative sample pairs as the input of the classifier to obtain the predicted label, so as to build a vulnerability detection model;

[0086] 3.1) According to the control flow, control dependency and data dependency edges in the control flow graph CFG and the program dependency graph PDG, the corresponding attention weight coefficients of each node in the joint graph are calculated:

[0087]

[0088] Where N e is the set of adjacent nodes of node i under edge e (including node i itself), For node i and node j (j∈N e ), || represents the concatenation operation, exp() is an exponential function with the natural constant e as the base, is the transpose of a learnable weight vector, It can be considered as the correlation between node i and node j of edge e. The larger the value, the stronger the correlation between the two nodes, and the larger the weight coefficient;

[0089] 3.2) The graph-level code representation layer consists of multiple graph attention layers. One graph attention layer corresponds to one edge type. According to the attention weight coefficient obtained in step 3.1), the graph attention layer updates the node vector representation in the joint graph:

[0090]

[0091] In the formula, is the vector representation of the node, σ is the activation function, W e is the weight coefficient of edge e;

[0092] 3.3) Please refer to Figure 3 As shown in Figure 1, multiple graph attention layers are averaged to form a multi-head attention layer, and the final vector representation of the nodes in the graph is obtained:

[0093]

[0094] The multi-head attention mechanism concatenates or averages the features of each head to obtain a new feature representation;

[0095] 3.4) According to the different importance of nodes in the joint graph, the feature joint graph vector of the positive and negative sample pairs is obtained:

[0096]

[0097] Where AGG is the aggregation function, V G is the node set of the graph structure, α n is the attention weight of node n, expressed as:

[0098]

[0099] In the formula, is the transpose of the learnable weight vector;

[0100] 3.5) The original sample and the positive sample are regarded as similar pairs, and the original sample and the hard negative sample are regarded as dissimilar pairs;

[0101] 3.6) Use cosine similarity to minimize the loss. The specific contrastive learning loss is:

[0102]

[0103] Where, sim() is the cosine similarity function, N is the number of positive and negative sample pairs, h is the vector representation of the original sample image, and h + is the vector representation of the positive sample image, h - is the vector representation of the negative sample image, τ is the temperature parameter;

[0104] 3.7) Add a classifier, use the feature joint graph vector of the positive and negative sample pairs as the input of the classifier, and use the activation function softmax to train the classifier to obtain the predicted label. If the predicted label is 1, it is a vulnerable code, and if the label is 0, it is a non-vulnerable code.

[0105] 4) Extract vulnerable code from the vulnerability database NVD and the open source code library GitHub as the code to be tested, and use the vulnerability detection model to analyze whether the code to be tested contains vulnerabilities.

[0106] 4.1) Extract the vulnerability code from the vulnerability database NVD and the open source code library GitHub as the code to be tested, and execute step 2) on the code file to be tested to generate the corresponding code attribute graph;

[0107] 4.2) The code attribute graph obtained through code feature representation learning is used as the input of the vulnerability detection model to obtain the predicted label of the code to be tested;

[0108] 4.3) Output the predicted label. If the predicted label is 1, it is a vulnerable code, and if the predicted label is 0, it is a non-vulnerable code.

[0109] Example 2

[0110] Corresponding to the source code vulnerability detection method based on contrastive learning in Example 1, this Example 2 provides a source code vulnerability detection system based on contrastive learning, including a data set construction module, a feature learning module, a model construction module, and a vulnerability detection module, such as Figure 1 As shown;

[0111] The dataset construction module is used to collect vulnerability data from the vulnerability database NVD and the open source code library GitHub to build a vulnerability dataset, and perform data enhancement on the vulnerability dataset to form positive and negative sample pairs;

[0112] The dataset construction module also includes a collection unit, an extraction unit, and a generation unit;

[0113] The collection unit is used to collect vulnerability data from the vulnerability database NVD and the open source code library GitHub, and extract vulnerability files and corresponding vulnerability patch files;

[0114] The extraction unit is used to pre-process the extracted vulnerability files and corresponding vulnerability patch files, extract the vulnerability code in the vulnerability files and the patch code in the corresponding vulnerability patch files from the function level, remove redundant information, including header files, comments, and declared global parameters, and obtain a vulnerability data set;

[0115] The generation unit is used to select vulnerability codes from the vulnerability data set as hard negative samples, and the corresponding patch codes as original samples. The semantic preservation heuristic is applied to the original samples to change the grammatical structure of the original samples while keeping the semantics unchanged, thereby generating positive samples. The corresponding positive samples, original samples, and hard negative samples are combined as positive and negative sample pairs.

[0116] Among them, the semantic preservation heuristic rule is:

[0117] Rule 1: Introduce a new else branch, which introduces a new else branch statement into an existing if conditional statement, or reverses the condition in an if conditional statement;

[0118] Rule 2: Conditional expressions with AND, for if statements with AND containing conditions C1 and C2, are rewritten as nested if structures;

[0119] Rule 3: Introduce a guard outside the for loop and change the program containing the for loop statement to a new program structure by adding an if(True) statement outside the loop.

[0120] Rule 4: Introduce a guard in the outer layer of the while loop, and change the program containing the for loop statement into a new program structure by adding a while judgment statement in the outer layer of the loop;

[0121] Rule 5: Add a conditional break statement in the while loop. A while loop with a condition C1 can be changed into another infinite loop statement. A C1 conditional break statement needs to be added to the loop statement body.

[0122] The feature learning module is used to learn the code feature representation of positive and negative sample pairs, and convert the positive and negative sample pairs into a code attribute graph. The control flow graph CFG and program dependency graph PDG contained in the code attribute graph together form a joint graph;

[0123] The feature learning module also includes a representation unit, a sentence embedding unit, a node aggregation unit, and a graph embedding preparation unit;

[0124] The representation unit is used to perform code representation on all samples in each positive and negative sample pair, and generate a code property graph CPG for the positive and negative sample pairs through the tool Joern, where the code property graph CPG includes an abstract syntax tree AST, a control flow graph CFG, and a program dependency graph PDG;

[0125] See also Figure 2 As shown, the sentence embedding unit is used to segment the leaf nodes of the abstract syntax tree AST using byte pair encoding BPE, embed the obtained code tokens using Word2vec, and use the tree-based convolutional neural network TBCNN to generate sentence embedding vectors according to the hierarchical structure in the abstract syntax tree AST;

[0126] Among them, the execution process of the tree-based convolutional neural network TBCNN is as follows:

[0127] By sliding a triangular convolution kernel on the tree, the parent node representation is updated according to the child node representation. The update calculation formula is:

[0128]

[0129] In the formula, is the updated vector representation of the father node, H is the node set containing the father node itself in the triangular convolution kernel, is the vector representation of the nodes in the H set, σ represents the activation function, and W q and b represent the weight matrix and bias term respectively; the main components of TBCNN include: vector representation and encoding, tree-based convolution, dynamic pooling, a fully connected hidden layer and an output layer;

[0130] The node aggregation unit is used to perform maximum pooling through the tree aggregation layer to aggregate the vectors of all nodes of the tree into a single vector containing the corresponding statement semantics, and use this single vector as the initial feature vector of the corresponding statement node in the control flow graph CFG and the program dependency graph PDG:

[0131]

[0132] In the formula, h G is the vector representation set of nodes in CFG and PDG after the tree aggregation layer, TE() represents the tree embedding process, V G It is a node set of CFG, PDG, AST n is the AST whose root node is n;

[0133] The graph embedding preparation unit is used to combine the control flow graph CFG and the program dependency graph PDG into a joint graph, and input the control flow graph CFG, the program dependency graph PDG and the vectors of the statements in the joint graph into the graph-level code representation layer.

[0134] The model building module is used to embed the joint graph and apply contrastive learning to obtain the feature joint graph vector of the positive and negative sample pairs. At the same time, a classifier is added and the feature joint graph vector of the positive and negative sample pairs is used as the input of the classifier to obtain the predicted label, thereby building a vulnerability detection model.

[0135] The model building module also includes an attention weight calculation unit, a graph attention layer unit, a multi-head attention layer unit, a graph node aggregation unit, a comparison rule unit, a loss minimization unit, and a training classifier unit;

[0136] The attention weight calculation unit is used to calculate the corresponding attention weight coefficients of each node in the joint graph according to the control flow, control dependency and data dependency edges in the control flow graph CFG and the program dependency graph PDG:

[0137]

[0138] Where N e is the set of adjacent nodes of node i under edge e (including node i itself), For node i and node j (j∈N e ), || represents the concatenation operation, exp() is an exponential function with the natural constant e as the base, is the transpose of a learnable weight vector, It can be considered as the correlation between node i and node j of edge e. The larger the value, the stronger the correlation between the two nodes, and the larger the weight coefficient;

[0139] The graph attention layer unit is used to represent the layer according to the graph-level code. It consists of multiple graph attention layers. One graph attention layer corresponds to one type of edge. According to the attention weight coefficient, the graph attention layer updates the node vector representation in the joint graph:

[0140]

[0141] In the formula, is the vector representation of the node, σ is the activation function, W e is the weight coefficient of edge e;

[0142] The multi-head attention layer unit is used to average multiple graph attention layers to form a multi-head attention layer to obtain the final vector representation of the nodes in the graph:

[0143]

[0144] The multi-head attention mechanism concatenates or averages the features of each head to obtain a new feature representation;

[0145] The graph node aggregation unit is used to obtain the feature joint graph vector of the positive and negative sample pairs according to the different importance of the nodes in the joint graph:

[0146]

[0147] Where AGG is the aggregation function, V G is the node set of the graph structure, α n is the attention weight of node n, expressed as:

[0148]

[0149] In the formula, is the transpose of the learnable weight vector;

[0150] The contrast rule unit is used to treat the original sample and the positive sample as similar pairs, and the original sample and the hard negative sample as dissimilar pairs;

[0151] The loss minimization unit is used to minimize the loss using cosine similarity. The specific contrastive learning loss is:

[0152]

[0153] Where, sim() is the cosine similarity function, N is the number of positive and negative sample pairs, h is the vector representation of the original sample image, and h + is the vector representation of the positive sample image, h - is the vector representation of the negative sample image, τ is the temperature parameter;

[0154] The training classifier unit is used to add a classifier. The feature joint graph vector of the positive and negative sample pairs is used as the input of the classifier. The classifier is trained with the activation function softmax to obtain the predicted label. If the predicted label is 1, it is a vulnerable code, and if the label is 0, it is a non-vulnerable code.

[0155] The vulnerability detection module extracts vulnerability codes from the vulnerability database NVD and the open source code library GitHub as the code to be tested, and uses the vulnerability detection model to analyze whether the code to be tested contains vulnerabilities;

[0156] The vulnerability detection module also includes a graph generation unit, a detection unit, and an output unit;

[0157] The graph generation unit is used to extract vulnerability codes from the vulnerability database NVD and the open source code library GitHub as the code to be tested, and execute the feature learning module on the code file to be tested to generate the corresponding code attribute graph;

[0158] The detection unit is used to use the code attribute graph obtained through code feature representation learning as the input of the vulnerability detection model to obtain the predicted label of the code to be tested;

[0159] The output unit is used to output a predicted label. If the predicted label is 1, it is a vulnerable code, and if the predicted label is 0, it is a non-vulnerable code.

Claims

1. A source code vulnerability detection method based on contrastive learning, characterized in that: The following steps are involved: 1) Collect vulnerability data from the vulnerability database NVD and the open source code library GitHub to build a vulnerability dataset, and perform data enhancement on the vulnerability dataset to form positive and negative sample pairs; 2) Perform code feature representation learning on positive and negative sample pairs, and convert the positive and negative sample pairs into code attribute graphs. The control flow graph CFG and program dependency graph PDG contained in the code attribute graph together form a joint graph; Step 2) specifically includes the following steps: 2.1) Perform code representation on all samples in each positive and negative sample pair, and generate a code property graph CPG for the positive and negative sample pairs through the tool Joern, where the code property graph CPG contains the abstract syntax tree AST, control flow graph CFG and program dependency graph PDG; 2.2) Use byte pair encoding (BPE) to segment the leaf nodes of the abstract syntax tree AST, use Word2vec to embed the obtained code tokens, and use the tree-based convolutional neural network (TBCNN) to generate sentence embedding vectors according to the hierarchical structure in the abstract syntax tree AST; 2.3) After the tree aggregation layer, perform maximum pooling to aggregate the vectors of all nodes of the tree into a single vector containing the semantics of the corresponding statement, and use this single vector as the initial feature vector of the corresponding statement node in the control flow graph CFG and the program dependency graph PDG; 2.4) Combine the control flow graph CFG and the program dependency graph PDG into a joint graph, and input the control flow graph CFG, the program dependency graph PDG and the vectors of the statements in the joint graph into the graph-level code representation layer; 3) Embed the joint graph and apply contrastive learning to obtain the feature joint graph vector of the positive and negative sample pairs. At the same time, add a classifier and use the feature joint graph vector of the positive and negative sample pairs as the input of the classifier to obtain the predicted label, so as to build a vulnerability detection model; Step 3) specifically includes the following steps: 3.1) Calculate the corresponding attention weight coefficient of each node in the joint graph based on the control flow, control dependency and data dependency edges in the control flow graph CFG and the program dependency graph PDG; 3.2) The graph-level code representation layer consists of multiple graph attention layers. One graph attention layer corresponds to one edge type. According to the attention weight coefficient, the graph attention layer updates the node vector representation in the joint graph. 3.3) Multiple graph attention layers are averaged to form a multi-head attention layer to obtain the final vector representation of the nodes in the graph; 3.4) According to the different importance of nodes in the joint graph, the feature joint graph vector of the positive and negative sample pairs is obtained; 3.5) The original sample and the positive sample are regarded as similar pairs, and the original sample and the hard negative sample are regarded as dissimilar pairs; 3.6) Use cosine similarity to minimize the loss; 3.7) Add a classifier, use the feature joint graph vector of the positive and negative sample pairs as the input of the classifier, and use the activation function softmax to train the classifier to obtain the predicted label. If the predicted label is 1, it is a vulnerable code, and if the label is 0, it is a non-vulnerable code; 4) Extract vulnerable code from the vulnerability database NVD and the open source code library GitHub as the code to be tested, and use the vulnerability detection model to analyze whether the code to be tested contains vulnerabilities.

2. The source code vulnerability detection method based on contrastive learning according to claim 1 is characterized in that: Step 1) specifically includes the following steps: 1.1) Collect vulnerability data from the vulnerability database NVD and the open source code library GitHub, extract vulnerability files and corresponding vulnerability patch files; 1.2) Preprocess the extracted vulnerability files and corresponding vulnerability patch files, extract the vulnerability code in the vulnerability files and the patch code in the corresponding vulnerability patch files from the function level, remove redundant information, including header files, comments, and declared global parameters, and obtain the vulnerability dataset; 1.3) Select vulnerability codes from the vulnerability dataset as hard negative samples and the corresponding patch codes as original samples. Apply semantic preservation heuristics to the original samples to change the grammatical structure of the original samples while keeping the semantics unchanged, thereby generating positive samples. The corresponding positive samples, original samples, and hard negative samples are combined as positive and negative sample pairs.

3. The source code vulnerability detection method based on contrastive learning according to claim 1 is characterized in that: Step 4) specifically includes the following steps: 4.1) Extract the vulnerability code from the vulnerability database NVD and the open source code library GitHub as the code to be tested, and execute step 2) on the code file to generate the corresponding code attribute graph; 4.2) The code attribute graph obtained through code feature representation learning is used as the input of the vulnerability detection model to obtain the predicted label of the code to be tested; 4.3) Output the predicted label. If the predicted label is 1, it is a vulnerable code, and if the predicted label is 0, it is a non-vulnerable code.

4. A source code vulnerability detection system based on contrastive learning using the method of claim 1, characterized in that: Contains data set construction module, feature learning module, model construction module, and vulnerability detection module; The dataset construction module is used to collect vulnerability data from the vulnerability database NVD and the open source code library GitHub to build a vulnerability dataset, and perform data enhancement on the vulnerability dataset to form positive and negative sample pairs; The feature learning module is used to learn the code feature representation of positive and negative sample pairs, and convert the positive and negative sample pairs into a code attribute graph. The control flow graph CFG and program dependency graph PDG contained in the code attribute graph together form a joint graph; The model building module is used to embed the joint graph and apply contrastive learning to obtain the feature joint graph vector of the positive and negative sample pairs. At the same time, a classifier is added and the feature joint graph vector of the positive and negative sample pairs is used as the input of the classifier to obtain the predicted label, thereby building a vulnerability detection model. The vulnerability detection module is used to extract vulnerability code from the vulnerability database NVD and the open source code library GitHub as the code to be tested, and use the vulnerability detection model to analyze whether the code to be tested contains vulnerabilities.

5. The source code vulnerability detection system based on contrastive learning according to claim 4 is characterized in that: The dataset construction module also includes a collection unit, an extraction unit, and a generation unit; The collection unit is used to collect vulnerability data from the vulnerability database NVD and the open source code library GitHub, and extract vulnerability files and corresponding vulnerability patch files; The extraction unit is used to pre-process the extracted vulnerability files and corresponding vulnerability patch files, extract the vulnerability code in the vulnerability files and the patch code in the corresponding vulnerability patch files from the function level, remove redundant information, including header files, comments, and declared global parameters, and obtain a vulnerability data set; The generation unit is used to select vulnerability codes from the vulnerability dataset as hard negative samples and the corresponding patch codes as original samples. The semantic preservation heuristic applied to the original samples changes the grammatical structure of the original samples while keeping the semantics unchanged, thereby generating positive samples. The corresponding positive samples, original samples and hard negative samples are combined as positive and negative sample pairs.

6. The source code vulnerability detection system based on contrastive learning according to claim 5, characterized in that: The feature learning module also includes a representation unit, a sentence embedding unit, a node aggregation unit, and a graph embedding preparation unit; The representation unit is used to perform code representation on all samples in each positive and negative sample pair, and generate a code property graph CPG for the positive and negative sample pairs through the tool Joern, where the code property graph CPG includes an abstract syntax tree AST, a control flow graph CFG, and a program dependency graph PDG; The sentence embedding unit is used to segment the leaf nodes of the abstract syntax tree AST using byte pair encoding BPE, embed the obtained code tokens using Word2vec, and use the tree-based convolutional neural network TBCNN to generate sentence embedding vectors according to the hierarchical structure in the abstract syntax tree AST; The node aggregation unit is used to perform maximum pooling through the tree aggregation layer to aggregate the vectors of all nodes of the tree into a single vector containing the semantics of the corresponding statement, and use the single vector as the initial feature vector of the corresponding statement node in the control flow graph CFG and the program dependency graph PDG; The graph embedding preparation unit is used to combine the control flow graph CFG and the program dependency graph PDG into a joint graph, and input the control flow graph CFG, the program dependency graph PDG and the vectors of the statements in the joint graph into the graph-level code representation layer.

7. The source code vulnerability detection system based on contrastive learning according to claim 6, characterized in that: The model building module also includes an attention weight calculation unit, a graph attention layer unit, a multi-head attention layer unit, a graph node aggregation unit, a comparison rule unit, a loss minimization unit, and a training classifier unit; The attention weight calculation unit is used to calculate the corresponding attention weight coefficient of each node in the joint graph according to the three edges of control flow, control dependency and data dependency in the control flow graph CFG and the program dependency graph PDG; The graph attention layer unit is used to represent the layer according to the graph-level code, which consists of multiple graph attention layers. One graph attention layer corresponds to one edge type. According to the attention weight coefficient, the graph attention layer updates the node vector representation in the joint graph. The multi-head attention layer unit is used to average multiple graph attention layers to form a multi-head attention layer to obtain the final vector representation of the nodes in the graph; The graph node aggregation unit is used to obtain the feature joint graph vector of the positive and negative sample pairs according to the different importance of the nodes in the joint graph; The contrast rule unit is used to treat the original sample and the positive sample as similar pairs, and the original sample and the hard negative sample as dissimilar pairs; The loss minimization unit is used to minimize the loss using cosine similarity; The training classifier unit is used to add a classifier. The feature joint graph vector of the positive and negative sample pairs is used as the input of the classifier. The classifier is trained with the activation function softmax to obtain the predicted label. If the predicted label is 1, it is a vulnerable code, and if the label is 0, it is a non-vulnerable code.

8. The source code vulnerability detection system based on contrastive learning according to claim 4, characterized in that: The vulnerability detection module also includes a graph generation unit, a detection unit, and an output unit; The graph generation unit is used to extract vulnerability codes from the vulnerability database NVD and the open source code library GitHub as the code to be tested, and execute the feature learning module on the code file to be tested to generate the corresponding code attribute graph; The detection unit is used to use the code attribute graph obtained through code feature representation learning as the input of the vulnerability detection model to obtain the predicted label of the code to be tested; The output unit is used to output a predicted label. If the predicted label is 1, it is a vulnerable code, and if the predicted label is 0, it is a non-vulnerable code.

Images