Statement-level code vulnerability detection method based on fusion feature prompt tuning

Abstract

The invention discloses a statement-level code vulnerability detection method based on fusion feature prompt tuning, which comprises the following steps of: firstly, analyzing a source code by utilizing a static analysis tool, and generating a slice code segment by utilizing a program slicing technology; then, combining the statement-level hard prompt with the statement-level soft prompt to form a fusion feature prompt; then, fusing the slice features and the statement features, and embedding the features into the model; and finally, realizing statement-level vulnerability detection by utilizing fusion feature prompt tuning and a pre-training model. According to the method, the prompt template is designed according to the characteristics of the vulnerability statements, pre-training knowledge can be fully utilized, and the structure and attribute information of vulnerability codes can be learned, so that the model can more accurately distinguish the vulnerability statements from non-vulnerability statements, and the possibility of false alarms is reduced. In addition, according to the method, the understanding ability of the pre-training model on the context can be fully utilized, and fine grammar and semantic information in codes can be captured.

Description

Technical Field

[0001] The present invention relates to a vulnerability detection method, and in particular to a statement-level code vulnerability detection method based on fusion feature prompt tuning. Background Art

[0002] Vulnerability detection is an important process of identifying potential security vulnerabilities in software through various techniques and tools. As the complexity of software systems increases, the difficulty and importance of vulnerability detection are also increasing. Vulnerability detection methods based on machine learning do not require manual establishment of rules, but still require manual feature extraction. Vulnerability detection methods based on deep learning can handle complex code structures and large-scale data. Their application in vulnerability detection has undergone significant development, from the initial convolutional neural network (CNN) and long short-term memory network (LSTM) methods to the introduction of graph neural network (GNN), which further improves the ability to capture code structure and dependencies. However, these methods still face problems such as large data requirements, high training costs, and poor interpretability. In recent years, the rise of pre-trained models (such as CodeBERT and GraphCodeBERT) has significantly improved the automatic feature extraction and generalization capabilities of vulnerability detection by performing self-supervised learning and contrastive learning on large-scale code bases. Pre-trained models are trained using large-scale unlabeled data, reducing the reliance on manual feature engineering, and have strong context understanding and transfer learning capabilities.

[0003] However, there is a gap between the pre-training and fine-tuning processes of these pre-trained models. For example, pre-trained models such as CodeBERT and CodeT5 are usually pre-trained with the masked language modeling (MLM) objective. The input of MLM is a mixture of code snippets and natural language text, and the model is trained to predict randomly masked input tokens. However, when the model is fine-tuned on downstream tasks, such as vulnerability detection, the input of the model only involves source code without natural language, and the training objective is also changed from a generation problem to a classification problem. Therefore, the input and objectives between the pre-training task and the fine-tuning task are inconsistent, making it difficult for the knowledge of the pre-trained model to be fully mined, resulting in only suboptimal results for downstream tasks.

[0004] Prompt Tuning optimizes the performance of pre-trained models by designing and adjusting input prompts, and has shown great potential in the fields of code analysis and vulnerability detection. Compared with traditional fine-tuning, prompt tuning reduces computing resources and time requirements, effectively utilizes pre-trained knowledge, and improves flexibility. At present, there are few studies on vulnerability detection using prompt tuning. One is Gao et al.'s study on the effectiveness of prompt tuning in code intelligence tasks (Chaozheng Wang, Yuanhang Yang, Cuiyun Gao, Yun Peng, Hongyu Zhang, Michael R. Lyu: No More Fine-Tuning? An Experimental Evaluation of Prompt Tuning in Code Intelligence. ESEC / SIGSOFT FSE 2022: 382-394). This study explores the effectiveness of prompt tuning in code intelligence tasks and compares it with traditional fine-tuning methods. The research mainly focuses on three core tasks: function-level vulnerability detection, code summary generation, and code translation. The pre-trained models used include CodeBERT and CodeT5. The other is the ProRLearn method proposed by Ren et al. (Zilong Ren, Xiaolin Ju, Xiang Chen, Hao Shen: ProRLearn: Boosting Prompt Tuning-based Vulnerability Detection by Reinforcement Learning. 2024), which explores the combination of prompt tuning and reinforcement learning techniques to improve the performance of vulnerability detection models. This method greatly improves the performance of pre-trained models in vulnerability detection tasks by dynamically optimizing prompts. However, these two methods are currently limited to function-level vulnerability detection, and there is still a certain gap with statement-level vulnerability location. In addition, designing more effective prompt templates and optimization strategies is also worthy of further study.

[0005] Therefore, the current challenges of statement-level vulnerability detection mainly lie in:

[0006] 1. There is a gap between the pre-training and fine-tuning process of the training model. How to fully tap the knowledge of the pre-training model and adapt the model to the input characteristics of the vulnerability code?

[0007] 2. Statement-level vulnerability detection requires the model to understand deeper semantic information and be able to capture subtle differences between statements. It also requires the model to accurately detect and locate each code statement and identify small but potentially serious vulnerabilities.

[0008] 3. How to choose a suitable pre-trained language model that requires rich code corpus and strong semantic understanding ability to adapt to sentence-level vulnerability detection tasks. Summary of the invention

[0009] The present invention provides a statement-level code vulnerability detection method based on fusion feature prompt tuning. The method designs prompt templates according to the characteristics of the vulnerability statement, can make full use of pre-trained knowledge, learn the structure and attribute information of the vulnerability code, and enable the model to more accurately distinguish between vulnerability statements and non-vulnerability statements, reducing the possibility of false positives. In addition, the method can make full use of the pre-trained model's ability to understand context and capture subtle grammatical and semantic information in the code, which is crucial for accurately understanding and detecting statement-level vulnerabilities, because vulnerabilities are often embedded in complex code logic and context.

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

[0011] A statement-level code vulnerability detection method based on fusion feature prompt tuning includes the following steps:

[0012] Step 1: Use static analysis tools to parse the source code, generate an abstract syntax tree (AST) and a program dependency graph (PDG), and then use program slicing technology to generate sliced ​​code segments;

[0013] Step 2: Combine the hard prompts and soft prompts at the sentence level to form a fused feature prompt;

[0014] Step 3: Generate the type vector representation and semantic vector representation of the sentence, and weightedly fuse the two to form a comprehensive feature vector representation of the sentence;

[0015] Step 4: Convert the slice code segment vulnerability label [0,1] to a vector representation [f1,f2] as the slice code segment label vector, where f1 indicates that the slice code segment has no vulnerability and f2 indicates that the slice code segment has a vulnerability;

[0016] Step 5: Design task-related label sets for non-vulnerable sentences and vulnerability sentences respectively, using two label words to represent each category. The label set for vulnerability sentences is {Terrible, Unsatisfactory}, and the label set for non-vulnerable sentences is {Great, Satisfactory}.

[0017] Step 6: Use a supervised learning method to concatenate the slice code segment label vector and the statement comprehensive feature vector to obtain the final feature vector representation of each statement, which is sent to the CodeGeeX model. The result is forward calculated through the fully connected layer and the softmax layer. The cross entropy loss function is calculated using the category label information of each statement, and the network parameters are adjusted according to the error back propagation.

[0018] Step 7: Use the trained CodeGeeX model to perform vulnerability detection on the test set, and analyze the vulnerability detection results through vulnerability identification strategy and vulnerability location strategy, so as to achieve statement-level vulnerability identification and vulnerability location.

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

[0020] 1. This invention proposes a statement-level code vulnerability detection method based on fusion feature prompt tuning for statement-level vulnerability detection tasks, combining hard prompts and soft prompts at the statement level to improve the flexibility and adaptability of the model. Hard prompts guide the model to focus on specific tasks and key features through clear natural language prompts, and enhance the interpretability of the results; soft prompts use a learnable continuous differentiable vector as the prompt content, and splice it with the original input vector as a new input. The learnable vector reduces the burden of manually designing prompt templates, and can be optimized through gradients to quickly adapt to vulnerability detection tasks.

[0021] 2. The present invention comprehensively utilizes the slice code segment label vulnerability label and statement type features to realize the combination of slice code segment label features and statement-level features, thereby enhancing the model's ability to understand the code. The slice code segment label features provide high-level contextual information, while the statement type features refine the specific structure of the code, allowing for fine-grained control and guidance of the model, helping the model to better understand the context and capture nuances, so that the model can more comprehensively capture vulnerability-related information. This multi-level, multi-dimensional feature fusion makes vulnerability detection more accurate and reliable.

[0022] 3. The present invention can realize statement-level vulnerability identification and location. By integrating the features of code slices and statements, the syntax and semantic information of each code statement can be deeply analyzed. Compared with function-level or file-level detection methods, the model can accurately detect and locate each code statement, identifying small or potentially serious vulnerabilities. For example, problems such as boundary condition errors or uninitialized variables may only exist in certain specific statements. In contrast, coarse-grained detection (such as function-level or slice-level) may ignore these details. In addition, statement-level detection can provide precise location and detailed context information to help developers quickly locate and fix problems. In contrast, coarse-grained detection can only provide a warning over a larger range, and developers need to spend more time to locate specific problems.

[0023] 4. The present invention applies prompt tuning to the pre-trained model CodeGeeX to achieve fine-grained vulnerability detection tasks. CodeGeeX is a large-scale pre-trained programming language model based on Transformer. It is an autoregressive decoder generated from left to right, taking code or natural language identifiers (tokens) as input and predicting the probability distribution of the next identifier. This pre-training enables it to perform well in processing code-related tasks. In the vulnerability detection task, this ability enables it to perform deep semantic analysis of code snippets, understand the structure and logical relationships of the code, and identify complex vulnerability patterns, rather than just matching based on surface features. BRIEF DESCRIPTION OF THE DRAWINGS

[0024] Figure 1 It is a schematic diagram of the overall process of the statement-level code vulnerability detection method based on fusion feature prompt tuning of the present invention.

[0025] Figure 2 It is the vulnerability code.

[0026] Figure 3 It is a flowchart of the construction of fusion feature prompts of vulnerability code. DETAILED DESCRIPTION

[0027] 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.

[0028] The present invention provides a statement-level code vulnerability detection method based on fusion feature prompt tuning. First, the input code and vulnerability candidate syntax are matched, and the vulnerability candidate key points are extracted. These key points are used as slicing criteria to generate corresponding slice code segments using program slicing technology. This step helps to accurately extract code context information related to the vulnerability candidate key points. Soft prompts and hard prompts are combined at the statement level. Specifically, natural language prompt statements are designed for the code, task objectives are clarified, continuous vectors are used to replace part of the input, and these vectors are optimized through training. Then, slice features and statement features are fused, and these features are embedded into the model to help the model focus on specific types of statements and understand the code logic more accurately. Finally, fine-tuning is performed on the annotated data set to optimize the model parameters and prompt vectors. A supervised learning method is used to optimize the model by comparing the model output with the annotation results and using the loss function to improve the model's vulnerability detection capability and accuracy. Figure 1 As shown, the specific steps of the method are as follows:

[0029] Step 1: Use static analysis tools to parse the source code, generate an abstract syntax tree (AST) and a program dependency graph (PDG), and then use program slicing technology to generate sliced ​​code segments. The specific steps are as follows:

[0030] Step 11: Use static analysis tools to parse the source code, generate AST and PDG, and extract statement type information.

[0031] Step 12: According to the candidate vulnerability syntax, traverse the AST, match and extract the four key vulnerability nodes that meet the rules (API function call, pointer reference, array reference, arithmetic expression).

[0032] Step 13: Analyze the data dependencies in PDG, use the vulnerability candidate key nodes as slicing criteria, use program slicing technology to generate forward slices and backward slices, and merge the two into slice code segments.

[0033] Step 2: Combine the sentence-level hard prompts and soft prompts to form a fused feature prompt. The specific steps are as follows:

[0034] Step 21: Reconstruct the input hard prompts, which is specifically implemented by adding natural language prompts after each code statement, that is, [STATEMENT] Review the statements: [MASK], to reconstruct the model input. Each word in these prompts has a clear and understandable meaning. The purpose is to elicit the specific task knowledge learned during the pre-training process during the tuning phase. Through clear prompt content, the model can better understand the intention and context of each statement.

[0035] Step 22: Construct soft prompt input, which is specifically implemented by taking a learnable continuous vector as the prompt content and concatenating it with the original input vector as the new input. Unlike hard prompts, the template words of soft prompts are no longer natural language, but a learnable continuous vector. These vectors can be optimized through gradients and quickly adapt to vulnerability detection tasks.

[0036] Step 23: Combine the soft prompts with the hard prompts. Specifically, the non-critical natural language parts in the hard prompts are replaced by learnable continuous vectors (for example, [STATEMENT], Review, and [MASK] are clear prompts and are therefore retained) to form a new fused feature prompt, namely [STATEMENT]Review[α1][α2]:[MASK], which aims to reduce the burden of manually selecting prompt templates. The soft prompt vector and model parameters are adjusted by optimizing the loss function to guide the model to focus on the vulnerability detection task and better adapt to different input features.

[0037] Step 3: Generate the type vector representation and semantic vector representation of the sentence, and weightedly fuse the two to form a comprehensive feature vector representation of the sentence. The specific steps are as follows:

[0038] Step 31: According to the syntax structure of the code, the code statements are divided into five types of statements: Function, IdentifierDeclStatement, Condition, ExpressionStatement and Return. Function means defining a function, including function name, parameter list and function body. In the code, identifying function definition helps to understand the structure and modularization of the code. By analyzing function definition, the scope and calling relationship of the function can be determined. IdentifierDeclStatement means declaring a variable and assigning an initial value to it. Identifying variable declaration helps to track the life cycle and data flow of the variable, and helps to find vulnerabilities such as uninitialized variable use and variable range overflow. Condition means conditional statement, which controls the execution path of the program. It determines which branch to execute based on the truth of the condition. The analysis of conditional statement helps to understand the logical branch and control flow of the program, and helps to detect logical vulnerabilities and conditional competition. ExpressionStatement is an expression statement, which is usually used for assignment, function call or operation, which helps to understand the interaction between variables and functions, and helps to find arithmetic overflow and unhandled vulnerabilities. Return is a return statement. Analyzing the return statement helps to understand the output behavior of the function and helps to find whether the type and range of the return value are correct.

[0039] Step 32: Generate one-hot encoding vectors [t1, t2, t3, t4, t5] for each sentence type and use the pre-trained model to generate the semantic vectors [v1, v2, …, v n ], where n is the number of lines of code, and the statement type feature and the statement semantic vector are given learnable weights ω1 and ω2 to generate a weighted fused statement comprehensive feature vector: s i =ω1*t j +ω2*v i , i∈[1,n],j∈[1,5], where t j Represents the statement v in the slice i The corresponding statement type vector.

[0040] Step 4: Convert the slice code segment vulnerability label [0,1] into a vector representation [f1,f2] as the slice code segment label vector, where f1 indicates that the slice code segment has no vulnerability and f2 indicates that the slice code segment has a vulnerability.

[0041] Step 5: Design task-related label sets for non-vulnerable sentences and vulnerability sentences respectively, and use two label words to represent each category. These label sets provide a clear classification basis for the model. The label set for vulnerability sentences is {Terrible, Unsatisfactory}, and the label set for non-vulnerable sentences is {Great, Satisfactory}.

[0042] Step 6: Use the supervised learning method to concatenate the slice code segment label vector and the statement comprehensive feature vector to obtain the final feature vector representation of each statement, which is sent to the CodeGeeX model. The forward calculation results are performed through the fully connected layer and the softmax layer. The cross entropy loss function is calculated using the category label information of each statement, and the network parameters are adjusted according to the error back propagation. The specific steps are as follows:

[0043] Step 61: Slice the code segment label vector [f1, f2] and the statement comprehensive feature vector s i =ω1*t j +ω2*v i Concatenate and get the final vector representation of the sentence, denoted as z = [f k ,s1,s2,…,s n ], f k The label vector corresponding to the slice code snippet.

[0044] Step 62: Send the vector into the CodeGeeX model, pass it through the fully connected layer and the softmax layer, and use the formula Get the forward calculation result, where σ(z) i represents the Softmax value of the i-th sentence, z i represents the i-th component of the input vector z, and N represents the dimension of the input vector.

[0045] Step 63: Use the category label information of each sentence to calculate the cross entropy loss function, the formula is Among them, N represents the number of samples, C represents the number of categories, and y i,c represents the true label of the i-th sample in the c-th category, p i,c Represents the predicted probability of the i-th sample in the c-th class. The loss function is optimized to adjust the soft prompt vectors α1, α2, weight coefficients ω1, ω2 and model parameters.

[0046] Step 7: Use the trained CodeGeeX model to perform vulnerability detection on the test set, and analyze the vulnerability detection results through vulnerability identification strategy and vulnerability location strategy, so as to achieve statement-level vulnerability identification and vulnerability location.

[0047] Example:

[0048] Taking the vulnerability code of CVE-2022-25258 as an example, the role of fusion feature prompt fine-tuning and vulnerability detection results are explained. The code slicing part is as follows Figure 2 As shown in the figure, the process of constructing fusion feature hints is as follows Figure 3 CVE-2022-25258 is a vulnerability in the Linux kernel USB Gadget subsystem. Due to the lack of proper validation of operating system descriptor interface requests, these requests may involve array index out of bounds or null pointer dereference, thereby causing memory corruption.

[0049] The interface variable is derived from w_value (line 5) and later used in the usb_configuration->interface array in the count_ext_prop function (line 13). Since the size of the c->interface array is MAX_CONFIG_INTERFACES (16) and the interface variable is not validated in both case USB_RECIP_INTERFACE (line 2) and the called function (line 9), this allows the attacker to index the c->interface array outside of its actual boundaries. In addition, the count_ext_prop function lacks validation, which leads to a NULL pointer dereference when the retrieved usb_function pointer is NULL. It should break if the value of the interface variable is greater than or equal to MAX_CONFIG_INTERFACES or when the retrieved usb_function pointer is NULL. Therefore, line 13 has a vulnerability of type CWE-119: Improper Restriction of Operations within the Bounds of a Memory Buffer.

[0050] The method proposed in the present invention can first provide slice-level features for the code by fusing feature prompts. This step enables the model to understand the overall structure and logic of the code by providing high-level contextual information. Next, statement-level features are provided for each statement to refine the specific structure of the code, allowing the model to perform fine-grained control and guidance on the code. For example, for the vulnerable statement on line 13, its statement type is ExpressionStatement. ExpressionStatement represents a statement containing an expression. This statement type may have vulnerabilities such as SQL injection, buffer overflow, null pointer dereference, resource leakage, etc. Statement-level feature prompts can guide the model to focus on analyzing whether the statement has the above-mentioned types of vulnerabilities. The fusion of slice-level features and statement-level features realizes multi-level and multi-dimensional feature synthesis, enabling the model to capture information related to vulnerabilities more comprehensively.

[0051] In addition, the method of combining hard and soft prompts at the statement level improves the flexibility and adaptability of the model, especially in dealing with complex contexts. CWE-119 type vulnerabilities (such as buffer overflows) are often hidden in complex code logic. The present invention can accurately detect and locate each code statement, identify small and potentially serious vulnerabilities, and thus improve the accuracy and reliability of overall vulnerability detection.

Claims

1. A statement-level code vulnerability detection method based on fusion feature prompt tuning, characterized in that The method comprises the following steps: Step 1: Use static analysis tools to parse the source code, generate an abstract syntax tree AST and a program dependency graph PDG, and then use program slicing technology to generate sliced ​​code segments; Step 2: Combine the hard prompts and soft prompts at the sentence level to form a fused feature prompt; Step 3: Generate the type vector representation and semantic vector representation of the sentence, and weightedly fuse the two to form a comprehensive feature vector representation of the sentence; Step 4: Convert the slice code segment vulnerability label [0,1] to a vector representation [f1,f2] as the slice code segment label vector, where f1 indicates that the slice code segment has no vulnerability and f2 indicates that the slice code segment has a vulnerability; Step 5: Design task-related label sets for non-vulnerable sentences and vulnerability sentences respectively, using two label words to represent each category. The label set for vulnerability sentences is {Terrible, Unsatisfactory}, and the label set for non-vulnerable sentences is {Great, Satisfactory}. Step 6: Use a supervised learning method to concatenate the slice code segment label vector and the statement comprehensive feature vector to obtain the final feature vector representation of each statement, which is sent to the CodeGeeX model. The result is forward calculated through the fully connected layer and the softmax layer. The cross entropy loss function is calculated using the category label information of each statement, and the network parameters are adjusted according to the error back propagation. Step 7: Use the trained CodeGeeX model to perform vulnerability detection on the test set, and analyze the vulnerability detection results through vulnerability identification strategy and vulnerability location strategy, so as to achieve statement-level vulnerability identification and vulnerability location.

2. The statement-level code vulnerability detection method based on fusion feature prompt tuning according to claim 1 is characterized in that The specific steps of step 1 are as follows: Step 11: Use static analysis tools to parse the source code, generate AST and PDG, and extract statement type information; Step 12: According to the candidate vulnerability syntax, traverse the AST, match and extract the four key vulnerability nodes that meet the rules: API function call, pointer reference, array reference, and arithmetic expression; Step 13: Analyze the data dependencies in PDG, use the vulnerability candidate key nodes as slicing criteria, use program slicing technology to generate forward slices and backward slices, and merge the two into slice code segments.

3. The statement-level code vulnerability detection method based on fusion feature prompt tuning according to claim 1 is characterized in that The specific steps of step 2 are as follows: Step 21: Reconstruct the input hard prompts and add natural language prompts after each code statement to reconstruct the model input; Step 22: Construct soft prompt input, take a learnable continuous vector as the prompt content, and concatenate it with the original input vector as the new input; Step 23: Combine the soft prompts with the hard prompts, and use the learnable continuous vector to replace the non-critical natural language part of the hard prompts to form a new fused feature prompt.

4. The statement-level code vulnerability detection method based on fusion feature prompt tuning according to claim 1 is characterized in that The specific steps of step 3 are as follows: Step 31: Classify the code statements into five types according to the grammatical structure of the code: Function, IdentifierDeclStatement, Condition, ExpressionStatement and Return; Step 32: Generate one-hot encoding vectors [t1, t2, t3, t4, t5] for each sentence type and use the pre-trained model to generate the semantic vectors [v1, v2, ..., v n ], assign learnable weights ω1 and ω2 to the sentence type feature and the sentence semantic vector, and generate a weighted fused sentence comprehensive feature vector: s i =ω1*t j +ω2*v i , i∈[1,n], j∈[1,5], where n is the number of lines of code, t j Represents the statement v in the slice i The corresponding statement type vector.

5. The statement-level code vulnerability detection method based on fusion feature prompt tuning according to claim 1 is characterized in that The specific steps of step 6 are as follows: Step 61: Slice the code segment label vector [f1, f2] and the statement comprehensive feature vector s i =ω1*t j +ω2*v i Concatenate and get the final vector representation of the sentence, denoted as z = [f k , s1, s2, ..., s n ], f k Represents the label vector corresponding to the slice code segment, f1 means that the slice code segment has no vulnerability, and f2 means that the slice code segment has a vulnerability; Step 62: Send the vector into the CodeGeeX model, pass it through the fully connected layer and the softmax layer, and use the formula Get the forward calculation result, where σ(z) i represents the Softmax value of the i-th sentence, z i represents the i-th component of the input vector z, and N represents the dimension of the input vector; Step 63: Use the category label information of each sentence to calculate the cross entropy loss function, the formula is Among them, N represents the number of samples, C represents the number of categories, and y i，c represents the true label of the i-th sample in the c-th category, p i，c Represents the predicted probability of the i-th sample in the c-th category, and optimizes the loss function to adjust the soft prompt vectors α1, α2, weight coefficients ω1, ω2 and model parameters.

Images