A Vulnerability Detection Algorithm Based on Transformer Model

Abstract

In today’s Internet background and the rapid development of computer science and technology, new software is born every day, whether it is on the computer or mobile phone and on the hardware. In order to meet people’s various daily needs, developers need to continuously develop new software and firmware. The software development process requires the reuse of shared codes and the realization of the middle-station module codes. These reusable codes can save developers’ development time and improve efficiency. The code of the middle-station model is highly complex, and the vulnerabilities hidden in it are not easy to be discovered. A large number of vulnerabilities are inevitably introduced, which leads to immeasurable losses in downstream task modules. In order to enable these middle-station codes to better serve downstream tasks and discover the vulnerabilities hidden in them in time, it is first necessary to extract the defined software method body from the source code. We build an abstract syntax tree for the method to form a statement set; then, the variable names, function names, and strings in the method are replaced. Each statement in the code is given a number to construct a node set. The dependency between functions and variables includes data dependency and control dependency extraction and the node set itself as the input feature of the model. This paper uses Transformer model to model the sequence information. Transformer model can make the information of each node in the sequence fully interact. Based on the Transformer model, this paper further attempts to add the attention structure to improve the probability of detecting vulnerabilities. In the final experimental results, the model can detect vulnerabilities in the code with an accuracy of 95.04% and a recall rate of 88.89%, which also proves that transformer can accurately detect vulnerabilities in the sequence.