Vulnerability mining method and device based on multi-task learning

Abstract

The application discloses a vulnerability mining method and device based on multi-task learning, and the method comprises the following steps: based on the abstract syntax tree, the control flow graph and the program dependency graph of the source code, a code attribute graph is constructed; the tree structure and the graph structure in the code attribute graph are separated; the nodes in the tree structure are encoded to obtain the initial node vector representation of the abstract syntax tree; the initial node vector representation of the abstract syntax tree, the tree structure and the graph structure are respectively input into a vulnerability recognition neural network model and a vulnerability positioning neural network model to obtain a vulnerability recognition result and a vulnerability positioning result. In order to solve the problems that the accuracy of the existing learning-based source code vulnerability mining method needs to be improved and the detection granularity is coarse, the application proposes a vulnerability mining method based on multi-task learning. The application can accurately identify the vulnerability and position the vulnerability at the same time, and improve the security vulnerability mining ability of the vulnerability mining personnel.

Description

Technical Field

[0001] This invention belongs to the field of computer technology and relates to a vulnerability discovery method and apparatus based on multi-task learning. Background Technology

[0002] In recent years, software security trends have become increasingly severe, making the assurance of software security ever more crucial. Vulnerabilities, as one of the key factors affecting software security, deserve sufficient attention. According to data released by the U.S. National Vulnerability Database, the number of vulnerabilities released annually was less than 8,000 before 2017, while from 2017 to 2020, the number of vulnerabilities released annually exceeded 14,000, and the trend continues to rise. At the same time, software is becoming increasingly large and complex, with code volumes reaching millions of pages, and the vulnerabilities exhibiting diverse and complex characteristics, making manual identification of vulnerabilities impractical. In this context, many studies have attempted to explore automated and effective vulnerability discovery methods.

[0003] Source code vulnerability discovery is a common automated vulnerability discovery method based on static analysis, which directly discovers vulnerabilities based on information in the source code. Existing source code vulnerability discovery methods can be mainly divided into code similarity-based methods and pattern-based methods. The former discovers vulnerabilities by comparing known vulnerable code, and can only find reproducible vulnerabilities caused by code cloning. The latter characterizes vulnerabilities based on vulnerability patterns defined by experts or learned, thereby achieving vulnerability discovery. Pattern-based methods can be further subdivided into rule-based methods and learning-based methods. The former relies on experts to define vulnerability rules and then matches these rules in different ways to discover vulnerabilities. The latter uses machine learning techniques to learn vulnerability patterns from data, thereby predicting vulnerabilities in unknown code based on these patterns, and thus discovering vulnerabilities. Because learning-based methods can effectively solve the problems of relying on expert knowledge and difficulty in detecting zero-day vulnerabilities, this type of method has been widely studied. However, existing learning-based source code vulnerability discovery methods have certain limitations in code semantic modeling and detection granularity, which prevents them from simultaneously meeting the requirements of high accuracy and high practicality. On the one hand, existing methods do not comprehensively model the semantic information of the code, making it difficult to cover the semantic features of different types of vulnerabilities, resulting in decreased accuracy when facing different types of vulnerabilities. On the other hand, existing methods have a coarse detection granularity and cannot provide clear vulnerability locations, making subsequent manual verification difficult and requiring improvement in practicality. These problems prevent existing learning-based source code vulnerability discovery methods from effectively uncovering vulnerabilities. Summary of the Invention

[0004] The purpose of this invention is to provide a vulnerability discovery method and apparatus based on multi-task learning. This method models and learns various semantic information of source code by constructing an encoder, calculates the semantic vector representation of source code, and jointly trains the two tasks of vulnerability identification and vulnerability localization by sharing encoder parameters. This enables accurate identification and localization of vulnerabilities, thereby improving the security vulnerability discovery capabilities of vulnerability researchers.

[0005] To achieve the above objectives, the present invention adopts the following technical solution:

[0006] A vulnerability discovery method based on multi-task learning, comprising the following steps:

[0007] Construct a code property graph based on the abstract syntax tree, control flow graph, and program dependency graph of the source code;

[0008] Separate the tree structure from the graph structure in the code attribute graph;

[0009] Encode the nodes in the tree structure to obtain the initial node vector representation of the abstract syntax tree;

[0010] The initial node vector representation, tree structure, and graph structure of the abstract syntax tree are input into the vulnerability identification neural network model to obtain the vulnerability identification result; wherein, the vulnerability identification neural network model includes: a first code semantic encoder and a first classification module;

[0011] The first code semantic encoder is used to calculate the source code vector representation based on the initial node vector representation, tree structure, and graph structure of the abstract syntax tree;

[0012] The first classification module is used to calculate the vulnerability identification result based on the source code vector representation;

[0013] The initial node vector representation, tree structure, and graph structure of the abstract syntax tree are input into the vulnerability localization neural network model to obtain the vulnerability localization result; wherein, the vulnerability localization neural network model includes: a second code semantic encoder, an attention layer, and a second classification module;

[0014] The second code semantic encoder is used to compute the final vector representation of nodes in the graph structure based on the initial node vector representation of the abstract syntax tree, the tree structure, and the graph structure;

[0015] The attention layer is used to assign attention weights to the final vector representation of each node in the graph structure and to perform attention calculations.

[0016] The second classification module is used to calculate the vulnerability location result based on the attention calculation result.

[0017] Furthermore, the code property graph is constructed through the following steps:

[0018] 1) Generate the attribute graph of the abstract syntax tree. , where the node set Nodes in For nodes and edge sets of an abstract syntax tree The edge in For the edges of the abstract syntax tree, the function Side Marked as an edge of the abstract syntax tree, function For nodes Assign type attributes to code statements and code attributes to nodes;

[0019] 2) Generate the attribute graph of the control flow graph. , where the node set Nodes in An edge set is a node representing a statement and a predicate in an abstract syntax tree. The edge in Corresponding to the jump edges in the control flow graph, the function For the edge Assign jump condition flags;

[0020] 3) Generate the attribute graph of the program dependency graph. , where the node set Nodes in For nodes and edge sets of an abstract syntax tree The edge in Corresponding to the dependency edges in the program dependency graph, the function For the edge Assign dependency tags, which include: control dependencies or data dependencies, functions For each edge marked as a data dependency Indicates the corresponding symbol on which it depends, or for each edge marked as a control dependency. Indicates the state of the predicate that controls the dependencies;

[0021] 4) Combined Attribute Diagram Attribute diagram With attribute graph Obtain the code property map ,in , , , .

[0022] Furthermore, the separated code attribute graph , among which , , The set of nodes in the abstract syntax tree representing the i-th statement. Let represent the set of edges in the abstract syntax tree of the i-th statement. and Let i be the root node of the abstract syntax tree for the i-th statement. It equals the number of statements in the source code.

[0023] Furthermore, the encoding of nodes in the tree structure includes:

[0024] 1) Use the PACE algorithm to vectorize the code statement type attributes to obtain the code statement type representation;

[0025] 2) Use the word2vec algorithm to vectorize the node code attributes to obtain the node code representation;

[0026] 3) Concatenate the code statement type representation and the node code representation to obtain the initial node vector representation of the abstract syntax tree.

[0027] Further, the source code vector representation is calculated through the following steps:

[0028] 1) Update the vector representations of all nodes in the tree structure based on the initial vector representation of the abstract syntax tree;

[0029] 2) Aggregate the vector representations of all nodes in the updated tree structure to obtain the statement vector representation;

[0030] 3) Initialize the vector representation of the nodes in the graph structure according to the statement vector representation, and update the vector representation of the nodes in the graph structure using graph-based convolution to obtain the final vector representation of the nodes in the graph structure;

[0031] 4) The final vector representation of the nodes in the aggregate graph structure is used to obtain the vector representation of the source code.

[0032] Furthermore, the first classification module includes a fully connected layer and a softmax layer, wherein the fully connected layer is used to perform linear transformation and nonlinear mapping on the vector representation of the program source code, and the softmax layer is used to perform binary classification based on the calculation results of the fully connected layer to obtain vulnerability identification results.

[0033] Furthermore, the vulnerability identification neural network model and the vulnerability localization neural network model are trained through the following steps:

[0034] 1) The parameters of the tree embedding module, tree pooling module, and graph embedding module in the first code semantic encoder and the second code semantic encoder are shared;

[0035] 2) Cross-entropy loss is used as the loss for vulnerability identification task and vulnerability location task, respectively, and the average value of the loss for vulnerability identification task and vulnerability location task is defined as the joint loss;

[0036] 3) Use the Adam optimizer to optimize the joint loss and train the vulnerability identification neural network model and the vulnerability localization neural network model.

[0037] A vulnerability discovery device based on multi-task learning includes:

[0038] Build modules are used to construct code property graphs based on source code abstract syntax trees, control flow graphs, and program dependency graphs;

[0039] The separation module is used to separate the tree structure from the graph structure in the code property graph;

[0040] The encoding module is used to encode the nodes in the tree structure to obtain the initial node vector representation of the abstract syntax tree;

[0041] The identification module is used to input the initial node vector representation, tree structure, and graph structure of the abstract syntax tree into the vulnerability identification neural network model to obtain the vulnerability identification result; wherein, the vulnerability identification neural network model includes: a first code semantic encoder and a first classification module;

[0042] The first code semantic encoder is used to calculate the source code vector representation based on the initial node vector representation, tree structure, and graph structure of the abstract syntax tree;

[0043] The first classification module is used to calculate the vulnerability identification result based on the source code vector representation;

[0044] The localization module is used to input the initial node vector representation, tree structure, and graph structure of the abstract syntax tree into the vulnerability localization neural network model to obtain the vulnerability localization result; wherein, the vulnerability localization neural network model includes: a second code semantic encoder, an attention layer, and a second classification module;

[0045] The second code semantic encoder is used to compute the final vector representation of nodes in the graph structure based on the initial node vector representation of the abstract syntax tree, the tree structure, and the graph structure;

[0046] The attention layer is used to assign attention weights to the final vector representation of each node in the graph structure and to perform attention calculations.

[0047] The second classification module is used to calculate the vulnerability location result based on the attention calculation result.

[0048] A storage medium storing a computer program, wherein the computer program is configured to execute the method described above at runtime.

[0049] An electronic device includes a memory and a processor, wherein the memory stores a program for performing the methods described above.

[0050] Compared with existing technologies, this invention addresses the shortcomings of current learning-based source code vulnerability discovery methods, which suffer from low accuracy and coarse detection granularity. It proposes a multi-task learning-based vulnerability discovery method. This invention constructs an encoder to model and learn various semantic information from the source code, calculates the semantic vector representation of the source code, and jointly trains two tasks—vulnerability identification and vulnerability localization—while sharing encoder parameters. This enables accurate vulnerability identification and localization, thereby improving the security vulnerability discovery capabilities of vulnerability hunters. Attached Figure Description

[0051] Figure 1 This is a flowchart of a vulnerability discovery method based on multi-task learning.

[0052] Figure 2 This is a schematic diagram of the final generated code attribute graph.

[0053] Figure 3 This is a structural diagram of a vulnerability identification neural network model.

[0054] Figure 4 This is a structural diagram of a neural network model for vulnerability localization. Detailed Implementation

[0055] The technical solutions in the embodiments of the present invention will now be clearly and completely described with reference to the accompanying drawings. Obviously, the described embodiments are merely specific embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0056] The vulnerability discovery method of the present invention has the following overall process: Figure 1 As shown, the main steps include:

[0057] Generate a Code Property Graph (CPG) from the source code and preprocess it to generate a code property graph with separate tree and graph structures, such as... Figure 2 As shown, CPG is a combined data structure that integrates an Abstract Syntax Tree (AST), a Control Flow Graph (CFG), and a Program Dependency Graph (PDG). The AST reflects the syntactic structure within a statement, the CFG reflects the control flow between statements, and the PDG reflects the data and control dependencies between statements. CPG can effectively cover the semantic information required for vulnerability detection.

[0058] The features of the AST nodes are encoded to obtain the initial vector representation of the AST nodes. The initial vector representation is a vector used to represent the semantic information of the nodes. This initial vector representation contains two fields: a type field and a code field. The type field is vectorized using PACE, and the code field is vectorized using word2vec. Then, the vectorized encodings of the two fields are concatenated to obtain the final initial vector representation.

[0059] Build a neural network model for vulnerability identification, such as Figure 2 As shown, the vulnerability identification neural network model consists of a code semantic encoder, a fully connected layer, and a softmax layer. The code semantic encoder comprises a tree embedding module, a tree pooling module, a graph embedding module, and a graph pooling module. The tree embedding and tree pooling modules learn the vector representation of the AST in the CPG to obtain a statement-level vector representation, which is used to initialize the vector representations of the CFG and PDG. The graph embedding and graph pooling modules learn the vector representations of the CFG and PDG in the CPG to obtain the vector representation of the source code. The fully connected layer and softmax layer then perform a binary classification of whether the source code contains vulnerabilities based on the vector representation of the source code.

[0060] Construct a neural network model for vulnerability localization, such as Figure 3 As shown, the vulnerability localization neural network model consists of a code semantic encoder, an attention layer, a fully connected layer, and a softmax layer. The code semantic encoder differs from the one in 3) in that it lacks a graph pooling module. The attention layer learns different attention weights for different nodes. The fully connected layer and softmax layer perform binary classification to determine whether a corresponding line of code is a vulnerability location based on the vector representation of each graph node.

[0061] The initial vector representation of the AST is used as input to the vulnerability identification neural network model and the vulnerability localization neural network model, and the two models are jointly trained to achieve multi-task learning.

[0062] Use a trained vulnerability identification neural network model to predict whether a given source code contains a vulnerability, and use a trained vulnerability localization neural network model to predict the specific line number where the vulnerability is located in the given source code.

[0063] In one example, step 1) includes the following steps:

[0064] Generate an AST from the program source code and transform it into a property graph, denoted as . Specifically, assuming the attribute graph is... , where the node set The nodes in the abstract syntax tree are given by the nodes of the original abstract syntax tree. The edges in the function are given by the edges of the original abstract syntax tree. Furthermore, the function... Label each edge as an AST edge. (Function) Assign a `type` attribute and a `code` attribute to each node. The value of the `type` attribute is a string, corresponding to the statement type of the code represented by that node, for example, " CallExpression " represents a function call statement, ConditionExpression The colon represents a conditional statement. The value of the `code` attribute is also a string, corresponding to the code represented by that node.

[0065] Generate a CFG from the program source code and transform it into a property graph, denoted as . Specifically, assuming the attribute graph is... , where the node set yes A subset of , which corresponds to the nodes in the AST representing statements and predicates. Furthermore, the edge labeling function... From the set of symbols Each edge is assigned a label to indicate the conditions for control flow graph transitions.

[0066] Generate a PDG from the program source code and transform it into an attribute graph, denoted as . Specifically, assuming the attribute graph is... , where the node set , The edges in the graph correspond to the edges of the original program dependency graph. Furthermore, the edge labeling function... From the set of symbols Each edge is assigned a label to indicate either a control dependency or a data dependency. (Function) Assign each data dependency edge symbol Attributes are assigned to indicate the corresponding symbols that are depended upon, and each control dependency edge is assigned an attribute. condition Attributes to indicate the state of predicates that control dependencies, such as true or false .

[0067] Will and Combined into CPG, denoted as Specifically, assuming the code attribute diagram is... ,in , , ,and .

[0068] The tree structure and graph structure in the code property graph are separated, where the tree structure refers to the AST and the graph structure refers to the CFG and PDG. After separation ,in , , This represents the set of nodes in the AST of the i-th statement. This represents the set of edges in the AST of the i-th statement. and It is the root node of the AST of the i-th statement. This is equal to the number of statements in the program's source code. The above statement means that each statement in the program's source code has exactly one node in the graph structure, and the edges of the graph structure... This represents the control flow, data dependencies, and control dependencies between different statements. Each node in the graph structure corresponds to a tree structure, representing the syntactic structure within that statement.

[0069] In one example, step 2) includes the following steps:

[0070] The PACE algorithm is used to vectorize the type field. Assuming the string to be encoded is S, the encoding result of the PACE algorithm is calculated using the following formula: ,in The constituent characters of S, . It is a well-known one-hot coding algorithm in the field.

[0071] The `code` field is vectorized using the word2vec algorithm. A sequence of words (tokens) is extracted from the program's source code, a word2vec model is trained, and the vectors output by this model are used as the token encodings. Since the code represented by a node can consist of multiple tokens, in this case, the average of the encoded vectors of each token can be used as the encoded vector for that code field. For example, The feature vector is the average of the feature vectors of the 8 tokens, that is... .

[0072] The vectorized codes of the two fields are concatenated to obtain the final initial vector representation of the AST. The initial feature vector of node n. It can be formally expressed as ,in, Represents vector concatenation. and These represent the feature vectors of the type field and the node field, respectively. The feature vector representing the token. K is the token contained in the code field.

[0073] In one example, the specific technical solution for the vulnerability identification neural network model in step 3) is as follows:

[0074] The vulnerability identification neural network model consists of a code semantic encoder, a fully connected layer, and a softmax layer.

[0075] The code semantic encoder consists of a tree embedding module, a tree pooling module, a graph embedding module, and a graph pooling module. The tree embedding and tree pooling modules are responsible for learning the vector representations of the AST in the CPG to obtain statement-level vector representations. The graph embedding and graph pooling modules are responsible for learning the vector representations of the CFG and PDG in the CPG to obtain the vector representations of the source code.

[0076] The tree embedding module in the code semantic encoder updates the vector representations of all nodes in the tree structure using tree-based convolutions. During each convolution, the vector representation of the parent node is updated based on the vector representations of its child nodes in the subtree. Assume the vector representation of the parent node is... The vector representation of the child nodes is Then the new vector representation of the parent node Calculated by the following formula ,in Let be the weight matrix for node i, and b be the bias term. This is the activation function.

[0077] The tree pooling module in the code semantic encoder aggregates the vector representations of all nodes in the tree structure into a vector representation of the entire tree structure, and uses this as the vector representation of the statement represented by the current tree result. Specifically, the tree pooling module extracts and concatenates the maximum values ​​of each dimension of all nodes in the tree structure to obtain an aggregated vector representation, which is used to initialize the vector representation of the graph structure.

[0078] The graph embedding module in the code semantic encoder updates the vector representations of all nodes in the graph structure using graph-based convolutions. During each convolution, the vector representation of the parent node is updated based on the vector representations of its neighboring nodes. Suppose the vector representation of a node is... The vector representation of its one-hop neighbor is Then the new vector representation of the parent node Calculated by the following formula ,in This is the weight matrix. The attention weights between the central node and the i-node. This is the activation function.

[0079] The graph pooling module in the code semantic encoder aggregates the vector representations of all nodes in the graph structure into a vector representation of the entire graph structure, and uses this as the vector representation of the program source code represented by the current graph structure. Specifically, the graph pooling module extracts and concatenates the maximum values ​​of each dimension of all nodes in the graph structure to obtain the aggregated vector representation.

[0080] Fully connected layers perform linear transformations and non-linear mappings on the vector representation of the program source code. Assuming the vector representation of the program source code output by the code semantic encoder is x, then the output of the fully connected layer is... ,in Here is the weight matrix, and b is the bias term. This is the activation function.

[0081] The softmax layer performs binary classification based on the calculation results of the fully connected layer. Assuming the output of the fully connected layer is logits, the output of the softmax layer is... The final binary classification result is the category containing the larger softmax value. The classification result indicates whether a vulnerability exists in the program's source code.

[0082] In one example, in step 4), the specific technical solution for the vulnerability localization neural network model is as follows:

[0083] The vulnerability localization neural network model consists of a code semantic encoder, an attention layer, a fully connected layer, and a softmax layer.

[0084] The code semantic encoder consists of a tree embedding module, a tree pooling module, and a graph embedding module. The structure of these three modules is the same as that in the vulnerability identification neural network model.

[0085] The attention layer assigns different attention weights to different statements. When classifying statement x, statement x is used as the query term, and the attention weights of statement x and all other semantics are calculated.

[0086] Fully connected layers perform linear transformations and non-linear mappings on the vector representation of the program source code. Assuming the vector representation of the program source code output by the code semantic encoder is x, then the output of the fully connected layer is... ,in Here is the weight matrix, and b is the bias term. This is the activation function.

[0087] The softmax layer performs binary classification based on the calculation results of the fully connected layer. Assuming the output of the fully connected layer is logits, the output of the softmax layer is... The final binary classification result is the category containing the larger softmax value. The classification result indicates whether the current statement is a vulnerability location.

[0088] In one example, the specific technical solution for joint training in step 5) is as follows:

[0089] Parameters of the tree embedding module, tree pooling module, and graph embedding module in the shared semantic encoder.

[0090] Cross-entropy loss is used as the loss for vulnerability identification and vulnerability localization tasks, respectively, and the average of the two losses is defined as the joint loss.

[0091] The Adam optimizer is used to optimize the joint loss described above, and the neural network model is trained.

[0092] In one example, a trained vulnerability identification neural network model is used to predict whether a given source code contains a vulnerability, and a trained vulnerability localization neural network model is used to predict the specific line number where the vulnerability is located in the given source code.

[0093] The above embodiments are only used to illustrate the technical solutions of the present invention and are not intended to limit them. Those skilled in the art can modify or make equivalent substitutions to the technical solutions of the present invention without departing from the spirit and scope of the present invention. The scope of protection of the present invention should be determined by the claims.

Claims

1. A vulnerability discovery method based on multi-task learning, comprising the following steps: Construct a code property graph based on the abstract syntax tree, control flow graph, and program dependency graph of the source code; Separate the tree structure from the graph structure in the code attribute graph; Encode the nodes in the tree structure to obtain the initial node vector representation of the abstract syntax tree; The initial node vector representation, tree structure, and graph structure of the abstract syntax tree are input into the vulnerability identification neural network model to obtain the vulnerability identification result; wherein, the vulnerability identification neural network model includes: a first code semantic encoder and a first classification module; The first code semantic encoder is used to calculate the source code vector representation based on the initial node vector representation, tree structure, and graph structure of the abstract syntax tree; The first classification module is used to calculate the vulnerability identification result based on the source code vector representation; The initial node vector representation, tree structure, and graph structure of the abstract syntax tree are input into the vulnerability localization neural network model to obtain the vulnerability localization result; wherein, the vulnerability localization neural network model includes: a second code semantic encoder, an attention layer, and a second classification module; The second code semantic encoder is used to compute the final vector representation of nodes in the graph structure based on the initial node vector representation of the abstract syntax tree, the tree structure, and the graph structure; The attention layer is used to assign attention weights to the final vector representation of each node in the graph structure and to perform attention calculations. The second classification module is used to calculate the vulnerability location result based on the attention calculation result.

2. The method as described in claim 1, characterized in that, Construct a code property graph using the following steps: 1) Generate the attribute graph G of the abstract syntax tree A =(V A E A , λ A μ A ), where the node set V A node v in A For nodes of an abstract syntax tree, the edge set E A edge e in A For the edges of an abstract syntax tree, the function λ A edge e A Marked as an edge in the abstract syntax tree, function μ A For node v A Assign type attributes to code statements and code attributes to nodes; 2) Generate the attribute graph G of the control flow graph C =(V C E C , λ C ), where the node set V C node v in C For nodes representing statements and predicates in an abstract syntax tree, the edge set E C edge e in C Corresponding to the jump edges in the control flow graph, the function λ C For edge e C Assign jump condition flags; 3) Generate the attribute graph G of the program dependency graph. P =(V P E P , λ P μ P ), where the node set V P node v in P For nodes of an abstract syntax tree, the edge set E P edge e in P Corresponding to the dependency edges in the program dependency graph, the function λ P For edge e P Assign dependency tags, which include: control dependencies or data dependencies, and function μ. P For each edge e marked as a data dependency P Indicates the corresponding symbol on which it depends, or for each edge marked as a control dependency. P Indicates the state of the predicate that controls the dependencies; 4) Combined Attribute Diagram G A Attribute graph G C With attribute graph G P The resulting code attribute graph is G = (V, E, λ, μ), where V = V A E = E A ∪E C ∪E P , λ=λ A ∪λ C ∪λ P , μ=μ A ∪μ P .

3. The method as described in claim 1, characterized in that, Separated code attribute graph G′=(V T E T V G E G ,λ,μ), where The set of nodes in the abstract syntax tree representing the i-th statement. Let represent the set of edges in the abstract syntax tree of the i-th statement. and Let |V be the root node of the abstract syntax tree for the i-th statement. G | equals the number of statements in the source code.

4. The method as described in claim 1, characterized in that, The encoding of nodes in the tree structure includes: 1) Use the PACE algorithm to vectorize the code statement type attributes to obtain the code statement type representation; 2) Use the word2vec algorithm to vectorize the node code attributes to obtain the node code representation; 3) Concatenate the code statement type representation and the node code representation to obtain the initial node vector representation of the abstract syntax tree.

5. The method as described in claim 1, characterized in that, The source code vector representation is calculated using the following steps: 1) Update the vector representations of all nodes in the tree structure based on the initial vector representation of the abstract syntax tree; 2) Aggregate the vector representations of all nodes in the updated tree structure to obtain the statement vector representation; 3) Initialize the vector representation of the nodes in the graph structure according to the statement vector representation, and update the vector representation of the nodes in the graph structure using graph-based convolution to obtain the final vector representation of the nodes in the graph structure; 4) The final vector representation of the nodes in the aggregate graph structure is used to obtain the vector representation of the source code.

6. The method as described in claim 1, characterized in that, The first classification module includes a fully connected layer and a softmax layer, wherein the fully connected layer is used to perform linear transformation and nonlinear mapping on the vector representation of the program source code, and the softmax layer is used to perform binary classification based on the calculation results of the fully connected layer to obtain vulnerability identification results.

7. The method according to any one of claims 1 to 6, characterized in that, The following steps are used to train the vulnerability identification neural network model and the vulnerability localization neural network model: 1) The parameters of the tree embedding module, tree pooling module, and graph embedding module in the first code semantic encoder and the second code semantic encoder are shared; 2) Use cross-entropy loss as the loss for vulnerability identification task and vulnerability location task respectively, and define the average of the loss for vulnerability identification task and vulnerability location task as the joint loss; 3) Use the Adam optimizer to optimize the joint loss and train the vulnerability identification neural network model and the vulnerability localization neural network model.

8. A vulnerability discovery device based on multi-task learning, comprising: Build modules are used to construct code property graphs based on source code abstract syntax trees, control flow graphs, and program dependency graphs; The separation module is used to separate the tree structure from the graph structure in the code property graph; The encoding module is used to encode the nodes in the tree structure to obtain the initial node vector representation of the abstract syntax tree; The identification module is used to input the initial node vector representation, tree structure, and graph structure of the abstract syntax tree into the vulnerability identification neural network model to obtain the vulnerability identification result; wherein, the vulnerability identification neural network model includes: a first code semantic encoder and a first classification module; The first code semantic encoder is used to calculate the source code vector representation based on the initial node vector representation, tree structure, and graph structure of the abstract syntax tree; The first classification module is used to calculate the vulnerability identification result based on the source code vector representation; The localization module is used to input the initial node vector representation, tree structure, and graph structure of the abstract syntax tree into the vulnerability localization neural network model to obtain the vulnerability localization result; wherein, the vulnerability localization neural network model includes: a second code semantic encoder, an attention layer, and a second classification module; The second code semantic encoder is used to compute the final vector representation of nodes in the graph structure based on the initial node vector representation of the abstract syntax tree, the tree structure, and the graph structure; The attention layer is used to assign attention weights to the final vector representation of each node in the graph structure and to perform attention calculations. The second classification module is used to calculate the vulnerability location result based on the attention calculation result.

9. A storage medium storing a computer program, wherein, The computer program is configured to execute any of the methods described in claims 1-7 at runtime.

10. A computer device, characterized in that, It includes a memory and a processor, wherein the memory stores a computer program and the processor is configured to run the computer program to perform the method as described in any one of claims 1-7.

Images