Source code vulnerability detection method and system based on adaptive graph neural network
By preprocessing and learning features from source code using an adaptive graph neural network, and combining CodeBERT and GCN modules, the problem of low efficiency and high false positive rate of existing methods in large-scale complex software is solved, and more efficient vulnerability detection is achieved.
Patent Information
- Application Number
- CN202411312472.0
- Authority / Receiving Office
- CN · China
- Patent Type
- Applications(China)
- Current Assignee / Owner
- Filing Date
- 2024-09-20
- Publication Date
- 2025-01-07
AI Technical Summary
Existing static analysis and dynamic testing methods are inefficient and have a high false positive rate in large-scale complex software. Traditional token sequence-based methods cannot preserve the structural information of the source code, and graph-based representation methods are insufficient in capturing complex structures and semantic relationships.
An adaptive graph neural network-based approach is adopted to preprocess the source code, extract the code attribute graph, use CodeBERT for vector representation, combine an adaptive attention convolution module and a GCN module to learn the global and local features of the graph, and perform vulnerability detection through a fully connected layer.
It improves the accuracy and efficiency of software vulnerability detection, can capture the syntax and semantic information of the code more comprehensively, and enhances the model's ability to understand vulnerability characteristics.
Smart Images

Figure CN119272275A_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the technical field of software vulnerability detection, and in particular to a source code vulnerability detection method and system based on an adaptive graph neural network. Background Art
[0002] Software vulnerabilities refer to defects or weaknesses in software systems that can be exploited by attackers. As software systems become increasingly complex, vulnerability detection becomes more and more challenging. Traditional static analysis and dynamic testing methods are often inefficient and have high false positive rates when dealing with large-scale complex software. In recent years, vulnerability detection methods based on deep learning have shown great potential, but how to effectively represent and learn the grammatical and semantic features of code remains a key issue.
[0003] In recent years, due to the significant advantages of deep learning in modeling capabilities and feature learning, it has begun to be widely used in the field of source code vulnerability detection to help human experts build vulnerability features and assist in vulnerability detection. Source code vulnerability detection methods based on deep learning are classified according to the representation method of the source code in the preprocessing stage, which can include token sequence-based and graph-based representation methods.
[0004] Although token-based code representation is close to the semantics of the source code, it cannot retain important structural information such as control flow and data flow, which limits the effectiveness of vulnerability detection. Current preprocessing methods have gradually strengthened the embodiment of structural information in token sequences, and improved the quality of code representation through various tools and methods. In contrast, methods based on graph representation can more comprehensively capture the complex structure and semantic relationships of source code, and are suitable for deep learning models to learn vulnerability features in programs. Therefore, the present invention chooses to represent the source code based on graph representation to ensure accurate detection of real vulnerabilities. Summary of the invention
[0005] The purpose of the present invention is to solve the problems existing in the prior art and propose a source code vulnerability detection method and system based on an adaptive graph neural network, which improves the accuracy and efficiency of software vulnerability detection.
[0006] To achieve the above purpose, the technical solution adopted is:
[0007] A source code vulnerability detection method based on adaptive graph neural network, comprising:
[0008] Preprocess the source code;
[0009] Extracting code property graph from normalized code using parser;
[0010] Use CodeBERT to convert code attribute graphs into learnable vector representations;
[0011] Calculate the centrality index of each node in the code attribute graph and concatenate it with the node feature vector;
[0012] Combine the adaptive attention convolution module and the GCN module to learn the global and local features of the graph;
[0013] The features learned by the adaptive attention convolution module and the GCN module are integrated into the fully connected layer to output the final vulnerability detection results.
[0014] According to the source code vulnerability detection method based on the adaptive graph neural network of the present invention, further, preprocessing the source code includes: removing irrelevant information in the source code; mapping the customized variable names and function names into a unified representation.
[0015] According to the source code vulnerability detection method based on adaptive graph neural network of the present invention, further, using a parser to extract a code property graph from the standardized code includes: first, using the parser to convert the source code into an abstract syntax tree AST, and constructing a control flow graph CFG based on the control flow information extracted from the AST; then analyzing the data dependency relationship in the code to generate a data dependency graph DDG; analyzing the function call relationship in the code to generate a call dependency graph CDG, and finally fusing AST, CFG, DDG and CDG to obtain a code property graph.
[0016] According to the source code vulnerability detection method based on the adaptive graph neural network of the present invention, further, using CodeBERT to convert the code attribute graph into a learnable vector representation includes:
[0017] First, the preprocessed code line content is used as the node label, and the label of each node is extracted as a sentence. Then, the sentence is segmented using CodeBERT's tokenizer, and the sentence is decomposed into a list of tokens. Finally, the list of tokens is input into the CodeBERT model to obtain the vector representation of each node label. In addition, when vectorizing the edges of the code attribute graph, the edge features are directly encoded, and a four-dimensional vector is generated for each edge in the order of AST, CFG, DDG, and CDG. A value of 1 indicates that this type of edge exists between the corresponding two nodes.
[0018] According to the source code vulnerability detection method based on the adaptive graph neural network of the present invention, further, the centrality index of each node in the code attribute graph is calculated and concatenated with the node feature vector, including:
[0019] The centrality index of each node in the code attribute graph is calculated using the eigenvector centrality function, Katz centrality function and betweenness centrality function respectively. The three results are concatenated to obtain the centrality measurement vector of the node; and the encoded node label eigenvector is connected with the centrality measurement vector to obtain the final node feature matrix.
[0020] According to the source code vulnerability detection method based on an adaptive graph neural network of the present invention, further, the adaptive attention convolution module first performs a linear transformation on the updated graph features of the input; then calculates the attention weights through the full_attention_conv function to obtain the weighted features; finally, the original input is added to the weighted features through the residual connection.
[0021] According to the source code vulnerability detection method based on adaptive graph neural network of the present invention, further, in the process of node feature propagation using GCN module, multiple hidden layers are added in the middle layer to deepen the features. After each hidden layer receives the node features from the previous layer, it performs layer-by-layer convolution operations on the node features and the features of its neighboring nodes, thereby generating a new node feature representation.
[0022] According to the source code vulnerability detection method based on the adaptive graph neural network of the present invention, further, the output of the adaptive attention convolution module and the GCN module are combined by weighted sum fusion or splicing fusion, and the weighted sum fusion formula is:
[0023] x=ω×x GCN +(1-ω)×x 自适应注意力卷积
[0024] Splicing fusion formula:
[0025] x=concat(x GCN ,x 自适应注意力卷积 )
[0026] Among them, ω is the fusion weight, x represents the concatenated feature vector, and x GCN represents the features extracted by the GCN module, x 自适应注意力卷积 It represents the features extracted by the adaptive attention convolution module, and concat is the concatenation function.
[0027] According to the source code vulnerability detection method based on the adaptive graph neural network of the present invention, further, the outputs of the adaptive attention convolution module and the GCN module are fused, and all node information in the graph is integrated into a global feature vector of a fixed length through global average pooling, and then binary classification is performed through a fully connected layer to determine whether there is a vulnerability in the source code.
[0028] Furthermore, the present invention also provides a source code vulnerability detection system based on an adaptive graph neural network, comprising:
[0029] Code normalization module, used to preprocess source code;
[0030] A CPG graph extraction module, used to extract a code property graph from the normalized code using a parser;
[0031] A graph vector embedding module, which uses CodeBERT to convert code attribute graphs into learnable vector representations;
[0032] The graph feature extraction module is used to calculate the centrality index of each node in the code attribute graph and concatenate it with the node feature vector;
[0033] The graph neural network training module is used to combine the adaptive attention convolution module and the GCN module to learn the global and local features of the graph;
[0034] The vulnerability detection module is used to fuse the features learned by the adaptive attention convolution module and the GCN module into the fully connected layer and output the final vulnerability detection results.
[0035] The beneficial effects achieved by adopting the above technical solution are:
[0036] The present invention proposes a source code vulnerability detection method based on an adaptive graph neural network. The method is based on the powerful representation ability of the Code Property Graph (CPG) and the excellent learning ability of the adaptive graph neural network. First, the source code is represented by using CPG, and the CPG nodes are vectorized and represented in combination with CodeBERT to fully capture the syntax and semantic information of the code. Subsequently, graph centrality analysis is introduced to extract high-level graph node features to capture graph node features that may be related to vulnerability features and enhance the model's understanding of the overall structure of the code. Finally, an adaptive graph neural network is designed, which effectively integrates the global dependency capture capability of the Transformer adaptive attention mechanism and the local structure analysis capability of the Graph Convolutional Networks (GCN) to achieve comprehensive learning of code features. The present invention improves the accuracy and efficiency of software vulnerability detection. BRIEF DESCRIPTION OF THE DRAWINGS
[0037] In order to more clearly illustrate the technical solutions of the embodiments of the present invention, the following briefly introduces the drawings of the embodiments of the present invention, wherein the drawings are only used to illustrate some embodiments of the present invention, but not to limit all embodiments of the present invention thereto.
[0038] Figure 1It is a schematic diagram of the process of a source code vulnerability detection method based on an adaptive graph neural network according to an embodiment of the present invention;
[0039] Figure 2 is an example diagram of the code standardization process of an embodiment of the present invention;
[0040] Figure 3 is a CPG graph conversion process diagram of an embodiment of the present invention;
[0041] Figure 4 It is the comparison result of 8 vulnerability detection methods in the embodiment of the present invention on four indicators;
[0042] Figure 5 is the ablation experiment result of the embodiment of the present invention;
[0043] Figure 6 is a thermal diagram of an ablation experiment according to an embodiment of the present invention. DETAILED DESCRIPTION
[0044] The following will be combined with the drawings of specific embodiments of the present invention to clearly and completely describe the exemplary scheme of the embodiment of the present invention. Unless otherwise defined, the technical terms or scientific terms used in the present invention should be the common meanings understood by people with ordinary skills in the field.
[0045] This embodiment discloses a source code vulnerability detection method based on an adaptive graph neural network. The method process is as follows: Figure 1 As shown, the following steps are included:
[0046] Step S101: pre-process the source code, remove irrelevant information, and unify variable names and function names.
[0047] In software source code vulnerability detection, the rich semantic information of the source code can help accurately express the vulnerability characteristics. This case solution focuses on detecting vulnerabilities directly from the source code. In this process, noise information such as comments in the source code, complex and non-standard variable and function naming, etc. will affect the model's ability to understand code features and vulnerability semantics. In order to reduce the impact of these irrelevant information on code semantics, a code normalization preprocessing step is introduced.
[0048] For example, Figure 2 The code shown is from the real vulnerability dataset Big-Vul. In order to obtain a simplified representation of the real code, the code normalization process includes the following three specific steps. Figure 2 shown.
[0049] Step a: Remove annotation information
[0050] During the code normalization process, the comments in the source code are removed to ensure that only executable code is further processed in the subsequent process.
[0051] Step b: Unify variable names
[0052] The normalization operation involves mapping custom variable names to a unified format.
[0053] Step c: Unify function names
[0054] Similarly, custom function names will also be mapped to a unified function name format.
[0055] The preprocessing step of code normalization attempts to eliminate the semantic noise introduced by different naming styles and conventions in the code by mapping variable and function names to a unified representation, so that the model can focus more on the structure and logic of the code and understand the vulnerability characteristics hidden in the specific pattern of the code, rather than being disturbed by diverse naming, thereby improving the model's ability to understand the vulnerability code pattern. At the same time, consistent variable and function name representation makes the training data more semantically unified, allowing the deep learning model to learn in a more consistent data environment, which helps improve the generalization ability of the model.
[0056] In addition, the standardized code is more concise in structure, reducing the burden of the model to handle different naming methods. This not only improves the efficiency of model training, but also reduces the complexity of the model, helping to speed up the training process and reduce the consumption of computing resources.
[0057] Step S102: Utilize a parser to extract a Code Property Graph (CPG) from the normalized code.
[0058] The code property graph is a comprehensive code representation method that can simultaneously reflect the semantics, syntax, data flow, and control flow information of the code. During the CPG generation process, first, the source code is converted into an abstract syntax tree AST using a parser, and a control flow graph CFG is constructed based on the control flow information extracted from the AST, including control structures such as conditional statements and loops in the code. Then, the data dependencies in the code are analyzed to generate a data dependency graph DDG, which shows the definition and usage of variables, as well as the flow path of data in the code. The function call relationship in the code is analyzed to generate a call dependency graph CDG. Finally, AST, CFG, DDG, and CDG are merged to obtain a CPG graph. The CPG graphical representation of the above code example is as follows: Figure 3 As shown, each node represents a line of normalized code, and edges of different colors represent different dependencies between two lines of code.
[0059] This case solution uses the open source static code analysis tool Joern to extract the CPG graph. The Joern tool can complete code analysis and CPG graph generation without pre-compilation, supports multiple programming languages, and is suitable for code to be tested of any granularity, greatly reducing the cost of use during the conversion process.
[0060] Step S103: Use CodeBERT to convert the code attribute graph into a learnable vector representation.
[0061] The nodes in the code attribute graph converted to dot format represent the code lines of the function. The normalized code line content as the label attribute of the node needs to be embedded and represented before it can be used as the input of the subsequent model training process. Therefore, vectorizing the nodes is a key preprocessing step before model training. This case solution uses the pre-trained CodeBERT model to embed the nodes to capture the semantic information of the label.
[0062] Specifically, we first extract the label of each node as a sentence, then use CodeBERT’s tokenizer to tokenize the sentence and decompose it into a list of tokens. Finally, we input the list of tokens into the CodeBERT model to obtain the vector representation of each node label.
[0063] For the edge features of the CPG graph, it involves the process of converting a multigraph into a directed graph. The NetworkX library is used to read graph data stored in the DOT format. Since the original data is a multigraph, it needs to be converted into a directed graph first. This conversion process involves merging multiple edges and merging the attributes of the edges into the attributes of a single edge. The edge attributes include four label attributes: AST, CFG, DDG, and CDG. Since in the CPG multigraph, the edge between two points is not unique, that is, an edge from point 1 to point 2 may be an edge with label = AST, which indicates a grammatical relationship, and also an edge with label = DDG, which indicates a data flow relationship. Therefore, for the encoding of the edge features, a four-dimensional vector is generated for each edge in the order of AST, CFG, DDG, and CDG, and a 1 indicates that there is an edge of this type between the corresponding two nodes.
[0064] Step S104: Calculate the centrality index of each node in the code attribute graph and concatenate it with the node feature vector of step S103.
[0065] In addition to the label attribute characteristics of nodes and edges themselves, studying the graph characteristics of the code CPG graph structure itself is also a key research topic. The centrality index of the graph evaluates the role of key nodes in the function based on the importance and influence of the nodes in the network. In order to capture the graph node characteristics that may be related to vulnerability characteristics, the following three centrality indicators are used:
[0066] (1) Eigenvector centrality: In a vulnerable function, nodes with high eigenvector centrality are usually nodes that are connected to many other important nodes. Such nodes may represent key execution paths or control flows in the vulnerable function. Therefore, eigenvector centrality can help identify potential key nodes in the vulnerable function.
[0067] (2) Katz centrality: Katz centrality considers the importance of a node through all paths, and the longer the path, the less important the impact. Therefore, Katz centrality can help find key nodes in the vulnerability function that are far away in the network but still have significant impact.
[0068] (3) Betweenness centrality: During vulnerability detection, nodes with high betweenness centrality may be nodes on the critical path connecting important modules or key functions in the vulnerable function. Attackers may use these nodes to spread malicious code or affect the execution process of the vulnerable function. Therefore, betweenness centrality can help identify important nodes related to the exploitation or introduction of the vulnerability.
[0069] This step will calculate the centrality index of each node in the CPG graph and embed the node centrality measurement vector. At the same time, for each node, the encoded node label feature vector is connected with the node centrality measurement vector to obtain the final node feature matrix. This step is summarized as Algorithm 1 as follows:
[0070]
[0071]
[0072] Considering the actual detection effect and performance, the dimension of each node is set to 700. At the same time, since the centrality analysis will consume a lot of computing resources, in order to balance the computing resources, the tolerance parameter of the eigenvector centrality calculation function is set to 0.001, and the attenuation factor of the Katz centrality calculation function is set to 0.2.
[0073] Step S105: Combine the adaptive attention convolution module and the GCN module to learn the global and local features of the graph.
[0074] Although traditional graph neural network models can process graph structured data, they have limitations when processing long-distance dependencies and global information. On the other hand, the Transformer model performs well in processing sequence data and capturing long-distance dependencies, but does not work well when directly applied to graph structured data. Therefore, this case combines the Transformer and GCN graph deep learning methods to design a graph neural network training module.
[0075] The core idea of the graph neural network training module is to combine the advantages of traditional GCN local feature extraction and Transformer to capture long-distance dependencies, effectively learning and representing the structural features of the code. This module mainly consists of two sub-modules: the adaptive attention convolution module and the GCN module.
[0076] The adaptive attention convolution module is the core component of the graph neural network training module. The adaptive attention convolution module first performs a linear transformation on the input graph features (i.e., node features), and the definition formula (1) is as follows:
[0077] q=W q (x),k=W k (x),v=W v (x) (1)
[0078] Where q is the query vector, k is the key vector, v is the value vector, x is the input node feature matrix, and W q is the query weight matrix, W k is the key weight matrix, W v is the value weight matrix.
[0079] Then, the attention weight is calculated by the full_attention_conv function. The calculation process involves all nodes, so that each node can interact with all other nodes in the graph, not just directly connected nodes. After calculation, the attention weight is dynamically adjusted, and the model redistributes the attention level of different nodes. The full_attention_conv function formula is shown in (2):
[0080]
[0081] Where Q is the query matrix, K is the key matrix, V is the value matrix, and d k is the dimension of the key vector and softmax is the softmax function applied to each row of the matrix.
[0082] At the same time, the adaptive attention convolution module adopts a multi-head attention mechanism and stacks multiple adaptive attention convolution layers, allowing information to propagate at a deeper level, further enhancing the model's ability to capture different types of long-distance dependencies. Finally, the adaptive attention convolution module enhances the adaptive ability through residual connections, allowing the model to maintain attention to the original input in the deep network while learning more complex feature representations. The formula is as follows:
[0083] output=α×attention_output+(1-α)×input (3)
[0084] Among them, output is the final output of the adaptive attention convolution layer, attention_output is the output calculated by the self-attention mechanism, and input is the original input of the adaptive attention convolution layer. α is an adjustable parameter used to balance the attention output and the original input.
[0085] Through the above adaptive attention allocation mechanism, while retaining the long-distance dependency of the entire feature graph, attention is redistributed to different areas of the image according to the graph features, making it possible to achieve more fine-grained local feature analysis in significant areas. At the same time, coarse-grained analysis in insignificant areas will also improve the efficiency of model operation. For example, when processing a code snippet that may have a buffer overflow vulnerability, the adaptive attention convolution layer may pay more attention to memory allocation and access-related nodes, and pay less attention to other less relevant code parts. This adaptive attention allocation mechanism enables the model to capture potential vulnerability patterns more accurately.
[0086] At the same time, the GCN module, as a supplement to the adaptive attention convolution module, focuses on extracting local graph structure features. In practical applications, GCN may be better at identifying certain types of vulnerabilities, such as those that are mainly determined by local code structures. By aggregating information from neighboring nodes, GCN can effectively capture local dependencies in the code, which is particularly important for detecting vulnerabilities with strong locality, such as uninitialized variables and null pointer references.
[0087] It works by updating the representation of the central node by aggregating information from neighboring nodes:
[0088]
[0089] in, is the normalized adjacency matrix, H (I) is the node feature of layer I, W (I) is a learnable weight matrix and σ is a nonlinear activation function.
[0090] In the GCN design of this module, after the first initialization layer receives the initial node features and graph structure information (edge index and edge weight), a graph convolution layer is used to perform a convolution operation on the input features, and the convolved features are batch normalized. Then, the normalized features are nonlinearly activated to introduce nonlinearity, so that the model can represent more complex feature relationships. At the same time, Dropout is applied to the activated features to prevent overfitting.
[0091] Then, multiple hidden layers are added to the middle layer to deepen the features and extract deeper features layer by layer. After each hidden layer receives the node features from the previous layer, it performs convolution operations on the node features and the features of its neighboring nodes layer by layer. At the same time, batch normalization and activation operations are performed after each layer of convolution to ensure the stability of the features in multi-layer transmission. Dropout is also applied to randomly inactivate the features.
[0092] Finally, the output layer receives the node features from the last hidden layer, and uses the last layer of GCNConv to map the features to the output channels to obtain the local features of the graph processed by GCN.
[0093] Step S106, integrate the features learned by the adaptive attention convolution module and the GCN module into the fully connected layer, and output the final vulnerability detection result.
[0094] The core idea of the vulnerability detection module is to fuse the features extracted by the graph feature neural network to achieve vulnerability detection. The innovation of this module lies in its flexible feature fusion strategy and graph-level pooling operation.
[0095] Depending on the needs of specific tasks, you can choose to use weighted sum fusion or splicing fusion to combine the outputs of the adaptive attention convolution module and the GCN module. Weighted sum fusion formula:
[0096] x=ω×x GCN +(1-ω)×x 自适应注意力卷积 (5)
[0097] Splicing fusion formula:
[0098] x=concat(x GCN ,x 自适应注意力卷积 ) (6)
[0099] Among them, ω is the fusion weight, x represents the concatenated feature vector, and x GCN represents the features extracted by the GCN module, x 自适应注意力卷积 It represents the features extracted by the adaptive attention convolution module, and concat is the concatenation function.
[0100] This design enables the model to adaptively adjust feature fusion strategies in different types of vulnerability detection tasks. For example, for complex vulnerability types that require consideration of both global and local information, splicing fusion may perform better; while for certain vulnerability types that mainly rely on global or local information, weighted sum fusion may be more suitable.
[0101] Since the lengths of the source codes used for training and testing are different, the sizes of the generated CPG graphs are different, and the lengths of the graph feature vectors are not uniform, global average pooling is used here to integrate the information of all nodes in the graph into a global feature vector of fixed length, so that the model is invariant to the scale changes of the graph, thereby improving the generalization and robustness of the model.
[0102] After feature fusion, due to the different lengths of the source code used for training and testing, the sizes of the generated CPG graphs are different, and the length of the graph feature vector is not uniform. Here, global average pooling is used to integrate the information of all nodes in the graph into a global feature vector of fixed length, so that the model is invariant to the scale change of the graph, improving the generalization ability and robustness of the model. The formula of global average pooling is as follows:
[0103]
[0104] Among them, GAP(d) represents the global average pooling result of the d-th feature map of the input tensor, H is the height of the feature map, W is the width of the feature map, and x(i,j,d) represents the value of the d-th feature map at position (i,j) in the input tensor.
[0105] Through global average pooling, the model comprehensively considers the information of all nodes to form a comprehensive code snippet representation.
[0106] This representation is particularly suitable for vulnerability types that require consideration of the overall code structure and semantics, such as security vulnerabilities caused by logical errors or design flaws. Finally, the module performs binary classification through a fully connected layer, mapping the graph-level features to a two-dimensional output space to indicate the probability of a vulnerability.
[0107] Corresponding to the above method, this embodiment also proposes a source code vulnerability detection system based on an adaptive graph neural network, comprising:
[0108] Code normalization module, used to preprocess source code.
[0109] The CPG graph extraction module is used to extract the code property graph from the normalized code using the parser.
[0110] A graph vector embedding module for converting code attribute graphs into learnable vector representations using CodeBERT.
[0111] The graph feature extraction module is used to calculate the centrality index of each node in the code attribute graph and concatenate it with the node feature vector.
[0112] The graph neural network training module is used to combine the adaptive attention convolution module and the GCN module to learn the global and local features of the graph.
[0113] The vulnerability detection module is used to fuse the features learned by the adaptive attention convolution module and the GCN module into the fully connected layer and output the final vulnerability detection results.
[0114] The following is a comprehensive experimental evaluation on a real vulnerability dataset to verify the effectiveness and superiority of the proposed method. First, the experimental preparation is explained; second, a hyperparameter optimization experiment is conducted to improve the model performance to the best state; at the same time, a comparative analysis of vulnerability detection results is conducted to verify the superiority of the proposed method; finally, an ablation experiment is conducted to verify the importance of evaluation components and analyze the role of different components in the method.
[0115] ①Dataset
[0116] The vulnerability data comes from the vulnerability dataset Big-Vul. The Big-Vul dataset is a real vulnerability dataset of C / C++ code, which covers 3754 code vulnerabilities, involving 91 different vulnerability types, and a total of 11834 vulnerable functions. At the same time, in order to facilitate subsequent performance comparisons with other methods in the field, and to balance the number of vulnerability data and non-vulnerability data, another vulnerability dataset SARD is used for non-vulnerability functions, and 10,000 non-vulnerability functions are selected. SARD contains a large number of constructed artificial vulnerability functions and non-vulnerability functions, and is currently one of the largest and most commonly used datasets.
[0117] Note: The Big-Vul dataset crawls data from the public Common Vulnerabilities and Exposures (CVE) database and CVE-related source code repositories based on a crawler method, marking the functions before the patch as vulnerable functions and the functions after the patch as non-vulnerable functions. However, in the actual dataset process, it was found that some functions were not associated with specific CVE numbers. It was believed that this part of the data needed to be filtered because it could not be labeled and lacked authenticity. Therefore, 9984 vulnerable functions in the Big-Vul dataset were finally used.
[0118] ② Experimental environment and configuration
[0119] The experiment was conducted on a Linux server equipped with 128GB memory, 16-core Intel Xeon processor and 16GB video memory NVIDIA RTX A4000 series graphics card. At the same time, the dataset was randomly divided into training set, validation set and test set according to 8:1:1. The code was developed based on PyTorch, Numpy and other libraries.
[0120] ③Evaluation criteria
[0121] To evaluate the effectiveness of the vulnerability detection model, the following criteria are used to evaluate the method:
[0122] True positive (TP): the number of samples correctly predicted to have vulnerabilities;
[0123] True negative (TN): the number of samples correctly predicted as having no vulnerabilities;
[0124] False positive (FP): the number of samples that are incorrectly predicted to have vulnerabilities;
[0125] False negative (FN): the number of samples that are incorrectly predicted to be free of vulnerabilities;
[0126] Accuracy: accuracy = (TP + TN) / (TP + TN + FP + FN);
[0127] Precision: precision = TP / (TP+FP);
[0128] Recall: recall = TP / (TP+FN);
[0129] ·F1 score: F1=2*(precision*recall) / (precision+recall).
[0130] ④ Experimental results and analysis
[0131] 1. Hyperparameter tuning experiment
[0132] In order to obtain the best model detection performance, we first conduct hyperparameter tuning experiments on the neural network parameters and use the grid search method to tune the hyperparameters. During the tuning process, the loss function uses the cross entropy loss function and the optimizer is Adam. The parameters, ranges, and gradients for tuning are shown in the following table:
[0133] Table 1 Hyperparameters and their value ranges, gradients
[0134]
[0135]
[0136] Among them, graph_weight is a parameter of the model that balances the final contribution of the adaptive attention convolution module and the GCN module to the feature vector. When graph_weight is 0, the output vector of the adaptive attention convolution module is fully used, and when graph_weight is 1, the output vector of GCN is fully used. At the same time, since some parameters are related to the computational efficiency of the model, time cost is also an important reference in the process of hyperparameter tuning.
[0137] Finally, according to the results of the grid search method and the comprehensive time consumption, the final parameters used are: GCN hidden_channels = 128, GCN num_layers = 3, GCN dropout 0.5, adaptive attention convolution layer hidden_channels = 64, adaptive attention convolution layer num_layers = 2, adaptive attention convolution layer dropout = 5, adaptive attention convolution layer num_heads = 3, adaptive attention convolution layer alpha = 0.5, adaptive attention convolution layer graph_weight = 0.8, adaptive attention convolution layer aggregate = add, Adam learning rate is 1e -3 .
[0138] 2. Vulnerability detection experiment
[0139] The proposed method is compared with several advanced vulnerability detection tools, including three rule-based commercial vulnerability detection tools Checkmark, FlawFinder and RATS, and four deep learning-based vulnerability detection methods LineVul, LineVD, DeepLineDP and VulSniper. The experimental results are shown in Figure 4 .
[0140] Rule-based methods are vulnerability detection schemes that directly match source code based on vulnerability-related patterns defined by cybersecurity experts. Currently, commonly used vulnerability-related patterns are often related to various numerical indicators of source code, such as cyclomatic complexity, McCabe, complexity indicators, and code change metrics. Vulnerability detection methods that use software metrics are based on the following assumptions: complex code is difficult to understand, and therefore difficult to maintain and test, and is therefore more likely to contain vulnerabilities. Similarly, frequently modified code is more prone to errors and is therefore more likely to have vulnerabilities. Although these metrics have long been used to measure software quality and predict defects, they do not directly correspond to the concept of code security. In fact, vulnerable code blocks are not always complex or large. Therefore, machine learning models based only on these metrics usually perform mediocre in vulnerability detection. As shown in the precision, recall, and F1 scores in the table, the detection effects of Checkmark, FlawFinder, and RATS are not ideal. From the perspective of code auditing, factors such as code complexity can be considered as part of the vulnerability, but they are usually not the decisive factor in determining whether the code has vulnerabilities. Different from the rule-based method, this method combines the Transformer adaptive attention mechanism and the GCN model to automatically learn the vulnerability patterns in the training samples, achieving lower missed alarm rate and false alarm rate.
[0141] The method of representing source code based on token sequence is to decompose the source code into token sequences according to specific word segmentation rules and then train the model. During the code compilation process, the source code is processed by the compiler, and different tokens are separated from the continuous character stream. These tokens are the smallest grammatical units of the code. From the perspective of the compiler, the code is a series of tokens with specific meanings. Therefore, some studies are based on an assumption: extracting lexical relations and semantic information in the source code through specific rules and converting the source code into a sequence of tokens is a reasonable abstract representation of the source code. Among the comparative methods, LineVul uses token sequences to represent source code. LineVul first converts the source code into a token sequence through lexical analysis, and then uses the pre-trained BERT model for vector embedding and vulnerability detection. Because LineVul regards the source code as plain text similar to natural language and directly uses BERT for prediction, it lacks consideration of the rich structural information of the code language and loses a lot of key information in the code, resulting in poor detection results.
[0142] The method based on graph representation of source code is to process the source code into different graph representations of the code and then conduct model training. Compared with token sequence representation, graph representation is a more abstract and complex representation form. Token sequence representation cannot retain the complex structure of the source code. The semantics of the representation depends on the single order in the final token sequence. Therefore, it can only represent linear semantics. Graph representation can represent many-to-many connections between vertices (i.e., statements or tokens), and can reflect the complex structural relationship between elements in the program. At the same time, due to the variety of graph representations of code, different graph types have different representation capabilities for graphs. Among them, DeepLineDP and LineVD are two representative methods of using AST tree representation and PDG graph to represent source code, respectively. However, since AST tree tends to reflect the grammatical structure of source code, and current vulnerabilities are often more closely related to source code semantic information, the detection effect of DeepLineDP is not ideal. Although PDG graph can represent various dependencies in source code, it lacks the description of grammatical relationships and node (i.e., statement or token) semantics. Therefore, LineVD detection effect is better than DeepLineDP, but there is still room for improvement. Similarly, VulSniper also uses CPG to represent source code and designs a neural network combined with an attention module to learn feature weights for vulnerability detection. However, its simple attention module cannot capture the ability of long-distance dependencies, nor can it allocate attention weights to different parts of the graph according to graph features, ignoring important local features and global graph features in CPG, resulting in poor detection results for VulSniper.
[0143] Experimental results show that the F1 score of this method on the vulnerability dataset is 82.9, which is 13.6%-49.9% higher than other vulnerability detection methods, showing obvious advantages.
[0144] 3. Ablation experiment
[0145] In order to evaluate the contribution of each component in the method, a series of ablation experiments are designed. The experimental groups of ablation experiment design are as follows:
[0146] Baseline: Node embedding using word2vec + traditional GCN + no centrality analysis;
[0147] Control group 1: using word2vec for node embedding + traditional GCN + centrality analysis;
[0148] Control group 2: CodeBERT for node embedding + traditional GCN + no centrality analysis;
[0149] Control group 3: word2vec node embedding + model combining adaptive attention convolution and GCN + no centrality analysis;
[0150] Control group 4: CodeBERT for node embedding + traditional GCN + centrality analysis;
[0151] Control group 5: CodeBERT for node embedding + model combining adaptive attention convolution and GCN + no centrality analysis;
[0152] Control group 6: word2vec node embedding + model combining adaptive attention convolution and GCN + centrality analysis;
[0153] Full version: Node embedding using CodeBERT + model combining adaptive attention convolution and GCN + with centrality analysis.
[0154] The experimental results are as follows Figure 5 shown.
[0155] In order to more intuitively show the impact of different component combinations on the final experimental performance, the ablation experiment results are visualized as a heat map, where each cell represents the difference between the performance indicator under a specific configuration and the baseline (such as accuracy, F1 score, loss, AUC, etc.). The horizontal axis represents various performance indicators, and the vertical axis represents the control group. Figure 6 shown.
[0156] From the heat map, we can observe that the four performance indicators are improved when using CodeBERT to embed graph nodes compared to Word2Vec. Analysis shows that Word2Vec is designed for word embedding tasks in the field of natural language processing, and compared with natural language, code language has more contextual semantic associations, including function calls, variable scopes, etc. Therefore, the model pre-trained on large-scale code language is obviously better than Word2Vec based only on word co-occurrence in embedding code graph nodes due to its strong context perception and generalization capabilities. At the same time, the change in embedding method has a relatively small performance improvement compared to the change in model. The reason may be that the complex edge feature information in CPG already contains certain structural semantics, thereby reducing the model's dependence on high-level semantic understanding.
[0157] At the same time, whether the global features of the graph can be added to the nodes through centrality analysis to complement the local features of the nodes and provide additional insights into graph features also plays a key role in improving performance. The presence or absence of centrality analysis also brought 1.5 and 1.5 improvements in accuracy and precision to the model respectively. Therefore, providing each node with its importance information in the entire graph, identifying the key nodes of the CPG graph, and highlighting the different importance of nodes in the graph structure are beneficial to the expression of the overall characteristics of the code. The performance improvement brought by centrality analysis is close to the improvement brought by the selection of different node embedding methods, which is much smaller than the improvement brought by model changes. This may be because the overall implementation of the neural network includes GCN, and GCN already has a certain ability to capture graph structure information. The information provided by centrality analysis may have some overlap with the features captured by the GCN model as the baseline method, resulting in limited performance improvement. Similarly, the performance improvement of the control group of the model combining adaptive attention convolution and GCN and centrality analysis is slightly lower than that of the other control groups that also use the two components. This may also be due to the fact that the long-distance dependency capture ability brought by the introduction of Transform partially overlaps with the global features brought by centrality analysis.
[0158] Finally, as can be seen from the heat map, the components of the model that combine adaptive attention convolution and GCN have the greatest impact on the overall performance. The accuracy and precision increased by 7.5 and 8.5 respectively. Therefore, it is effective to use the graph features to reflect the importance of different nodes and edges in the graph, and to use the adaptive attention allocation mechanism to perform different granularity analysis and feature extraction on the significant and insignificant areas in the graph. At the same time, it is proved that it is feasible to use the strong local structure analysis and feature extraction capabilities of GCN to supplement the Transformer's ability to capture global relationships but lack local structure analysis. The relative maximum improvement in performance brought about by the optimization of the model also shows that in the vulnerability detection or source code analysis tasks based on deep learning, the appropriate model design and implementation are the most important factors in all steps.
[0159] Finally, it should be noted that the above-described embodiments are only specific implementations of the present invention, which are used to illustrate the technical solutions of the present invention, rather than to limit them. The protection scope of the present invention is not limited thereto. Although the present invention is described in detail with reference to the above-described embodiments, ordinary technicians in the field should understand that any technician familiar with the technical field can still modify the technical solutions recorded in the above-described embodiments within the technical scope disclosed by the present invention, or can easily think of changes, or make equivalent replacements for some of the technical features therein; and these modifications, changes or replacements do not make the essence of the corresponding technical solutions deviate from the spirit and scope of the technical solutions of the embodiments of the present invention, and should be included in the protection scope of the present invention. Therefore, the protection scope of the present invention shall be subject to the protection scope of the claims.
Claims
1. A source code vulnerability detection method based on adaptive graph neural network, characterized in that: Include: Preprocess the source code; Extracting code property graph from normalized code using parser; Use CodeBERT to convert code attribute graphs into learnable vector representations; Calculate the centrality index of each node in the code attribute graph and concatenate it with the node feature vector; Combine the adaptive attention convolution module and the GCN module to learn the global and local features of the graph; The features learned by the adaptive attention convolution module and the GCN module are integrated into the fully connected layer to output the final vulnerability detection results.
2. The source code vulnerability detection method based on adaptive graph neural network according to claim 1 is characterized in that: Preprocessing the source code includes: removing irrelevant information in the source code; mapping the customized variable names and function names into a unified representation.
3. The source code vulnerability detection method based on adaptive graph neural network according to claim 1 is characterized in that: Using a parser to extract a code property graph from normalized code includes: first, using a parser to convert the source code into an abstract syntax tree AST, and constructing a control flow graph CFG based on the control flow information extracted from the AST; then analyzing the data dependency relationship in the code to generate a data dependency graph DDG; analyzing the function call relationship in the code to generate a call dependency graph CDG; and finally, fusing AST, CFG, DDG and CDG to obtain a code property graph.
4. The source code vulnerability detection method based on adaptive graph neural network according to claim 3 is characterized in that: Using CodeBERT to convert code attribute graphs into learnable vector representations includes: First, the preprocessed code line content is used as the node label, and the label of each node is extracted as a sentence. Then, the sentence is segmented using CodeBERT's tokenizer, and the sentence is decomposed into a list of tokens. Finally, the list of tokens is input into the CodeBERT model to obtain the vector representation of each node label. In addition, when vectorizing the edges of the code attribute graph, the edge features are directly encoded, and a four-dimensional vector is generated for each edge in the order of AST, CFG, DDG, and CDG. A value of 1 indicates that this type of edge exists between the corresponding two nodes.
5. The source code vulnerability detection method based on adaptive graph neural network according to claim 4 is characterized in that: Calculate the centrality index of each node in the code attribute graph and concatenate it with the node feature vector, including: The centrality index of each node in the code attribute graph is calculated using the eigenvector centrality function, Katz centrality function and betweenness centrality function respectively. The three results are concatenated to obtain the centrality measurement vector of the node; and the encoded node label eigenvector is connected with the centrality measurement vector to obtain the final node feature matrix.
6. The source code vulnerability detection method based on adaptive graph neural network according to claim 1 is characterized in that: The adaptive attention convolution module first linearly transforms the updated graph features of the input; then calculates the attention weights through the full_attention_conv function to obtain the weighted features; finally, the original input is added to the weighted features through the residual connection.
7. The source code vulnerability detection method based on adaptive graph neural network according to claim 1 is characterized in that: In the process of using the GCN module to propagate node features, multiple hidden layers are added to the middle layer for feature deepening. After each hidden layer receives the node features from the previous layer, it performs layer-by-layer convolution operations on the node features and the features of its neighboring nodes to generate new node feature representations.
8. The source code vulnerability detection method based on adaptive graph neural network according to claim 1 is characterized in that: The outputs of the adaptive attention convolution module and the GCN module are combined using weighted sum fusion or concatenated fusion. The weighted sum fusion formula is: x=ω×x GCN +(1-ω)×x 自适应注意力卷积 Splicing fusion formula: x=concat(x GCN ,x 自适应注意力卷积 ) Among them, ω is the fusion weight, x represents the concatenated feature vector, and x GCN represents the features extracted by the GCN module, x 自适应注意力卷积 It represents the features extracted by the adaptive attention convolution module, and concat is the concatenation function.
9. The source code vulnerability detection method based on adaptive graph neural network according to claim 1 is characterized in that: After fusing the outputs of the adaptive attention convolution module and the GCN module, all node information in the graph is integrated into a global feature vector of fixed length through global average pooling, and then binary classification is performed through a fully connected layer to determine whether there are vulnerabilities in the source code.
10. A source code vulnerability detection system based on adaptive graph neural network, characterized in that: Include: Code normalization module, used to preprocess source code; A CPG graph extraction module, used to extract a code property graph from the normalized code using a parser; A graph vector embedding module, which uses CodeBERT to convert code attribute graphs into learnable vector representations; The graph feature extraction module is used to calculate the centrality index of each node in the code attribute graph and concatenate it with the node feature vector; The graph neural network training module is used to combine the adaptive attention convolution module and the GCN module to learn the global and local features of the graph; The vulnerability detection module is used to fuse the features learned by the adaptive attention convolution module and the GCN module into the fully connected layer and output the final vulnerability detection results.
Citation Information
Cited By
Software supply chain risk monitoring and early warning method and device
CN119939607A
Software vulnerability detection method and device based on image and domain adaptation
CN120354419A
Code-level security vulnerability intelligent verification method, electronic equipment and storage medium
CN120671152A
Classification method and device based on graph neural network, computer equipment and medium
CN120707937A