A code vulnerability detection and line risk sorting method based on mixed features

Abstract

The application belongs to the field of network space security, and particularly relates to a code vulnerability detection and line risk sorting method based on mixed features; comprising: collecting a code data set and preprocessing the same; extracting a code mark vector; extracting a code statement vector; learning a global feature vector of a graph structure of the code by using a graph attention network and a pooling layer; learning a semantic global feature vector of the code; weighting and splicing the global feature vector of the graph structure and the global feature vector of the semantics to obtain a mixed feature vector; inputting the mixed feature vector into a full connection layer and a Softmax layer for training to obtain a trained code vulnerability detection model; sorting the code lines by risk by using an HA-IG method to obtain a risk sorting result of the code lines; and the application can better cope with different types and structures of codes and improve adaptability to diversified vulnerability scenarios.

Description

Technical Field

[0001] This invention belongs to the field of cyberspace security, specifically relating to a method for code vulnerability detection and line risk ranking based on hybrid features. Background Technology

[0002] With the widespread application of software systems across various fields, the security threats posed by software vulnerabilities are becoming increasingly serious. In particular, the frequent occurrence of security problems caused by software vulnerabilities has resulted in significant losses to society and the economy. Code vulnerability detection is a crucial technical means to ensure the security of software systems. Its purpose is to detect potential security vulnerabilities in the code during the code development phase or during software operation, and then to fix them. Traditional vulnerability detection methods mainly rely on techniques such as static analysis, dynamic analysis, and symbolic execution. However, these methods have shown certain limitations when dealing with large projects, complex dependencies, and novel vulnerability types.

[0003] With the rapid development of deep learning technology, deep learning-based code vulnerability detection methods have gradually gained attention. In recent years, the introduction of Natural Language Processing (NLP) technology and Graph Neural Networks (GNNs) has enhanced capabilities in code semantic understanding and dependency modeling. Models such as Long Short-Term Memory Networks (LSTM) and Gated Recurrent Units (GRUs) can learn rich contextual information from large-scale code corpora, effectively improving code representation capabilities. Simultaneously, Graph Neural Networks can model the dependencies between statements, functions, and variables in the code, capturing deeper structural information.

[0004] However, current deep learning-based code vulnerability detection methods still have room for improvement in some aspects. On the one hand, existing detection methods struggle to fully integrate features from different levels within the code, posing challenges in feature fusion. On the other hand, existing detection methods are typically coarse-grained, making it difficult to identify and rank high-risk lines of code. Therefore, how to effectively combine information from various dimensions of the code to more accurately detect vulnerabilities has become an urgent problem to be solved. Summary of the Invention

[0005] To address the shortcomings of existing technologies, this invention proposes a code vulnerability detection and line risk ranking method based on hybrid features. The method includes: inputting code into a code vulnerability detection model for training during the training process, and obtaining code vulnerability detection results and line risk ranking results during the testing process.

[0006] The method includes the following steps:

[0007] S1: Collect the code dataset and preprocess it to obtain the training set and test set;

[0008] S2: Use pre-trained Word2Vec to extract the embedding vectors of code tags;

[0009] S3: Extract the vectors of code statements by sequentially using bidirectional RG-LRU and multi-head Word Attention layers;

[0010] S4: Based on the vector of the code statements, the graph attention network GAT and the pooling layer are used in sequence to learn the global feature vector N1 of the graph structure of the code;

[0011] S5: Based on the vector of the code statements, the semantic global feature vector N2 of the code is learned by sequentially using bidirectional RG-LRU and multi-head Sentence Attention layers;

[0012] S6: Weighted concatenation of the global feature vector of the code graph structure with the global feature vector of the semantics to obtain a hybrid feature vector;

[0013] S7: Input the hybrid feature vector into the fully connected layer and the Softmax layer in sequence for training to obtain the trained code vulnerability detection model;

[0014] S8: During the testing process, the code is inspected using the above steps, and the HA-IG method proposed in this solution is used to rank the risk of the code lines, resulting in the risk ranking of the code lines.

[0015] Preferably, the preprocessing of the code dataset includes:

[0016] S11: Remove comments from the code to obtain clean code;

[0017] S12: Organize each statement in the code into a single line;

[0018] S13: Mark vulnerable code functions as 1, and mark patched and secure code functions as 0;

[0019] S14: Split the dataset into a training set and a test set.

[0020] Preferably, the process of extracting embedding vectors of code tags using pre-trained Word2Vec includes:

[0021] S21: Use the BPE tokenizer to split the code text into code tags;

[0022] S22: Use pre-trained Word2Vec to embed code tags as vector representations of code tags.

[0023] Preferably, the process of extracting vectors of code statements by sequentially using bidirectional RG-LRU and multi-head Word Attention layers includes:

[0024] S31: Input the embedding vector of the code tag into the bidirectional RG-LRU layer for learning to obtain the hidden state representation of the tag;

[0025] S32: Input the hidden state representation of the label into the multi-head Word Attention layer to obtain the attention score of the label;

[0026] S33: The hidden state representation of the label is weighted and summed with the attention scores of each head, and then divided by the number of attention heads to obtain the vector of code statements.

[0027] Preferably, the process of sequentially learning the global feature vector N1 of the graph structure using the graph attention network GAT and pooling layers includes:

[0028] S41: Use Joern to extract the data dependencies and control dependencies of the code to obtain an adjacency matrix containing the dependencies;

[0029] S42: Input the vector of code statements and the adjacency matrix into GAT to obtain the statement vector after learning the dependencies;

[0030] S43: Pass the statement vectors after learning the dependencies through an average pooling layer to obtain the global feature vector N1 of the code's graph structure.

[0031] Preferably, the process of learning the semantic global feature vector N2 of the code by sequentially using bidirectional RG-LRU and multi-head Sentence Attention layers includes:

[0032] S51: Input the vector of the code statement into the bidirectional RG-LRU layer for learning to obtain the hidden state representation of the statement;

[0033] S52: Input the hidden state representation of the sentence into the multi-head Sentence Attention layer to obtain the attention score of the sentence;

[0034] S53: The hidden state representation of the statement is weighted and summed with the attention scores of each head, and then divided by the number of attention heads to obtain the semantic global feature vector N2 of the code.

[0035] Preferably, during testing, the process of using the HA-IG method to prioritize lines of code based on risk includes:

[0036] S81: Sort the attention scores of the labels calculated in the multi-head Word Attention layer in descending order, take the top k options, and obtain the label risk table;

[0037] S82: Calculate the average value c of the integrated gradient for each label row by row;

[0038] S83: Calculate the risk score for a line of code using the following formula:

[0039]

[0040] Among them, risk i Let represent the risk score of the line of code, c represent the average ensemble gradient of all labels in that line, k represent the total number of elements in the label risk table, t represent the rank of the highest-ranking label in the label risk table for that line (t = k if the label does not appear in the table), and a represent the risk score of the line of code. si This represents the average attention score of this line of code in a multi-head Sentence Attention layer, calculated by the number of heads, where u represents the scaling factor;

[0041] S84: Sort the risk scores of the code lines in descending order to obtain the risk ranking table of the code lines.

[0042] Furthermore, the average value c of the ensemble gradient for each marker in the code line is calculated as follows:

[0043]

[0044] Among them, IG it Let x represent the ensemble gradient of the t-th code tag in the i-th line. it Let x′ be the vector representation of the t-th code tag in the i-th line. it represents the baseline input, F represents the model's prediction function, and m represents the total number of labels in the i-th row.

[0045] The beneficial effects of this invention are as follows:

[0046] 1. This invention proposes a novel vulnerability detection method (framework) that better combines the semantic information and structural dependency information of the code, making full use of their complementary information, ensuring the integrity of information during feature extraction, and not omitting key information, thereby improving the accuracy of vulnerability detection.

[0047] 2. This invention introduces a Real-Gated Linear Recurrent Unit, which overcomes the drawback of excessive parameters in Long Short-Term Memory networks and the instability and inaccuracy of Gated Recurrent Units in processing long-term dependent sequence data. It simultaneously maintains a smaller parameter set and the ability to process longer sequences.

[0048] 3. This invention improves the word attention and sentence attention in the hierarchical architecture into a multi-head form, thereby improving the overall robustness and generalization ability of the model.

[0049] 4. This invention proposes an improved method for risk ranking of lines of code—HA-IG. This method combines two attention scores and ensemble gradients proposed in this paper to rank lines of code by risk, overcoming the limitations of a single indicator and improving the accuracy of risk ranking. Attached Figure Description

[0050] Figure 1 This is a diagram illustrating the overall framework of the code vulnerability detection and line risk ranking method based on hybrid features in this invention. Detailed Implementation

[0051] The technical solutions of the embodiments of the present invention will be clearly and completely described below with reference to the accompanying drawings. Obviously, the described embodiments are only some embodiments of the present invention, and not all embodiments. Based on the embodiments of the present invention, all other embodiments obtained by those skilled in the art without creative effort are within the scope of protection of the present invention.

[0052] This invention proposes a method for code vulnerability detection and line risk ranking based on hybrid features, such as... Figure 1 As shown, the method includes the following:

[0053] During training, the code is input into the code vulnerability detection model for training. During testing, the code vulnerability detection results and the line risk ranking results of the code are obtained.

[0054] The method includes the following steps:

[0055] S1: Collect the code dataset and preprocess it to obtain the training set and test set.

[0056] S11: Remove comments from the code to obtain clean code.

[0057] Preferably, code datasets can be collected from NVD, SARD vulnerability databases, GitHub, etc., and irrelevant data such as comments can be removed to obtain clean code.

[0058] S12: Organize each statement in the code into the same line and split it according to the statement. In this solution, the code statement and the code line have the same meaning.

[0059] S13: Mark vulnerable code functions as 1, and mark patched and secure code functions as 0.

[0060] S14: Split the dataset into a training set and a test set.

[0061] The dataset is split into independent training and testing datasets. Preferably, the ratio of the training to testing datasets is 8:2.

[0062] S2: Use pre-trained Word2Vec to extract the embedding vectors of code tags.

[0063] S21: Use the BPE tokenizer to split the code text into code tags;

[0064] S22: Use pre-trained Word2Vec to embed code tags as vector representations of code tags.

[0065] S3: The vectors of code statements are extracted sequentially using a bidirectional Real-Gated Linear Recurrent Unit (RG-LRU) and a multi-head WordAttention layer.

[0066] S31: Input the embedding vector of the code tag into the bidirectional RG-LRU layer for learning to obtain the hidden state representation of the tag;

[0067] The RG-LRU cell is described as follows:

[0068] r t =σ(W r x t +b r )

[0069] i t =σ(W i x t +b i )

[0070]

[0071] Where σ is the sigmoid function, x t W represents the input at the current moment. r and W i Let b be the weight matrix. r and b i For the bias term, a = σ(Λ), Λ is a learnable parameter matrix, c is a scalar constant, ⊙ denotes the Hadamard product, and h t-1 h represents the hidden state at the previous time step. t This indicates the current hidden state and output.

[0072] S32: Input the hidden state representation of the label into the multi-head Word Attention layer to obtain the attention score of the label;

[0073] S33: The hidden state representation of the label is weighted and summed with the attention scores of each head, and then divided by the number of attention heads to obtain the vector of code statements.

[0074] The description of the multi-head Word Attention layer and the calculation of the code statement vectors are as follows:

[0075] u it =tanh(W w h it +b w )

[0076]

[0077] s i ′=∑ t α it h it

[0078] Where, α it W is the attention score for the t-th label in the i-th row. w Let b be the weight matrix. w h is the bias term. it The currently marked hidden state, u w s is a learnable context vector. i ' is the vector representation of a single-headed computation statement. This scheme adopts a multi-headed approach, that is, using multiple n... w W w and b w Simultaneously learning, the final code statement vector representation is as follows:

[0079]

[0080] Among them, s i is a vector representation of the code statements, and m is the number of attention heads.

[0081] S4: The global feature vector N1 of the graph structure of the code is learned by sequentially using the graph attention network GAT and the pooling layer.

[0082] S41: Use Joern to extract the data dependencies and control dependencies of the code to obtain an adjacency matrix containing the dependencies;

[0083] Joern is used to extract the data dependency graph and control dependency graph of code functions. Each statement is treated as a node, and an initial adjacency matrix is ​​created based on the number of nodes, with each element having a value of 0. For each edge with a dependency relationship, the corresponding position in the adjacency matrix is ​​marked with 1, and the rest are marked with 0. Based on the data dependency relationship or control dependency relationship, following the above steps, an adjacency matrix containing the dependencies can be obtained.

[0084] S42: Input the vector of code statements and the adjacency matrix into GAT to obtain the statement vector after learning the dependencies;

[0085] The GAT is described as follows:

[0086]

[0087]

[0088] Where l represents the current state, W is the node embedding vector of the current layer. (l) It is a learnable weight matrix, a (l) It is a learnable weight vector used to compute attention scores between nodes, || denotes vector concatenation, LeakyReLU is the activation function, and N i Let σ be the set of neighboring nodes of node i, and let σ be the sigmoid activation function.

[0089] S43: Pass the statement vectors after learning the dependencies through an average pooling layer to obtain the global feature vector N1 of the code's graph structure.

[0090] The global average pooling operation is represented as:

[0091]

[0092] Where N1 represents the global feature vector of the graph structure of the code, V is the set of nodes, and h v It is the feature vector of node v, where each node represents a code statement.

[0093] S5: The semantic global feature vector N2 of the code is learned by sequentially using bidirectional RG-LRU and multi-head Sentence Attention layers.

[0094] S51: Input the vector of the code statement into the bidirectional RG-LRU layer for learning to obtain the hidden state representation of the statement;

[0095] The RG-LRU cell is described in step S31.

[0096] S52: Input the hidden state representation of the sentence into the multi-head Sentence Attention layer to obtain the attention score of the sentence;

[0097] S53: The hidden state representation of the statement is weighted and summed with the attention scores of each head, and then divided by the number of attention heads to obtain the semantic global feature vector N2 of the code.

[0098] The description of the multi-head Sentence Attention layer and the calculation of the semantic global feature vector of the code are shown below:

[0099] u i =tanh(W s h i +b s )

[0100]

[0101] e′=∑ i α i h i

[0102] Where, α i Let W be the attention score for the i-th line of code. s Let b be the weight matrix. s For the bias term, h i This represents the hidden state of the current code statement, u s Let e ​​be a learnable context vector, and e′ be a semantic global feature vector representation computed by a single head. This scheme adopts a multi-head approach, that is, it uses multiple u... s W s and b s Simultaneous learning is performed, and the final semantic global feature vector is represented as follows:

[0103]

[0104] Where N2 is the semantic global feature vector of the code, and m is the number of attention heads.

[0105] S6: Weighted concatenation of the global feature vector of the code graph structure with the global feature vector of the semantics to obtain a hybrid feature vector.

[0106] The global feature vector N1 of the code graph structure and the global feature vector N2 of the semantics are weighted and concatenated. The weight parameters can be represented as (a, b). The result is a mixed feature vector of the code function [a*N1||b*N2]. The weight parameters can be adjusted arbitrarily according to the experimental conditions. The initial value can be set to (1, 1).

[0107] S7: Input the hybrid feature vector into the fully connected layer and the Softmax layer for training to obtain the trained code vulnerability detection model.

[0108] The fully connected layer is described as follows:

[0109] y = W T x+b

[0110] Where W is the weight matrix, b is the bias term, and x is the input vector.

[0111] The output of the Softmax layer is the vulnerability detection result of the code, which is specifically described as follows:

[0112]

[0113] Among them, z i is the i-th component of the input, and e is the base of the natural logarithm.

[0114] Preferably, during training, dropout is used for regularization, and the dropout rate can be set to 0.5.

[0115] The performance evaluation metrics for the model include precision, recall, accuracy, and the F1 score.

[0116]

[0117] Wherein, TP represents the number of samples where the true value is positive and the predicted value is positive; FP represents the number of samples where the true value is negative and the predicted value is positive; FN represents the number of samples where the true value is positive and the predicted value is negative; and TN represents the number of samples where the true value is negative and the predicted value is negative.

[0118] S8: During the testing process, the code is inspected using the above steps, and the HA-IG method proposed in this solution is used to rank the risk of the code lines, resulting in the risk ranking of the code lines.

[0119] S81: Sort the attention scores of the labels calculated in the multi-head Word Attention layer in descending order, take the top k options, and obtain the label risk table;

[0120] Attention score α it The calculation method is shown in step S3.

[0121] S82: Calculate the average value c of the ensemble gradient for each label row by row.

[0122]

[0123] Among them, IG it Let x represent the ensemble gradient of the t-th code tag in the i-th line. it Let x′ be the vector representation of the t-th code tag in the i-th line. it represents the baseline input (usually the zero vector or the mean), F represents the model's prediction function, and m represents the total number of labels in the i-th row.

[0124] S83: Calculate the risk score for a line of code using the following formula:

[0125]

[0126] Among them, risk i Let represent the risk score of the line of code, c represent the average ensemble gradient of all labels in that line, k represent the total number of elements in the label risk table, t represent the rank of the highest-ranking label in the label risk table for that line (t = k if the label does not appear in the table), and a represent the risk score of the line of code. si This represents the average attention score of this line of code in a multi-head Sentence Attention layer, calculated by the number of heads (the attention score α of the line of code). i The calculation is shown in step S5), where u represents the scaling factor (the initial value can be set to 1, and it can be increased as needed when there are too many lines of code);

[0127] S84: Sort the risk scores of the code lines in descending order to obtain the risk ranking table of the code lines.

[0128] After training, the code vulnerability detection model with optimal parameters is saved. The code to be detected is then obtained and input into the trained code vulnerability detection model to obtain the code vulnerability detection results and the line risk ranking results of the code.

[0129] The above-described embodiments further illustrate the purpose, technical solution, and advantages of the present invention. It should be understood that the above-described embodiments are merely preferred embodiments of the present invention and are not intended to limit the present invention. Any modifications, equivalent substitutions, improvements, etc., made to the present invention within the spirit and principles of the present invention should be included within the protection scope of the present invention.

Claims

1. A method for code vulnerability detection and line risk ranking based on hybrid features, characterized in that, include: During training, the code is input into the code vulnerability detection model for training, and during testing, the code vulnerability detection results and the line risk ranking results of the code are obtained; The method includes the following steps: S1: Collect the code dataset and preprocess it to obtain the training set and test set; S2: Use pre-trained Word2Vec to extract the embedding vectors of code tags; S3: Extract the vector of the code statement by sequentially using a bidirectional RG-LRU and a multi-head Word Attention layer; the process of extracting the vector of the code statement by sequentially using a bidirectional RG-LRU and a multi-head Word Attention layer includes: S31: Input the embedding vector of the code tag into the bidirectional RG-LRU layer for learning to obtain the hidden state representation of the tag; S32: Input the hidden state representation of the label into the multi-head Word Attention layer to obtain the attention score of the label; S33: The hidden state representation of the label is weighted and summed with the attention scores of each head, and then divided by the number of attention heads to obtain the vector of code statements; S4: Based on the vector of the code statements, the graph attention network GAT and the average pooling layer are used in sequence to learn the global feature vector N1 of the graph structure of the code; S5: Based on the vector of the code statements, the semantic global feature vector N2 of the code is learned by sequentially using bidirectional RG-LRU and multi-head Sentence Attention layers; S6: Weighted concatenation of the global feature vector of the code graph structure with the global feature vector of the semantics to obtain a hybrid feature vector; S7: Input the hybrid feature vector into the fully connected layer and the Softmax layer in sequence for training to obtain the trained code vulnerability detection model; S8: During testing, the code is inspected using the above steps, and the HA-IG method is used to rank the lines of code by risk, obtaining the risk ranking results. The process of ranking lines of code by risk using the HA-IG method during testing includes: S81: Sort the attention scores of the labels calculated in the multi-head Word Attention layer in descending order, take the top k options, and obtain the label risk table; S82: Calculate the average value c of the integrated gradient for each label row by row; S83: Calculate the risk score for a line of code using the following formula: Among them, risk i Let represent the risk score of the line of code, c represent the average ensemble gradient of all labels in that line, k represent the total number of elements in the label risk table, t represent the rank of the highest-ranking label in the label risk table for that line (t = k if the label does not appear in the table), and a represent the risk score of the line of code. si This represents the average attention score of this line of code in a multi-head Sentence Attention layer, calculated by the number of heads, where u represents the scaling factor. The average value c of the ensemble gradient for each marker in the code line is calculated as follows: Among them, IG it Let x represent the ensemble gradient of the t-th code tag in the i-th line. it Let x′ be the vector representation of the t-th code tag in the i-th line. it represents the baseline input, F represents the model's prediction function, and m represents the total number of labels in the i-th row; S84: Sort the risk scores of the code lines in descending order to obtain the risk ranking table of the code lines.

2. The method for code vulnerability detection and line risk ranking based on hybrid features according to claim 1, characterized in that, The preprocessing of the code dataset includes: S11: Remove comments from the code to obtain clean code; S12: Organize each statement in the code into a single line; S13: Mark vulnerable code functions as 1, and mark patched and secure code functions as 0; S14: Split the dataset into a training set and a test set.

3. The method for code vulnerability detection and line risk ranking based on hybrid features according to claim 1, characterized in that, The process of extracting embedding vectors of code tags using pre-trained Word2Vec includes: S21: Use the BPE tokenizer to split the code text into code tags; S22: Use pre-trained Word2Vec to embed code tags as vector representations of code tags.

4. The method for code vulnerability detection and line risk ranking based on hybrid features according to claim 1, characterized in that, The process of learning the global feature vector N1 of the graph structure of the code by sequentially using the graph attention network GAT and pooling layers includes: S41: Use Joern to extract the data dependencies and control dependencies of the code to obtain an adjacency matrix containing the dependencies; S42: Input the vector of code statements and the adjacency matrix into GAT to obtain the statement vector after learning the dependencies; S43: Pass the statement vectors after learning the dependencies through an average pooling layer to obtain the global feature vector N1 of the code's graph structure.

5. The method for code vulnerability detection and line risk ranking based on hybrid features according to claim 1, characterized in that, The process of learning the semantic global feature vector N2 of the code by sequentially using bidirectional RG-LRU and multi-head Sentence Attention layers includes: S51: Input the vector of the code statement into the bidirectional RG-LRU layer for learning to obtain the hidden state representation of the statement; S52: Input the hidden state representation of the sentence into the multi-head Sentence Attention layer to obtain the attention score of the sentence; S53: The hidden state representation of the statement is weighted and summed with the attention scores of each head, and then divided by the number of attention heads to obtain the semantic global feature vector N2 of the code.

Images