Vulnerability detection method based on pre-trained code language model and convolutional neural network

By combining the pre-trained code language model GraphCodeBERT and convolutional neural networks, and utilizing the embedding vectors of data propagation chains and token sequences, a vulnerability detection model is constructed. This solves the problem of difficulty in capturing source code semantics and structural information in existing technologies, and achieves more efficient vulnerability detection.

CN119442246BActive Publication Date: 2025-10-10SOUTH CHINA UNIV OF TECH
View PDF 2 Cites 0 Cited by

Patent Information

Application Number
CN202411459493.5
Authority / Receiving Office
CN · China
Patent Type
Patents(China)
Current Assignee / Owner
Filing Date
2024-10-18
Publication Date
2025-10-10
Estimated Expiration
2044-10-18

AI Technical Summary

Technical Problem

Existing vulnerability detection methods are difficult to effectively capture the semantic and structural information of source code, resulting in poor vulnerability detection results.

Method used

A combined method based on the pre-trained code language model GraphCodeBERT and convolutional neural network is adopted. Through the embedding vector input of the data propagation chain and the code token sequence, combined with the pre-trained code language model and convolutional neural network, a vulnerability detection model is constructed to learn the semantic and structural information of the source code.

Benefits of technology

It improves the accuracy and effectiveness of vulnerability detection, enhances the generalization ability of the model, and can better capture key vulnerability features and filter out irrelevant features.

✦ Generated by Eureka AI based on patent content.

Smart Images

  • Figure CN119442246B_ABST
    Figure CN119442246B_ABST
Patent Text Reader

Abstract

The application relates to the technical field of vulnerability detection, and relates to a vulnerability detection method based on a pre-trained code language model and a convolutional neural network, which comprises the following steps: S1, obtaining a vulnerability data set, preprocessing the vulnerability data set, and constructing a model training data set; S2, obtaining a data propagation chain and a code token sequence according to source code of the model training data set, embedding the data propagation chain and the token sequence to obtain an embedding vector corresponding to the source code; S3, constructing a vulnerability detection model, training the vulnerability detection model to obtain a trained vulnerability detection model; and S4, outputting, by the trained vulnerability detection model, whether the source code to be detected has a vulnerability. The application constructs a vulnerability detection model based on a pre-trained code language model and a convolutional neural network, adopts a data propagation chain and a code token sequence as model input, can learn semantic information and structural information of source code, captures key vulnerability features, and improves the accuracy and effectiveness of vulnerability detection.
Need to check novelty before this filing date? Find Prior Art

Description

TECHNICAL FIELD

[0001] The present application relates to the technical field of vulnerability detection, and particularly relates to a vulnerability detection method based on a pre-trained code language model and a convolutional neural network. BACKGROUND

[0002] Vulnerability detection is a process of identifying and discovering potential security vulnerabilities in software systems, which can be exploited by malicious users to cause unauthorized access, data leakage or other security problems. The purpose of vulnerability detection is to find and fix these vulnerabilities before attackers exploit them, ensuring the security and integrity of software systems.

[0003] Software vulnerability detection has a wide range of applications in various industries and scenarios, from development to deployment and operation, and plays a key role in ensuring the security, compliance and reliability of software products and services. For example, in the early development stage of software, using vulnerability detection tools can detect potential vulnerabilities, help identify and fix vulnerabilities before code release, and reduce the cost of later repairs; in the system maintenance of financial institutions and e-commerce platforms, frequent vulnerability detection is required to ensure the security of software and the safe transmission and storage of sensitive data. Through continuous vulnerability detection, enterprises can effectively prevent security incidents and protect their own and user interests.

[0004] As the complexity of software increases, vulnerabilities in source code become more hidden and difficult to find. Traditional vulnerability detection methods sometimes fail to find deep code vulnerabilities. Therefore, people have begun to apply deep learning to vulnerability detection methods in order to seek better vulnerability detection results. Among these methods, sequence-based and graph-based methods have achieved some results, but they still have shortcomings and only perform moderately. Sequence-based methods focus on learning semantic information of source code, while ignoring rich structural information in source code. Graph-based methods can capture structural information in code well, but they often cannot distinguish subtle programming patterns from code structure graphs and extract long-distance dependent semantic information. In view of the shortcomings of these two types of methods, a vulnerability detection method that can learn semantic information and structural information of source code at the same time is needed to improve the effectiveness of vulnerability detection. SUMMARY

[0005] To solve the technical problems existing in the prior art, the present application provides a detection method based on a pre-trained code language model and a convolutional neural network, which integrates the advantages of pre-trained models and convolutional neural networks, so that the detection model can fully learn the semantic information and structural information of source code, capture key vulnerability features and filter irrelevant features, improve the accuracy of vulnerability detection, and improve the effectiveness of vulnerability detection.

[0006] The purpose of the present invention can be achieved by taking the following technical solutions:

[0007] The vulnerability detection method based on the pre-trained code language model and convolutional neural network includes the following steps:

[0008] S1. Obtain a vulnerability dataset, preprocess the vulnerability dataset, and construct a model training dataset;

[0009] S2. Obtain the data propagation chain and code token sequence based on the source code of the model training dataset, and embed the data propagation chain and token sequence to obtain the embedding vector corresponding to the source code;

[0010] S3. Build a vulnerability detection model based on the pre-trained code language model and convolutional neural network, and train the vulnerability detection model to obtain a trained vulnerability detection model;

[0011] S4. Preprocess the source code to be detected to obtain the embedding vector corresponding to the source code, and output whether the source code to be detected has a vulnerability through the trained vulnerability detection model.

[0012] Specifically, step S1 includes: obtaining an open source vulnerability dataset, generating new samples between minority class samples in the vulnerability dataset by interpolation, increasing the number of minority class samples, and after the vulnerability dataset reaches class balance, randomly shuffling the vulnerability dataset into a training set and a test set according to a preset ratio to construct a model training dataset.

[0013] Specifically, obtaining the data propagation chain and code token sequence based on the source code of the model training dataset includes:

[0014] Use the Tree-sitter tool to convert the source code into an abstract syntax tree, mark the location of the leaf nodes of the AST, and form a variable set V;

[0015] For each variable in the variable set V, a directed edge is constructed to represent the propagation relationship between the variables;

[0016] Filter out the variables with no edges flowing into them (variable nodes with in-degree 0) and record them as the initial node set;

[0017] Taking each variable in the initial node set InitVars as the initial node, the breadth-first search algorithm is executed according to the edge relationship set E to construct a data propagation chain starting from the variable;

[0018] The BPE serialization tool in the pre-trained code language model GraphCodeBERT is used to serialize the source code into a token sequence, and tokens composed of only a single character and non-variables in the token sequence are filtered.

[0019] Specifically, the embedding of the data propagation chain and the token sequence obtains an embedding vector corresponding to the source code, comprising:

[0020] The data propagation chain and the token sequence are connected into a long sequence I, and the pre-trained code language model GraphCodeBERT is used to embed the long sequence I, so as to convert the long sequence I into a vector to obtain an embedding vector corresponding to the embedding vector of the source code; the long sequence I = {[CLS], C, [SEP], V}, wherein C is the token sequence converted from the source code and its position in the token sequence, V is a set of all variables and edge relationship positions extracted from the data propagation chain, [CLS] is a special token placed at the beginning of the sequence, and [SEP] is a unique token used to separate two different types of data.

[0021] Specifically, the vulnerability detection model comprises a pre-trained code language model module, a convolutional neural network module, and a Sigmoid activation function.

[0022] The pre-trained code language model module is a pre-trained code language model GraphCodeBERT, which comprises 12 layers of Transformer Encoder layers, and each layer of the Transformer Encoder comprises a feedforward neural network and a masked multi-head attention mechanism in structure; the pre-trained code language model GraphCodeBERT uses the embedding vector of the data propagation chain and the code token sequence as input, and the embedding vector H 0 After passing through the 12 layers of Transformer Encoder layers, a code feature vector is output.

[0023] The convolutional neural network module comprises an input layer, a hidden layer, and an output layer, wherein the hidden layer comprises a two-dimensional convolutional layer, five Net_Block small network layers, a maximum pooling layer, and a fully connected layer, the two-dimensional convolutional layer, the five Net_Block small network layers, the maximum pooling layer, and the fully connected layer are connected in sequence, the input layer of the convolutional neural network module is connected to the output of the pre-trained code language model GraphCodeBERT, the output layer is connected to the fully connected layer and outputs a one-dimensional vector, and a prediction probability is output through a Sigmoid activation function.

[0024] Specifically, the convolution kernel size of the two-dimensional convolution layer is 3×512; each Net_Block small network layer consists of a maximum pooling layer and two identical one-dimensional convolution layers, and the convolution kernel sizes of the maximum pooling layer and the one-dimensional convolution layer are 4×1 and 3×1 respectively; the output of the 5-layer Net_Block passes through a maximum pooling layer, and the convolution kernel size of the maximum pooling layer is x×1, where x corresponds to the size of the third dimension of the tensor output by the Net_Block small network layer.

[0025] Specifically, the vulnerability detection model is trained to obtain a trained vulnerability detection model, including: training the vulnerability detection model using the embedding vectors of the code in the model training data set and the true labels corresponding to the source code, and the loss function in the training process is a binary cross entropy function, and the cross entropy function formula is:

[0026]

[0027] Among them, N is the number of samples during training, i is the subscript of each sample, and y i is the true label of the sample, is the predicted output of the vulnerability detection model.

[0028] Specifically, step S4 includes:

[0029] The data propagation chain and token sequence are extracted from the source code to be detected, and the data propagation chain and token sequence are embedded to obtain the embedding vector corresponding to the source code. The embedding vector corresponding to the source code is input into the trained vulnerability detection model; the prediction probability is output through the Sigmoid activation function of the vulnerability detection model; a threshold probability is set. If the output prediction probability is greater than the threshold probability, the source code is judged to have a vulnerability, otherwise the source code is judged to have no vulnerability.

[0030] Compared with the prior art, the present invention has the following advantages and beneficial effects:

[0031] The present invention provides a vulnerability detection method based on a pre-trained code language model and a convolutional neural network. By extracting a data propagation chain and a source code token sequence from the source code as code representation, the data propagation chain represents the structural information of the source code, and the source code token sequence represents the semantic information of the source code. The two are embedded together as input variables of the vulnerability detection model, which helps the vulnerability detection model learn more comprehensive source code information.

[0032] By constructing a vulnerability detection model, a pre-trained code language model GraphCodeBERT is used as the first module of the vulnerability detection model. The pre-trained code language model uses a graph-guided attention mask mechanism, which can enable the Transformer to learn semantic information while having the ability to learn structural information. In addition, since GraphCodeBERT is a pre-trained model, using it as part of the vulnerability detection model can improve the generalization ability of the vulnerability detection model. A convolutional neural network model is used as the second module of the vulnerability detection model, which can capture more context information and key vulnerability features from the hidden layer state output by the pre-trained code language model GraphCodeBERT, effectively improving the performance of the vulnerability detection model and improving the accuracy and effectiveness of vulnerability detection. BRIEF DESCRIPTION OF DRAWINGS

[0033] In order to more clearly illustrate the technical solutions of the embodiments of the present application or the prior art, the drawings needed in the embodiments or prior art description will be briefly introduced below. Obviously, the drawings in the following description are only some embodiments of the present application, and those skilled in the art can obtain other drawings from the structures shown in the drawings without creative labor.

[0034] Figure 1 is a flowchart of the vulnerability detection method based on the pre-trained code language model and the convolutional neural network in the embodiments of the present application;

[0035] Figure 2 is a schematic diagram of the flow of obtaining the data propagation chain in the embodiments of the present application;

[0036] Figure 3 is a schematic diagram of the overall network architecture of the vulnerability detection method in the embodiments of the present application;

[0037] Figure 4 is a schematic diagram of the network architecture of the pre-trained code language model module in the embodiments of the present application;

[0038] Figure 5 is a schematic diagram of the network architecture of the convolutional neural network module in the embodiments of the present application. DETAILED DESCRIPTION

[0039] The technical solutions of the present application will be described in further detail below in conjunction with the drawings and embodiments. Obviously, the described embodiments are only part of the embodiments of the present application, not all embodiments. The embodiments of the present application are not limited thereto. Based on the embodiments in the present application, all other embodiments obtained by those skilled in the art without creative labor are within the scope of protection of the present application.

[0040] Example 1:

[0041] The present invention discloses a vulnerability detection method based on a pre-trained code language model and a convolutional neural network. The method adopts a data propagation chain and a code token sequence as model inputs, and combines the pre-trained code language model and the convolutional neural network model, so that the detection model can learn the semantic and structural information of the source code, thereby improving the accuracy and effectiveness of vulnerability detection.

[0042] like Figure 1 As shown, the vulnerability detection method based on the pre-trained code language model and convolutional neural network proposed in the present invention includes the following steps:

[0043] S1. Obtain the vulnerability dataset, preprocess the vulnerability dataset, and build a model training dataset.

[0044] Specifically, an open source vulnerability dataset is obtained, and new samples are generated between the minority class samples in the vulnerability dataset through interpolation to increase the number of minority class samples. After the vulnerability dataset reaches class balance, the vulnerability dataset is randomly shuffled and divided into a training set and a test set according to a preset ratio to construct a model training dataset.

[0045] In this embodiment, the widely used open source dataset REVEAL can be selected. The REVEAL dataset is a vulnerability dataset collected in the real world rather than automatically synthesized. It traces the source code of two open source projects, Linux Debian Kernel and Chromium, and collects 18,169 function source codes, of which 16,505 are vulnerability-free and 1,664 are vulnerable. Because the dataset has the problem of class sample imbalance, the SMOTE algorithm is used to generate new samples between minority class samples by interpolation to increase the number of minority class samples. After the dataset reaches class balance, it is randomly shuffled and divided into a training set and a test set in a ratio of 7:3.

[0046] S2. Obtain the data propagation chain and code token sequence based on the source code, and embed the data propagation chain and token sequence to obtain the embedding vector corresponding to the source code.

[0047] S21. Extract the data propagation chain from the source code. Figure 2 As shown in the figure, a schematic diagram of the process of obtaining a data propagation chain in an embodiment of the present invention is shown. Given a source code for calculating the difference between x and y, the source code is converted into an abstract syntax tree using the Tree-sitter tool, the positions of the leaf nodes of the AST are marked, and a variable set V is formed, which is recorded as V = {x 1 ,y 2 ,diff 3 ,…,y12 ,diff 13}. Among them, the Tree-sitter tool is a cross-platform code parser generation tool and incremental parsing library written mainly in Rust. For each variable in the variable set V, a directed edge is constructed to represent the propagation relationship between variables, for example, ε = <v i ,v j > indicates that the jth variable comes from the ith variable or is calculated from the ith variable, and the edge relationship set is denoted as E = { <x 1 ,x 5 >, <x 1 ,x 9 >,…, <diff 10 ,diff 13 Then, we filter out the variables without inflow edges (variable nodes with in-degree 0), and record them as the initial node set InitVars = {x 1 ,y 2 ,0 4 Then, using each variable in the initial node set InitVars as an initial node, a breadth-first search algorithm is performed based on the edge relationship set E to construct a data propagation chain starting from that variable. In particular, if a variable is a discrete element and not connected to any variable node, it is not considered a data propagation chain. Once all variables in the initial node set InitVars have been searched, the data propagation chain of the source code has been successfully constructed.

[0048] like Figure 2 As shown, the data propagation chain is recorded as For each pc i (i∈{1,2,…,n pc}),pc i =(V′,E′), where and For example, with variable y 2 The data propagation chain for the starting node is:

[0049] ({y 2 ,y 8 ,y 6 ,…,diff 10 ,diff 13},{ <y 2 ,y 8 >, <y 8 ,diff 7 >, <diff 7 ,diff 13 <, <diff 10 ,diff 13 >}).

[0050] S22, convert the source code into a token sequence, tokenize the source code using the Byte Pair Encoding (BPE) serialization tool pre-trained in the pre-trained code language model GraphCodeBERT to obtain a token sequence, and the obtained token sequence is denoted as "token" refers to a basic unit in text, which can be a word, a character, or a word or phrase after tokenization. In order to avoid the interference of irrelevant data noise, filter the tokens in the token sequence that are composed of only single characters and are not variables. Taking the source code in Figure 2 as an example, the corresponding token sequence is C={int, foo, …, return, diff}.

[0051] S23, connect the data propagation chain and the token sequence into a long sequence I, embed the long sequence I through the pre-trained code language model GraphCodeBERT, and convert the long sequence I into a vector to obtain an embedding vector corresponding to the embedding vector of the source code. Wherein, the sequence I={[CLS], C, [SEP], V}, wherein C is the token sequence converted from the source code and its position in the token sequence, V is a set of all variables and edge relationship positions extracted from the data propagation chain, [CLS] is a special token placed at the beginning of the sequence, and [SEP] is a unique token used to separate two different types of data.

[0052] In this embodiment, taking the source code of Figure 2 as an example, the generated long sequence is:

[0053] The pre-trained code language model GraphCodeBERT is used to embed this sequence, and after embedding the sequence I, it is converted from text form to vector to obtain an embedding vector H=[[0.1554, -0.0277, 0.0685, …, 0.0083, 0.0076], …, [0.0414, 0.0451, -0.0232, …, 0.0191, 0.0240]]. For each token embedded, it and its position in the source code are embedded together and connected as its corresponding output vector. Each variable in the data propagation chain is represented by a special embedding marker to indicate that the variable is a node in the data propagation chain.

[0054] S3, construct a vulnerability detection model based on the pre-trained code language model and the convolutional neural network, train the vulnerability detection model, and obtain the trained vulnerability detection model.

[0055] like Figure 3 As shown in FIG, a schematic diagram of the overall network architecture of the vulnerability detection method is shown. The vulnerability detection model includes a pre-trained code language model module, a convolutional neural network module, and a Sigmoid activation function. After the source code is pre-processed in step S1, it is processed by the pre-trained code language model module, and then the Sigmoid activation function is used to transform the output to between 0 and 1. That is, the output of the model is a decimal between 0 and 1, which represents the probability of whether the source code has a vulnerability.

[0056] like Figure 4 As shown in FIG, the network architecture diagram of the pre-trained code language model module in the embodiment of the present invention is shown. The pre-trained code language model module is the pre-trained code language model GraphCodeBERT. The pre-trained code language model GraphCodeBERT includes 12 layers of Transformer Encoder (encoder) layers. The Transformer Encoder of each layer is structurally consistent, mainly including a feedforward neural network and a masked multi-head attention mechanism. The data propagation chain and the embedding vector of the code token sequence obtained in step S1 are used as input, denoted as the embedding vector H 0 Embedding vector H 0 It will pass through the 12-layer Transformer Encoder layer and output the hidden layer state, which is the code feature vector H 12 ,get:

[0057] H n =transformer n (H n-1 ),n∈[1,12]; (1)

[0058] Among them, this formula (1) represents the calculation of the nested 12-layer transformer encoder. Each layer of transformer encoder takes the output of the previous layer of transformer encoder as input, and uses the embedding vector H generated in step S2 0 For example, after calculation, the hidden layer state H 12 =[[-0.5043,0.2987,0.3640,…,0.2221,0.2459,0.3616],…,[-0.0141,0.1988,0.2927,…,0.0441,-0.1409,0.2995]]. H 12 Represents the code feature vector obtained by learning the input vector using the 12-layer transformer encoder.

[0059] Since the pre-training code language model GraphCodeBERT introduces a mask attention matrix based on graph guidance in the learning process of the 12-layer Transformer Encoder, the vulnerability detection model can learn the data flow structure information of the source code and does not introduce too complex code structure information, while avoiding the learning of irrelevant information.

[0060] As shown in Figure 5 the specific logical structure diagram of the convolutional neural network module, the convolutional neural network module includes an input layer, a hidden layer and an output layer, wherein the hidden layer includes a two-dimensional convolutional layer, five Net_Block small network layers, a maximum pooling layer and a fully connected layer, the two-dimensional convolutional layer, the five Net_Block small network layers, the maximum pooling layer and the fully connected layer are connected in sequence, the input layer of the convolutional neural network module is connected with the output of the pre-training code language model GraphCodeBERT, the output layer is connected with the fully connected layer and outputs a one-dimensional vector, and a prediction probability is output using a Sigmoid activation function. Specifically, the output H 12 As the input of the convolutional neural network module, the hidden layer state first passes through a two-dimensional convolutional layer to achieve data dimension reduction, and the convolutional kernel size of the two-dimensional convolutional layer is (3x512). Then it passes through five layers of Net_Block small network layers, each of which is composed of a maximum pooling layer and two one-dimensional convolutional layers with the same size, and the convolutional kernel sizes of the maximum pooling layer and the one-dimensional convolutional layer are (4x1) and (3x1), respectively. Multiple layers of Net_Block help more context semantic information to be integrated into each embedding position. While deepening the number of Net_Block layers, a skip connection is used to avoid model degradation caused by network layer deepening and gradient disappearance. In order to prevent model overfitting and be able to retain key features, the output of the five layers of Net_Block will also pass through another maximum pooling layer, and the convolutional kernel size of this maximum pooling layer is (x1), where x corresponds to the size of the third dimension of the tensor output by the Net_Block layer. Finally, a one-dimensional vector O is output through a fully connected layer, and a prediction probability of 0-1 is output using a Sigmoid activation function:

[0061]

[0062] In the subsequent step, a threshold T is set, for example, set to 0.5. If the prediction probability output by the model is greater than 0.5, it is judged as having a vulnerability, otherwise it is judged as not having a vulnerability.

[0063] ​Specifically, the vulnerability detection model is trained to obtain a trained vulnerability detection model, including:

[0064] According to the operation in step S1, the data set source code is converted into an embedding vector of the code, and the embedding vector of the code in the model training data set and the corresponding real label of the source code are used to train the vulnerability detection model. The loss function in the training process adopts a binary cross-entropy function, and the cross-entropy function formula is:

[0065]

[0066] Where N is the number of samples during training, i is the index of each sample, y i is the real label of the sample, is the predicted output of the vulnerability detection model.

[0067] In addition, since the training of a deep learning model requires a large amount of training data to make the model converge, a public vulnerability data set is generally selected for training, and subsequent fine-tuning is performed using the own data set, so that the model is more suitable for a specific type of vulnerability and a specific scene.

[0068] S4, pre-processing the source code to be detected to obtain an embedding vector corresponding to the source code, and outputting whether the source code to be detected has a vulnerability through the trained vulnerability detection model.

[0069] Specifically, for a source code to be detected, first, the method described in step S1 is used for data preprocessing, and the data propagation chain and token sequence are extracted from the source code to be detected. The embedding vector corresponding to the source code is obtained by embedding the data propagation chain and the token sequence, and then the embedding vector corresponding to the source code is input into the trained vulnerability detection model in step S3. The vulnerability detection model will finally output a prediction probability of 0-1 through the Sigmoid activation function, which represents the probability of whether the source code to be detected has a vulnerability. Then set a threshold probability, if the prediction probability output by the vulnerability detection model is greater than the threshold probability, it is judged that there is a vulnerability, otherwise there is no vulnerability. For example, the threshold probability is set to 0.5, and if the prediction probability is greater than 0.5, the source code to be detected is considered to be a code with a vulnerability.

[0070] In summary, the present invention extracts data propagation chains and token sequences from the source code, embeds them and inputs them into the detection model as code representation. The pre-trained code language model GraphCodeBERT is used to learn the rich code structure information and semantic information contained in the data propagation chain and token sequence, and ignores irrelevant redundant information, so that the detection model can learn more comprehensive code features. Since GraphCodeBERT has a pre-training basis for six programming languages, this also greatly enhances the generalization ability of the detection model. The convolutional neural network model provided by the present invention takes the hidden layer state of GraphCodeBERT as input, and enhances the ability of the vulnerability detection model to extract contextual information and key vulnerability features in the code. The vulnerability detection model proposed by the present invention orthogonally combines the pre-trained code language model GraphCodeBERT and the convolutional neural network model, combining the advantages of both, so that the vulnerability detection model has good performance and high accuracy, and provides strong technical support for detecting vulnerabilities in the software development and maintenance process.

[0071] The above embodiments are preferred implementation modes of the present invention, but the implementation modes of the present invention are not limited to the above embodiments. Any other changes, modifications, substitutions, combinations, and simplifications that do not deviate from the spirit and principles of the present invention should be considered as equivalent replacement methods and are included in the scope of protection of the present invention.

Claims

1. A vulnerability detection method based on a pre-trained code language model and a convolutional neural network, characterized in that: The following steps are involved: S1. Obtain a vulnerability dataset, preprocess the vulnerability dataset, and construct a model training dataset; S2. Obtain the data propagation chain and code token sequence based on the source code of the model training dataset, and embed the data propagation chain and token sequence to obtain the embedding vector corresponding to the source code; The method of obtaining the data propagation chain and code token sequence based on the source code of the model training dataset includes: Use the Tree-sitter tool to convert the source code into an abstract syntax tree, mark the positions of the leaf nodes of the AST, and form a variable set V; For each variable in the variable set V, a directed edge is constructed to represent the propagation relationship between the variables; Filter the variables without edge inflow and record them as the initial node set; Taking each variable in the initial node set as the initial node, the breadth-first search algorithm is executed according to the edge relationship set E to construct a data propagation chain starting from the variable; Use the BPE serialization tool in the pre-trained code language model GraphCodeBERT to tokenize the source code to obtain a token sequence, and filter out tokens in the token sequence that only consist of single characters and are non-variable; The embedding of the data propagation chain and the token sequence to obtain the embedding vector corresponding to the source code includes: Connect the data propagation chain and token sequence into a long sequence I , through the pre-trained code language model GraphCodeBERT for long sequences I Embed the long sequence I Convert into a vector to get the embedding vector corresponding to the source code of the embedding vector; long sequence 𝐼={[𝐶𝐿𝑆],𝐶,[𝑆𝐸𝑃],𝑉}, where, C It is the token sequence converted from the source code and its position in the token sequence. V It is the set of all variables and edge relationship positions in the extracted data propagation chain. [𝐶𝐿𝑆] is a special token placed at the beginning of the sequence. [𝑆𝐸𝑃] is a unique token used to separate two different types of data. S3. Build a vulnerability detection model based on the pre-trained code language model and convolutional neural network, train the vulnerability detection model, and obtain a trained vulnerability detection model; The vulnerability detection model includes a pre-trained code language model module, a convolutional neural network module and a Sigmoid activation function; The pre-trained code language model module is a pre-trained code language model GraphCodeBERT, which includes 12 layers of Transformer Encoder layers. The Transformer Encoder of each layer structurally includes a feedforward neural network and a masked multi-head attention mechanism. The pre-trained code language model GraphCodeBERT uses the embedding vector of the data propagation chain and the code token sequence as input, and outputs a code feature vector after the embedding vector passes through the 12 layers of Transformer Encoder layers. The convolutional neural network module includes an input layer, a hidden layer, and an output layer, wherein the hidden layer includes a two-dimensional convolutional layer, a 5-layer Net_Block small network layer, a maximum pooling layer, and a fully connected layer. The two-dimensional convolutional layer, the 5-layer Net_Block small network layer, the maximum pooling layer, and the fully connected layer are connected in sequence. The input layer of the convolutional neural network module is connected to the output of the pre-trained code language model GraphCodeBERT, and the output layer is connected to the fully connected layer and outputs a 1-dimensional vector, and the predicted probability is output through the Sigmoid activation function; The convolution kernel size of the two-dimensional convolution layer is 3 x 512; each Net_Block small network layer consists of a maximum pooling layer and two identical one-dimensional convolution layers, and the convolution kernel sizes of the maximum pooling layer and the one-dimensional convolution layer are 4 x 1 and 3 x 1 respectively; the output of the 5-layer Net_Block passes through a maximum pooling layer, and the convolution kernel size of the maximum pooling layer is 𝑥 x1, where 𝑥 corresponds to the size of the third dimension of the tensor output by the Net_Block small network layer; S4. Preprocess the source code to be detected to obtain the embedding vector corresponding to the source code, and output whether the source code to be detected has a vulnerability through the trained vulnerability detection model.

2. The vulnerability detection method based on a pre-trained code language model and a convolutional neural network according to claim 1, characterized in that: The step S1 includes: obtaining an open source vulnerability dataset, generating new samples between minority class samples in the vulnerability dataset by interpolation to increase the number of minority class samples, and after the vulnerability dataset reaches class balance, randomly shuffling the vulnerability dataset into a training set and a test set according to a preset ratio to construct a model training dataset.

3. The vulnerability detection method based on a pre-trained code language model and a convolutional neural network according to claim 1, characterized in that: The vulnerability detection model is trained to obtain a trained vulnerability detection model, including: training the vulnerability detection model using the embedding vectors of the code in the model training data set and the real labels corresponding to the source code, and the loss function in the training process is a cross entropy function of binary classification, and the cross entropy function formula is: ; in, N is the number of samples during training, 𝑖 is the subscript of each sample, and y 𝑖 is the true label of the sample, is the predicted output of the vulnerability detection model.

4. The vulnerability detection method based on a pre-trained code language model and a convolutional neural network according to claim 1, characterized in that: The step S4 comprises: The data propagation chain and token sequence are extracted from the source code to be detected, and the data propagation chain and token sequence are embedded to obtain the embedding vector corresponding to the source code. The embedding vector corresponding to the source code is input into the trained vulnerability detection model; the prediction probability is output through the Sigmoid activation function of the vulnerability detection model; a threshold probability is set. If the output prediction probability is greater than the threshold probability, the source code is judged to have a vulnerability, otherwise the source code is judged to have no vulnerability.

Citation Information

Patent Citations

  • Model construction method for software vulnerability mining and software vulnerability mining method

    CN117171757A

  • Transform attention mechanism and 2D-3D feature cross fusion-based drug response prediction model

    CN117912590A