Marking Key Segment Of Program Input Via Attention Mechanism

Key segment of a program input is the specific part of the input that has significant affect on the execution of target function. Marking key segment plays an important role in software security analysis. Traditional dynamic analysis methods can not mark the key segments correctly because of control flow dependency problem. The root cause of such problem is that implicit flow analysis method cannot cover all the behavior of the code fragment in a branch, especially when the code snippet contains unexpected jump behavior. The neural network can learn to fit the behavior of the program with proper training data. In this paper, we introduce the attention based neural network to mark the key segments of program input accurately and efficiently. We propose an attention based two-parts network structure and map program inputs into the target code execution by such network. Then we propose a two-step training method to train our network to calculate the importance of each input component on the execution of target function. Finally, we mark the key segments by statistical analysis method. We implement such method and develop a key segment marking tool AttentionMark. Experiments on four real-world software show that AttentionMark outperforms NeuralTaint and traditional dynamic analysis tool in key segment marking.