A vulnerability detection method, device and storage medium

Abstract

The present invention discloses a vulnerability detection method, device and storage medium, the method comprising: obtaining data to be detected; the data to be detected is assembly data included in a binary program file; using a preset detection model to identify the data to be detected, and determining vulnerability data in the data to be detected; wherein the preset detection model is obtained based on BERT neural network, graph neural network and classifier training; the BERT neural network is used to identify assembly instruction features in CFG; the graph neural network is used to identify graph structure features of the CFG; the classifier is used to identify vulnerability data based on assembly instruction features and graph structure features.

Description

Technical Field

[0001] The present invention relates to the field of computers, and in particular to a vulnerability detection method, device and storage medium. Background Art

[0002] Existing technical solutions for vulnerability detection include: source code review, fuzz testing, dynamic and static analysis, etc.

[0003] Source code review is a vulnerability detection technology that is performed after obtaining the software source code. Usually, security technology experts are required to check the software source code. Through source code review, memory leaks, stack overflows, and other security risks can be found.

[0004] Fuzz testing is a software testing technique that inputs automatically or semi-automatically generated random data into a program to monitor program anomalies, such as crashes, assertion failures, etc., to detect possible program errors.

[0005] Static analysis refers to scanning the program code without running the code, using lexical analysis, syntax analysis, control flow, data flow analysis and other techniques to find various defects in the program. Program static analysis can help software developers and quality assurance personnel find structural errors, security vulnerabilities and other problems in the code, thereby ensuring the overall quality of the software.

[0006] Dynamic analysis refers to the analysis of a program by executing it on a processor. In order for dynamic analysis to be effective, the target program must be executed with enough test inputs to cover almost all possible outputs.

[0007] The above solutions have different defects respectively, specifically:

[0008] Source code review: Security experts or developers are required to conduct a detailed review of the source code. The time developers spend on each function almost doubles, and for some very large projects, a detailed review of the source code is almost an impossible task.

[0009] Fuzz testing: Fuzz testing is currently the most widely used way to detect vulnerabilities, but its implementation is very slow, usually measured in months, and is limited to vulnerabilities that cause crashes.

[0010] Static analysis: Very slow, at least measured in days, requires security experts, is limited to given rules, and is difficult to scale.

[0011] Dynamic analysis: Some tools can be used to assist in execution, but it is usually difficult to implement and has limited application because covering all possible outputs is a very difficult task. Summary of the invention

[0012] In view of this, the main purpose of the present invention is to provide a vulnerability detection method, device and storage medium.

[0013] To achieve the above object, the technical solution of the present invention is achieved as follows:

[0014] An embodiment of the present invention provides a vulnerability detection method, the method comprising:

[0015] Acquire data to be detected; the data to be detected is assembly data included in a binary program file;

[0016] Using a preset detection model to identify the data to be detected, and determining vulnerability data in the data to be detected;

[0017] Wherein, the preset detection model is obtained based on BERT neural network, graph neural network and classifier training;

[0018] The BERT neural network is used to identify assembly instruction features in CFG;

[0019] The graph neural network is used to identify the graph structure features of the CFG;

[0020] The classifier is used to identify vulnerability data based on assembly instruction features and graph structure features.

[0021] Furthermore, the method also includes: training a BERT neural network, training a graph neural network, and training a classifier; or, the method includes: jointly training a BERT neural network, a graph neural network, and a classifier.

[0022] Furthermore, the training of the BERT neural network includes:

[0023] Acquire a training data set; the training data set includes: at least one CFG and a label corresponding to each CFG; the label indicates whether it is a vulnerability; the CFG includes: multiple CFG nodes and edges between CFG nodes; the CFG nodes are assembly instructions, and the edges are conversion instructions; the CFG is obtained based on using a disassembler to split the program into determined functions;

[0024] After converting the tokens of the assembly instructions in the CFG node into one-hot encoding, the tokens are input into the BERT neural network and trained using MLM as a task. The trained BERT neural network is then used as part of the preset detection model.

[0025] Among them, training the BERT neural network includes: initial training and secondary training;

[0026] The initial training includes: in the BERT neural network, the input layer randomly masks a preset proportion of inputs, and the output layer predicts the original inputs to perform initial training on the network;

[0027] The secondary training includes: after obtaining the features of each CFG node, performing the prediction task of the next CFG node, taking the upstream and downstream nodes together with the randomly constructed CFG nodes as input to predict whether the paired CFG nodes do have upstream and downstream relationships.

[0028] Furthermore, the training graph neural network includes:

[0029] Obtain a training data set; the training data set includes: at least one CFG and a label corresponding to each CFG; the label represents whether it is a vulnerability; wherein each CFG node has a defined node embedding; the embedding represents a vector;

[0030] According to each CFG in the training data set and the label corresponding to the CFG, the graph neural network is trained to obtain a trained graph neural network.

[0031] Furthermore, the training classifier includes:

[0032] According to the assembly instruction features, graph structure features and labels corresponding to the CFGs obtained by processing each CFG in the training data set, a fully connected neural network is input, the fully connected neural network is trained, and a trained fully connected neural network is obtained as a classifier.

[0033] Furthermore, the classifier adopts an encoder model in which multi-layer perceptron (MLP) layers are stacked, and a softmax activation function is used for classification training; the trained classifier can be used to determine the location of vulnerability data.

[0034] Further, after obtaining the data to be detected, the method further includes: converting the data to be detected into at least one CFG to be detected;

[0035] Accordingly, the using of a preset detection model to identify the data to be detected and determine the vulnerability data in the data to be detected includes:

[0036] The at least one CFG to be detected is identified using a preset detection model to determine vulnerability data in the data to be detected.

[0037] An embodiment of the present invention provides a vulnerability detection device, the device comprising:

[0038] An acquisition module, used for acquiring data to be detected; the data to be detected is assembly data included in a binary program file;

[0039] A detection module, used to identify the data to be detected by using a preset detection model, and determine the vulnerability data in the data to be detected;

[0040] Wherein, the preset detection model is obtained based on BERT neural network, graph neural network and classifier training;

[0041] The BERT neural network is used to identify assembly instruction features in CFG;

[0042] The graph neural network is used to identify the graph structure features of the CFG;

[0043] The classifier is used to identify vulnerability data based on assembly instruction features and graph structure features.

[0044] An embodiment of the present invention provides a vulnerability detection device, including a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor implements the steps of any one of the above vulnerability detection methods when executing the program.

[0045] An embodiment of the present invention further provides a computer-readable storage medium having a computer program stored thereon, and when the computer program is executed by a processor, the steps of any one of the above vulnerability detection methods are implemented.

[0046] A vulnerability detection method provided by an embodiment of the present invention includes: obtaining data to be detected; the data to be detected is assembly data included in a binary program file; using a preset detection model to identify the data to be detected, and determining vulnerability data in the data to be detected; wherein the preset detection model is obtained based on BERT neural network, graph neural network and classifier training; the BERT neural network is used to identify assembly instruction features in CFG; the graph neural network is used to identify graph structure features of the CFG; the classifier is used to identify vulnerability data based on assembly instruction features and graph structure features; in this way, fast, efficient and accurate vulnerability detection is achieved. BRIEF DESCRIPTION OF THE DRAWINGS

[0047] Figure 1 A schematic diagram of a vulnerability detection method provided by an embodiment of the present invention;

[0048] Figure 2 A schematic diagram of a vocabulary provided by an embodiment of the present invention;

[0049] Figure 3 A schematic diagram of a DUGNN model architecture provided by an embodiment of the present invention;

[0050] Figure 4 A schematic diagram of a vulnerability detection device provided by an embodiment of the present invention;

[0051] Figure 5 A schematic diagram of another vulnerability detection device provided by an embodiment of the present invention. DETAILED DESCRIPTION

[0052] The method provided by an embodiment of the present invention obtains data to be detected; the data to be detected is assembly data included in a binary program file; a preset detection model is used to identify the data to be detected, and vulnerability data in the data to be detected is determined; wherein the preset detection model is obtained based on BERT neural network, graph neural network and classifier training; the BERT neural network is used to identify assembly instruction features in CFG; the graph neural network is used to identify graph structure features of the CFG; the classifier is used to identify vulnerability data based on assembly instruction features and graph structure features.

[0053] The present invention will be further described in detail below in conjunction with the embodiments.

[0054] Figure 1 A schematic diagram of a vulnerability detection method provided by an embodiment of the present invention; Figure 1 As shown, it is applied to electronic devices such as computers, servers, etc.; the method includes:

[0055] Step 101, obtaining data to be detected; the data to be detected is the assembly data included in the binary program file;

[0056] Step 102: using a preset detection model to identify the data to be detected, and determining vulnerability data in the data to be detected;

[0057] Wherein, the preset detection model is obtained based on BERT neural network, graph neural network and classifier training;

[0058] The BERT neural network is used to identify assembly instruction features in CFG;

[0059] The graph neural network is used to identify the graph structure features of the CFG;

[0060] The classifier is used to identify vulnerability data based on assembly instruction features and graph structure features.

[0061] Thus, it can be identified whether the corresponding data to be detected is vulnerability data.

[0062] Specifically, the method also includes: training a BERT neural network, training a graph neural network, and training a classifier; or, the method includes: jointly training a BERT neural network, a graph neural network, and a classifier.

[0063] That is, the three can be trained at the same time, or they can be trained separately.

[0064] Specifically, the training of the BERT neural network includes:

[0065] Acquire a training data set; the training data set includes: at least one CFG and a label corresponding to each CFG; the label indicates whether it is a vulnerability; the CFG includes: multiple CFG nodes and edges between CFG nodes; the CFG nodes are assembly instructions, and the edges are conversion instructions; the CFG is obtained based on using a disassembler to split the program into determined functions;

[0066] After converting the tokens of the assembly instructions in the CFG node into one-hot encoding, they are input into the BERT neural network and trained using the Masked Language Modeling (MLM) as a task. The trained BERT neural network is used as part of the preset detection model.

[0067] Among them, training the BERT neural network includes: initial training and secondary training;

[0068] The initial training includes: in the BERT neural network, the input layer randomly masks a preset proportion of inputs, and the output layer predicts the original inputs to perform initial training on the network;

[0069] The secondary training includes: after obtaining the features of each CFG node, performing the prediction task of the next CFG node, taking the upstream and downstream nodes together with the randomly constructed CFG nodes as input to predict whether the paired CFG nodes do have upstream and downstream relationships.

[0070] Specifically, the training graph neural network includes:

[0071] Obtain a training data set; the training data set includes: at least one CFG and a label corresponding to each CFG; the label represents whether it is a vulnerability; wherein each CFG node has a defined node embedding;

[0072] According to each CFG in the training data set and the label corresponding to the CFG, the graph neural network is trained to obtain a trained graph neural network.

[0073] Here, the embedding represents a vector; that is, each CFG node corresponds to a defined node vector.

[0074] Specifically, the training classifier includes:

[0075] According to the assembly instruction features, graph structure features and labels corresponding to the CFGs identified from the training data set, a fully connected neural network is input to train the fully connected neural network to obtain a trained fully connected neural network as a classifier.

[0076] That is, a fully connected neural network is used as the classifier to be trained, and classification training is performed on the classifier to be trained based on the labeled vulnerability dataset;

[0077] The obtained classifier after classification training is used as the classifier in the detection model.

[0078] The classifier adopts an encoder model in which a multi-layer perceptron MLP layer is stacked, and a softmax activation function is used for classification training; the trained classifier can be used to determine the location of vulnerability data.

[0079] Specifically, after acquiring the data to be detected, the method further includes: converting the data to be detected into at least one CFG to be detected;

[0080] Accordingly, the using of a preset detection model to identify the data to be detected and determine the vulnerability data in the data to be detected includes:

[0081] The at least one CFG to be detected is identified using a preset detection model to determine vulnerability data in the data to be detected.

[0082] The following is a detailed description of the training data set:

[0083] The computer (subject) can obtain a binary program and disassemble it. The disassembler can be used to split the program into many functions, that is, a control flow graph (CFG) is created for each function.

[0084] CFG is a special graph in which consistent assembly instructions form nodes and conversion instructions are edges (that is, the CFG includes: multiple CFG nodes and edges between CFG nodes; the CFG nodes are consistent assembly instructions (hereinafter referred to as assembly instructions), and the edges are conversion instructions). Consistent assembly instructions are instructions that do not change the execution flow (for example, add, mov), while conversion instructions are instructions that change the execution flow (for example, jmp, jme, jne, call, ret).

[0085] Here, each function in the binary is represented by a CFG, and each CFG is assigned a label (there are two types of labels: vulnerability and no vulnerability). If there is a vulnerability, its corresponding Common Weakness Enumeration (CWE) label is marked. The training data set consists of various open source programs with known vulnerabilities. The accuracy of the data set can also be verified by fuzz testing. Among them, there are at least 100 examples in each CWE label category.

[0086] The CFG includes a plurality of CFG nodes, each of which is an assembly instruction; that is, the nodes in the CFG are assembly instructions, and then the assembly instructions are split to obtain a plurality of tokens;

[0087] That is to say, the assembly instructions can be split into smaller units, and this step can be called tokenization. Each smaller unit is called a token. All the instructions corresponding to the assembler programs in the node need to be tokenized. Assume that the instruction has exactly k (k=3) operands, the extra operands will be pruned, and the missing operands will be filled with empty tokens. One operand is represented by m (m=13) tokens.

[0088] The following is an explanation of the graph neural network:

[0089] The neural network should be able to learn assembly instructions before learning binary vulnerabilities. In order to properly learn assembly instructions in the deep learning model, preprocessing is required; preprocessing includes:

[0090] A vocabulary is constructed for the assembly instructions by a suitable morphological analyzer, which extracts library functions, opcodes, registers and hexadecimal values ​​from the assembly instructions and treats them as words.

[0091] Assembly instructions usually have a fixed syntax that uses 1 opcode and 3 operands. The tokenizer represents each operand as 13 fields. Therefore, the tokenized instruction length is 40 (3 operands, each represented as 13 fields, and one field for the operator).

[0092] Here, the set of all possible tokens is called a vocabulary (e.g. Figure 2), let the vocabulary size be Vt; the number of words in this vocabulary is much smaller than the number of word sets used in general NLP. Since assembly instructions usually have a fixed syntax, which means that it uses one opcode and three operands, a vector of fixed dimension size can be created to represent assembly instructions. When applied, the assembly instructions are read by the parser and an appropriate vector representation is constructed for the syntax. Each opcode and operand is represented by a vector to facilitate subsequent processing. This process is called "instruction embedding".

[0093] What needs to be explained is that in neural networks, embedding is a special kind of vector, so embedding is definitely a vector, but a vector is not necessarily an embedding. The embedding obtained by the instruction embedding process is a vector.

[0094] The library functions, opcodes, registers, and hexadecimal values ​​are described below:

[0095] Library function is a way to put functions into a library for others to use. The method is to compile some commonly used functions and put them into a file for different people to call. When calling, just add the file name of the function using #include<>. It is usually put into the lib file.

[0096] The opcode refers to the part of the instruction or field (usually represented by code) specified in the computer program to perform the operation. In fact, it is the assembly instruction number, which is used to tell the CPU which instruction needs to be executed.

[0097] The function of a register is to store binary code. It is composed of flip-flops with storage function. One flip-flop can store 1 bit of binary code, so a register that stores n bits of binary code needs to be composed of n flip-flops.

[0098] It is represented by numbers 0 to 9 and letters A to F (or a to f), where A to F represents 10 to 15, which are called hexadecimal numbers.

[0099] About Instruction Embedding. In order to obtain the embedding of tokenized instructions, the word2vec model is usually used; for example, the skip-gram version of the word2vec model is used, that is, given an input token, the context token is predicted. The model is trained on the assembler instruction dataset of the token in an unsupervised manner in the training dataset part.

[0100] The next step is to create token embeddings. The input to this step is the tokenization instructions. After training this model, a real-valued vector is obtained for each token, and these real-valued vectors are called embeddings; the above creation process is called "token embedding".

[0101] One-hot encoding is used to vectorize instructions. Here, one-hot encoding is a method of quantizing categorical data. The length of the vector generated by this method is equal to the size of the vocabulary Vt. If the data belongs to the i-th category, the i-th bit is marked as 1, and the other bits are 0. Thus, the tokenized instructions can be converted into vectors.

[0102] In order to obtain the most effective representation for detecting CWE, the BERT (Bidirectional Encoder Representations from Transformers) model is used to implement token embedding. Specifically, BERT is trained in an unsupervised manner (no labeled data is required) using MLM (Masked Language Model) and NSP (Next Sentence Prediction) tasks, where a sentence is a line in an assembly function obtained from a binary file, and the "unique hot" representation of the token is fed into the model to get embedded:

[0103]

[0104] v = BERT (v hot )∈R n×1

[0105] Where n=510 is the instruction embedding size, v is the instruction embedding, and Vt is the vocabulary size. As a result, a real-valued vector is obtained for each token.

[0106] Use BERT to learn the language features of assembly instructions in CFG nodes.

[0107] The specific scheme is as follows: convert the tokens of the assembly instructions in the CFG node into one-hot encoding, input it into the BERT neural network, and use MLM as the task.

[0108] In the BERT neural network, the input layer randomly masks a portion of the input (30%), and the output layer predicts the original input to perform initial training on the network. After obtaining the language features of the assembly instructions for each CFG node, the next node prediction task is performed, which is to put the upstream and downstream nodes together with some randomly constructed nodes as input to predict whether the paired nodes do have upstream and downstream relationships. The neural network is trained twice using this task. In this way, an encoder for instruction embedding can be obtained. This encoder accepts assembly instructions and inputs instruction embedding.

[0109] After that, the next step is to create a graph embedding. The following is an explanation of graph embedding:

[0110] In this step, a CFG graph is taken as input, where each CFG node has a defined node embedding. The output of this step is the embedding of the entire CFG graph. This embedding is further used to classify the entire graph (each graph is a function) as vulnerable or not vulnerable, so a graph encoder is used to generate this embedding.

[0111] The graph encoder is generated based on the DUGNN model. First, the node embedding is calculated; then, the node embedding and the graph adjacency matrix are fed into the DUGNN to obtain the graph embedding matrix; and the attention layer is used to convert it into a graph embedding vector.

[0112] like Figure 3 The following is a brief DUGNN model architecture. X-node feature matrix and A-adjencyMatrix represent input data, and Graph embedding represents output data. Attention represents the neural network layer. GCNIN and Pooling represent that the following blocks are repeated multiple times. Arrows represent data and operation flows.

[0113] The DUGNN model is defined as follows ( Figure 3 ). The DUGNN model uses graph convolutional layers, which we call gconv layers here, and M is the number of gconv layers in the model. The gconv layer is defined according to equation (1). Each gconv layer consists of N convolutional layers called gfc, which are defined in a nested manner in equations (2) and (3).

[0114]

[0115]

[0116]

[0117] We will use MLP to represent fully-connected layers.

[0118] To increase the model complexity of the graph encoder, high-order statistics of features are captured during the graph convolution (GCN) operation, similar to the graph capsule network. In addition to increasing the model complexity, there is also the problem of Laplacian smoothing in the GCN model. To overcome the smoothing problem, it is proposed to cascade the output of the intermediate encoding layer of the graph convolution (GCN) and feed it into the subsequent layers. In this way, the graph encoder can learn multi-scale smooth features of the nodes, which can avoid the problem of under-smoothing or over-smoothing in GCN-based models. In addition, it also has the side benefits of alleviating the vanishing gradient problem and enhancing feature propagation in deep networks.

[0119] To obtain graph embedding, we need to obtain the CFG node representation in the CFG graph as described above. The first m node instructions ("tokenized instructions") are vectorized using one-hot encoding into vectors:

[0120]

[0121] This formula represents the one-hot encoding of the tokenized instruction i to obtain a vector.

[0122]

[0123] The tokenized instructions in the same node are connected (concat) to obtain the instruction vector Vhot.

[0124] Then, apply the BERT model to get node embeddings:

[0125] X=BERT(v hot )

[0126] X∈R m×1

[0127] Here, m=510 is the size of the embedding. Hence, X is the embedding of a node in the graph.

[0128] In the above formula, k represents the node embedding size;

[0129] M represents the number of gconv layers in DUGNN;

[0130] N represents the number of gfc in the gconv layer;

[0131] P represents the polynomial exponent in the gconv layer;

[0132] X represents graph node embedding;

[0133] A represents the graph adjacency matrix;

[0134] n represents the number of nodes in the graph;

[0135] l represents the graph embedding size;

[0136] R represents the real number space;

[0137] X represents the resulting node embedding.

[0138] Below we will use g(A) to represent the graph filter function with adjacency matrix A and identity matrix I:

[0139] g(A)=A+I

[0140] The graph embedding vgraph is defined as the output of the last convolution. The convolution is defined as F:

[0141] υ graph =F M (X, A)∈R l×n

[0142] During a convolution operation, the output of an intermediate layer is fed into subsequent layers:

[0143] F l (X, A) = concat(F l-1 (X, A), selu(g conv(F l-1 (X, A), A)))

[0144] Each convolution gconv layer is applied to X 1 , X 2 , ...X P , gfc is computed for all possible combinations of i and p, where i∈{1,...,N} and p∈{1,...,P}.

[0145]

[0146] The gfc layer is defined as:

[0147]

[0148]

[0149] Using the attention layer MLP attention Put v graph ∈R l×|nodes| Convert to v′ graph ∈R l ×1 :

[0150] w attention =softmax(MLP attention (v graph ))

[0151] w attention ∈R n×1 ,

[0152] Among them, w attention is the attention score.

[0153] Sum the attention scores of the nodes to get the graph vector embedding:

[0154] v′ graph =sum(v graph ·w attention )∈R l×1

[0155] Further explanation about the classifier is as follows.

[0156] A multi-layer perceptron (MLP) layer is stacked on the encoder model, and a softmax activation function is used for classification:

[0157]

[0158] Select nodes that are more likely to be vulnerability locations based on the attention score:

[0159] node=argmax(|w attention |)

[0160] Argmax is to find the node with the largest absolute value of the attention score.

[0161] The description of classification training is as follows:

[0162] Considering that the loss in the model consists of two parts: CELoss (cross entropy loss) of the label and BCELoss (binary cross entropy loss) of the attention layer.

[0163] If the node is where the vulnerability occurs, BCELoss is applied to A one-hot vector where 1 appears at the index of the vulnerability location and 0 at other index locations. Otherwise,

[0164] is a position where all Vector.

[0165] That is, the following loss function is used for classification training:

[0166]

[0167] Among them, class_frac and attention_frac correspond to the scores of CELoss and BCELoss losses respectively. and They are the labels and predictions of nodes with vulnerabilities, respectively.

[0168] In summary, in the method provided in the embodiment of the present invention, the training of a preset detection model (i.e., a neural network) includes the following steps:

[0169] First, train the BERT model on unlabeled assembly instructions;

[0170] Then, train the DUGNN model on a huge training dataset whose labels may not be accurate enough;

[0171] Next, continue training it on a small training dataset with accurate labels;

[0172] Finally, the model is checked on the Juliet dataset (synthetic data not suitable for training).

[0173] The network structure and weights of the trained neural assembly language encoder and graph neural network encoder are fixed (the weights are not updated during the back propagation process) as the backbone network encoder of CFG; then, a fully connected neural network is connected as a classifier, and classification training is performed based on the labeled vulnerability data set. During the training process, the parameters of the neural network are updated to train the neural network. The trained network can be used to detect vulnerabilities.

[0174] Figure 4 A schematic diagram of the structure of a vulnerability detection device provided by an embodiment of the present invention; Figure 4 As shown, the device comprises:

[0175] An acquisition module, used for acquiring data to be detected; the data to be detected is assembly data included in a binary program file;

[0176] A detection module, used to identify the data to be detected by using a preset detection model, and determine the vulnerability data in the data to be detected;

[0177] Wherein, the preset detection model is obtained based on BERT neural network, graph neural network and classifier training;

[0178] The BERT neural network is used to identify assembly instruction features in CFG;

[0179] The graph neural network is used to identify the graph structure features of the CFG;

[0180] The classifier is used to identify vulnerability data based on assembly instruction features and graph structure features.

[0181] Specifically, the device also includes: a preprocessing module, used for training a BERT neural network, training a graph neural network, and training a classifier; or, used for jointly training a BERT neural network, a graph neural network, and a classifier.

[0182] Specifically, the preprocessing module is used to obtain a training data set; the training data set includes: at least one CFG and a label corresponding to each CFG; the label represents whether it is a vulnerability; the CFG includes: multiple CFG nodes and edges between CFG nodes; the CFG nodes are assembly instructions, and the edges are conversion instructions; the CFG is obtained based on using a disassembler to split the program into determined functions;

[0183] After converting the tokens of the assembly instructions in the CFG node into one-hot encoding, the tokens are input into the BERT neural network and trained using MLM as a task. The trained BERT neural network is then used as part of the preset detection model.

[0184] Among them, training the BERT neural network includes: initial training and secondary training;

[0185] The initial training includes: in the BERT neural network, the input layer randomly masks a preset proportion of inputs, and the output layer predicts the original inputs to perform initial training on the network;

[0186] The secondary training includes: after obtaining the features of each CFG node, performing the prediction task of the next CFG node, taking the upstream and downstream nodes together with the randomly constructed CFG nodes as input to predict whether the paired CFG nodes do have upstream and downstream relationships.

[0187] Specifically, the preprocessing module is used to obtain a training data set; the training data set includes: at least one CFG and a label corresponding to each CFG; the label represents whether it is a vulnerability; wherein each CFG node has a defined node embedding; the embedding represents a vector;

[0188] According to each CFG in the training data set and the label corresponding to the CFG, the graph neural network is trained to obtain a trained graph neural network.

[0189] Specifically, the preprocessing module is used to input a fully connected neural network according to the assembly instruction features, graph structure features and labels corresponding to the CFGs obtained by processing each CFG in the training data set, train the fully connected neural network, and obtain the trained fully connected neural network as a classifier.

[0190] Specifically, the classifier adopts an encoder model in which a multi-layer perceptron MLP layer is stacked, and a softmax activation function is used for classification training; the trained classifier can be used to determine the location of vulnerability data.

[0191] Specifically, the acquisition module is used to convert the data to be detected into at least one CFG to be detected after acquiring the data to be detected;

[0192] Correspondingly, the detection module is used to use a preset detection model to identify the at least one CFG to be detected and determine the vulnerability data in the data to be detected.

[0193] It should be noted that: the vulnerability detection device provided in the above embodiment only uses the division of the above program modules as an example when implementing the corresponding vulnerability detection method. In actual applications, the above processing can be assigned to different program modules as needed, that is, the internal structure of the electronic device is divided into different program modules to complete all or part of the processing described above. In addition, the device provided in the above embodiment and the embodiment of the corresponding method belong to the same concept, and its specific implementation process is detailed in the method embodiment, which will not be repeated here.

[0194] Figure 5 A schematic diagram of the structure of a communication device provided by an embodiment of the present invention, such as Figure 5 As shown, the communication device 50 includes: a processor 501 and a memory 502 for storing a computer program that can be run on the processor; when the processor 501 is used to run the computer program, it executes: obtaining data to be detected; the data to be detected is assembly data included in a binary program file; using a preset detection model to identify the data to be detected, and determining vulnerability data in the data to be detected; wherein the preset detection model is obtained based on BERT neural network, graph neural network and classifier training; the BERT neural network is used to identify assembly instruction features in CFG; the graph neural network is used to identify graph structure features of the CFG; the classifier is used to identify vulnerability data based on assembly instruction features and graph structure features.

[0195] The relevant operations have been described in the above embodiments and will not be repeated here.

[0196] In actual application, the communication device 50 may further include: at least one network interface 503. The various components in the communication device 50 are coupled together via a bus system 504. It is understood that the bus system 504 is used to realize the connection and communication between these components. In addition to the data bus, the bus system 504 also includes a power bus, a control bus, and a status signal bus. However, for the sake of clarity, Figure 5 In the figure, various buses are labeled as bus system 504. There may be at least one processor 501. The network interface 503 is used for wired or wireless communication between the communication device 50 and other devices. The memory 502 is used to store various types of data to support the operation of the communication device 50.

[0197] An embodiment of the present invention also provides a computer-readable storage medium having a computer program stored thereon; when the computer program is executed by a processor, the following steps are performed: obtaining data to be detected; the data to be detected is assembly data included in a binary program file; using a preset detection model to identify the data to be detected, and determining vulnerability data in the data to be detected; wherein the preset detection model is obtained by training based on a BERT neural network, a graph neural network and a classifier; the BERT neural network is used to identify assembly instruction features in a CFG; the graph neural network is used to identify graph structure features of the CFG; the classifier is used to identify vulnerability data based on assembly instruction features and graph structure features.

[0198] The relevant operations have been described in the above embodiments and will not be repeated here.

[0199] The above is only a specific embodiment of the present invention, but the protection scope of the present invention is not limited thereto. Any person skilled in the art can easily think of changes or substitutions within the technical scope disclosed by the present invention, which should be included in the protection scope of the present invention. Therefore, the protection scope of the present invention should be based on the protection scope of the claims.

Claims

1. A vulnerability detection method, characterized in that: The method comprises: Acquire data to be detected; the data to be detected is assembly data included in a binary program file; Using a preset detection model to identify the data to be detected, and determining vulnerability data in the data to be detected; Wherein, the preset detection model is obtained based on BERT neural network, graph neural network and classifier training; The BERT neural network is used to identify assembly instruction features in CFG; The graph neural network is used to identify the graph structure features of the CFG; The classifier is used to identify vulnerability data based on assembly instruction features and graph structure features; The method further includes: training a BERT neural network, training a graph neural network, and training a classifier; or, the method includes: jointly training a BERT neural network, a graph neural network, and a classifier; The training of the BERT neural network includes: Acquire a training data set; the training data set includes: at least one CFG and a label corresponding to each CFG; the label indicates whether it is a vulnerability; the CFG includes: multiple CFG nodes and edges between CFG nodes; the CFG nodes are assembly instructions, and the edges are conversion instructions; the CFG is obtained based on using a disassembler to split the program into determined functions; After converting the tokens of the assembly instructions in the CFG node into one-hot encoding, the tokens are input into the BERT neural network and trained using MLM as a task. The trained BERT neural network is then used as part of the preset detection model. Among them, training the BERT neural network includes: initial training and secondary training; The initial training includes: in the BERT neural network, the input layer randomly masks a preset proportion of inputs, and the output layer predicts the original inputs to perform initial training on the network; The secondary training includes: after obtaining the features of each CFG node, performing the prediction task of the next CFG node, taking the upstream and downstream nodes together with the randomly constructed CFG nodes as input to predict whether the paired CFG nodes do have upstream and downstream relationships.

2. The method according to claim 1, characterized in that The training graph neural network comprises: Obtain a training data set; the training data set includes: at least one CFG and a label corresponding to each CFG; the label represents whether it is a vulnerability; wherein each CFG node has a defined node embedding; the embedding represents a vector; According to each CFG in the training data set and the label corresponding to the CFG, the graph neural network is trained to obtain a trained graph neural network.

3. The method according to claim 2, characterized in that The training classifier comprises: According to the assembly instruction features, graph structure features and labels corresponding to the CFGs obtained by processing each CFG in the training data set, a fully connected neural network is input, the fully connected neural network is trained, and a trained fully connected neural network is obtained as a classifier.

4. The method according to claim 3, characterized in that The classifier adopts an encoder model in which a multi-layer perceptron MLP layer is stacked, and a softmax activation function is used for classification training; the trained classifier can be used to determine the location of vulnerability data.

5. The method according to claim 1, characterized in that After acquiring the data to be detected, the method further includes: converting the data to be detected into at least one CFG to be detected; Accordingly, the using of a preset detection model to identify the data to be detected and determine the vulnerability data in the data to be detected includes: The at least one CFG to be detected is identified using a preset detection model to determine vulnerability data in the data to be detected.

6. A vulnerability detection device, characterized in that: The device comprises: An acquisition module, used for acquiring data to be detected; the data to be detected is assembly data included in a binary program file; A detection module, used to identify the data to be detected by using a preset detection model, and determine the vulnerability data in the data to be detected; Wherein, the preset detection model is obtained based on BERT neural network, graph neural network and classifier training; The BERT neural network is used to identify assembly instruction features in CFG; The graph neural network is used to identify the graph structure features of the CFG; The classifier is used to identify vulnerability data based on assembly instruction features and graph structure features; Also includes: training a BERT neural network, training a graph neural network, and training a classifier; or, includes: jointly training a BERT neural network, a graph neural network, and a classifier; The training of the BERT neural network includes: Acquire a training data set; the training data set includes: at least one CFG and a label corresponding to each CFG; the label indicates whether it is a vulnerability; the CFG includes: multiple CFG nodes and edges between CFG nodes; the CFG nodes are assembly instructions, and the edges are conversion instructions; the CFG is obtained based on using a disassembler to split the program into determined functions; After converting the tokens of the assembly instructions in the CFG node into one-hot encoding, the tokens are input into the BERT neural network and trained using MLM as a task. The trained BERT neural network is then used as part of the preset detection model. Among them, training the BERT neural network includes: initial training and secondary training; The initial training includes: in the BERT neural network, the input layer randomly masks a preset proportion of inputs, and the output layer predicts the original inputs to perform initial training on the network; The secondary training includes: after obtaining the features of each CFG node, performing the prediction task of the next CFG node, taking the upstream and downstream nodes together with the randomly constructed CFG nodes as input to predict whether the paired CFG nodes do have upstream and downstream relationships.

7. A vulnerability detection 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 steps of the method according to any one of claims 1 to 5 are implemented.

8. A computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 5 are implemented.

Images