CN111753303B - A multi-granularity code vulnerability detection method based on deep learning and reinforcement learning - Google Patents

Abstract

The invention discloses a multi-granularity code vulnerability detection method based on deep learning and reinforcement learning, which comprises the following steps: 1) Analyzing the source code to obtain an intermediate code representation corresponding to the code; 2) Slicing the intermediate code to obtain a code segment smaller than the source program; 3) Converting an input code segment into a low-dimensional continuous real-valued vector by using a code segment representation method; 4) Inputting the vector representation of the code segment into a coarse-grained code vulnerability detection model based on deep learning, and judging whether the code segment contains defects; 5) And constructing a fine-grained code vulnerability detection model based on reinforcement learning, and predicting code lines which specifically cause vulnerabilities in code segments containing defects. The invention provides a complete multi-granularity code vulnerability detection framework, applies reinforcement learning to the field of fine-granularity code vulnerability detection for the first time, and provides a new code segmentation representation learning model to fully utilize semantic information of a program, thereby improving the accuracy and the practicability of vulnerability detection.

Description

A multi-granularity code vulnerability detection method based on deep learning and reinforcement learning

technical field

The invention relates to a code loophole detection method, in particular to a multi-granularity code loophole detection method based on deep learning and reinforcement learning technology.

Background technique

Software vulnerabilities refer to defects in software during its life cycle, and these defects may be exploited by criminals to bypass system access control, illegally steal higher privileges and manipulate the system arbitrarily, such as triggering privileged commands and accessing sensitive information , impersonate identity, monitor system operation, etc. If security-related vulnerabilities cannot be identified and repaired in a timely manner, the vulnerabilities will be easily exploited by malicious attackers, causing the system to be invaded and resulting in unreliable system operation results, or serious security issues such as arbitrary command execution and arbitrary file reading. question.

Code analysis is the main method to check and discover the inherent defects and utilization methods in software codes, and has always been a research hotspot in the field of information security and software security. However, as the existing software systems become more and more complex, the frequency of vulnerabilities and hackers' attack methods continue to increase, the traditional vulnerability detection tools based on predefined rules can no longer meet the needs of modern software development, and more and more researchers have begun to pay attention to Code vulnerability detection method based on machine learning and deep learning. Machine learning-based vulnerability detection methods rely on experts to manually define code characteristics (for example, software complexity metrics (software complexity metrics), function calls (function calls), code changes (code changes) and system calls (system calls)), and then use Machine learning models automatically classify vulnerable and non-vulnerable code. However, because the definition of code features is highly subjective, it is usually only applicable to specific projects and has poor generalization ability. Moreover, the granularity of the code input to the machine learning model is usually coarse, and it is impossible to determine the exact location of the vulnerable code line. Vulnerability detection methods based on deep learning do not require experts to manually define features, and can automatically generate vulnerability patterns from a large amount of historical data. Automatically generate transitions and significantly improve the effectiveness of vulnerability detection. However, related research on this method has just started, most of which focus on coarser-grained code vulnerability detection, such as detecting code vulnerabilities at the function level or file level, while related research on the "vulnerability structure" of code is Very lacking. The vulnerability structure not only enables the detection tool to judge whether the code contains a vulnerability, but also points out the specific form of the vulnerability in the code and the location where the vulnerability occurs. In order to make the vulnerability detection method based on deep learning better applied in practice, it is necessary to study the structure of code vulnerabilities.

The literature VulDeeLocator (Z.Li, D.Zou, S.Xu, Z.Chen, Y.Zhu, and H.Jin, VulDeeLocator: A Deep Learning-based Fine-grained Vulnerability Detector, arXivpreprint arXiv:2001.02350, 2020) is currently The only document that can be retrieved is based on deep learning to achieve fine-grained vulnerability location at the sentence level. This method adds the location information of the vulnerability sentence to the deep neural network model training, so that the model can focus on the vulnerability-related information during training. statement to implement statement-level fine-grained vulnerability detection. However, from the experimental results, VulDeeLocator is not significantly better than traditional rule-based vulnerability detection tools in fine-grained detection. This is because the network structure of this method cannot fully capture the semantic information of the program and the code structure related to the vulnerability. .

Contents of the invention

The purpose of the present invention is to provide a multi-granularity code vulnerability detection method based on deep learning and reinforcement learning, which can not only detect defective software modules (for example, functions) at a coarser granularity, but also At the granular level, the code statements that may cause vulnerabilities in the module are located, that is, multi-granularity code vulnerability detection is realized. In addition, since the accuracy of the vulnerability detection model depends on the accuracy of the model input, that is, the accuracy of the code fragment representation. Aiming at the problem that the previous token-based code representation methods fail to make full use of code semantic information, the present invention proposes a new code segment representation learning method (Staged Code Representation Learning). The method first learns a vector representation of each statement in the code, and then learns a vector representation of the entire program based on the vector representation of each statement. By separating the vector representation of the statement and the vector representation of the program, the learned code vector can capture the subtler structural (syntactic) differences and more complex semantic information of the program.

The purpose of the present invention is achieved through the following technical solutions:

A multi-granularity code vulnerability detection method based on deep learning and reinforcement learning. First, source code is parsed to obtain an intermediate code representation corresponding to the code. Compared with source code representation, using intermediate code representation can capture more program control flow and variable definition-use information. Second, take the key points that may cause vulnerabilities as the slicing standard, slice the intermediate code, and obtain code segments (code gadgets) that are smaller than the source program, so as to reduce the length of the model input sequence and avoid loophole-related statements. impact of important information. Thirdly, the code segmentation representation learning model proposed by the present invention is used to convert the input code segment into a low-dimensional continuous real-valued vector. Then the vector representation of the code segment is input into the coarse-grained code vulnerability detection model based on deep learning, and it is judged whether the code segment corresponding to the input vector contains a vulnerability. Finally, if the coarse-grained model detects that the input code segment contains a vulnerability, the fine-grained detection model based on reinforcement learning will continue to make the next step of judgment, that is, to find out the possible code lines that cause the vulnerability.

Specifically include the following steps:

Step 1: Use the Clang tool to statically analyze the source program to obtain the intermediate code representation of the program;

Step 2: Extract the key points that may cause vulnerabilities, generate slicing standards, then slice the intermediate code, and merge the forward slicing and backward slicing to obtain the code segment of the program;

Step 3: Use the code segmentation representation learning method to represent the code segment as a low-dimensional continuous real-valued vector;

Step 4: Input the vector representation of the code segment into the coarse-grained code vulnerability detection model based on deep learning to determine whether the code segment contains vulnerabilities;

Step 5: Construct a fine-grained vulnerability detection model based on reinforcement learning to predict the code line that specifically causes the vulnerability in the code segment containing the defect.

Compared with the prior art, the present invention has the following advantages:

1. The present invention proposes a new multi-granularity code vulnerability detection method and framework. Compared with the existing method that can only detect vulnerabilities at a coarser granularity level, its advantage is that it can complete coarse-grained vulnerability detection. Code module detection can also realize the location of fine-grained code vulnerability statements, which improves the practicability of data-driven vulnerability detection methods and the interpretability of model prediction results.

2. The present invention proposes a novel code segmentation representation learning method. Compared with the existing token-based code representation method, its advantage is that it can make full use of the local and global semantic information of the program and improve the generated The accuracy of the code representation vector, which in turn improves the vulnerability detection ability of the model.

3. The present invention proposes for the first time the application of reinforcement learning technology to fine-grained code vulnerability detection tasks, by continuously trying various possible combinations of vulnerability statements on training data instances, and using the number of vulnerability statements contained in each combination as an adjustment model Behavior signals are fed back to the Policy module, and through continuous signal accumulation, the model can automatically learn the code structure related to the vulnerability.

Description of drawings

FIG. 1 is an overall flow chart of the multi-granularity code vulnerability detection method proposed by the present invention.

Figure 2 is the statement node corresponding to the key point of the function call in the syntax tree.

Fig. 3 is the statement node corresponding to the key point of array definition in the syntax tree.

Fig. 4 is the statement node corresponding to the pointer definition key point in the syntax tree.

Fig. 5 is the statement node corresponding to the key point of the assignment expression in the syntax tree.

Detailed ways

The technical solution of the present invention will be further described below in conjunction with the accompanying drawings, but it is not limited thereto. Any modification or equivalent replacement of the technical solution of the present invention without departing from the spirit and scope of the technical solution of the present invention should be covered by the present invention. within the scope of protection.

This embodiment implements coarse-grained and fine-grained code vulnerability identification and vulnerability statement location based on deep learning technology and reinforcement learning technology, respectively. First, the source code is parsed to obtain the intermediate code representation (such as LLVM IR) corresponding to the code. Second, the intermediate code is sliced using keypoints that may cause vulnerabilities as the slicing standard to obtain code segments (code gadgets) that are smaller than the source program. Thirdly, the code segmentation representation learning method is used to convert the input code segment into a low-dimensional continuous real-valued vector. Then input the vector representation of the code segment into the coarse-grained vulnerability detection model based on deep learning, and judge whether the code segment corresponding to the input vector contains a vulnerability. Finally, if the coarse-grained model detects that the input code segment contains a vulnerability, the fine-grained vulnerability detection model based on reinforcement learning will continue to make the next step of judgment, that is, find out the line of code that causes the vulnerability.

As shown in Figure 1, the specific steps are as follows:

Step 1: Use the Clang tool to statically analyze the source program to obtain the intermediate code representation of the program.

Step 2: Extract the key points that may cause vulnerabilities, generate slicing standards, then slice the intermediate code, and merge the forward slicing and backward slicing to obtain the code segment of the program. The specific steps are as follows:

Step 21: using program analysis technology to parse the program into an abstract syntax tree form;

Step 22: Traverse the generated abstract syntax tree through the node matching algorithm, and find four syntax tree nodes that may cause code vulnerabilities: (1) function call statement node (ie Callee), as shown in Figure 2; (2) array definition statement Node (i.e. IdentifierDeclStatement, and the statement contains "[" and "]" characters), as shown in Figure 3; (3) pointer definition statement node (i.e., IdentifierDeclStatement, and the statement contains "*" characters), as shown in Figure 4 (4) expression statement node (i.e. ExpressionStatement), as shown in Figure 5;

Step 23: Filter the above four types of nodes extracted in the program, and select the syntax tree nodes that meet the conditions as the key points that may cause vulnerabilities, for example: if the identifier (identify) corresponding to the "Callee" node is in the predefined In the list of library functions that may cause vulnerabilities, the parent node (CallExpression) corresponding to the "Callee" node is the statement S in the slice standard;

Step 24: Extract the variables involved from the statement S as V in the slice standard, combine S and V to finally obtain the slice standard (S, V) for extracting slices;

Step 25: Use the program analysis technology to parse the code segment into a program dependency graph, perform forward and backward slice analysis on the program dependency graph according to the slice standard generated in step 24, and combine the forward slice and the backward slice to obtain four possibilities Program slices related to key points that cause vulnerabilities;

Step 26: Convert the program slice of the source program into the program slice of the intermediate code according to the corresponding relationship between the source program and the intermediate code.

Step 3: Use the code segmentation representation learning method to represent the code segment as a low-dimensional continuous real-valued vector. The specific steps are as follows:

Step 31: Split the code segment in units of statements;

Step 32: Construct a CNN-based statement encoding network (Statement Encode Network, SENet), the schematic diagram of the model is shown in the upper middle part of Figure 1;

Step 33: splitting each statement obtained in step 31 into a sequence of tokens with a space as a separator, and outputting a vector representation of the statement as a statement encoding network;

Step 34: Build an LSTM-based program encoding network (Program Encode Network, PENet), the schematic diagram of the model is shown in the upper right part of Figure 1;

Step 35: Take the vector representation of each sentence contained in the code segment in step 33 as the input of the program encoding network, and output the vector representation of the hidden layer of the last time step as the vector representation of the code segment.

Step 4: Input the vector representation of the code segment into the coarse-grained code vulnerability detection model based on deep learning to determine whether the code segment contains vulnerabilities. The specific steps are as follows:

Step 41: Construct a coarse-grained code vulnerability detection model (DetectorNetwork, DNet) based on a fully connected single hidden layer. The schematic diagram of the model is shown in the middle and lower part of Figure 1;

Step 42: take the vector representation of the code segment as the input of the model, and output the probability that the code segment contains a vulnerability;

Step 43: Use the predicted probability and the real label as the input of the cross-entropy loss function, calculate the predicted error, and use the backpropagation algorithm to update the code segmentation representation learning model (ie SENet based on CNN and PENet based on LSTM) And the parameters of the coarse-grained vulnerability detection model (DNet).

Step 5: Construct a fine-grained vulnerability detection model (Policy Network, PNet) based on reinforcement learning, and predict the code line that specifically causes the vulnerability in the code segment containing the defect. The specific steps are as follows:

Step 51: Build a fine-grained code vulnerability prediction model PNet based on reinforcement learning, the schematic diagram of the model is shown on the right side of Figure 1;

Step 52: splicing the vector representation of the current program statement and the context vector representation of the statement as the state (state) representation of the reinforcement learning at time step t;

Step 53: According to the state representation at time t, predict the action that the agent (agent or policy) may take. If the action is "relevant", it means that the input sentence at time t may cause code loopholes statement, if the action is "irrelevant", it means that the current statement will not cause code vulnerabilities, and the same action prediction is made for each statement in the code segment to generate an action sequence;

计算奖励(reward)，其中U是步骤53中预测的动作序列中和漏洞相关的代码行，V是真实的漏洞所在的代码行；Step 54: According to the formula

Calculate the reward (reward), where U is the code line related to the vulnerability in the predicted action sequence in step 53, and V is the code line where the real vulnerability is located;

Step 55: Update the parameters of the model PNet according to the classic REINFORCE algorithm and the Policy gradient algorithm, so that PNet can automatically learn the code structure related to the vulnerability.

Example:

Taking the specific vulnerability examples in the data set Software Assurance Reference Dataset (SARD) as an example, the detection process of the multi-granularity code vulnerability detection method based on deep learning and reinforcement learning proposed by the present invention is analyzed. The contents of the four source code files involved in the vulnerability examples are shown in Table 1 to Table 4 respectively. First execute steps 1 and 2 of the specific embodiment of the present invention, convert the source code into intermediate code, and use the dangerous function memset that may cause loopholes as a key point to generate program slices for the intermediate code, as shown in Table 5, in In this slice code, the actual vulnerable statements are lines 10, 11, 34 and 35. Then execute steps 3 and 4 to learn the vector representation of each line of code in the program and the vector representation of the entire program, and use the vector representation of the program as the input vector of the coarse-grained vulnerability detection model (DNet), and the model prediction result is 1, that is Slices contain holes. Finally, step 5 is executed, and the vector representation of each statement in the program is used as the input of the fine-grained vulnerability detection model (PNet), and the predicted action (action) sequence is output as shown in Table 6, where the number 0 represents that the corresponding code statement does not contain a vulnerability , the number 1 represents that the corresponding code statement may contain a vulnerability. Since the position index of the number 1 in Table 6 is also 10, 11, 34 and 35, the fine-grained vulnerability detection model accurately identifies the specific vulnerability location. It can be seen from the above examples that the method proposed by the present invention not only realizes coarse-grained vulnerability code detection, but also realizes locating specific code statements that may cause vulnerabilities.

Table 1 CWE124_Buffer_Underwrite__char_declare_memmove_53a.c

Table 2 CWE124_Buffer_Underwrite__char_declare_memmove_53b.c

Table 3 CWE124_Buffer_Underwrite__char_declare_memmove_53c.c

Table 4 CWE124_Buffer_Underwrite__char_declare_memmove_53d.c

Table 5 uses memset as the key point to generate program slices for intermediate codes

Table 6 Action sequence predicted by the fine-grained vulnerability detection model for Table 5

Claims (1)

1. A multi-granularity code vulnerability detection method based on deep learning and reinforcement learning, characterized in that said method comprises the steps: Step 1: Use the Clang tool to statically analyze the source program to obtain the intermediate code representation of the program; Step 2: Extract the key points that may cause vulnerabilities, generate slicing standards, then slice the intermediate code, and merge the forward slicing and backward slicing to obtain the code segment of the program. The specific steps are as follows: Step 21: using program analysis technology to parse the program into an abstract syntax tree form; Step 22: traverse the generated abstract syntax tree through the node matching algorithm, and find four syntax tree nodes that may cause code loopholes: (1) function call statement node; (2) array definition statement node; (3) pointer definition statement node; (4) Expression statement node; Step 23: Filter the above four types of nodes extracted in the program, and select the syntax tree nodes that meet the conditions as the key points that may cause vulnerabilities; Step 24: Extract the variables involved from the statement S as V in the slice standard, combine S and V to finally obtain the slice standard (S, V) for extracting slices; Step 25: Use the program analysis technology to parse the code segment into a program dependency graph, perform forward and backward slice analysis on the program dependency graph according to the slice standard generated in step 24, and combine the forward slice and the backward slice to obtain four possibilities Program slices related to key points that cause vulnerabilities; Step 26: Convert the program slice of the source program into the program slice of the intermediate code according to the corresponding relationship between the source program and the intermediate code; Step 3: Use the code segmentation representation learning method to represent the code segment as a low-dimensional continuous real-valued vector. The specific steps are as follows: Step 31: Split the code segment in units of statements; Step 32: Construct a CNN-based sentence encoding network SENet; Step 33: splitting each statement obtained in step 31 into a token sequence with a space as a delimiter as the input of the statement encoding network, and outputting a vector representation of the statement; Step 34: Construct PENet, an LSTM-based program encoding network; Step 35: using the vector representation of each statement contained in the code segment in step 33 as the input of the program encoding network, outputting the vector representation of the hidden layer of the last time step as the vector representation of the code segment; Step 4: Input the vector representation of the code segment into the coarse-grained code vulnerability detection model based on deep learning to determine whether the code segment contains vulnerabilities. The specific steps are as follows: Step 41: Build a coarse-grained code vulnerability detection model DNet based on a fully connected single hidden layer; Step 42: take the vector representation of the code segment as the input of the model, and output the probability that the code segment contains a vulnerability; Step 43: Use the predicted probability and the real label as the input of the cross-entropy loss function, calculate the predicted error, and use the backpropagation algorithm to update the parameters of the code segmentation representation learning model and the coarse-grained vulnerability detection model DNet. The learning model includes the parameters based on SENet implemented by CNN and PENet based on LSTM; Step 5: Construct a fine-grained vulnerability detection model based on reinforcement learning, and predict the code line that specifically causes the vulnerability in the code segment containing the defect. The specific steps are as follows: Step 51: Construct PNet, a fine-grained code vulnerability prediction model based on reinforcement learning; Step 52: Concatenate the vector representation of the current program statement and the context vector representation of the statement as the state state representation of the reinforcement learning at time step t; Step 53: According to the state representation at time t, predict the action taken by the agent or policy. If the action is "relevant", it means that the input sentence at time t is a statement that causes code loopholes. If the action is "not relevant irrelevant", it means that the current statement will not cause code loopholes, the same action prediction is made for each statement in the code segment, and an action sequence is generated; Step 54: According to the formula

Calculate the reward reward, where U is the code line related to the vulnerability in the predicted action sequence in step 53, and V is the code line where the real vulnerability is located; Step 55: Update the parameters of the model PNet according to the classic REINFORCE algorithm and the Policy gradient algorithm, so that the model PNet can automatically learn the code structure related to the vulnerability.

Images