Devign: Effective Vulnerability Identification by Learning Comprehensive Program Semantics via Graph Neural Networks

Introduction

The number of software vulnerabilities has been increasing rapidly recently, either reported publicly through CVE (Common Vulnerabilities and Exposures) or discovered internally in proprietary code. In particular, the prevalence of open-source libraries not only accounts for the increment, but also propagates impact. These vulnerabilities, mostly caused by insecure code, can be exploited to attack software systems and cause substantial damages financially and socially.

Vulnerability identification is a crucial yet challenging problem in security. Beside the classic approaches such as static analysis , dynamic analysis and symbolic execution, a number of advances have been made in applying machine learning as a complementary approach. In these early methods , features or patterns hand-crafted by human experts are taken as inputs by machine learning algorithms to detect vulnerabilities. However, the root causes of vulnerabilities varies by types of weaknesses and libraries, making it impractical to characterize all vulnerabilities in numerous libraries with the hand-crafted features.

To improve usability of the existing approaches and avoid the intense labor of human experts on feature extraction, recent works investigate the potential of deep neural networks on a more automated way of vulnerability identification . However, all of these works have major limitations in learning comprehensive program semantics to characterize vulnerabilities of high diversity and complexity in real source code. First, in terms of learning approaches, they either treat the source code as a flat sequence, which is similar to natural languages or represent it with only partial information. However, source code is actually more structural and logical than natural languages and has heterogeneous aspects of representation such as Abstract Syntax Tree (AST), data flow, control flow and etc. Moreover, vulnerabilities are sometimes subtle flaws that require comprehensive investigation from multiple dimensions of semantics. Therefore, the drawbacks in the design of previous works limit their potentiality to cover various vulnerabilities. Second, in terms of training data, part of the data in is labeled by static analyzers, which introduced high percentage of false positives that are not real vulnerabilities. Another part like , are simple artificial code (even with “good” or “bad” inside the code to distinguish the vulnerable code and non-vulnerable code) that are far beyond the complexity of real code .

To this end, we propose a novel graph neural network based model with composite programming representation for factual vulnerability data. This allows us to encode a full set of classical programming code semantics to capture various vulnerability characteristics. A key innovation is a new Conv module which takes as input a graph’s heterogeneous node features from gated recurrent units. The Conv module hierarchically chooses more coarse features via leveraging the traditional convolutional and dense layers for graph level classification. Moreover, to both testify the potential of the composite programming embedding for source code and the proposed graph neural network model for the challenging task of vulnerability identification, we compiled manually labeled data sets from 4 popular and diversified libraries in C programming language. We name this model Devign (Deep Vulnerability Identification via Graph Neural Networks).

In the composite code representation, with ASTs as the backbone, we explicitly encode the program control and data dependency at different levels into a joint graph of heterogeneous edges with each type denoting the connection regarding to the corresponding representation. The comprehensive representation, not considered in previous works, facilitates to capture as extensive types and patterns of vulnerabilities as possible, and enables to learn better node representation through graph neural networks.

We propose the gated graph neural network model with the Conv module for graph-level classification. The Conv module learns hierarchically from the node features to capture the higher level of representations for graph-level classification tasks.

We implement Devign, and evaluate its effectiveness through manually labeled data sets (cost around 600 man-hours) collected from the 4 popular C libraries. We make two datasets public together with more details (https://sites.google.com/view/devign). The results show that Devign achieves an average 10.51% higher accuracy and 8.68% F1 score than baseline methods. Meanwhile, the Conv module brings an average 4.66% accuracy and 6.37% F1 gain. We apply Devign to 40 latest CVEs collected from the 4 projects and get 74.11% accuracy, manifesting its usability of discovering new vulnerabilities.

The Devign Model

Vulnerability patterns manually crafted with the code property graphs, integrating all syntax and dependency semantics, have been proved to be one of the most effective approaches to detect software vulnerabilities. Inspired by this, we designed Devign to automate the above process on code property graphs to learn vulnerable patterns using graph neural networks . Devign architecture shown in Figure 1, which includes the three sequential components: 1) Graph Embedding Layer of Composite Code Semantics, which encodes the raw source code of a function into a joint graph structure with comprehensive program semantics. 2) Gated Graph Recurrent Layers, which learn the features of nodes through aggregating and passing information on neighboring nodes in graphs. 3) The Conv module that extracts meaningful node representation for graph-level prediction.

where $\mathcal{L}(\cdot)$ is the cross entropy loss function, $\omega(\cdot)$ is a regularization, and $\lambda$ is an adjustable weight.

2 Graph Embedding Layer of Composite Code Semantics

As illustrated in Figure 1, the graph embedding layer $EMB$ is a mapping from the function code $c_{i}$ to graph data structures as the input of the model, i.e.,

In this section, we describe the motivation and method on why and how to utilize the classical code representations to embed the code into a composite graph for feature learning.

In program analysis, various representations of the program are utilized to manifest deeper semantics behind the textual code, where classic concepts include ASTs, control flow, and data flow graphs that capture the syntactic and semantic relationships among the different tokens of the source code. Majority of vulnerabilities such as memory leak are too subtle to be spotted without a joint consideration of the composite code semantics . For example, it is reported that ASTs alone can be used to find only insecure arguments . By combining ASTs with control flow graphs, it enables to cover two more types of vulnerabilities, i.e., resource leaks and some use-after-free vulnerabilities. By further integrating the three code graphs, it is possible to describe most types except two that need extra external information (i.e., race condition depending on runtime properties and design errors that are hard to model without details on the intended design of a program)

Though manually crafted the vulnerability templates in the form of graph traversals, it conveyed the key insight and proved the feasibility to learn a broader range of vulnerability patterns through integrating properties of ASTs, control flow graphs and data flow graphs into a joint data structure. Beside the three classical code structures, we also take the natural sequence of source code into consideration, since the recent advance on deep learning based vulnerability detection has demonstrated its effectiveness . It can complement the classical representations because its unique flat structure captures the relationships of code tokens in a ‘human-readable’ fashion.

2.2 Graph Embedding of Code

Next we briefly introduce each type of the code representations and how we represent various subgraphs into one joint graph, following a code example of integer overflow as in Figure 2(a) and its graph representation as shown in Figure 2(b).

AST is an ordered tree representation structure of source code. Usually, it is the first-step representation used by code parsers to understand the fundamental structure of the program and to examine syntactic errors. Hence, it forms the basis for the generation of many other code representations and the node set of AST $V^{ast}$ includes all the nodes of the rest three code representations used in this paper. Starting from the root node, the codes are broken down into code blocks, statements, declaration, expressions etc., and finally into the primary tokens that form the leaf nodes. The major AST nodes are shown in Figure 2. All the boxes are AST nodes, with specific codes in the first line and node type annotated. The blue boxes are leaf nodes of AST and purple arrows represent the child-parent AST relations.

CFG describes all paths that might be traversed through a program during its execution. The path alternatives are determined by conditional statements, e.g., if, for, and switch statements. In CFGs, nodes denote statements and conditions, and they are connected by directed edges to indicate the transfer of control. The CFG edges are highlighted with green dashed arrows in Figure 2. Particularly, the flow starts from the entry and ends at the exit, and two different paths derive at the if statements.

DFG tracks the usage of variables throughout the CFG. Data flow is variable oriented and any data flow involves the access or modification of certain variables. A DFG edge represents the subsequent access or modification onto the same variables. It is illustrated by orange double arrows in Figure 2 and with the involved variables annotated over the edge. For example, the parameter $b$ is used in both the if condition and the assignment statement.

In order to encode the natural sequential order of the source code, we use NCS edges to connect neighboring code tokens in the ASTs. The main benefit with such encoding is to reserve the programming logic reflected by the sequence of source code. The NCS edges are denoted by red arrows in Figure 2, connect all the leaf nodes of the AST.

Consequently, a function $c_{i}$ can be denoted by a joint graph $g$ with the four types of subgraphs (or $4$ types of edges) sharing the same set of nodes $V=V^{ast}$. As shown in Figure (2), every node $v\in V$ has two attributes, Code and Type. Code contains the source code represented by $v$, and the type of $v$ denotes the type attribute. The initial node representation $x_{v}$ shall reflect the two attributes. Hence, we encode Code by using a pre-trained word2vec model with the code corpus built on the whole source code files in the projects, and Type by label encoding. We concatenate the two encodings together as the initial node representation $x_{v}$.

3 Gated Graph Recurrent Layers

The key idea of graph neural networks is to embed node representation from local neighborhoods through the neighborhood aggregation. Based on the different techniques for aggregating neighborhood information, there are graph convolutional networks , GraphSAGE , gated graph recurrent networks and their variants. We chose the gated graph recurrent network to learn the node embedding, because it allows to go deeper than the other two and is more suitable for our data with both semantics and graph structures .

where $AGG(\cdot)$ denotes an aggregation function that could be one of the functions $\{MEAN,MAX,SUM,CONCAT\}$ to aggregate the information from different edge types to compute the next time-step node embedding $h^{(t)}$. We use the $SUM$ function in the implementation. The above propagation procedure iterates over $T$ time steps, and the state vectors at the last time step $H_{i}^{(T)}=\{h_{j}^{(T)}\}_{j=1}^{m}$ is the final node representation matrix for the node set $V$.

4 The Conv Layer

The generated node features from the gated graph recurrent layers can be used as input to any prediction layer, e.g., for node or link or graph-level prediction, and then the whole model can be trained in an end-to-end fashion. In our problem, we require to perform the task of graph-level classification to determine whether a function $c_{i}$ is vulnerable or not. The standard approach to graph classification is gathering all these generated node embeddings globally, e.g., using a linear weighted summation to flatly adding up all the embeddings as shown in Eq (5),

Thus, we design the Conv module to select sets of nodes and features that are relevant to the current graph-level task. Previous works in proposed to use a SortPooling layer after the graph convolution layers to sort the node features in a consistent node order for graphs without fixed ordering, so that traditional neural networks can be added after it and trained to extract useful features characterizing the rich information encoded in graph. In our problem, each code representation graph has its own predefined order and connection of nodes encoded in the adjacent matrix, and the node features are learned through gated recurrent graph layers instead of graph convolution networks that requires to sort the node features from different channels. Therefore, we directly apply 1-D convolution and dense neural networks to learn features relevant to the graph-level task for more effective predictionWe also tried LSTMs and BiLSTMs (with and without attention mechanisms) on the sorted nodes in AST order, however, the convolution networks work best overall.. We define $\sigma(\cdot)$ as a 1-D convolutional layer with maxpooling, then

Let $l$ be the number of convolutional layers applied, then the Conv module, can be expressed as

where we firstly apply traditional 1-D convolutional and dense layers respectively on the concatenation $[H_{i}^{(T)},x_{i}]$ and the final node features $H_{i}^{(T)}$, followed by a pairwise multiplication on the two outputs, then an average aggregation on the resulted vector, and at last make a prediction.

Evaluation

We evaluate the benefits of Devign against a number of state-of-the-art vulnerability discovery methods, with the goal of understanding the following questions:

How does our Devign compare to the other learning based vulnerability identification methods?

How does our Conv module powered Devign compare to the Ggrn with the flat summation in Eq (5) for the graph-level classification task?

Can Devign learn from each type of the code representations (e.g., a single-edged graph with one type of information)? And how do the Devign models with the composite graphs (e.g., all types of code representations) compare to each of the single-edged graphs?

Can Devign have a better performance compare to some static analyzers in the real scenario where the dataset is imbalanced with an extremely low percentage of vulnerable functions?

How does Devign perform on the latest vulnerabilities reported publicly through CVEs?

It is never trivial to obtain high-quality data sets of vulnerable functions due to the demand of qualified expertise. We noticed that despite released data sets of vulnerable functions, the labels are generated by statistic analyzers which are not accurate. Other potential datasets used in are not available. In this work, supported by our industrial partners, we invested a team of security to collect and label the data from scratch. Besides raw function collection, we need to generate graph representations for each function and initial representations for each node in a graph. We describe the detailed procedures below.

Raw Data Gathering To test the capability of Devign in learning vulnerability patterns, we evaluate on manually-labeled functions collected from 4 large C-language open-source projects that are popular among developers and diversified in functionality, i.e., Linux Kernel, QEMU, Wireshark, and FFmpeg.

To facilitate and ensure the quality of data labelling, we started by collecting security-related commits which we would label as vulnerability-fix commits or non-vulnerability fix commits, and then extracted vulnerable or non-vulnerable functions directly from the labeled commits. The vulnerability-fix commits (VFCs) are commits that fix potential vulnerabilities, from which we can extract vulnerable functions from the source code of versions previous to the revision made in the commits. The non-vulnerability-fix commits (non-VFCs) are commits that do not fix any vulnerability, similarly from which we can extract non-vulnerable functions from the source code before the modification. We adopted the approach proposed in to collect the commits. It consists of the following two steps. 1) Commits Filtering. Since only a tiny part of commits are vulnerability related, we exclude the security-unrelated commits whose messages are not matched by a set of security-related keywords such as DoS and injection. The rest, more likely security-related, are left for manual labelling. 2) Manual Labelling. A team of four professional security researchers spent totally 600 man-hours to perform a two round data labelling and cross-verification.

Given a VFC or non-CFC, based on the modified functions, we extract the source code of these functions before the commit is applied, and assign the labels accordingly.

Graph Generation We make use of the open-source code analysis platform for C/C++ based on code property graphs, Joern , to extract ASTs and CFGs for all functions in our data sets. Due to some inner compile errors and exceptions in Joern, we can only obtain ASTs and CFGs for part of functions. We filter out these functions without ASTs and CFGs or with oblivious errors in ASTs and CFGs. Since the original DFGs edges are labeled with the variables involved, which tremendously increases the number of the types of edges and meanwhile complicates embedded graphs, we substitute the DFGs with three other relations, LastRead (DFG_R), LastWrite (DFG_W), and ComputedFrom (DFG_C) , to make it more adaptive for the graph embedding. DFG_R represents the immediately last read of each occurrence of the variable. Each occurrence can be directly recognized from the leaf nodes of ASTs. DFG_W represents the immediately last write of each occurrence of variables. Similarly, we makes these annotations to the leaf node variables. DFG_C determines the sources of a variable. In an assignment statement, the left-hand-side (lhs) variable is assigned with a new value by the right-hand-side (rhs) expression. DFG_C captures such relations between the lhs variable and each of the rhs variable.

Further, we remove functions with node size greater than 500 for computational efficiency, which accounts for 15%. We summarize the statistics of the data sets in Table 1.

2 Baseline Methods

In the performance comparison, we compare Devign with the state-of-the-art machine-learning-based vulnerability prediction methods, as well as the gated graph recurrent network (Ggrn) that used the linearly weighted summation for classification.

: We collect totally 4 complexity metrics and 11 vulnerability metrics for each function using Joern, and utilize Xgboost for classification. Here we did not use the proposed binning and ranking method because it was not learning based, but a heuristic designed to rank the likelihood of being vulnerable. We search the best parameters via Bayes Optimization .

: It treats the source code as natural languages and input the tokenized code into bidirectional LSTMs with initial embeddings trained via Word2vec. Here we implemented a 3-layer bidirectional for the best performance.

It is an improved version of with the attention mechanism .

: Similar to , it takes source code as natural languages and utilizes the bag of words to get the initial embeddings of code tokens, and then feeds them to CNNs to learn.

3 Performance Evaluation

Devign Configuration In the embedding layer, the dimension of word2vec for the initial node representation is 100. In the gated graph recurrent layer, we set the the dimension of hidden states as 200, and number of time steps as 6. For the Conv parameters of Devign, we apply (1, 3) filter with ReLU activation function for the first convolution layer which is followed by a max pooling layer with (1, 3) filter and (1, 2) stride, and (1, 1) filter for the second convolution layer with a max pooling layer with (2, 2) filter and (1, 2) stride. We use the Adam optimizer with learning rate 0.0001 and batch size 128, and $L2$ regularization to avoid overfitting. We randomly shuffle each dataset and split 75% for the training and the rest 25% for validation. We train our model on Nvidia Graphics Tesla M40 and P40, with 100-epoch patience for early stopping.

Results Analysis We use accuracy and F1 score to measure performance. Table 3 summarizes all the experiment results. First, we analyze the results regarding Q1, the performance of Devign with other learning based methods. From the results on baseline methods, Ggrn and Devign with composite code representations, we can see that both Ggrn and Devign significantly outperform the baseline methods in all the data sets. Especially, compared to all the baseline methods, the relative accuracy gain by Devign is averagely 10.51%, at least 8.54% on the QEMU dataset. Devign (Composite) outperforms the 4 baseline methods in terms of F1 score as well, i.e., the relative gain of F1 score is 8.68% on the average and the minimum relative gains on each dataset (Linux Kernel, QEMU, Wirshark, FFmpeg and Combined) are 2.31%, 11.80%, 6.65%, 4.04% and 4.61% respectively. As Linux follows best practices of coding style, the F1 score 84.97 by Devign is highest among all datasets. Hence, Devign with comprehensive semantics encoded in graphs performs significantly better than the state-of-the-art vulnerability identification methods.

Next, we investigate the answer to Q2 about the performance gain of Devign against Ggrn. We first look at the score with the composite code representation. It demonstrates that, in all the data sets, Devign reaches higher accuracy (an average of 3.23%) than Ggrn, where the highest accuracy gain is 5.12% on the FFmpeg data set. Also Devign gets higher F1, an average of 3.92% than Ggrn, where the highest F1 gain is 6.82 % on the QEMU data set. Meanwhile, we look at the score with each single code representation, from which, we get similar conclusion that generally Devign significantly outperforms Ggrn, where the maximum accuracy gain is 9.21% for the DFG_W edge and the maximum F1 gain is 17.13% for the DFG_C. Overall the average accuracy and F1 gain by Devign compared with Ggrn are 4.66%, 6.37% among all cases, which indicates the Conv module extracts more related nodes and features for graph-level prediction.

Then we check the results for Q3 to answer whether Devign can learn different types of code representation and the performance on composite graphs. Surprisingly we find that the results learned from single-edged graphs are quite encouraging in both of Ggrn and Devign. For Ggrn, we find that the accuracy in some specific types of edges is even slightly higher than that in the composite graph, e.g., both CFG and NCS graphs have better results on the FFmpeg and combined data set. For Devign, in terms of accuracy, except the Linux data set, the composite graph representation is overall superior to any single-edged graph with the gain ranging from 0.11% to 3.75%. In terms of F1 score, the improvement brought by composite graph compared with the single-edged graphs is averagely 2.69%, ranging from 0.4% to 7.88% in the Devign in all tests. In summary, composite graphs help Devign to learn better prediction models than single-edged graphs.

To answer Q4 on comparison with static analyzers on the real imbalanced dataset, we randomly sampled the test data to create imbalanced datasets with 10% vulnerable functions according to a large industrial scale analysis . We compare with the well-known open-source static analyzers Cppcheck, Flawfinder, and a commercial tool CXXX which we hide the name for legal concern. The results are shown in Table , where our approach outperforms significantly all static analyzers with an F1 score 27.99 higher. Meanwhile, static analyzers tend to miss most vulnerable functions and have high false positives, e.g., Cppcheck found 0 vulnerability in 3 out of the 4 single project datasets.

Finally to answer Q5 on the latest exposed vulnerabilities, we scrape the latest 10 CVEs of each project respectively to check whether Devign can be potentially applied to identify zero-day vulnerabilities. Based on commit fix of the 40 CVEs, we totally get 112 vulnerable functions. We input these functions into the trained Devign model and achieve an average accuracy of 74.11%, which manifests Devign’s potentiality of discovering new vulnerabilities in practical applications.

Related Work

The success of deep learning has inspired the researchers to apply it for more automated solutions to vulnerability discovery on source code . The recent works treat source code as flat natural language sequences, and explore the potential of natural language process techniques in vulnerability detection. For instance, built models upon LSTM/BiLSTM neural networks, while proposed to use the CNNs instead.

To overcome the limitations of the aforementioned models on expressing logic and structures in code, a number of works have attempted to probe more structural neural networks such as tree structures or graph structures for various tasks. For instance, proposed to generate logical formulas for program verification through gated graph recurrent networks, and aimed at prediction of variable names and variable miss-usage. proposed Gemini for binary code similarity detection, where functions in binary code are represented by attributed control flow graphs and input Structure2vec for learning graph embedding. Different from all these works, our work targeted at vulnerability identification, and incorporated comprehensive code representations to express as many types of vulnerabilities as possible. Beside, our work adopt gated graph recurrent layers in to consider semantics of nodes (e.g., node annotations) as well as the structural features, both of which are important in vulnerability identification. Structure2vec focuses primarily on learning structural features. Compared with that applies gated graph recurrent network for variable prediction, we explicitly incorporate control flow graph into the composite graph and propose the Conv module for efficient graph-level classification.

Conclusion and Future Work

We introduce a novel vulnerability identification model Devign that is able to encode a source-code function into a joint graph structure from multiple syntax and semantic representations and then leverage the composite graph representation to effectively learn to discover vulnerable code. It achieved a new state of the art on machine-learning-based vulnerable function discovery on real open-source projects. Interesting future works include efficient learning from big functions via integrating program slicing, applying the learnt model to detect vulnerabilities cross projects, and generating human-readable or explainable vulnerability assessment.

References