Source code vulnerability detection method and device based on code language model and code attribute graph, electronic equipment and storage medium

Abstract

The invention discloses a code language model and code attribute graph-based source code vulnerability detection method, which comprises the following steps of: analyzing a to-be-detected source code function by using a tool to generate an initial code attribute graph, and converting a to-be-detected source code into a natural code sequence; associating the natural code sequence with the initial code attribute graph to obtain an associated code attribute graph; performing one-hot coding on the node type to obtain a node type code; inputting the associated code attribute graph into a CodeBERT model, and outputting a node content vector and a sequence representation of a function; the node type codes and the node content vectors are connected to obtain joint initial features of the nodes; inputting the joint initial features of the nodes into a GGNN model, and outputting graph-level features; and performing linear interpolation according to the sequence representation and the graph-level features of the function to obtain a vulnerability prediction result. According to the method, the code vulnerabilities can be detected more comprehensively and accurately, and the detection capability of the vulnerabilities in a complex code structure is improved.

Description

Technical Field

[0001] The present invention relates to the technical field of code vulnerability detection, and in particular to a source code vulnerability detection method, device, electronic device and storage medium based on a code language model and a code attribute graph. Background Art

[0002] Software vulnerability detection technology is a major concern in the software industry and network security. In recent years, popular deep learning-based code vulnerability detection methods can be roughly divided into two categories: sequence-based methods and graph-based methods.

[0003] Research based on deep sequence models generally slices and normalizes the original source code, and then parses it into a series of tokens. Subsequently, these tokens are converted into suitable vectors for neural network processing. Models based on RNN and Transformer are used to learn the context information sequence in the token and make the final defect prediction. Works such as RNN-based VulDeePecker and SySeVR introduce lexical analysis to turn the source code into more fine-grained code snippets. Transformer-based methods usually ignore code slicing and normalization strategies and choose to directly tokenize the source code.

[0004] In the prior art, although sequence models (such as RNN, LSTM, GRU and CNN, etc.) can learn the contextual information of the code when processing source code, they are obviously insufficient in capturing the program's hierarchy, execution flow, and data and control dependencies. Existing methods often use code sequence models or code graph models alone, failing to fully utilize the advantages of both, resulting in limited vulnerability detection performance. For example, a single sequence model may ignore code structure information, while a single graph model may lack a deep understanding of code semantics. Summary of the invention

[0005] In order to solve the above problems existing in the prior art, the present invention provides a source code vulnerability detection method, device, electronic device and storage medium based on a code language model and a code attribute graph. The technical problem to be solved by the present invention is achieved through the following technical solutions:

[0006] A first aspect of an embodiment of the present invention provides a source code vulnerability detection method based on a code language model and a code property graph, comprising the following steps:

[0007] Using the Joern tool to parse the source code function to be tested to generate an initial code property graph, and at the same time converting the source code to be tested into a natural code sequence;

[0008] Associating the natural code sequence with the initial code attribute graph to obtain an associated code attribute graph; wherein the node content in the associated code attribute graph includes the initial node content of the initial code attribute graph and the statement set of the natural code sequence;

[0009] Performing one-hot encoding on the node type of each node in the association code attribute graph to obtain a node type code;

[0010] Input the associated code attribute graph into the trained CodeBERT model, and output the sequence representation of node content vectors and functions;

[0011] The node type code and the node content vector are connected to obtain a joint initial feature of the node; wherein, during the training process of the CodeBERT model and the GGNN model, the CodeBERT model and the GGNN model are trained simultaneously;

[0012] Inputting the joint initial features of the nodes into the trained GGNN model, and outputting graph-level features;

[0013] Linear interpolation is performed according to the sequence representation of the function and the graph-level features to obtain a vulnerability prediction result.

[0014] In one embodiment of the present invention, the node type is encoded as:

[0015] T v =onehot(v.type())

[0016] Among them, T v represents the node type code, v represents the node in the associated code attribute graph, and v.type() represents the node type of each node in the associated code attribute graph.

[0017] In one embodiment of the present invention, the sequence of the node content vector and function is expressed as:

[0018] C v ,S v =CodeBERT(v.fragment(),C)

[0019] Among them, C v Represents the node content vector, S v Represents a sequence representation of a function.

[0020] 4. The method according to claim 3, characterized in that the joint initial features of the nodes are:

[0021] x v =concat(T v ,C v )

[0022] Among them, x v Represents the joint initial features of the nodes.

[0023] In one embodiment of the present invention, the step of inputting the joint initial features of the nodes into the trained GGNN model and outputting graph-level features includes:

[0024] Inputting the joint initial features of all nodes and the edge features of the associated code attribute graph into the trained GGNN model to obtain the hidden state representation of the node;

[0025] Aggregate the hidden state representations and output graph-level features.

[0026] In one embodiment of the present invention, the hidden state of the node represents for:

[0027]

[0028] Wherein, E represents the edge feature of the associated code attribute graph;

[0029] The graph-level feature x g for:

[0030]

[0031] In one embodiment of the present invention, it also includes:

[0032] During the training process of the CodeBERT model and the GGNN model, the parameters of the CodeBERT model and the GGNN model are optimized through the cross entropy loss function based on the output of the CodeBERT model and the output of the GGNN model.

[0033] A second aspect of an embodiment of the present invention provides a source code vulnerability detection device based on a code language model and a code attribute graph, comprising:

[0034] A first feature extraction module is used to parse the source code function to be tested using the Joern tool to generate an initial code attribute graph, and convert the source code to be tested into a natural code sequence;

[0035] A second feature extraction module is used to associate the natural code sequence with the initial code attribute graph to obtain an associated code attribute graph; wherein the node content in the associated code attribute graph includes the initial node content of the initial code attribute graph and the statement set of the natural code sequence;

[0036] An encoding module, used for performing one-hot encoding on the node type of each node in the associated code attribute graph to obtain a node type code;

[0037] A trained CodeBERT model is used to input the associated code attribute graph into the trained CodeBERT model, and output a sequence representation of a node content vector and a function;

[0038] A third feature extraction module is used to connect the node type code and the node content vector to obtain a joint initial feature of the node; wherein, during the training process of the CodeBERT model and the GGNN model, the CodeBERT model and the GGNN model are trained simultaneously;

[0039] The trained GGNN model is used to input the joint initial features of the nodes into the trained GGNN model to output graph-level features;

[0040] The prediction module is used to perform linear interpolation according to the sequence representation of the function and the graph-level features to obtain a vulnerability prediction result.

[0041] A third aspect of an embodiment of the present invention provides an electronic device, comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein when the processor executes the program, a source code vulnerability detection method based on a code language model and a code property graph provided by the first aspect of an embodiment of the present invention is implemented.

[0042] A fourth aspect of an embodiment of the present invention provides a computer-readable storage medium having a computer program stored thereon. When the computer program is executed by a processor, a source code vulnerability detection method based on a code language model and a code property graph provided by the first aspect of an embodiment of the present invention is implemented.

[0043] Beneficial effects of the present invention:

[0044] The present invention combines the pre-trained code language model CodeBERT with the code property graph (CPG). Through CPG, information such as AST, CFG and PDG are integrated to provide a unified and comprehensive representation for the code. At the same time, the pre-trained code model is used to extract local semantic features as CPG node embedding, thereby fusing the context information of the code sequence with the structural information in the code property graph, and detecting code vulnerabilities more comprehensively and accurately. At the same time, the GGNN with a gating mechanism is introduced. By assigning a gated recurrent unit (GRU) to each node, it can effectively integrate the information of adjacent nodes when updating node embedding, better aggregate syntax, control flow and data flow information, thereby accurately capturing the dependencies in the code property graph and improving the detection ability of vulnerabilities in complex code structures. The joint training method simultaneously optimizes the parameters of the pre-trained code model (such as CodeBERT) and the gated code GNN (such as GGNN), so that the model can simultaneously learn the semantic information of the code sequence and the structural information of the code property graph, thereby improving the vulnerability detection performance and giving full play to the complementary advantages of the two models.

[0045] Other features and advantages of the present invention will be described in the following description, and partly become apparent from the description, or understood by practicing the present invention. The purpose and other advantages of the present invention can be realized and obtained by the structures particularly pointed out in the written description, claims, and drawings.

[0046] The technical solution of the present invention is further described in detail below through the accompanying drawings and embodiments. BRIEF DESCRIPTION OF THE DRAWINGS

[0047] The accompanying drawings are used to provide a further understanding of the present invention and constitute a part of the specification. Together with the embodiments of the present invention, they are used to explain the present invention and do not constitute a limitation of the present invention. In the accompanying drawings:

[0048] Figure 1 A schematic diagram of a flow chart of a source code vulnerability detection method based on a code language model and a code attribute graph provided by an embodiment of the present invention;

[0049] Figure 2 It is a schematic diagram of a source code vulnerability detection device framework based on a code language model and a code attribute graph;

[0050] Figure 3 A schematic diagram of an initial code attribute graph for training source code provided by an embodiment of the present invention;

[0051] Figure 4 A schematic diagram of the joint initial feature extraction process of nodes and the embedding process of the GGNN model provided in an embodiment of the present invention;

[0052] Figure 5The present invention is a schematic diagram of a source code vulnerability detection device based on a code language model and a code attribute graph. DETAILED DESCRIPTION

[0053] The present invention is further described in detail below with reference to specific embodiments, but the embodiments of the present invention are not limited thereto.

[0054] like Figure 1 As shown, a first aspect of an embodiment of the present invention provides a source code vulnerability detection method based on a code language model and a code property graph, comprising the following steps:

[0055] Step 11, using the Joern tool to parse the source code function to be tested to generate an initial code property graph, and at the same time convert the source code to be tested into a natural code sequence.

[0056] Step 12, associate the natural code sequence with the initial code attribute graph to obtain an associated code attribute graph.

[0057] The node content in the associated code attribute graph includes the initial node content of the initial code attribute graph and the statement set of the natural code sequence.

[0058] Step 13, one-hot encode the node type of each node in the associated code attribute graph to obtain the node type code.

[0059] Step 14: Input the associated code attribute graph into the trained CodeBERT model and output the sequence representation of the node content vector and function.

[0060] Step 15, connect the node type code and the node content vector to obtain the joint initial feature of the node.

[0061] Among them, during the training process of the CodeBERT model and the GGNN model, the CodeBERT model and the GGNN model are trained at the same time.

[0062] Step 16: Input the joint initial features of the nodes into the trained GGNN model and output graph-level features.

[0063] Step 17: Perform linear interpolation based on the sequence representation of the function and the graph-level features to obtain the vulnerability prediction result.

[0064] In this embodiment, a code property graph (CPG) is introduced, and a pre-trained code language model is designed to be combined with CPG. CPG integrates the abstract syntax tree (AST), control flow graph (CFG), and program dependency graph (PDG), and merges various code attributes (including syntax, flow control, and data dependency) into a unified graph structure. The pre-trained code language model CodeBERT extracts local semantic features of the code and embeds them as nodes in CPG.

[0065] GGNN uses gating mechanisms (such as GRU units) to update node embeddings. Through message passing and aggregation mechanisms, each node can receive information from neighboring nodes and effectively capture the dependency information between nodes during the iteration process, including syntax, control flow, and data flow information. The strategy of jointly training CodeBERT and GGNN modules optimizes the parameters of both through the cross-entropy loss across code graph nodes, realizing the synergistic utilization of the contextual information of the code sequence (provided by CodeBERT) and the relational information in the code attribute graph (captured by GGNN). In addition, by using the pre-trained CodeBERT as an auxiliary classifier and linearly interpolating the two prediction results, more flexible information fusion is achieved. The interpolation parameters are adjusted according to different task requirements so that the model can balance the advantages of different models, further improve the model performance, and enhance the accuracy and generalization ability of vulnerability detection.

[0066] like Figure 2 As shown, a second aspect of an embodiment of the present invention provides a source code vulnerability detection method based on a code language model and a code property graph, comprising the following steps:

[0067] Step 21, start model training:

[0068] Step 211, using the Joern tool to parse the source code functions in the original data set to generate a training initial code attribute graph, and at the same time convert the training set source code into a training natural code sequence.

[0069] For each original dataset, the Joern tool is used to parse the source code into a code property graph (CPG). In this process, functions with graph generation errors are strictly excluded to ensure data quality. Then the dataset is divided into training set, validation set and test set according to experimental requirements. For example, when studying the dependence of the model on the scale of training data, training data is selected from the original dataset in different proportions (such as 1%, 5%, 10%, etc.), while ensuring the randomness and representativeness of the validation set and test set.

[0070] Among them, the original function-level source code is converted into a fixed-length feature vector containing semantic and syntactic structure information. This conversion prepares a suitable data format for the effective processing of the subsequent graph neural network model and code language model. To achieve this goal, two special code representation strategies are adopted for the graph neural network and code sequence language model in this embodiment.

[0071] For code graph representation: Use the open source code analysis tool Joern to parse the source code and generate a code property graph, such as Figure 3This code property graph provides a unified and concise representation that combines control flow and data flow with an abstract syntax tree and dependency graph. Functions with errors are strictly excluded during graph generation to ensure data quality.

[0072] Each node v in CPG has two attributes: node type and original code snippet. Node type information includes type identifiers such as IdentifierDeclType and keywords or operators such as int, char, for. For the representation of nodes, the function is decomposed into a statement set C = c1, c2, ..., c n , each c i Mapped to node v in CPG i .

[0073] For code sequence representation: follow the method proposed in Devign by converting function-level code into natural code sequence. This method aligns code serialization with the natural order of the source code, thereby preserving the logical order of the code.

[0074] Step 212: for the training set, the training natural code sequence is associated with the training initial code attribute graph to obtain a training associated code attribute graph.

[0075] The node contents in the training associated code attribute graph include the initial node contents of the training initial code attribute graph and the statement set of the training natural code sequence.

[0076] In the process of generating the code property graph, functions are transformed into comprehensive graphs containing various types of nodes (such as variables and function calls) and edges (including control flow and data flow), which convey different types of information. As the core of the code property graph, the AST (Abstract Syntax Tree) captures syntactic information and models the hierarchical function structure to outline the composition of the syntax and code. However, since the abstract syntax tree mainly provides static representation and lacks the ability to infer the dynamic behavior of the program, to solve this problem, the code property graph adds other types of edges to represent data flow and control flow, thereby enriching the graph and gaining insight into the execution context and dependencies between code segments. This integration builds a more comprehensive understanding of both the static structure and dynamic behavior of the program.

[0077] Step 213, one-hot encode the node type of each node in the training associated code attribute graph to obtain a training node type code.

[0078] Step 214, input the training association code attribute graph into the CodeBERT model, and output the sequence representation of the training node content vector and the training function.

[0079] Step 215, connect the training node type code and the training node content vector to obtain the joint initial features of the training node, which is the process of feature extraction, such as Figure 4 Schematic diagram of the node feature extraction process shown in .

[0080] Step 216, input the joint initial features of the training nodes into the GGNN model, and output the training graph-level features.

[0081] We use Gated Graph Neural Network (GGNN) to explore code property graphs, using its advanced capabilities to discern information flow patterns across nodes, thereby revealing structural insights related to code properties. Gated Graph Neural Network takes the feature vectors of all nodes along the graph edge as input. For a given embedded graph, Gated Graph Neural Network assigns a Gated Recurrent Unit (GRU) to each node. Figure 4 As shown in , this GRU updates the current vertex embedding by integrating the embeddings of all neighboring nodes. Finally, all vertex embeddings are aggregated into a vector to represent the entire code attribute graph.

[0082] The architecture of GGNN is mainly based on the graph structure, aiming to process and learn the information of nodes in the graph and the dependencies between nodes. i (V,X,A), where V represents the node set, X represents the node features, and A represents the adjacency relationship (i.e., the connection between nodes). GGNN creates a node v for each node. j ∈V assigns a gated recurrent unit (GRU) to implement information processing.

[0083] Neighborhood information aggregation: In GGNN, the update of a node depends on the information of its neighboring nodes. First, the connection relationship between nodes is defined by the adjacency matrix A. At each time step t, node v j It aggregates information from its neighboring nodes. Specifically, through the formula:

[0084]

[0085] To achieve information aggregation, is node v j The hidden state at time t-1 (initialized as node feature x j and fill with zeros), b is the bias vector, A (v,g) is the adjacency matrix associated with node v and graph g, which determines which nodes v are neighbors and the weights of information transfer. This process enables nodes to obtain contextual information about their neighborhood, thereby better understanding their position and role in the graph.

[0086] Gated recurrent unit (GRU) update: GRU is used to update the hidden state of the node. GRU is a special recurrent neural network (RNN) unit that can effectively process long sequence data and avoid the gradient disappearance or explosion problem in traditional RNN. The GRU in GGNN is updated through two gates (update gate and reset gate ) to control the flow and update of information.

[0087] The calculation formula of the update gate is:

[0088]

[0089] The calculation formula for the reset gate is:

[0090]

[0091] Where W z , U z , W r , U r is a learnable weight matrix, σ is the sigmoid activation function, and AGG is an aggregation function (the SUM function is often used in practical applications) used to integrate information from different neighbor nodes and different types of edges.

[0092] The candidate hidden state is calculated as:

[0093]

[0094] in represents element-wise multiplication. Finally, the hidden state of node v at time t is Update by formula:

[0095]

[0096] This formula determines how much old state information to retain and how much new candidate state information to accept based on the value of the update gate, thereby achieving gradual update of node status and information fusion.

[0097] Graph-level representation and prediction: After multiple time steps of information transmission and node updates, GGNN needs to summarize the information of the entire graph for final prediction.

[0098]

[0099] Sum the final hidden states of all nodes to get a vector representation of the entire graph This vector summarizes the information of all nodes in the graph and can reflect the overall characteristics and structural information of the graph.

[0100] Step 217, linear interpolation is performed according to the sequence representation of the training function and the training graph-level features to obtain the training vulnerability prediction result.

[0101] In this step, CodeBERT is used to generate node embeddings for the gated graph neural network, thereby implicitly combining CodeBERT and the gated graph neural network. Here, the advantages of pre-training and graph-based methods are further explicitly combined by utilizing interpolated predictions. Specifically, an auxiliary classifier is introduced, which obtains the prediction result of CodeBERT Softmax(W Embedding) by inputting the code embedding Embedding into a dense layer with softmax activation. Finally, the predictions of GGNN and CodeBERT are linearly interpolated, which is expressed as follows:

[0102] Pred=λPred GGNN +(1-λ)×Softmax(W·Embedding)

[0103] The parameter λ controls the trade-off between the two. λ=1 means that only the full Vul-LMGNN model is used, while λ=0 means that only the CodeBERT module is adjusted. When λ is in the range of (0,1), it allows for a balanced integration of predictions from both models. The fine-tuned CodeBERT model adjusts and optimizes the input graph of the gated graph neural network. The subsequent interpolated predictions promote the appropriate trade-off between the graph model and the sequence model detection, achieving excellent detection results. Code embedding is also the sequence representation of the training function.

[0104] Step 218, optimize the parameters of the CodeBERT model and the GGNN model through the cross entropy loss function based on the output of the CodeBERT model and the output of the GGNN model, and return to execute step 214 until the cross entropy loss function is minimized and the training is completed.

[0105] Among them, the expression of the cross entropy loss function L is:

[0106]

[0107] Among them, M represents the number of categories, y ic is a binary indicator (0 or 1) indicating whether the class label c is the correct classification of observation i. During this training process, CodeBERT updates the node embedding at each iteration, thereby gradually improving the complementary advantages of CodeBERT and the gated graph neural network GGNN. Here,

[0108]

[0109] α(·)=MAXPOOL(Relu(CONV(·)))

[0110] Among them, α(·) defines a layer of convolution kernel maximum pooling, which is used to activate the output features, x i Represents the training node content vector.

[0111] The code flow example of model training is as follows:

[0112]

[0113]

[0114] Step 22: After the trained model is obtained, it is verified and tested through the verification set and the test set. The verification and testing process is the same as the steps of step 213 to step 217.

[0115] Step 23, using the trained model to test the source code data set, the detection process is the same as the data processing process of steps 211 to 217 above, the only difference is whether the data processing object is the training set or the source code to be tested and whether it contains the real label of the training set, including the following steps:

[0116] Step 231: Use the Joern tool to parse the source code function to generate an initial code property graph, and convert the source code to be tested into a natural code sequence.

[0117] In this step, the method is the same as that in the above-mentioned model training step 211.

[0118] Step 232, associate the natural code sequence with the initial code attribute graph to obtain an associated code attribute graph.

[0119] The node content in the associated code attribute graph includes the initial node content of the initial code attribute graph and the statement set of the natural code sequence.

[0120] The method of obtaining the associated code attribute graph in this step is the same as that in the above step 212. Here, the associated code attribute graph is:

[0121] (V,E)=extract_code_property_graph(C)

[0122] Where, for each node v∈V, E is a set of edge features, where the node content of node v includes the initial node content of the initial code attribute graph and the statement set of the natural code sequence. The statement set of the function C=c1,c2,...,c n , each c i Mapped to node v in the associated code property graph i .

[0123] Step 233, perform one-hot encoding on the node type of each node in the associated code attribute graph to obtain the node type code:

[0124] T v =onehot(v.type())

[0125] Among them, T v Represents the node type encoding, v represents the node in the associated code attribute graph, and v.type() represents the node type of each node in the associated code attribute graph.

[0126] Step 234, input the associated code attribute graph into the trained CodeBERT model, and output the sequence representation of the node content vector and function:

[0127] C v ,S v =CodeBERT(v.fragment(),C)

[0128] Among them, C v Represents the node content vector, S v Represents the sequence representation of the function by the CodeBERT model.

[0129] Step 235, connect the node type code and the node content vector to obtain the joint initial feature x of the node v .here,

[0130] x v =concat(T v ,C v ).

[0131] Step 236, input the joint initial features of the nodes into the trained GGNN model, and output graph-level features.

[0132] Specifically, the joint initial features x of all nodes v Input into the gated graph neural network (GGNN) to get the hidden state representation of the node And aggregate these hidden states Get the graph-level feature representation: graph-level feature x g .

[0133] Here, the letter representation of each data is different from the letter in step 216, which is a simplified representation of the letter and processing process in step 216, but the corresponding meaning is the same, for example, wait.

[0134] Step 237, perform linear interpolation based on the sequence representation of the function and the graph-level features to obtain the vulnerability prediction result.

[0135] Specifically, the output S of the CodeBERT model is v Input the dense layer with softmax activation and get the prediction result Softmax(W·S v ), graph-level features x g As the prediction result of GGNN, the two perform linear interpolation prediction to obtain the final vulnerability prediction result.

[0136] Here, according to the above linear interpolation formula Pred = λPred GGNN +(1-λ)×Softmax(W·Embedding), the vulnerability prediction result is expressed as:

[0137] Pred=λx g +(1-λ)×Softmax(W·S v ).

[0138] The key points of the technical contents of the present invention are:

[0139] 1. Integrate multiple code representation methods

[0140] 1. Construct a code property graph (CPG): Integrate multiple code attributes such as the abstract syntax tree (AST), control flow graph (CFG), and program dependency graph (PDG) into a unified code property graph structure to fully capture information such as code syntax, control flow, and data dependencies, providing richer context for vulnerability detection. For example, in the code property graph of a function, nodes and edges can clearly present various structures and dependencies within the function, allowing the model to better understand the code logic.

[0141] 2. Combining code sequences with graph structures: By combining pre-trained code language models (such as CodeBERT) with code attribute graphs, the semantic information in the code sequences is retained, and the graph structure's ability to model code structures and dependencies is utilized, achieving complementary advantages between the two.

[0142] 2. Propose Gated Code GNN

[0143] 1. Effectively aggregate information: Use gating mechanisms (such as GRU units) to effectively aggregate information such as syntax, control flow, and data flow in the code attribute graph, so that nodes can selectively retain and fuse information from neighboring nodes when updating their states, thereby better capturing the complex dependencies and structural features in the code.

[0144] 2. Handling long-distance dependencies: Compared with traditional graph neural networks, gated graph neural networks can better handle long-distance dependencies between nodes in graph structures, overcoming the limitations of traditional graph neural networks in handling contextual relationships of distantly connected nodes, and helping to more comprehensively analyze the overall logic of the code.

[0145] 3. Joint training and interpolation prediction strategy

[0146] 1. Joint training mechanism: Innovatively jointly train CodeBERT and gated graph neural network modules, and optimize the parameters of both through cross-entropy loss across code graph nodes, so that the model can simultaneously learn the contextual information of the code sequence and the relational information in the graph structure, thereby improving the accuracy of vulnerability detection.

[0147] 2. Interpolation prediction method: An auxiliary classifier is introduced to perform linear interpolation on the prediction results of CodeBERT and gated graph neural network, and the contribution of the two in the final prediction is flexibly adjusted according to the interpolation parameter λ to further improve the model performance and realize the explicit fusion of the advantages of pre-trained models and graph-based methods.

[0148] The vulnerability detection method of the present invention has the following beneficial effects:

[0149] 1. More comprehensive code information capture

[0150] 1. Integrate multiple code attributes: Sequence-based methods in the prior art often ignore key code attributes such as syntax, control flow, and data dependency. Graph-based methods may underestimate code semantics and have difficulty capturing long-distance context information. The code property graph (CPG) constructed by the present invention integrates the abstract syntax tree (AST), control flow graph (CFG), and program dependency graph (PDG) to fully present the static structure and dynamic behavior of the code. For example, when analyzing a function, CPG can simultaneously display the syntax structure, execution order, and data dependency, allowing the model to have a deeper understanding of the code logic and provide rich information for vulnerability detection.

[0151] 2. Combining sequence and graph information: Extract local semantic features of the code through pre-trained code language models (such as CodeBERT) and embed them as CPG nodes, and then use the gated code graph neural network (Gated Code GNN) to process CPG to achieve the fusion of code sequence semantic information and graph structure information, overcome the limitations of a single method, such as when detecting vulnerabilities, it can comprehensively consider the code context and structural relationship to improve detection accuracy.

[0152] 2. Efficient model training and optimization

[0153] 1. Joint training strategy: A method for jointly training CodeBERT and Gated Code GNN modules is proposed. The parameters of both are optimized simultaneously through cross-entropy loss across code graph nodes, so that the model can collaboratively learn information in code sequences and graph structures, and enhance the ability to identify vulnerability patterns. For example, during the training process, CodeBERT continuously optimizes node embedding, enhances the complementary advantages with Gated Code GNN, and improves the overall model performance.

[0154] 2. Improved node embedding initialization: Use pre-trained CodeBERT to initialize CPG node embedding. Compared with the traditional method of training static word embedding models, it uses large-scale pre-training to obtain prior knowledge. By fine-tuning CodeBERT on the target dataset, it can better adapt to specific domain data and improve the accuracy of node vector representation, thereby improving the model's effectiveness in vulnerability detection.

[0155] 3. Flexible prediction and performance improvement

[0156] 1. Interpolation prediction mechanism: An auxiliary classifier is introduced to perform linear interpolation on the prediction results of Vul-LMGNN and CodeBERT. The contributions of the two are balanced by adjusting the interpolation parameter λ to achieve an explicit fusion of the advantages of the pre-trained model and the graph-based method. The model prediction strategy can be flexibly adjusted in different data sets and tasks to improve vulnerability detection performance.

[0157] 2. Significant performance advantages: Experiments on four real-world vulnerability datasets show that the Vul-LMGNN model of the present invention outperforms existing advanced methods in terms of accuracy, precision, recall and F1 score, especially on small-scale datasets, where the F1 score is improved by about 10%, proving the high efficiency and superiority of the model in actual vulnerability detection.

[0158] The effect of the invention is further described below through experimental data:

[0159]

[0160] Table 1. Performance metrics of various models on datasets with specific CWES

[0161] As can be seen from Table 1, overall, Vul-LMGNN (the present invention) shows obvious advantages on the dataset containing specific CWE, and its comprehensive performance is better than other comparison models. This is due to its unique architecture design. By combining the pre-trained code language model with the code attribute graph and using GGNN for information fusion, it can more effectively capture the semantic and structural information of the code, thereby improving the accuracy of vulnerability detection.

[0162] Transformer-based models also perform well on some data sets, but their performance is not stable enough. This may be because although they can learn some contextual information of the code, they have limitations in dealing with complex code structures and dependencies.

[0163] The GNN-based models TextGCN and Devign (AST) perform poorly when dealing with specific CWE vulnerability detection tasks, and may need to further optimize their graph structure representation and node feature extraction methods to better meet the needs of vulnerability detection.

[0164]

[0165] Table 2. Performance indicators of various models on datasets without specific CWES

[0166] As shown in Table 2, Vul-LMGNN also shows good performance on the dataset that does not distinguish specific CWE types, and it is better than other comparison models in terms of accuracy, precision, recall, and F1 score. This further proves the effectiveness and generalization ability of the Vul-LMGNN model, which can adapt to different types of vulnerability detection tasks.

[0167] The performance of Transformer-based models on these datasets varies. Some models perform well on some datasets, but overall they are not as stable and excellent as Vul-LMGNN. This may be because they are not adaptable enough to datasets with unclear vulnerability types and cannot fully utilize the information in the code to make accurate judgments.

[0168] The GNN-based model TextGCN and the CNN-based model TextCNN perform relatively poorly on datasets that do not distinguish between CWE types, indicating that they may not be able to effectively capture the key features in the code when processing such data, and need to improve their feature extraction and model architecture.

[0169] The Devign (AST) model still has a large gap with Vul-LMGNN on the data set that does not distinguish CWE types, indicating that there are deficiencies in its model design or data processing method, and further optimization is needed to improve vulnerability detection performance.

[0170] like Figure 5 As shown, the third aspect of the embodiment of the present invention provides a source code vulnerability detection device based on a code language model and a code attribute graph, comprising:

[0171] A first feature extraction module 51 is used to parse the source code function to be tested using the Joern tool to generate an initial code attribute graph, and convert the source code to be tested into a natural code sequence;

[0172] The second feature extraction module 52 is used to associate the natural code sequence with the initial code attribute graph to obtain an associated code attribute graph; wherein the node content in the associated code attribute graph includes the initial node content of the initial code attribute graph and the statement set of the natural code sequence;

[0173] The encoding module 53 is used to perform one-hot encoding on the node type of each node in the associated code attribute graph to obtain a node type code;

[0174] The trained CodeBERT model 54 is used to input the associated code attribute graph into the trained CodeBERT model and output a sequence representation of a node content vector and a function;

[0175] The third feature extraction module 55 is used to connect the node type code and the node content vector to obtain the joint initial feature of the node; wherein, during the training process of the CodeBERT model and the GGNN model, the CodeBERT model and the GGNN model are trained simultaneously;

[0176] The trained GGNN model 56 is used to input the joint initial features of the nodes into the trained GGNN model and output graph-level features;

[0177] The prediction module 57 is used to perform linear interpolation according to the sequence representation of the function and the graph-level features to obtain the vulnerability prediction result.

[0178] In one embodiment of the present invention, the node type is encoded as:

[0179] T v =onehot(v.type())

[0180] Among them, T v Represents the node type encoding, v represents the node in the associated code attribute graph, and v.type() represents the node type of each node in the associated code attribute graph.

[0181] In one embodiment of the present invention, the sequence of node content vectors and functions is represented as:

[0182] C v ,S v =CodeBERT(v.fragment(),C)

[0183] Among them, c v Represents the node content vector, S v Represents a sequence representation of a function.

[0184] In one embodiment of the present invention, the joint initial features of the nodes are:

[0185] xv =concat(T v ,C v )

[0186] Among them, x v Represents the joint initial features of the nodes.

[0187] In one embodiment of the present invention, the joint initial features of the nodes are input into the trained GGNN model, and the graph-level features are output, including:

[0188] The joint initial features of all nodes and the edge features of the associated code attribute graph are input into the trained GGNN model to obtain the hidden state representation of the node;

[0189] Aggregate hidden state representations and output graph-level features.

[0190] In one embodiment of the present invention, the hidden state of a node represents for:

[0191]

[0192] Among them, E represents the edge feature of the associated code attribute graph;

[0193] Graph-level features x g for:

[0194]

[0195] In one embodiment of the present invention, it also includes:

[0196] During the training process of the CodeBERT model and the GGNN model, the parameters of the CodeBERT model and the GGNN model are optimized through the cross entropy loss function based on the output of the CodeBERT model and the output of the GGNN model.

[0197] A fourth aspect of an embodiment of the present invention provides an electronic device, comprising a memory, a processor, and a computer program stored in the memory and executable on the processor. When the processor executes the program, a source code vulnerability detection method based on a code language model and a code property graph provided in the above-mentioned embodiment of the present invention is implemented.

[0198] The fifth aspect of an embodiment of the present invention further provides a computer-readable storage medium on which a computer program is stored. When the computer program is executed by a processor, the steps of a source code vulnerability detection method based on a code language model and a code property graph provided in the above-mentioned embodiment of the present invention are implemented.

[0199] The memory may include a random access memory (RAM) or a non-volatile memory (NVM), such as at least one disk memory. Optionally, the memory may also be at least one storage device located away from the aforementioned processor.

[0200] The above-mentioned processor can be a general-purpose processor, including a central processing unit (CPU), a network processor (NP), etc.; it can also be a digital signal processor (DSP), an application specific integrated circuit (ASIC), a field programmable gate array (FPGA) or other programmable logic devices, discrete gate or transistor logic devices, discrete hardware devices.

[0201] The method provided in the embodiment of the present invention can be applied to electronic devices. Specifically, the electronic device can be: a desktop computer, a portable computer, an intelligent mobile terminal, a server, etc. This is not limited here, and any electronic device that can implement the present invention belongs to the protection scope of the present invention.

[0202] As for the device / electronic device embodiment, since it is basically similar to the method embodiment, the description is relatively simple, and the relevant parts can be referred to the partial description of the method embodiment.

[0203] The present invention is described with reference to flowcharts and / or block diagrams of methods, devices (apparatus), and computer program products according to embodiments of the present invention. It should be understood that each process and / or block in the flowchart and / or block diagram, as well as the combination of processes and / or blocks in the flowchart and / or block diagram, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, a special-purpose computer, an embedded processor, or other programmable data processing device to generate a machine, so that the instructions executed by the processor of the computer or other programmable data processing device generate instructions for implementing the processes in the flowchart and / or block diagram. Figure 1 A process or multiple processes and / or boxes Figure 1 A device that provides the functions specified in a block or multiple blocks.

[0204] These computer program instructions may also be stored in a computer-readable memory capable of directing a computer or other programmable data processing device to operate in a specific manner, so that the instructions stored in the computer-readable memory produce an article of manufacture comprising an instruction device, which implements the process Figure 1 A process or multiple processes and / or boxes Figure 1 A function specified in one or more boxes.

[0205] These computer program instructions can also be loaded onto a computer or other programmable data processing device so that a series of operating steps are executed on the computer or other programmable device to produce a computer-implemented process, thereby providing instructions for implementing the process. Figure 1 A process or multiple processes and / or boxes Figure 1 The steps for the functions specified in one or more boxes.

[0206] Obviously, those skilled in the art can make various changes and modifications to the present invention without departing from the spirit and scope of the present invention. Thus, if these modifications and variations of the present invention fall within the scope of the claims of the present invention and their equivalents, the present invention is also intended to include these modifications and variations.

Claims

1. A source code vulnerability detection method based on a code language model and a code attribute graph, characterized in that: The following steps are involved: Use the Joem tool to parse the source code function to generate an initial code property graph, and convert the source code to be tested into a natural code sequence; Associating the natural code sequence with the initial code attribute graph to obtain an associated code attribute graph; wherein the node content in the associated code attribute graph includes the initial node content of the initial code attribute graph and the statement set of the natural code sequence; Performing one-hot encoding on the node type of each node in the association code attribute graph to obtain a node type code; Input the associated code attribute graph into the trained CodeBERT model, and output the sequence representation of node content vectors and functions; The node type code and the node content vector are connected to obtain a joint initial feature of the node; wherein, during the training process of the CodeBERT model and the GGNN model, the CodeBERT model and the GGNN model are trained simultaneously; Inputting the joint initial features of the nodes into the trained GGNN model, and outputting graph-level features; Linear interpolation is performed according to the sequence representation of the function and the graph-level features to obtain a vulnerability prediction result.

2. The method according to claim 1, characterized in that The node type is encoded as: T v =onehot(v.type()) Among them, T v represents the node type code, v represents the node in the associated code attribute graph, and v.type() represents the node type of each node in the associated code attribute graph.

3. The method according to claim 1, characterized in that The sequence of node content vectors and functions is represented as: C v ,S v =CodeBERT(v.fragment(),C) Among them, C v Represents the node content vector, S v Represents a sequence representation of a function.

4. The method according to claim 3, characterized in that The joint initial features of the nodes are: x v =concat(T v ,C v ) Among them, x v Represents the joint initial features of the nodes.

5. The method according to claim 1, characterized in that The step of inputting the joint initial features of the nodes into the trained GGNN model and outputting graph-level features includes: Inputting the joint initial features of all nodes and the edge features of the associated code attribute graph into the trained GGNN model to obtain the hidden state representation of the node; Aggregate the hidden state representations and output graph-level features.

6. The method according to claim 5, characterized in that The hidden state of the node is represented by for: Wherein, E represents the edge feature of the associated code attribute graph; The graph-level feature x g for:

7. The method according to claim 5, characterized in that Also includes: During the training process of the CodeBERT model and the GGNN model, the parameters of the CodeBERT model and the GGNN model are optimized through the cross entropy loss function based on the output of the CodeBERT model and the output of the GGNN model.

8. A source code vulnerability detection device based on a code language model and a code attribute graph, characterized in that: include: A first feature extraction module is used to parse the source code function to be tested using the Joern tool to generate an initial code attribute graph, and convert the source code to be tested into a natural code sequence; A second feature extraction module is used to associate the natural code sequence with the initial code attribute graph to obtain an associated code attribute graph; wherein the node content in the associated code attribute graph includes the initial node content of the initial code attribute graph and the statement set of the natural code sequence; An encoding module, used for performing one-hot encoding on the node type of each node in the associated code attribute graph to obtain a node type code; A trained CodeBERT model is used to input the associated code attribute graph into the trained CodeBERT model, and output a sequence representation of a node content vector and a function; A third feature extraction module is used to connect the node type code and the node content vector to obtain a joint initial feature of the node; wherein, during the training process of the CodeBERT model and the GGNN model, the CodeBERT model and the GGNN model are trained simultaneously; The trained GGNN model is used to input the joint initial features of the nodes into the trained GGNN model to output graph-level features; The prediction module is used to perform linear interpolation according to the sequence representation of the function and the graph-level features to obtain a vulnerability prediction result.

9. An electronic device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, characterized in that: When the processor executes the program, the source code vulnerability detection method based on the code language model and the code property graph is implemented as described in any one of claims 1 to 7.

10. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the source code vulnerability detection method based on a code language model and a code property graph according to any one of claims 1 to 7 is implemented.

Images