Modern Vulnerability Detection: Using GNNs to Find Subtle Bugs

For over 20 years, static application security testing (SAST) has been the foundation of secure coding. However, beneath the surface, many legacy SAST tools still operate using basic techniques such as regular expressions and lexical pattern matching; essentially, sophisticated versions of the Unix command grep. As a result, most SAST tools suffer from what I call “false positive fatigue.” These tools report every occurrence of a strcpy() (or similar) regardless of whether the buffer is mathematically proven to be safe.

This article explores an innovative method for detecting vulnerabilities using graph neural networks (GNNs). In contrast to viewing source code as a linear string of characters, GNNs represent code as a structured graph of logical and data-flow structures. As such, we can now develop models that understand how a user’s input at line 10 in the code ultimately relates to a database query at line 50, even when variable names are changed three times between those two points in the code.

Why Traditional Tools Fail

Traditional SAST tools fail to identify vulnerabilities due to their flat representation of code. For example, consider this Python code snippet:

def get_user_data(user_input):

sanitized = clean_input(user_input)

# ... 50 lines of complex logic ...

cursor.execute("SELECT * FROM users WHERE name = " + sanitized)

A Regex-based tool identifies the SQL injection threat due to the string concatenation in a SQL query. The tool does not recognize the clean_input() function because it cannot understand the data flow across the function boundaries. However, a GNN can model the path that the variable follows. Therefore, if the path includes a sanitization node, the GNN can learn to classify it as “safe.”

Transforming Code to Graphs: Understanding the CPG

To utilize GNNs, we first convert source code into a math-friendly format known as a Code Property Graph (CPG). The CPG merges three classic graph types:

• Abstract Syntax Tree (AST) – the syntax tree of the source code (loops, if statements)
• Control Flow Graph (CFG) – the order in which the source code executes (Path A vs. Path B)
• Program Dependence Graph (PDG) – the dependencies between variables (data flow)

Therefore, the CPG is a rich graph where the nodes represent code elements (variable declarations, operator assignments), and the edges represent relationships between those code elements (“calls,” “defines,” “depends on”).

Tools: Generating Graphs With Joern

Fortunately, you do not need to write a parser from scratch to create a CPG. Joern is an open-source tool that parses C/C++, Java, and Python and generates CPGs.

# Install Joern (Linux/Mac)

./joern-install.sh

# Convert a source file into a CPG

joern-parse --output cpg.bin ./my_vulnerable_app/

You can then export the generated graph to a format your neural network can read (e.g., a CSV file of nodes and edges).

The GNN Model: Learning “Shapes” of Bugs

After generating a graph, you can feed it into a graph neural network (GNN). Similar to a standard neural network that expects a static-sized image or text vector, a GNN utilizes a technique called “message passing.”

• Node embedding: Every node (for example, x = 5) receives a vector to represent its node type (assignment) and content.
• Message passing: Every node communicates with its neighboring nodes. The x variable node sends a message to the if (x > 0) node.
• Aggregation: After 3-4 iterations of passing messages, every node “knows” about its immediate neighborhood. The SQL Execute node “knows” it was connected to a User Input node four hops away.

Implementation Example (PyTorch Geometric)

Below is a simplified version of a GNN model to classify vulnerabilities using the popular PyTorch Geometric library.

import torch
from torch_geometric.nn import GCNConv, global_mean_pool
import torch.nn.functional as F

class VulnDetectorGNN(torch.nn.Module):

    def __init__(self, num_node_features, num_classes):
        super(VulnDetectorGNN, self).__init__()

        # Graph Convolutional Layers
        self.conv1 = GCNConv(num_node_features, 64)
        self.conv2 = GCNConv(64, 64)
        self.conv3 = GCNConv(64, 64)

        # Final classifier
        self.linear = torch.nn.Linear(64, num_classes)

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch

        # 1. Message Passing Layers (with ReLU activation)
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = self.conv2(x, edge_index)
        x = F.relu(x)
        x = self.conv3(x, edge_index)

        # 2. Global Pooling (Aggregate all nodes into 1 graph vector)
        x = global_mean_pool(x, batch)

        # 3. Classifier (Safe vs. Vulnerable)
        x = F.dropout(x, p=0.5, training=self.training)
        x = self.linear(x)

        return F.log_softmax(x, dim=1)

# Create Model
model = VulnDetectorGNN(num_node_features=100, num_classes=2)
print(model)

Dataset: CodeXGLUE

You cannot train a model without data. The CodeXGLUE dataset provided by Microsoft is widely accepted within the industry as the standard dataset for developing models that can detect vulnerabilities. CodeXGLUE contains thousands of C/C++ functions labeled as either “Vulnerable” or “Safe” based upon actual CVEs found in well-known open-source projects such as FFmpeg and QEMU.

Training tip: Due to the fact that real-world vulnerabilities occur very rarely (possibly 1 out of every 1,000 functions), you must balance your dataset (over-sample the vulnerable functions), or your model will simply guess “Safe” 99.9% of the time and claim to have a high degree of accuracy.

Application of GNNs in DevSecOps

While engineers should not discard their legacy SAST tools, such as Fortify or Checkmarx, GNNs are best utilized as a second opinion.

• Triage assistant: Run standard SAST on your application. Take the top 500 “High” findings and pass them through a GNN model trained specifically on your codebase’s history of “False Positive” vulnerabilities.
• Filter: If the GNN indicates that the finding appears to resemble one of the 500 “False Positive” vulnerabilities that were identified previously, then mark it as low-priority.
• Custom rules: Utilize GNNs to discover vulnerabilities that cannot be identified by regex-based rules, including complex logic vulnerabilities or missing authorization checks spanning multiple files.

Conclusion

The future of vulnerability detection will be driven by semantics (the meaning behind the code) rather than syntax (how the code is written). By modeling source code as a graph, we can better capture the author’s intent. Although GNNs consume more computational resources than a simple regex-based script, their reduced false positive rates make them a valuable addition to the current array of application security tools.