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

Abstract

The application discloses a binary slice-level fast vulnerability detection method based on an xLSTM model, which adopts a segmented code representation architecture and a novel xLSTM model. The segmented code representation architecture is composed of an instruction encoding network and a program encoding network. By combining the instruction encoding network and the program encoding network, fine-grained semantic information and global structural information of the code can be comprehensively captured, thereby enhancing the accuracy of the vulnerability detection model. In addition, the method adopts a novel xLSTM model. The xLSTM introduces an exponential gating mechanism. Compared with the traditional Sigmoid gating, the exponential gating can provide more dynamic information filtering capability, which helps to improve the memory and forgetting process. At the same time, the xLSTM introduces additional normalization and stabilization steps in the gating calculation, which improves the stability of the model. In addition, the xLSTM introduces matrix memory, which allows parallel processing and improves storage capacity.

Description

TECHNICAL FIELD

[0001] The application belongs to the field of software vulnerability detection, and relates to a binary program vulnerability detection method, in particular to a binary program code vulnerability detection method based on an xLSTM model and a segmented code representation architecture. BACKGROUND

[0002] With the rapid development of information technology, Internet of Things (IoT) devices have been deeply integrated into many aspects of daily life. These devices mostly distribute their software in the form of binary programs, however, with the large-scale deployment of IoT devices, the security problems of underlying software have become increasingly prominent. Software vulnerabilities refer to defects in software that may be exploited by malicious individuals during its life cycle, which will lead to serious consequences such as privacy leakage, unauthorized access, and ransom. 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, taking advantage of its powerful feature extraction and pattern recognition capabilities to deeply mine vulnerability patterns in code, thereby significantly improving the accuracy and efficiency of vulnerability identification.

[0003] Due to the lack of readability of binary code, 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 for word embedding, then use deep neural networks to perform representation learning on the embedded assembly code, extract vulnerability features, and send them to a classifier for vulnerability classification. Due to the highly sequential nature of binary code, recurrent neural networks (RNN) are the ideal choice for most binary code vulnerability detection methods. However, due to the long sequence length of assembly code, RNN is prone to gradient vanishing or gradient explosion problems, which hinders the model's ability to capture long-range dependencies and affects its understanding and parsing ability of complex code logic. Although long short-term memory networks (LSTM) and gated recurrent units (GRU) are widely considered to alleviate these problems, they also face challenges when dealing with long assembly code. In recent years, the Transformer architecture has achieved remarkable success in natural language processing tasks, effectively capturing context information through self-attention mechanisms, and improving the model's performance in text generation, translation, sentiment analysis, and other tasks. Therefore, many researchers have applied the Transformer architecture to the binary code vulnerability detection task. However, the computational complexity of the Transformer architecture is O(n 2), and due to the fact that the assembly instruction sequence can be quite long, this leads to the problem of excessive consumption of computing resources and long processing time of the vulnerability detection model based on the Transformer architecture.

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

[0005] In view of the fact that the current binary program slice-level vulnerability detection method based on the traditional RNN model is difficult to capture 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 occupation, the present application provides a binary slice-level fast vulnerability detection method based on the xLSTM model. The method adopts a segmented code representation architecture and a novel xLSTM model. The segmented code representation architecture is composed 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 the tokens within the instruction. The program encoding network focuses on extracting the context information of the code segment, thereby improving the understanding and analysis ability of the program semantics. By combining the instruction encoding network and the program encoding network, the fine-grained semantic information and the global structure information of the code can be comprehensively captured, thereby enhancing the accuracy of the vulnerability detection model. In addition, the method adopts a novel xLSTM model, which can effectively learn long-distance code semantics for the long sequence characteristics of assembly code, improve the accuracy of the model, and reduce the resource consumption of the model. The xLSTM introduces an exponential gating mechanism, which can provide more dynamic information filtering capability compared with the traditional Sigmoid gating, and is helpful to improve the memory and forgetting process. At the same time, the xLSTM introduces an additional normalization and stabilization step in the gating calculation, which improves the stability of the model. In addition, the xLSTM introduces a matrix memory, which allows parallel processing and improves the storage capacity.

[0006] The purpose of the present application is achieved by the following technical solutions:

[0007] A binary slice-level fast vulnerability detection method based on the xLSTM model, comprising the following steps:

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

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

[0010] Step 3: generating assembly code slices according to the assembly code slicing criteria, and marking the vulnerability code slices.

[0011] Step 4: word embedding of tokens in the assembly code slice using Word2Vec to obtain the initial vector representation of each statement in the assembly code slice, then input it into the instruction encoding network composed of BiGRU model to obtain the hidden vector representation of each instruction, and finally input the hidden vector representation of each statement into the program encoding network composed of xLSTM model to obtain the hidden vector representation of the assembly code slice;

[0012] Step 5: input the hidden vector representation of the assembly code slice obtained in step 4 into the classifier composed of full connection layer (FCN) and SoftMax, calculate the cross-entropy loss according to the output result of the classifier and the actual label of the slice, and update the parameters of the segmented code representation network composed of the instruction encoding network and the program encoding network and the classifier network through back propagation until the model training is completed;

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

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

[0015] 1. The xLSTM model is first introduced into the field of binary program vulnerability detection. 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, capture long-term dependencies and deep semantic information in the assembly instruction sequence with long context characteristics, and thus identify 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, making its computational complexity lower than that of the Transformer architecture, so it is more suitable for implementing binary program vulnerability detection tasks in resource-limited situations such as Internet of Things devices.

[0016] 2. The segmented code representation architecture is introduced into the binary vulnerability detection task. The segmented code representation architecture is composed 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 context information of the code segment, thereby improving the understanding and analysis ability of the overall semantics of the program. By combining code semantic information of different granularities, it significantly improves the accuracy and robustness of binary program vulnerability detection. BRIEF DESCRIPTION OF DRAWINGS

[0017] Figure 1 Figure 1 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 with reference to the accompanying drawings, but is not limited thereto. Any modification or equivalent replacement of the technical solution of the present invention that does not depart from the spirit and scope of the technical solution of the present invention should be included in the scope of protection 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 annotated. 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 learned feature vector representation of the 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 IDA Pro 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 IDA Pro to generate the control flow graph (CFG) of the assembly code.

[0024] Step 3: Generate assembly code slices according to the assembly code slice guidelines 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: Word embedding of tokens in the assembly code slice using Word2Vec to obtain the initial vector representation of each statement in the assembly code slice, then input it into the instruction encoding network composed of BiGRU model to obtain the hidden vector representation of each instruction, and finally input the hidden vector representation of each statement into the program encoding network composed of xLSTM model to obtain the hidden vector representation of the assembly code slice, the specific steps are as follows:

[0028] Step 4.1: Tokenize the assembly code segment obtained in step 3, and use Word2Vec to perform word embedding on the tokens in the assembly code slice to obtain the initial vector representation T1, T2,..., T n of each assembly instruction token, 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: Input the initial vector representation T1, T2,..., T n obtained in step 4.1 into the BiGRU model to capture the local semantic features of the instructions, 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] where BiGRU() represents a bidirectional gated recurrent neural network, and FI i represents the embedding representation of the ith assembly instruction.

[0032] Step 4.3: Input the hidden vector sequence FI1, FI2,..., FI n obtained in step 4.2 into the xLSTM model to capture the global semantic features of the slice, and obtain the hidden vector representation of each assembly code segment, the specific formula calculation is as follows:

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

[0034] where xLSTM() represents an extended long short-term memory network, and FA i is the embedding representation of the ith assembly code segment.

[0035] Step 5: The hidden vector representation of the assembly code slice obtained in step 4 is input into a classifier composed of a fully connected layer (FCN) and SoftMax, and the cross-entropy loss is calculated according to the output result of the classifier and the actual label of the slice. The parameters of the segmented code representation network composed of the instruction encoding network and the program encoding network and the classifier network are updated by backpropagation until the model training is completed. The specific steps are as follows:

[0036] Step 5.1: The hidden vector representation of the assembly code slice obtained in step 4 is input into a fully connected layer, and then passed through a SoftMax function to convert the input real vector into a probability distribution to obtain the prediction result.

[0037] Step 5.2: Calculate the cross-entropy loss according to the prediction result of step 5.1 and the label information of the assembly code slice. Specifically, use the label value marked in step 3. When the label value is 0, indicating no vulnerability, it is initialized to [1, 0], and when the label value is 1, indicating a vulnerability, it is initialized to the vector [0, 1], and the cross-entropy loss is calculated with the prediction result given by the segmented code representation network.

[0038] Step 5.3: Adjust the parameters of the segmented code representation network and the classifier network according to the cross-entropy loss by backpropagation until the loss value no longer decreases, and the training is completed.

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

[0040] Embodiment:

[0041] As shown in Figure 2 , the binary program code slice is shown in Figure 2 (a), and the assembly code slice is shown in Figure 2 (b). The number before each binary program statement in (a) represents the position of the statement in the original program, and the number before each assembly instruction in (b) represents the position of the instruction in the original assembly code. Figure 2 (b) is Figure 2 (a) the corresponding assembly code slice. After the original binary program code is disassembled into assembly code using the disassembly tool IDAPro, the assembly code slice can be obtained according to the assembly code slice criterion, as shown in Figure 2The vulnerability type of the vulnerability assembly code slice example is Stack Overflow (stack memory overflow), specifically, the vulnerability occurs at line 37, because the size and boundary of the target memory are not fully checked when initializing the data, which may modify the memory address at rbp-184. Next, mov BYTE PTR[rax], 0 attempts to write data to the memory location pointed to by rax, which will cause a stack overflow if rax has been incorrectly modified to an address outside the allocated range, which may trigger a program crash or be exploited by an attacker to execute arbitrary code.

[0042] As shown in Figure 1 The embodiment first tokenizes the assembly code slice, then inputs each token obtained by tokenization into the Word2Vec model to obtain the initial vector representation of the token, then inputs the obtained token embedding into the instruction encoding network to obtain the embedding representation of each assembly instruction, next inputs the instruction embedding representation into the program encoding network to obtain the embedding representation of the assembly code slice, and finally inputs the embedding representation of the assembly code slice into the full connection layer, and then through SoftMax, converts the input real number vector into a binary classification probability to obtain the prediction result of the slice. The method of the present application can accurately identify the vulnerability in the assembly code slice.

Claims

1. A binary slice-level fast vulnerability detection method based on the xLSTM model, characterized by The method comprises the following steps: Step 1: Use the decompiler tool IDA Pro 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 slice to obtain the initial vector representation of each statement in the assembly code slice. Then, feed it into the instruction encoding network composed of the BiGRU model to obtain the hidden vector representation of each instruction. Finally, feed the hidden vector representation of each statement into the program encoding network composed of the xLSTM model to obtain the hidden vector representation of the assembly code slice. Step 5: The hidden vector representation of the assembly code slice obtained in step 4 is fed into a classifier consisting of a fully connected layer and a SoftMax layer. The cross entropy loss is calculated based on the output of the classifier and the actual label of the slice. Backpropagation is then performed to update the parameters of the segmented code representation network and the classifier network, which are composed of the instruction encoding network and the program encoding network, 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 slice 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: Extract the vulnerability candidate key points of the assembly code according to the assembly code slicing criteria, obtain the bidirectional slicing assembly instruction set of the assembly code by traversing the control flow graph obtained in step 2 forward and backward, 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 without 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 embed the tokens of the assembly code segments using Word2Vec 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: The hidden vector sequence FI1, FI2, ..., FI obtained in step 4.2 is 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 backpropagation 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 indicates no vulnerability, and its initial vector is [1,0]. When the label value is 1, it indicates a vulnerability, it is initialized to a vector [0,1], and the cross entropy loss is calculated with the prediction result given by the segmented code representation network.

Images