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

Abstract

This invention discloses a statement-level code vulnerability detection method based on fused feature prompt tuning. First, static analysis tools are used to parse source code, and program slicing technology is used to generate sliced ​​code segments. Statement-level hard and soft prompts are then combined to form fused feature prompts. Slice features and statement features are then fused and embedded into a model. Finally, statement-level vulnerability detection is achieved using fused feature prompt tuning and a pre-trained model. This method designs prompt templates based on the characteristics of vulnerable statements, fully leveraging pre-trained knowledge to learn the structure and attribute information of vulnerable code. This allows the model to more accurately distinguish between vulnerable and non-vulnerable statements, reducing the likelihood of false positives. Furthermore, this method can fully leverage the pre-trained model's ability to understand context, capturing subtle grammatical and semantic information in the code.

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 a crucial process for identifying potential security vulnerabilities in software through various techniques and tools. As software systems become increasingly complex, vulnerability detection becomes increasingly challenging and important. Machine learning-based vulnerability detection methods eliminate the need for manual rule-building, but still require manual feature extraction. Deep learning-based vulnerability detection methods are capable of handling complex code structures and large amounts of data. Their application in vulnerability detection has undergone significant development, evolving from initial approaches such as convolutional neural networks (CNNs) and long short-term memory networks (LSTMs) to the introduction of graph neural networks (GNNs), further enhancing their ability to capture code structure and dependencies. However, these methods still face challenges 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 and contrastive learning on large code bases. Pre-trained models utilize large amounts of unlabeled data for training, reducing reliance on manual feature engineering and possessing powerful contextual 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 using 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 also changes 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 to fully mine the knowledge of the pre-trained model, 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 to traditional fine-tuning, prompt tuning reduces computational resources and time requirements, effectively leverages pre-trained knowledge, and improves flexibility. Currently, there is limited research on using prompt tuning for vulnerability detection. One study, by Gao et al., examined the effectiveness of prompt tuning for 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 explored the effectiveness of prompt tuning for code intelligence tasks and compared it with traditional fine-tuning methods. The study focused on three core tasks: function-level vulnerability detection, code summary generation, and code translation, using pre-trained models such as CodeBERT and CodeT5. Another paper, ProRLearn (Zilong Ren, Xiaolin Ju, Xiang Chen, Hao Shen: ProRLearn: Boosting Prompt Tuning-based Vulnerability Detection by Reinforcement Learning, 2024), proposed by Ren et al., explores combining prompt tuning and reinforcement learning techniques to improve the performance of vulnerability detection models. This method significantly improves the performance of pre-trained models in vulnerability detection tasks by dynamically optimizing prompts. However, both methods are currently limited to function-level vulnerability detection and still lag behind statement-level vulnerability localization. Furthermore, designing more effective prompt templates and optimization strategies warrants further research.

[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 processes of the training model. How to fully exploit the knowledge of the pre-trained model and adapt the model to the input characteristics of the vulnerable 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 capabilities to adapt to statement-level vulnerability detection tasks. Summary of the Invention

[0009] This paper provides a statement-level code vulnerability detection method based on fused feature prompt tuning. This method designs prompt templates based on the characteristics of vulnerable statements. This method fully leverages pre-trained knowledge to learn the structure and attributes of vulnerable code, enabling the model to more accurately distinguish between vulnerable and non-vulnerable statements and reducing the likelihood of false positives. Furthermore, this method leverages the pre-trained model's ability to understand context, capturing subtle syntactic and semantic information in the code. This is crucial for accurately understanding and detecting statement-level vulnerabilities, as vulnerabilities are often embedded in complex code logic and context.

[0010] The purpose 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 sentence-level hard and soft prompts to form fused feature prompts;

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

[0015] Step 4: Convert the vulnerability label of the code segment [0, 1] into a vector representation [f1, f2] as the code segment label vector, where f1 indicates that the code segment has no vulnerability and f2 indicates that the 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: Using a supervised learning method, concatenate the sliced ​​code segment label vector and the statement comprehensive feature vector to obtain the final feature vector representation of each statement. This vector is fed into the CodeGeeX model, and the results are 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 based on the error backpropagation.

[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 and vulnerability location strategies 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 paper proposes a statement-level code vulnerability detection method based on fusion feature prompt tuning for statement-level vulnerability detection tasks. This method combines statement-level hard and soft prompts to improve the model's flexibility and adaptability. Hard prompts use clear natural language prompts to guide the model's focus on specific tasks and key features, enhancing the interpretability of the results. Soft prompts use a learnable continuous differentiable vector as the prompt content, concatenating it with the original input vector as a new input. Learnable vectors reduce the burden of manually designing prompt templates and can be optimized via gradients, allowing for rapid adaptation to vulnerability detection tasks.

[0021] 2. This invention comprehensively utilizes sliced ​​code segment labels, vulnerability labels, and statement type features to achieve a combination of sliced ​​code segment label features and statement-level features, enhancing the model's ability to understand code. Sliced ​​code segment label features provide high-level contextual information, while statement type features refine the specific structure of the code, allowing for fine-grained control and guidance of the model, helping the model better understand context and capture nuances, enabling the model to 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 minor 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. This invention applies prompt tuning to the pre-trained model CodeGeeX to achieve fine-grained vulnerability detection. CodeGeeX is a large-scale pre-trained programming language model based on Transformer. It is a left-to-right generative autoregressive decoder that takes code or natural language tokens as input and predicts the probability distribution of the next token. This pre-training enables it to perform excellent performance on code-related tasks. In vulnerability detection tasks, this capability enables deep semantic analysis of code snippets, understanding the structure and logical relationships of the code, and identifying complex vulnerability patterns, rather than simply 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's 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 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.

[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 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 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 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 a 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 sentence-level hard and soft prompts to form fused feature prompts. The specific steps are as follows:

[0034] Step 21: Reconstruct the input hard prompts. This is achieved by adding natural language prompts after each code statement, namely [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 task-specific knowledge learned during pre-training during the tuning phase. With clear prompt content, the model can better understand the intent and context of each statement.

[0035] Step 22: Construct the soft hint input. This is achieved by concatenating a learnable continuous vector as the hint content with the original input vector as the new input. Unlike hard hints, the template word for soft hints is no longer natural language, but a learnable continuous vector. These vectors can be optimized using gradients to quickly adapt to vulnerability detection tasks.

[0036] Step 23: Combine soft prompts with hard prompts. Specifically, use a learnable continuous vector to replace the non-critical natural language part of the hard prompt (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. By optimizing the loss function, the soft prompt vector and model parameters are adjusted 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: Based on the code's grammatical structure, code statements are classified into five types: Function, IdentifierDeclStatement, Condition, ExpressionStatement, and Return. Function defines a function, including its name, parameter list, and body. Identifying function definitions in code helps understand the structure and modularity of the code. By analyzing function definitions, you can determine the function's scope and calling relationships. IdentifierDeclStatement declares a variable and assigns it an initial value. Identifying variable declarations helps track the variable's lifecycle and data flow, helping to discover vulnerabilities such as uninitialized variable usage and variable range overflow. Condition represents a conditional statement, which controls the program's execution path and determines which branch to execute based on the condition's truth. Analyzing conditional statements helps understand the program's logical branches and control flow, helping to detect logical vulnerabilities and conditional races. ExpressionStatement is an expression statement, typically used for assignments, function calls, or calculations. It helps understand the interaction between variables and functions and helps discover arithmetic overflows and unhandled vulnerabilities. Return represents a return statement. Analyzing return statements helps understand the function's output behavior and helps determine whether the return value's type and range 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 sentence type features and the sentence semantic vector are given learnable weights ω1 and ω2 to generate the weighted fusion sentence 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 vulnerability label of the slice code segment [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-specific label sets for non-vulnerable and vulnerabilities, using two labels to represent each category. These label sets provide a clear classification basis for the model. The label set for vulnerabilities is {Terrible, Unsatisfactory}, and the label set for non-vulnerabilities is {Great, Satisfactory}.

[0042] Step 6: Using a supervised learning method, concatenate the slice code segment label vector and the statement comprehensive feature vector to obtain the final feature vector representation of each statement. This vector is fed into the CodeGeeX model, and the results are 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 based on 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 Splicing, get the final vector representation of the sentence, denoted as z = [f k ,s1,s2,…,s n ], f k Indicates the label vector corresponding to the slice code segment.

[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. Optimize the loss function 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 and vulnerability location strategies 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's 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 is 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 case USB_RECIP_INTERFACE (line 2) or in the called function (line 9), this allows an attacker to index the c->interface array outside of its actual bounds. Furthermore, the count_ext_prop function lacks validation, which results in a NULL pointer dereference when the retrieved usb_function pointer is NULL. 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, it should break. Therefore, line 13 is vulnerable to 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 type of statement may have vulnerabilities such as SQL injection, buffer overflow, null pointer dereference, and resource leakage. 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 integration, enabling the model to more comprehensively capture information related to vulnerabilities.

[0051] Furthermore, the combination of statement-level hard and soft prompts improves the model's flexibility and adaptability, particularly in handling complex contexts. CWE-119 vulnerabilities (such as buffer overflows) are often hidden within complex code logic. This method can accurately detect and locate each code statement, identifying small but potentially serious vulnerabilities, thereby improving the accuracy and reliability of overall vulnerability detection.

Claims

1. A statement-level code vulnerability detection method based on fusion feature prompt tuning, characterized by 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 sentence-level hard and soft prompts to form fused feature prompts. The specific steps are as follows: Step 21: Reconstruct the input hard prompts by adding natural language prompts after each code statement to reconstruct the model input; Step 22: Construct soft prompt input, taking a learnable continuous vector as the prompt content and concatenating it with the original input vector as the new input; Step 23: Combine the soft prompt with the hard prompt and use a learnable continuous vector to replace the non-critical natural language part of the hard prompt to form a new 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. The specific steps 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 for each sentence type , use the pre-trained model to generate the semantic vector of the sentence , giving learnable weights to sentence type features and sentence semantic vectors and , generate the weighted fusion sentence comprehensive feature vector: , where n is the number of lines of code, Represents statements in a slice The corresponding statement type vector; Step 4: Convert the slice code segment vulnerability label [0,1] into a vector representation , as the slice code segment label vector, where Indicates that the slice code segment has no vulnerabilities. Indicates that there is a vulnerability in the slice code segment; 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: Using a supervised learning method, concatenate the slice code segment label vector and the statement comprehensive feature vector to obtain the final feature vector representation of each statement. This vector is fed into the CodeGeeX model, and the results are 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 based on the error back propagation. The specific steps are as follows: Step 61: Slice the code snippet label vector and sentence comprehensive feature vector Splicing, get the final vector representation of the sentence, denoted as , Represents the label vector corresponding to the slice code segment, Indicates that the slice code segment has no vulnerabilities. Indicates that there is a vulnerability in the slice code segment; 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 Indicates the The Softmax value of the sentence, represents the first A quantity, 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 , where N represents the number of samples and C represents the number of categories. Indicates the The sample in The true label of the class, Indicates the The sample in The predicted probability of the class, the optimized loss function adjusts the soft prompt vector , weight coefficient and model parameters; Step 7: Use the trained CodeGeeX model to perform vulnerability detection on the test set, and analyze the vulnerability detection results through vulnerability identification and vulnerability location strategies 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: Based on 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.

Images