Binary slice-level rapid vulnerability detection method based on xLSTM model

Abstract

The invention discloses a binary slice-level rapid vulnerability detection method based on an xLSTM model. According to the method, a segmented code representation architecture and a novel xLSTM model are adopted. The segmented code representation architecture is composed of an instruction coding network and a program coding network, and fine-grained semantic information and global structure information of codes can be comprehensively captured by combining the instruction coding network and the program coding network, so that the accuracy of a vulnerability detection model is enhanced. Besides, a novel xLSTM model is adopted in the method, an index gating mechanism is introduced into xLSTM, and compared with traditional Sigmoid gating, index gating can provide more dynamic information filtering capacity, and the memory and forgetting process can be improved. Meanwhile, the xLSTM introduces additional normalization and stabilization steps in gating calculation, so that the stability of the model is improved. Furthermore, the xLSTM introduces matrix memory, which allows parallel processing and improves storage capacity.

Description

Technical Field

[0001] The present invention belongs to the field of software vulnerability detection, and relates to a binary program vulnerability detection method, and specifically to a binary program code vulnerability detection method based on an xLSTM model and a segmented code representation architecture. Background Art

[0002] With the rapid development of information technology, Internet of Things (IoT) devices have been deeply integrated into many aspects of daily life. Most of these devices distribute their software in the form of binary programs. However, with the large-scale deployment of IoT devices, the security issues of basic software have become increasingly prominent. Software vulnerabilities refer to defects in software that may be exploited by criminals during its life cycle. Once exploited, they will lead to serious consequences such as privacy leakage, illegal authorization, and extortion. Software vulnerability detection technology is an important method to reduce security risks and improve software security. In recent years, many researchers have applied deep learning technology to the field of software vulnerability detection. With its powerful feature extraction and pattern recognition capabilities, they have deeply mined vulnerability patterns in the code, thereby significantly improving the accuracy and efficiency of vulnerability identification.

[0003] Since binary code is not readable, deep learning-based vulnerability detection methods usually use decompilation tools to convert binary code into assembly code for more in-depth analysis. After obtaining the corresponding assembly code, these methods treat the assembly code as natural language text and use natural language processing techniques such as Word2Vec to embed words. Then, deep neural networks are used to learn the representation of the embedded assembly code, extract vulnerability features from it, and feed it into the classifier for vulnerability classification. Due to the highly serialized characteristics of binary code, recurrent neural networks (RNNs) have become an ideal choice for most binary code vulnerability detection methods. However, due to the long sequence length of assembly code, RNNs are prone to gradient vanishing or gradient exploding problems, which hinders the model's ability to capture long-distance dependencies, thereby affecting its ability to understand and parse complex code logic. Although long short-term memory networks (LSTMs) and gated recurrent units (GRUs) are widely believed to alleviate these problems, they also face challenges when processing long assembly codes. In recent years, the Transformer architecture has achieved remarkable success in natural language processing tasks. It effectively captures contextual information through the self-attention mechanism, improving the performance of the model in a variety of tasks such as text generation, translation, and sentiment analysis. Therefore, many researchers have applied the Transformer architecture to binary code vulnerability detection tasks. However, the computational complexity of the Transformer architecture is O(n 2), and because the assembly instruction sequence may be quite lengthy, the vulnerability detection model based on the Transformer architecture suffers from excessive consumption of computing resources and long processing time.

[0004] Currently, no binary program slice-level vulnerability detection method based on the xLSTM model has been retrieved. Summary of the invention

[0005] The current binary program slice-level vulnerability detection method based on the traditional RNN model is difficult to capture the long-distance code semantics, and the binary program slice-level vulnerability detection method based on the Transformer architecture has the disadvantages of slow calculation time and high memory usage. The present invention provides a binary slice-level fast vulnerability detection method based on the xLSTM model. The method adopts a segmented code representation architecture and a new xLSTM model. The segmented code representation architecture consists of an instruction encoding network and a program encoding network. The instruction encoding network focuses on mapping the semantics of a single instruction or operation to a vector space, thereby capturing the semantic relationship between tokens inside the instruction; the program encoding network focuses on extracting the contextual information of the code fragment, thereby improving the understanding and analysis capabilities of the program semantics. By combining the instruction encoding network and the program encoding network, the fine-grained semantic information and global structural information of the code can be fully captured, thereby enhancing the accuracy of the vulnerability detection model. In addition, the method adopts a novel xLSTM model, which can more effectively learn long-distance code semantics for the long sequence characteristics of assembly code, improve the accuracy of the model, and reduce the resource overhead of the model. xLSTM introduces an exponential gating mechanism. Compared with traditional Sigmoid gating, exponential gating can provide more dynamic information filtering capabilities, which helps improve the memory and forgetting process. At the same time, xLSTM introduces additional normalization and stabilization steps in the gating calculation, which improves the stability of the model. In addition, xLSTM introduces matrix memory, which allows parallel processing and improves storage capacity.

[0006] The objective of the present invention is achieved through the following technical solutions:

[0007] A binary slice-level rapid vulnerability detection method based on an xLSTM model comprises the following steps:

[0008] Step 1: Use the decompilation tool IDAPro to automatically detect and identify the architecture and instruction set of the input file, and convert the binary program code into assembly code;

[0009] Step 2: Use the disassembly view of the decompilation tool IDAPro to generate the control flow graph (CFG) of the assembly code;

[0010] Step 3: Generate assembly code slices according to the assembly code slice criteria and mark the vulnerable code slices;

[0011] Step 4: Use Word2Vec to embed the tokens in the assembly code slices to obtain the initial vector representation of each statement in the assembly code slices, and then send it to the instruction encoding network composed of the BiGRU model to obtain the hidden vector representation of each instruction. Finally, the hidden vector representation of each statement is sent to the program encoding network composed of the xLSTM model to obtain the hidden vector representation of the assembly code slices.

[0012] Step 5: The hidden vector representation of the assembly code slice obtained in step 4 is sent to the classifier composed of the fully connected layer (FCN) and SoftMax. The cross entropy loss is calculated based on the output of the classifier and the actual label of the slice. The parameters of the segmented code representation network and the classifier network composed of the instruction encoding network and the program encoding network are updated by back propagation until the model training is completed.

[0013] Step 6: Use the segmented code representation network and classifier network trained in step 5 to perform vulnerability detection on the assembly code slices to be tested.

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

[0015] 1. This invention introduces the xLSTM model into the field of binary program vulnerability detection for the first time. The exponential gating and residual connection mechanisms introduced by xLSTM enable it to quickly adapt to new data and adjust its memory state, maintain or even enhance the gradient flow of the entire network, and capture long-term dependencies and deep semantic information in assembly instruction sequences with long context characteristics, thereby identifying complete vulnerability patterns in the code. Compared with the Transformer architecture, xLSTM does not rely on the self-attention mechanism and has a highly parallelized architecture, which makes its computational complexity lower than the Transformer architecture. Therefore, it is more suitable for binary program vulnerability detection tasks in resource-limited situations such as IoT devices.

[0016] 2. The present invention introduces a segmented code representation architecture into the binary vulnerability detection task. The segmented code representation architecture consists of an instruction encoding network and a program encoding network. The former can capture the local semantic information of a single instruction, and the latter can extract the contextual information of a code fragment, thereby improving the ability to understand and analyze the overall semantics of the program. By combining code semantic information of different granularities, the accuracy and robustness of binary program vulnerability detection can be significantly improved. BRIEF DESCRIPTION OF THE DRAWINGS

[0017] Figure 1 The figure is a schematic diagram of the overall process of the binary slice-level fast vulnerability detection method based on the xLSTM model.

[0018] Figure 2 Example diagram of assembly code slices generated for the present invention.

[0019] Figure 3 Schematic diagram of the segmented code representation learning architecture of the present invention. DETAILED DESCRIPTION

[0020] The technical solution of the present invention is further described below in conjunction with the accompanying drawings, but 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 included in the protection scope of the present invention.

[0021] The present invention provides a binary slice-level rapid vulnerability detection method based on the xLSTM model. First, the binary program code to be detected is converted into assembly code using the decompilation tool IDAPro. Secondly, the control flow graph (CFG) of the assembly code is generated using the disassembly view of the decompilation tool IDAPro. Then, assembly code slices are generated according to the slicing rules of the assembly code and the slice data are labeled. Next, the slices of the assembly code are segmented and embedded, and the obtained initial embedding is sent to the segmented code representation architecture for representation learning to obtain the hidden vector representation of the assembly code slice (such as Figure 3 Finally, the feature vector representation of the learned assembly slice is sent to the classifier composed of the fully connected layer (FCN) and SoftMax to determine whether the program under test contains vulnerabilities. Figure 1 As shown, the specific steps of the method are as follows:

[0022] Step 1: Use the decompilation tool IDAPro to automatically detect and identify the architecture and instruction set of the input file, and convert the binary program code into assembly code.

[0023] Step 2: Use the disassembly view of the decompilation tool IDAPro to generate the control flow graph (CFG) of the assembly code.

[0024] Step 3: Generate assembly code slices according to the assembly code slice criteria and mark the vulnerable code slices. The specific steps are as follows:

[0025] Step 3.1: According to the assembly code slicing criteria, extract the vulnerability candidate key points of the assembly code, obtain the bidirectional slicing assembly instruction set of the assembly code by forward and backward traversing the control flow graph obtained in step 2, and generate the slicing code segment of the assembly code.

[0026] Step 3.2: Mark the assembly code slices obtained in step 3.1. The assembly code slices containing vulnerabilities are marked as "1", and the assembly code slices not containing vulnerabilities are marked as "0".

[0027] Step 4: Use Word2Vec to embed the tokens in the assembly code slices to obtain the initial vector representation of each statement in the assembly code slices, and then send it to the instruction encoding network composed of the BiGRU model to obtain the hidden vector representation of each instruction. Finally, the hidden vector representation of each statement is sent to the program encoding network composed of the xLSTM model to obtain the hidden vector representation of the assembly code slice. The specific steps are as follows:

[0028] Step 4.1: Segment the assembly code segments obtained in step 3, and use Word2Vec to embed the tokens of the assembly code segments to obtain the initial vector representation T1, T2, ..., T of each assembly instruction token. n , where n is the number of tokens in each assembly instruction, T n is the embedding representation of the nth token.

[0029] Step 4.2: Denote the initial vectors obtained in step 4.1 as T1, T2, ..., T n Input into the BiGRU model to capture the local semantic features of the instruction, thereby obtaining the hidden vector representation of each assembly instruction. The specific calculation formula is as follows:

[0030] FI i =BiGRU(T1, T2, ..., T n )

[0031] Among them, BiGRU() represents a bidirectional gated recurrent neural network, FI i Represents the embedded representation of the i-th assembly instruction.

[0032] Step 4.3: Substitute the hidden vector sequence FI1, FI2, ..., FI obtained in step 4.2 n The global semantic features of the slice are captured by the xLSTM model, and the hidden vector representation of each assembly slice code segment is obtained. The specific calculation formula is as follows:

[0033] FA i =xLSTM(FI1, FI2, ..., FI n )

[0034] Among them, xLSTM() represents the extended long short-term memory network, FA i is the embedding representation of the i-th assembly slice code segment.

[0035] Step 5: The hidden vector representation of the assembly code slice obtained in step 4 is sent to the classifier composed of the fully connected layer (FCN) and SoftMax. The cross entropy loss is calculated based on the output of the classifier and the actual label of the slice. Back propagation updates the parameters of the segment code representation network and the classifier network composed of the instruction encoding network and the program encoding network until the model training is completed. The specific steps are as follows:

[0036] Step 5.1: Send the hidden vector representation of the assembly slice code segment obtained in step 4 to the fully connected layer, and then pass it through the SoftMax function to convert the input real number vector into a probability distribution to obtain the prediction result;

[0037] Step 5.2: Calculate the cross entropy loss based on the prediction results of step 5.1 and the label information of the assembly slice code segment. Specifically, use the label value marked in step 3. When the label value is 0, it means no vulnerability. Initialize it as a vector [1,0]. When the label value is 1, it means there is a vulnerability. Initialize it as a vector [0,1] and calculate the cross entropy loss with the prediction results given by the segmented code representation network.

[0038] Step 5.3: Adjust the parameters of the segmentation code representation network and the classifier network according to the cross entropy loss back propagation until the loss value no longer decreases, and the training ends.

[0039] Step 6: Use the segmented code representation network and classifier network trained in step 5 to perform vulnerability detection on the assembly code slices to be tested.

[0040] Example:

[0041] like Figure 2 As shown, Figure 2 (a) is a binary program code slice. The number before each binary program statement is the position of the statement in the original program. Figure 2 (b) is an assembly code slice. The number before each assembly instruction is the position of the instruction in the original assembly code. Figure 2 (b) Figure 2 (a) Corresponding assembly code slices. After decompiling the original binary program code into assembly code using the decompilation tool IDAPro, assembly code slices can be obtained according to the assembly code slice criteria, such as Figure 2As shown. The vulnerability type of this vulnerable assembly code slice example is Stack Overflow. Specifically, the vulnerability occurs on line 37 because the size and boundaries of the target memory are not fully checked when initializing the data, which may result in the modification of the memory address at rbp-184. Next, mov BYTE PTR[rax],0 attempts to write data to the memory location pointed to by rax. If rax has been incorrectly modified to an address outside the allocated range, this will cause a stack overflow, which may cause the program to crash or be exploited by an attacker to execute arbitrary code.

[0042] like Figure 1 As shown, this embodiment first performs word segmentation on the assembly code slice, and then inputs each token obtained by the word segmentation into the Word2Vec model to obtain the token initial vector representation, and then embeds the obtained token into the instruction encoding network to obtain the embedding representation of each assembly instruction, and then inputs the instruction embedding representation into the program encoding network to obtain the embedding representation of the assembly slice code segment, and finally sends the embedding representation of the assembly slice code segment to the fully connected layer, and then passes through SoftMax to convert the input real number vector into a binary classification probability to obtain the prediction result of the slice. The method of the present invention can accurately identify the vulnerabilities in the assembly code slice.

Claims

1. A binary slice-level fast vulnerability detection method based on the xLSTM model, characterized in that The method comprises the following steps: Step 1: Use the decompilation tool IDAPro to automatically detect and identify the architecture and instruction set of the input file, and convert the binary program code into assembly code; Step 2: Use the disassembly view of the decompilation tool IDAPro to generate a control flow graph of the assembly code; Step 3: Generate assembly code slices according to the assembly code slice criteria and mark the vulnerable code slices; Step 4: Use Word2Vec to embed the tokens in the assembly code slices to obtain the initial vector representation of each statement in the assembly code slices, and then send it to the instruction encoding network composed of the BiGRU model to obtain the hidden vector representation of each instruction. Finally, the hidden vector representation of each statement is sent to the program encoding network composed of the xLSTM model to obtain the hidden vector representation of the assembly code slices. Step 5: The hidden vector representation of the assembly code slice obtained in step 4 is sent to the classifier composed of the fully connected layer and SoftMax. The cross entropy loss is calculated based on the output of the classifier and the actual label of the slice. The parameters of the segmented code representation network and the classifier network composed of the instruction encoding network and the program encoding network are updated by back propagation until the model training is completed. Step 6: Use the segmented code representation network and classifier network trained in step 5 to perform vulnerability detection on the assembly code slices to be tested.

2. The binary slice-level rapid vulnerability detection method based on the xLSTM model according to claim 1 is characterized in that The specific steps of step 3 are as follows: Step 3.1: According to the assembly code slicing criteria, extract the vulnerability candidate key points of the assembly code, obtain the bidirectional slicing assembly instruction set of the assembly code by forward and backward traversing the control flow graph obtained in step 2, and generate the slicing code segment of the assembly code; Step 3.2: Mark the assembly code slices obtained in step 3.

1. The assembly code slices containing vulnerabilities are marked as "1", and the assembly code slices not containing vulnerabilities are marked as "0".

3. The binary slice-level rapid vulnerability detection method based on the xLSTM model according to claim 1 is characterized in that The specific steps of step 4 are as follows: Step 4.1: Segment the assembly code segments obtained in step 3, and use Word2Vec to embed the tokens of the assembly code segments to obtain the initial vector representation T1, T2, ..., T of each assembly instruction token. n , where n is the number of tokens in each assembly instruction, T n is the embedding representation of the nth token; Step 4.2: Denote the initial vectors obtained in step 4.1 as T1, T2, ..., T n Input into the BiGRU model to capture the local semantic features of the instruction, thereby obtaining the hidden vector representation of each assembly instruction. The specific calculation formula is as follows: BE i =BiGRU(T1,T2,...,T n ) Among them, BiGRU() represents a bidirectional gated recurrent neural network, FI i represents the embedded representation of the i-th assembly instruction; Step 4.3: Substitute the hidden vector sequence FI1, FI2, ..., FI obtained in step 4.2 n The global semantic features of the slice are captured by the xLSTM model, and the hidden vector representation of each assembly slice code segment is obtained. The specific calculation formula is as follows: FA i =xLSTM(FI1,FI2,...,FI n ) Among them, xLSTM() represents the extended long short-term memory network, FA i is the embedding representation of the i-th assembly slice code segment.

4. The binary slice-level rapid vulnerability detection method based on the xLSTM model according to claim 1 is characterized in that The specific steps of step 5 are as follows: Step 5.1: Send the hidden vector representation of the assembly slice code segment obtained in step 4 to the fully connected layer, and then pass it through the SoftMax function to convert the input real number vector into a probability distribution to obtain the prediction result; Step 5.2: Calculate the cross entropy loss based on the prediction results of step 5.1 and the label information of the assembly slice code segment; Step 5.3: Adjust the parameters of the segmentation code representation network and the classifier network according to the cross entropy loss back propagation until the loss value no longer decreases, and the training ends.

5. The binary slice-level rapid vulnerability detection method based on the xLSTM model according to claim 4 is characterized in that The specific steps of step 5.2 are as follows: Using the label value marked in step 3, when the label value is 0, it means no vulnerability, it is initialized as a vector [1,0], when the label value is 1, it means there is a vulnerability, it is initialized as a vector [0,1], and the cross entropy loss is calculated with the prediction result given by the segmented code representation network.

Images