A code vulnerability detection method based on graph contrast learning

Abstract

The present invention introduces the idea of ​​graph comparative learning into the code vulnerability detection task, and proposes a code vulnerability detection method based on graph comparative learning, which can be divided into two stages: initial node representation acquisition and code vulnerability detection based on graph comparative learning. In the initial node representation acquisition, an abstract syntax tree is first generated according to the code, and then a variety of edges representing the relationship between nodes are added on the basis of the abstract syntax tree to obtain the code AST extended graph, and finally word2vec is used to generate the initial node representation. In the code vulnerability detection stage based on graph comparative learning, the final node and final graph representation are first generated by a graph encoder, and then the interaction of information between these representations is used to determine whether the code has vulnerabilities. The present invention uses an unsupervised method to train the model, so code vulnerabilities can be detected without data labels. The comparative experiments of 6 vulnerability detections on the open source dataset SARD containing code vulnerabilities in multiple languages ​​verified the effectiveness of the present invention.

Description

Technical Field

[0001] The present invention belongs to the field of code vulnerability detection, and more specifically, provides a code vulnerability detection method based on graph contrastive learning (VDGCL). Background Art

[0002] 1) Code vulnerability detection

[0003] A code vulnerability refers to an exploitable code defect in the software source code, which is usually caused by errors in the design, development or configuration of the software. These defects may be exploited by criminals to bypass the system's access control and illegally steal higher permissions to manipulate the system at will, such as triggering privileged commands, accessing sensitive information, impersonating identities, eavesdropping on system operations, etc., thereby posing a serious threat to the confidentiality, integrity, availability, and access control of the system or its application data. Therefore, there is an urgent need to carry out systematic research in the field of vulnerability detection in order to efficiently and timely discover vulnerabilities in software systems, patch them in real time, and improve the security level of cyberspace.

[0004] Code vulnerability detection can be divided into traditional detection methods, machine learning-based methods, and deep learning-based methods. Traditional detection methods perform poorly when dealing with large and complex software systems and new vulnerabilities. They require a lot of manual participation, many vulnerabilities cannot be identified, and the false positive rate is high. Machine learning-based methods have achieved automation of code vulnerability detection to a certain extent, but they still require manual labor to extract data features and formulate vulnerability patterns. At the same time, they also face problems such as high false positive rates, inaccurate dataset information, and programming language diversity. Deep learning-based methods can automatically identify vulnerability features in code, but many methods convert the original source code into flat serialized data, ignoring the control flow features formed by jumps, loops, and judgments in the source code, and have low utilization of the logical and semantic information in the code. To better model complex code structures, the source code is abstracted into graph structure data based on its structural features, and then graph neural networks (GNNs) are used to detect vulnerabilities. However, the encoders designed by these methods usually use superimposed multi-layer GNNs, which cannot fully utilize the semantic information in the source code.

[0005] 2) Image Comparative Learning

[0006] The idea of ​​contrastive learning can be traced back to the 1990s. The specific method is to select anchor samples, positive samples, and negative samples, maximize the consistency between positive samples and anchor samples, and minimize the consistency between negative samples and anchor samples. It does not require the participation of data labels and has therefore been widely used in many fields. For example, in recommendation systems, CL4SRec encodes user interaction sequences, maximizes the consistency of the same user interaction sequence between different views, and minimizes the consistency of different user interaction sequences between different views. In natural language processing, VarCLR takes source code variable names as input, maximizes the consistency between similar variable names, and maximizes the consistency between dissimilar variable names. In computer vision, DIM uses the idea of ​​"maximizing the consistency between local information and global information" for contrastive learning, and DGI introduces this idea into the graph field, using the maximum consistency between local information (node ​​representation) and global information (graph representation) for contrastive learning operations, and has achieved good results in node classification, graph classification, and link prediction tasks. However, contrastive learning has not yet been used in the field of function-level vulnerability detection. Summary of the Invention

[0007] The purpose of this invention is to provide a code vulnerability detection method based on graph contrast learning. The goal is to train a model and detect code vulnerabilities without labels. This invention introduces graph contrast learning into the code vulnerability detection task for the first time, and learns graph data features by maximizing the consistency between local information and global information. The effectiveness of VDGCL is verified through experiments.

[0008] The specific technical solution for achieving the purpose of the present invention is:

[0009] A code vulnerability detection method based on graph contrastive learning can be divided into two stages: initial node representation acquisition and code vulnerability detection based on graph contrastive learning. The method includes the following steps:

[0010] The initial node represents acquisition (S1-S3):

[0011] S1: Build an Abstract Syntax Tree (AST) based on the source code: Slice the software code into functions or methods, and uniformly rename the variables in it. Unify the naming style of the variables in the code to avoid affecting the generated graph node representation. Finally, use the code analysis tool joern to generate an AST from the code with unified variable naming.

[0012] S2: Add multiple edges representing data flow and code jump relationships based on the AST obtained in S1: When constructing the AST extension graph of the code, the present invention intends to add 8 types of bidirectional edges to the previously obtained AST to enhance the logical information in the graph, thereby facilitating the information propagation of node relationships in the AST extension graph.

[0013] S3: Obtaining initial node representation: Based on the code information in the AST extended graph, initialization vector training is performed on the nodes in the graph. The present invention uses word2vec to initialize the node representation in the above AST extended graph.

[0014] Let the original view in contrastive learning be α, the expanded new view be β, and the AST expansion graphs in the two views are and The present invention uses and The final node represents H a 、H β And the final graph represents h a 、h β Conduct comparative learning.

[0015] Code vulnerability detection based on graph contrast learning (S4-S6):

[0016] S4: Expand the view:

[0017] Through the view expansion operation in graph comparative learning, personalized PageRank (Personalized PageRank, PPR) is used to expand the view according to the topological structure of the graph.

[0018] S5: Get the final node representation and graph representation:

[0019] a) Through the representation acquisition operation in graph contrast learning, the initial node representations and adjacency matrices of the two graphs are input into the GGNN encoder g θ (·), With the MLP encoder f ψ (·) Get the final node representation H of the two AST expansion graphs α 、H β .

[0020] b) Through the representation acquisition operation in graph contrast learning, the GGNN encoder g θ (·), The output of (·) is input to the pooling function POOLING and the MLP encoder f ψ (·) In the two AST expansion graphs, the final graph representation h α 、h β .

[0021] S6: Maximizing consistency

[0022] Through the consistency maximization operation in graph contrast learning, the present invention uses Graph Barlow Twins to calculate the consistency between different representations. The specific formula is: That is to maximize H α and h β The consistency between them, while maximizing H β and h α The consistency between them, N is the number of nodes in the AST expansion graph, and Node v i In H α and H β The representation vector in .

[0023] The beneficial effects of the present invention include:

[0024] 1) Targeted at the specific application of software code vulnerability detection, the graph comparative learning method and architecture are specifically designed. This method performs comparative learning operations based on the consistency of local and global information across different views of the graph, making better use of graph data than multi-layer neural networks.

[0025] 2) Using a variety of vulnerabilities in real scenarios, comparative experiments were conducted from multiple aspects. The results show that the proposed unsupervised code vulnerability detection method VDGCL is even more effective than the current most advanced supervised learning method in most cases, proving the effectiveness of graph comparative learning in code vulnerability detection. BRIEF DESCRIPTION OF THE DRAWINGS

[0026] Figure 1 is a flow chart of the present invention;

[0027] Figure 2 Code examples and corresponding AST expansion diagrams;

[0028] Figure 3 This is a flowchart for code vulnerability detection. DETAILED DESCRIPTION

[0029] The present invention is further described in detail with reference to the following specific examples and accompanying drawings. The processes, conditions, experimental methods, and the like for implementing the present invention, except for those specifically mentioned below, are common knowledge and common general knowledge in the art. Those skilled in the art can derive other drawings and other implementation methods based on these drawings without inventive effort. Figure 1 The flowchart of the present invention specifically includes the following steps:

[0030] The initial node indicates that the acquisition steps are S1-S3. Figure 2 (a) is a code example. Figure 2 (b) is based on Figure 2 (a) AST expansion graph generated by the code example.

[0031] S1: Construct AST based on source code: Slice the software code in units of functions or methods, and uniformly rename the variables therein, unify the naming style of the variables in the code to avoid affecting the generated graph node representation, and finally generate AST using the code after the variables are uniformly named. AST is an abstract representation of the grammatical structure of the program source code, which describes the logical structure inside the code in a tree form. Each node on the AST represents a structure in the source code, and non-leaf nodes represent syntax nodes (syntax nodes), such as non-terminal symbols such as if keywords and function declaration keywords in the programming language syntax. Leaf nodes represent semantic tags (syntax tokens) in the programming language syntax, such as terminal symbols such as identifier names and constant values. The present invention uses the code analysis tool joem as an AST generation tool. After the import command imports the code source file, the cpg.method(class name).plotDotAst command is executed to generate AST.

[0032] S2: Add multiple edges representing data flow and code jump relationships: When constructing the AST extension graph of the code, the present invention intends to add the following 8 types of bidirectional edges to the AST obtained above to enhance the logical information in the graph, thereby facilitating the information propagation of node relationships in the AST extension graph.

[0033] 1) Data and control flow: After importing the source code file using the import command in joern, directly executing the cpg.method(class name).plotDotCdg and cpg.method(class name).plotDotDdg commands can generate two types of graphs: control dependency graph and data dependency graph.

[0034] 2) GuardedBy: Connects a variable node to the associated expression node, indicating that the variable's value is determined by the expression. This type of edge can be used to detect operand reversal errors in binary operations.

[0035] 3) Jump: If a jump statement occurs during program execution, edges are used to connect the keyword nodes of the conditional statement before and after the jump. GuardedBy edges and Jump edges can record different branches of control flow. These two edges can be used to detect "double-free memory" and "improper resource locking" vulnerabilities.

[0036] 4) ComputedFrom: Taking the statement var = expr as an example, the variable var is calculated from the expression expr. This invention connects the node var to all variables in the expression. ComputedFrom edges can capture the specific location of buffer or variable usage and can be used to detect "null pointer dereference" vulnerabilities.

[0037] 5) NextToken: In an AST, edges connecting parent and child nodes cannot record the order of leaf nodes. Adding NextToken edges connects each leaf node to its successor node. NextToken edges record the order in which variables, operands, and other nodes are used, and can be used to detect "resource release error" vulnerabilities.

[0038] 6) LastUse: The same variable in the code is connected with a LastUse edge in the order in which it appears in the code. This is used to record variable usage and helps identify "double-free memory" vulnerabilities.

[0039] 7) GuardedByNegation: Connects the Boolean expression in the conditional statement with the variable inside the conditional statement with an edge.

[0040] S3: It is necessary to perform initialization vector training on the nodes in the graph according to the code information in the AST expansion graph. The present invention uses word2vec to initialize the node representation of the above-mentioned AST expansion graph. First, the nodes of the AST expansion graph are read from the document, the keywords or symbols represented by the nodes are obtained, and then each node is mapped to a vector. In this process, similar keywords or symbols are closer after being mapped to the vector space. Finally, the vectors of all keywords and symbols in the statement where the variable is located are spliced ​​together. For example, the vectors representing the statement int count=0 are spliced ​​together as "int", "count", "=" and "0". The initial node representation of the code AST expansion graph can be obtained.

[0041] The code vulnerability detection steps based on graph contrast learning are S4-S6. Figure 3 Flowchart for detecting code vulnerabilities using graph comparative learning.

[0042] S4: Expanded View

[0043] The present invention uses Personalized PageRank (PPR) to expand the view based on the topological structure of the graph. The PPR algorithm expands the view as follows:

[0044] S PPR =p[I N -(1-p)D -1 / 2 AD-1 / 2 ] -1

[0045] S PPR Represented by the figure Expanded graph is the identity matrix, p is the probability of returning to the source node for each jump in the random walk, D and A are the original graph and The degree matrix and adjacency matrix of .

[0046] S5: Obtaining the final node representation and graph representation

[0047] a) Input the initial node representation and adjacency matrix of the two graphs into the GGNN encoder g θ (·), (·) with the MLP encoder f ψ (·) Get the final node representation H of the two AST expansion graphs α 、H β .

[0048] b) GGNN encoder g θ (·), The output of (·) is input to the pooling function POOLING and the MLP encoder f ψ (·) In the two AST expansion graphs, the final graph representation h α 、h β , the pooling process can be expressed as follows:

[0049]

[0050] Node v is the output of layer l i , || is the vector concatenation operation, L is the number of pooling layers, is the pooling function parameter, σ is the PReLU nonlinear activation function. Finally, the MLP encoder f φ (·) Adjust the graph representation dimension to be the same as the node dimension, and obtain the final graph representation h α and h β .

[0051] S6: Maximizing consistency

[0052] The present invention uses Graph Barlow Twins to calculate consistency, which requires calculating H α and h β The consistency of H β and h αConsistency. Graph Barlow Twins calculates the cross-correlation matrix of two matrices and makes the diagonal elements of the cross-correlation matrix as close to 1 as possible and the off-diagonal elements as close to 0 as possible. The element in row i and column j The calculation is as follows:

[0053]

[0054] b is the index of batch data, z (1) and z (2) For standard normal distribution data, the matrix H (1) and H (2) The consistency between them is calculated as follows:

[0055]

[0056] The node represents the dimension. The final loss function is calculated as follows:

[0057]

[0058] N is and The number of nodes in and For node v i Final node representations in views α and β.

[0059] Experimental verification:

[0060] In order to verify the versatility of this invention, five different typical code vulnerability detection methods were selected for the experiment:

[0061] DEEPBUGS uses existing non-vulnerable code to generate vulnerable code, generates representation vectors for each of the two codes, and uses a feed-forward neural network to detect vulnerabilities.

[0062] μVulDeePecker adds data flow and control flow information based on AST, integrates function-level and statement-level code vulnerabilities, and uses a building-block BLSTM network to detect vulnerabilities.

[0063] Li et al. adopted an unsupervised approach to convert source code into a minimal intermediate representation and used CNN to detect vulnerabilities.

[0064] Devign adds data flow, control flow, and code keyword sequences to the AST to construct an AST expansion graph. It also generates a node representation of the AST expansion graph and inputs the node representation into the GNN to detect code vulnerabilities.

[0065] Funded adds more edges to Devign’s AST extension graph, enriching the information contained in the AST extension graph and using GGNN to detect code vulnerabilities.

[0066] The operating system used in the experimental machine is Ubuntu 18.04.6 LTS, the processor is Intel (R) Xeon (R) CPU E5-2678 v3, the memory capacity is 128GB, and the graphics card is NVIDIA 2080Ti. The programming language used is Python 3.8.0, the deep learning development framework used is Tensorflow v2.5.0, and the C language code is converted into an abstract syntax tree using Joern, and then the code tree is converted into an AST expansion graph, in which the node vector dimension is 100. The present invention uses small batch stochastic gradient descent (SGD) and Adam algorithm to optimize the model, with a learning rate of 0.01, a dropout rate of 0.2, a GGNN layer number of 2, and a 5-fold cross validation method to evaluate the performance of all methods on the data set. In the methods using graph neural networks to detect vulnerabilities such as Devign, FUNDED and VDGCL, the maximum batch processing node number limit is set to 10,000, which means that the number of nodes per batch of data during model training is at most 10,000, and the number of nodes actually used in each training is not necessarily the same. The present invention uses the average value of 5 experimental results as the final result, and stops training when the loss is less than 0.005 or the number of training times reaches 100.

[0067] As can be seen from Table 1, the VDGCL model proposed in this invention achieves optimal or suboptimal results in most evaluation indicators. Detecting code vulnerabilities based on serialized data requires converting the code into vectors or matrices, and then inputting them into neural networks such as CNN, RNN, and LSTM to detect vulnerabilities. However, these neural networks do not utilize the side information of the graph data, and the node information cannot be aggregated and updated along the edges in the graph, resulting in low accuracy. Detecting code vulnerabilities based on graph data can utilize both node information and edge information in the graph, but the encoders used in these methods cannot interact local information with global information in the graph, resulting in limited improvement in detection accuracy. VDGCL achieved optimal or suboptimal results in most evaluation indicators for the six vulnerabilities in the SARD dataset, with accuracy and F1 Score leading other suboptimal methods by up to 3.6% and 4.3% respectively, demonstrating the effectiveness of graph comparative learning in code vulnerability detection.

[0068] Table 1 Comparison of experimental results of VDGCL and other methods

[0069]

[0070] The above description is only a detailed description of the preferred embodiments and principles of the present invention. For ordinary technicians in this field, based on the ideas provided by the present invention, there may be changes in the specific implementation methods, and these changes should also be considered as the scope of protection of the present invention.

Claims

1. A code vulnerability detection method based on graph contrastive learning, which is divided into two stages: initial node representation acquisition and code vulnerability detection based on graph contrastive learning; characterized in that: The following steps are involved: The initial node represents the acquisition S1: Build an abstract syntax tree (AST) using source code; S2: Add multiple edges based on AST to build the code into an AST extension graph; S3: Get the initial node representation vector of the keyword in the AST expansion graph; Code vulnerability detection based on graph contrast learning S4: Through the view expansion operation in graph contrast learning, using the personalized pagerank algorithm, the original AST is used to expand the graph Expand a new AST extension graph ; S5: Through the representation acquisition operation in graph contrast learning, the final node representation of the two AST expansion graphs is obtained using the GGNN encoder, MLP encoder and POOLING pooling function 、 And the final graph representation 、 ; S6: Maximize the consistency maximization operation in graph contrast learning and consistency while maximizing and The consistency of the model parameters is then optimized iteratively.

2. The code vulnerability detection method based on graph comparative learning according to claim 1 is characterized in that: The AST in step S1 is generated by joern.

3. The code vulnerability detection method based on graph comparative learning according to claim 1 is characterized in that: The AST expansion graph in step S2 is to add multiple edges based on the AST generated in S1.

4. The code vulnerability detection method based on graph comparative learning according to claim 1 is characterized in that: In step S3, the initial node representation vector is generated by word2vec.

5. The code vulnerability detection method based on graph comparative learning according to claim 1 is characterized in that: Step S4 expands a new AST extension graph , using personalized pagerank algorithm in the extended graph The topological structure is expanded by deleting some edges in the original AST expansion graph to expand the new AST expansion graph.

6. The code vulnerability detection method based on graph comparative learning according to claim 1 is characterized in that: In step S5, the GGNN encoder is used 、 With MLP encoder Get the final node representation of the two AST expansion graphs 、 ; At the same time, the GGNN encoder 、 The output result is input to the pooling function POOLING and MLP encoder In the figure, we get the final graph representation of the two AST expansion graphs 、 .

7. The code vulnerability detection method based on graph comparative learning according to claim 1 is characterized in that: The formula for calculating consistency is ,maximize and consistency while maximizing and consistency, where Graph Barlow Twins loss function to measure the consistency of two matrices, is the number of nodes in the AST expansion graph, and Node exist and The representation vector in .

Images