Code vulnerability detection and row risk sorting method based on mixed features

Abstract

The invention belongs to the field of network space security, and particularly relates to a code vulnerability detection and row risk sorting method based on mixed features. Comprising the steps that a code data set is collected and preprocessed; extracting a vector marked by the code; extracting vectors of the code statements; learning a global feature vector of a graph structure of the code by adopting a graph attention network and a pooling layer; learning a semantic global feature vector of the code; performing weighted splicing on 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; performing risk sorting on the code lines by using an HA-IG method to obtain a risk sorting result of the code lines; according to the method, codes of different types and structures can be better handled, and the adaptability to diversified vulnerability scenes is improved.

Description

Technical Field

[0001] The present invention belongs to the field of cyberspace security, and specifically relates to a code vulnerability detection and risk ranking method based on hybrid features. Background Art

[0002] With the widespread application of software systems in various fields, the security threats brought by software vulnerabilities are becoming increasingly serious. In particular, the frequent occurrence of security issues caused by software vulnerabilities has caused huge losses to society and the economy. Code vulnerability detection is an important technical means to ensure the security of software systems. Its purpose is to detect potential security vulnerabilities in the code during the code development stage or the software operation process, and then repair them. Traditional vulnerability detection methods mainly rely on static analysis, dynamic analysis, and symbolic execution technologies, but these methods have shown certain limitations when dealing with large projects, complex dependencies, and new types of vulnerabilities.

[0003] With the rapid development of deep learning technology, code vulnerability detection methods based on deep learning have gradually received attention. In recent years, the introduction of natural language processing (NLP) technology and graph neural networks (GNN) has enabled stronger capabilities in code semantic understanding and dependency modeling. Models such as long short-term memory networks (LSTM) and gated recurrent units (GRU) can learn rich contextual information from large-scale code corpora, effectively improving the representation ability of the code. At the same time, graph neural networks can model the dependencies between statements, functions, and variables in the code, capturing deeper structural information.

[0004] However, the current deep learning-based code vulnerability detection methods still have room for improvement in some aspects. On the one hand, it is difficult for existing detection methods to fully combine features at different levels in the code, and there are challenges in feature fusion. On the other hand, existing detection methods are usually coarse-grained and difficult to identify and sort code lines with higher risks. Therefore, how to effectively combine information from various dimensions of the code to more accurately detect vulnerabilities in the code has become an urgent problem to be solved. Summary of the invention

[0005] In view of the shortcomings of the prior art, the present invention proposes a code vulnerability detection and line risk ranking method based on hybrid features, the method comprising: inputting the code into a code vulnerability detection model for training during a training process, and obtaining the code vulnerability detection result and the line risk ranking result of the code during a test process;

[0006] The method comprises the following steps:

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

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

[0009] S3: Use bidirectional RG-LRU and multi-head Word Attention layers to extract the vector of the code sentence;

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

[0011] S5: Based on the vector of the code sentence, bidirectional RG-LRU and multi-head Sentence Attention layers are used to learn the semantic global feature vector N2 of the code;

[0012] S6: weighted concatenation of the global feature vector of the code graph structure and the global feature vector of the semantics to obtain a mixed feature vector;

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

[0014] S8: During the testing process, the code is detected using the above steps, and the risk ranking of the code lines is performed using the HA-IG method proposed in this solution to obtain the risk ranking results of the code lines.

[0015] Preferably, the process of preprocessing the code data set includes:

[0016] S11: Delete the comments in the code to get pure code;

[0017] S12: Arrange each statement in the code into the same line;

[0018] S13: mark the code function with vulnerability as 1, and mark the patch of the vulnerability code function and the safe code function as 0;

[0019] S14: Split the dataset into training and testing sets.

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

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

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

[0023] Preferably, the process of sequentially using the bidirectional RG-LRU and the multi-head Word Attention layer to extract the vector of the code sentence 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 marker into the multi-head Word Attention layer to obtain the marker's attention score;

[0026] S33: Take the weighted sum of the hidden state representation of the marker and the attention score of each head, and then divide it by the number of attention heads to obtain the vector of the code statement.

[0027] Preferably, the process of sequentially using the graph attention network GAT and the pooling layer to learn the global feature vector N1 of the graph structure of the code includes:

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

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

[0030] S43: The sentence vector after learning the dependency is passed through the average pooling layer to obtain the global feature vector N1 of the code graph structure.

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

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

[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: Take the weighted sum of the hidden state representation of the sentence and the attention scores of each head, and then divide it by the number of attention heads to obtain the semantic global feature vector N2 of the code.

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

[0036] S81: Arrange the attention scores of the tags calculated in the multi-head Word Attention layer in descending order, take the first k options, and obtain the tag risk table;

[0037] S82: Obtain the average value c of the integrated gradient of each marker row by row;

[0038] S83: Calculate the risk score of the code line according to the following formula:

[0039]

[0040] Among them, risk i represents the risk score of the code line, c represents the average integrated gradient of all tags in the line, k represents the total number of elements in the tag risk table, t represents the rank of the highest-ranking tag in the tag risk table for the line (if it does not appear in the table, t = k), a si It represents the average of the attention scores of the code line in the multi-head Sentence Attention layer according to the number of heads, and u represents the scaling factor;

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

[0042] Furthermore, the average value c of the integrated gradient of each mark of the code line is calculated as follows:

[0043]

[0044] Among them, IG it represents the integrated gradient of the tth code token in the i-th row, x it represents the vector representation of the tth code tag in the i-th row, x′ it represents the baseline input, F represents the prediction function of the model, and m represents the total number of tokens in the i-th row.

[0045] The beneficial effects of the present invention are:

[0046] 1. The present invention proposes a new vulnerability detection method (framework), which better combines the semantic information and structural dependency information of the code, makes full use of the complementary information of the two, ensures the integrity of information in the feature extraction process, and does not miss key information, thereby improving the accuracy of vulnerability detection.

[0047] 2. The present invention introduces the Real-Gated Linear Recurrent Unit, which can overcome the shortcomings of the Long Short-Term Memory network (LSM) with too many parameters and the Gated Recurrent Unit (Gated Recurrent Unit) with insufficient stability and accuracy in processing long-term dependent sequence data. At the same time, it maintains the ability to process fewer parameters and longer sequences.

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

[0049] 4. The present invention proposes an improved method for risk ranking of code lines - HA-IG. This method combines the two attention scores and integrated gradient in the article to rank the risks of code lines, overcomes the limitations of a single indicator, and improves the accuracy of risk ranking. BRIEF DESCRIPTION OF THE DRAWINGS

[0050] Figure 1 This is an overall framework diagram of the code vulnerability detection and risk ranking method based on hybrid features in the present invention. DETAILED DESCRIPTION

[0051] The following will be combined with the drawings in the embodiments of the present invention to clearly and completely describe the technical solutions in the embodiments of the present invention. Obviously, the described embodiments are only part of the embodiments of the present invention, not all of the embodiments. Based on the embodiments of the present invention, all other embodiments obtained by ordinary technicians in this field without creative work are within the scope of protection of the present invention.

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

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

[0054] The method comprises the following steps:

[0055] S1: Collect code datasets and preprocess them to obtain training sets and test sets.

[0056] S11: Delete the comments in the code to get pure code.

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

[0058] S12: Arrange each statement in the code into the same line and split it by statement. In this solution, code statements and code lines have the same meaning;

[0059] S13: Mark the code function with vulnerability as 1, and mark the vulnerability code function patch and safe code function as 0.

[0060] S14: Split the dataset into training and testing sets.

[0061] The data set is split into a training data set and a test data set that are independent of each other. Preferably, the ratio of the training data set to the test data set can be set to 8:2.

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

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

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

[0065] S3: The bidirectional Real-Gated Linear Recurrent Unit (RG-LRU) and multi-head WordAttention layer are used in turn to extract the vector of the code sentence.

[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 description of the RG-LRU unit is as follows:

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

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

[0070]

[0071] Among them, σ is the sigmoid function, x t represents the input at the current moment, W r and W i is the weight matrix, b r and b i is the bias term, a=σ(Λ), Λ is the learnable parameter matrix, c is a scalar constant, ⊙ represents the Hadamard product, h t-1 represents the hidden state at the previous moment, h t Represents the current hidden state and output.

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

[0073] S33: Take the weighted sum of the hidden state representation of the marker and the attention score of each head, and then divide it by the number of attention heads to obtain the vector of the code statement.

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

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

[0076]

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

[0078] Among them, α it is the attention score of the tth tag in the i-th row, W w is the weight matrix, b w is the bias term, h it is the hidden state of the current mark, u w is the learnable context vector, s i ' is the vector representation of the statement calculated by a single head. This solution adopts the form of multiple heads, that is, multiple n w , W w and b w Learning is performed simultaneously, and the final code sentence vector is represented as:

[0079]

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

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

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

[0083] Use Joern to extract the data dependency graph and control dependency graph of the code function, treat each statement as a node, and create an initial adjacency matrix based on the number of nodes, with the value of each element being 0. For each edge with a dependency relationship, mark the corresponding position of the adjacency matrix as 1, and the rest as 0. According to the data dependency relationship or the control dependency relationship, follow the above steps to obtain an adjacency matrix containing the dependency relationship.

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

[0085] The description of GAT is as follows:

[0086]

[0087]

[0088] Among them, l represents the current state, is the node embedding vector of the current layer, W (l) is a learnable weight matrix, a (l) is a learnable weight vector and is used to calculate the attention score between nodes, || represents the vector concatenation operation, LeakyReLU is the activation function, and N i is the set of neighbor nodes of node i, and σ is the sigmoid activation function.

[0089] S43: The sentence vector after learning the dependency is passed through the average pooling layer to obtain the global feature vector N1 of the code graph structure.

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

[0091]

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

[0093] S5: Bidirectional RG-LRU and multi-head Sentence Attention layers are used in turn to learn the semantic global feature vector N2 of the code.

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

[0095] The description of the RG-LRU unit is shown 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: Take the weighted sum of the hidden state representation of the sentence and the attention scores of each head, and then divide it 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 as follows:

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

[0100]

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

[0102] Among them, α i is the attention score of the i-th line of code, W s is the weight matrix, b s is the bias term, h i is the hidden state of the current code statement, u s is a learnable context vector, and e′ is a semantic global feature vector representation of a single-head calculation. This scheme adopts a multi-head form, that is, multiple u s , W s and b s Learning is performed simultaneously, and the final semantic global feature vector is expressed as:

[0103]

[0104] Among them, 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 and the global feature vector of the semantics to obtain a mixed 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 parameter can be expressed as (a, b). The mixed feature vector [a*N1||b*N2] of the code function is obtained by merging. The weight parameter can be adjusted arbitrarily according to the experimental situation. The initial value can be set to (1, 1).

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

[0108] The description of the fully connected layer is:

[0109] y=W T x+b

[0110] Among them, 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 expressed as:

[0112]

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

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

[0115] The performance evaluation indicators of the model include precision, recall, accuracy and comprehensive evaluation indicator F1:

[0116]

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

[0118] S8: During the testing process, the code is detected using the above steps, and the risk ranking of the code lines is performed using the HA-IG method proposed in this solution to obtain the risk ranking results of the code lines.

[0119] S81: Arrange the attention scores of the tags calculated in the multi-head Word Attention layer in descending order, take the first k options, and obtain the tag risk table;

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

[0121] S82: Obtain the average value c of the integrated gradient of each marker row by row:

[0122]

[0123] Among them, IG it represents the integrated gradient of the tth code token in the i-th row, x it represents the vector representation of the tth code tag in the i-th row, x′ it represents the baseline input (usually a zero vector or the mean), F represents the prediction function of the model, and m represents the total number of tokens in the i-th row.

[0124] S83: Calculate the risk score of the code line according to the following formula:

[0125]

[0126] Among them, risk i represents the risk score of the code line, c represents the average integrated gradient of all tags in the line, k represents the total number of elements in the tag risk table, t represents the rank of the highest-ranking tag in the tag risk table for the line (if it does not appear in the table, t = k), a si Indicates the average of the attention scores of the code line in the multi-head Sentence Attention layer according to the number of heads (the attention score of the code line α i The calculation is shown in step S5), u represents the scaling factor (the initial value can be set to 1, and can be increased according to the situation when there are too many lines of code);

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

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

[0129] The above embodiments further illustrate the purpose, technical solutions and advantages of the present invention in detail. It should be understood that the above embodiments are only preferred implementation modes 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 in the protection scope of the present invention.

Claims

1. A code vulnerability detection and risk ranking method based on hybrid features, characterized in that: include: During the training process, the code is input into the code vulnerability detection model for training, and during the testing process, the code vulnerability detection results and the line risk ranking results of the code are obtained; The method comprises the following steps: S1: Collect and preprocess the code dataset to obtain the training set and test set; S2: Use pre-trained Word2Vec to extract the embedding vector of the code token; S3: Use bidirectional RG-LRU and multi-head Word Attention layers to extract the vector of the code sentence; S4: Based on the vector of the code statement, the graph attention network GAT and the average pooling layer are used in turn to learn the global feature vector N1 of the code graph structure; S5: Based on the vector of the code sentence, bidirectional RG-LRU and multi-head Sentence Attention layers are used to learn the semantic global feature vector N2 of the code; S6: weighted concatenation of the global feature vector of the code graph structure and the global feature vector of the semantics to obtain a mixed feature vector; S7: Input the mixed feature vector into the fully connected layer and the Softmax layer in sequence for training to obtain a trained code vulnerability detection model; S8: During the test, the code is detected using the above steps, and the risk ranking of the code lines is performed using the HA-IG method to obtain the risk ranking results of the code lines.

2. According to the method of claim 1, the method is characterized in that: The process of preprocessing the code dataset includes: S11: Delete the comments in the code to get pure code; S12: Arrange each statement in the code into the same line; S13: mark the code function with vulnerability as 1, and mark the patch of the vulnerability code function and the safe code function as 0; S14: Split the dataset into training and testing sets.

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

4. According to the method of code vulnerability detection and risk ranking based on hybrid features in claim 1, it is characterized in that: The process of extracting the vector of the code sentence using the bidirectional RG-LRU and multi-head Word Attention layers in sequence 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 marker into the multi-head Word Attention layer to obtain the marker's attention score; S33: Take the weighted sum of the hidden state representation of the marker and the attention score of each head, and then divide it by the number of attention heads to obtain the vector of the code statement.

5. According to the method of code vulnerability detection and risk ranking based on hybrid features in claim 1, it is characterized in that: The process of sequentially using the graph attention network GAT and the pooling layer to learn the global feature vector N1 of the graph structure of the code includes: S41: Use Joern to extract the data dependency and control dependency of the code and obtain the adjacency matrix containing the dependency relationship; S42: Input the vector and adjacency matrix of the code statement into GAT to obtain the statement vector after learning the dependency relationship; S43: The sentence vector after learning the dependency is passed through the average pooling layer to obtain the global feature vector N1 of the code graph structure.

6. The method for code vulnerability detection and risk ranking based on hybrid features according to claim 1 is characterized in that: The process of learning the semantic global feature vector N2 of the code using the bidirectional RG-LRU and multi-head Sentence Attention layers in sequence includes: S51: input the vector of the code sentence into the bidirectional RG-LRU layer for learning to obtain the hidden state representation of the sentence; 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: Take the weighted sum of the hidden state representation of the sentence and the attention scores of each head, and then divide it by the number of attention heads to obtain the semantic global feature vector N2 of the code.

7. The method for code vulnerability detection and risk ranking based on hybrid features according to claim 1 is characterized in that: During testing, the process of risk ranking of code lines using the HA-IG method includes: S81: Arrange the attention scores of the tags calculated in the multi-head Word Attention layer in descending order, take the first k options, and obtain the tag risk table; S82: Obtain the average value c of the integrated gradient of each marker row by row; S83: Calculate the risk score of the code line according to the following formula: Among them, risk i represents the risk score of the code line, c represents the average integrated gradient of all tags in the line, k represents the total number of elements in the tag risk table, t represents the rank of the highest-ranking tag in the tag risk table for the line (if it does not appear in the table, t = k), a si It represents the average of the attention scores of the code line in the multi-head Sentence Attention layer according to the number of heads, and u represents the scaling factor; S84: Sort the risk scores of the code lines in descending order to obtain a risk ranking table of the code lines.

8. The method for code vulnerability detection and risk ranking based on hybrid features according to claim 7 is characterized in that: The average value c of the integrated gradient for each marker of the code line is calculated as follows: Among them, IG it represents the integrated gradient of the tth code token in the i-th row, x it represents the vector representation of the tth code tag in the i-th row, x′ it represents the baseline input, F represents the prediction function of the model, and m represents the total number of tokens in the i-th row.

Images