基于图神经网络的代码漏洞检测方法

本文提出了一种利用控制流图和图神经网络的代码特征提取方案,针对源代码的结构和语义信息进行建模。通过LLVM编译器和word2vec对代码进行预处理和向量化,然后使用图神经网络模型(包含多层、残差连接)对图结构数据进行训练,以提升漏洞检测的准确性。实验表明,这种方法在开源数据集上表现优越,特别是在多层网络中,残差连接方案优于直接连接。
摘要由CSDN通过智能技术生成

文章结构

1. 引言

1.1 漏洞现状

1.2 漏洞研究传统特点

现有的研究方案大多把源代码转化为文本序列,使用自然语言处理中常用的循环神经网络进行自动化特征提取,但这忽略了代码中由于跳转、循环、判断形成的控制流特征,往往会有较大的信息损失。如何提取源代码样本的特征信息,把源代码样本转换成适合后续神经网络训练的数据形式并最大化保留样本信息是关键问题之一。

1.3 文章贡献

1.3.1 设计了一种基于中间语言控制流图的代码特征提取方案

本文通过代码中间语言的控制流图提取代码样本的结构特征,使用控制流图基本块中的代码序列提取代码样本的语义特征,将代码样本转化为图结构数据,以解决序列样本丢失代码结构特征信息的问题。通过提取源代码的控制流图结构信息,同时使用基本块的中间语言提取代码语义信息,嵌入图节点中,将源代码转化为图结构数据样本。

1.3.2 使用图神经网络模型处理代码图结构样本
1.3.3 基于设计方案设计了实验方案并对结果进行了分析

2. 相关工作

2.1 图表示发展
2.2 深度学习图表示发展

3. 方案设计

3.1 代码预处理LLVM+chang编译器
3.2 生成图 LLVM分析器 + word2vec
3.3 图转向量 PYG
3.4 训练模型
3.3.1 多层

本文提出的图神经网络结构由多组图卷积和池化层组成的基本单元来完成在图级别上的端到端分类学习问题,

3.3.2 残差

多组基本单元之间可以直接串联连接,也可以通过残差方式连接。文
献[16]指出残差连接方式在多层网络上有更好的效果。

3.3.3 通过应用多层感知器(MLP)和Softmax 层

4. 实验设计和评估

3.1 数据集

本文采用的数据集是来源于美国国家标准技术研究所(NIST)的软件保障参考数据集(SARD)

3.2 参数设置

本文的模型测试环境为 CentOS 7.6 操作系统,Intel Xeon CPU E5-2630 处理器,180 GB 内存,GeForce GTX 1080 显卡,8 GB 显存。模型使用 GPU 显卡加速模型训练。
在图结构数据生成部分,word2vec 算法生成的词嵌入参数大小设定为 256,基本块初始化向量的大小即 num_node_features 为 256。

3.3 评估指标
3.4 结果分析
  • 直接连接方案在大于 2 层网络层之后出现了准确率的大幅下降,而残差网络连
    接方案在多层网络环境下具有更好的效果,因此本文选择了效果最好的 3 层残差连接网络。
  • 本文提出的基于控制流图的特征提取与图神经网络模型在开源数据集上的准确率
    均大幅高于静态扫描工具

实验设置

基本过程

  • 提出了一种基于图神经网络的代码漏洞检测方法,通过中间语言的控制流图特征,实现了函数级别的智能化代码漏洞检测。首先,将源代码编译为中间表示,进而提取其包含结构信息的控制流图,同时使用词向量嵌入算法初始化基本块向量提取代码语义信息;然后,完成拼接生成图结构样本数据,使用多层图神经网络对图结构数据特征进行模型训练和测试。采用开源漏洞样本数据集生成测试数据对所提方法进行了评估,结果显示该方法有效提高了漏洞检测能力。

文章创新点

套用新模型GNN结合多层、残差、MLP
使用控制流图表示方法
编码方式结构+语义 dot + word2vec pyg

论文资源链接

# GPF ## 一、GPF(Graph Processing Flow):利用神经网络处理问题的一般化流程 1、节点预表示:利用NE框架,直接获得全每个节点的Embedding; 2、正负样本采样:(1)单节点样本;(2)节点对样本; 3、抽取封闭子:可做类化处理,建立一种通用数据结构; 4、子特征融合:预表示、节点特征、全局特征、边特征; 5、网络配置:可以是输入、输出的网络;也可以是输入,分类/聚类结果输出的网络; 6、训练和测试; ## 二、主要文件: 1、graph.py:读入数据; 2、embeddings.py:预表示学习; 3、sample.py:采样; 4、subgraphs.py/s2vGraph.py:抽取子; 5、batchgraph.py:子特征融合; 6、classifier.py:网络配置; 7、parameters.py/until.py:参数配置/帮助文件; ## 三、使用 1、在parameters.py中配置相关参数(可默认); 2、在example/文件夹中运行相应的案例文件--包括链接预测、节点状态预测; 以链接预测为例: ### 1、导入配置参数 ```from parameters import parser, cmd_embed, cmd_opt``` ### 2、参数转换 ``` args = parser.parse_args() args.cuda = not args.noCuda and torch.cuda.is_available() torch.manual_seed(args.seed) if args.cuda: torch.cuda.manual_seed(args.seed) if args.hop != 'auto': args.hop = int(args.hop) if args.maxNodesPerHop is not None: args.maxNodesPerHop = int(args.maxNodesPerHop) ``` ### 3、读取数据 ``` g = graph.Graph() g.read_edgelist(filename=args.dataName, weighted=args.weighted, directed=args.directed) g.read_node_status(filename=args.labelName) ``` ### 4、获取全节点的Embedding ``` embed_args = cmd_embed.parse_args() embeddings = embeddings.learn_embeddings(g, embed_args) node_information = embeddings #print node_information ``` ### 5、正负节点采样 ``` train, train_status, test, test_status = sample.sample_single(g, args.testRatio, max_train_num=args.maxTrainNum) ``` ### 6、抽取节点对的封闭子 ``` net = until.nxG_to_mat(g) #print net train_graphs, test_graphs, max_n_label = subgraphs.singleSubgraphs(net, train, train_status, test, test_status, args.hop, args.maxNodesPerHop, node_information) print('# train: %d, # test: %d' % (len(train_graphs), len(test_graphs))) ``` ### 7、加载网络模型,并在classifier中配置相关参数 ``` cmd_args = cmd_opt.parse_args() cmd_args.feat_dim = max_n_label + 1 cmd_args.attr_dim = node_information.shape[1] cmd_args.latent_dim = [int(x) for x in cmd_args.latent_dim.split('-')] if len(cmd_args.latent_dim) == 1: cmd_args.latent_dim = cmd_args.latent_dim[0] model = classifier.Classifier(cmd_args) optimizer = optim.Adam(model.parameters(), lr=args.learningRate) ``` ### 8、训练和测试 ``` train_idxes = list(range(len(train_graphs))) best_loss = None for epoch in range(args.num_epochs): random.shuffle(train_idxes) model.train() avg_loss = loop_dataset(train_graphs, model, train_idxes, cmd_args.batch_size, optimizer=optimizer) print('\033[92maverage training of epoch %d: loss %.5f acc %.5f auc %.5f\033[0m' % (epoch, avg_loss[0], avg_loss[1], avg_loss[2])) model.eval() test_loss = loop_dataset(test_graphs, model, list(range(len(test_graphs))), cmd_args.batch_size) print('\033[93maverage test of epoch %d: loss %.5f acc %.5f auc %.5f\033[0m' % (epoch, test_loss[0], test_loss[1], test_loss[2])) ``` ### 9、运行结果 ``` average test of epoch 0: loss 0.62392 acc 0.71462 auc 0.72314 loss: 0.51711 acc: 0.80000: 100%|███████████████████████████████████| 76/76 [00:07<00:00, 10.09batch/s] average training of epoch 1: loss 0.54414 acc 0.76895 auc 0.77751 loss: 0.37699 acc: 0.79167: 100%|█████████████████████████████████████| 9/9 [00:00<00:00, 34.07batch/s] average test of epoch 1: loss 0.51981 acc 0.78538 auc 0.79709 loss: 0.43700 acc: 0.84000: 100%|███████████████████████████████████| 76/76 [00:07<00:00, 9.64batch/s] average training of epoch 2: loss 0.49896 acc 0.79184 auc 0.82246 loss: 0.63594 acc: 0.66667: 100%|█████████████████████████████████████| 9/9 [00:00<00:00, 28.62batch/s] average test of epoch 2: loss 0.48979 acc 0.79481 auc 0.83416 loss: 0.57502 acc: 0.76000: 100%|███████████████████████████████████| 76/76 [00:07<00:00, 9.70batch/s] average training of epoch 3: loss 0.50005 acc 0.77447 auc 0.79622 loss: 0.38903 acc: 0.75000: 100%|█████████████████████████████████████| 9/9 [00:00<00:00, 34.03batch/s] average test of epoch 3: loss 0.41463 acc 0.81132 auc 0.86523 loss: 0.54336 acc: 0.76000: 100%|███████████████████████████████████| 76/76 [00:07<00:00, 9.57batch/s] average training of epoch 4: loss 0.44815 acc 0.81711 auc 0.84530 loss: 0.44784 acc: 0.70833: 100%|█████████████████████████████████████| 9/9 [00:00<00:00, 28.62batch/s] average test of epoch 4: loss 0.48319 acc 0.81368 auc 0.84454 loss: 0.36999 acc: 0.88000: 100%|███████████████████████████████████| 76/76 [00:07<00:00, 10.17batch/s] average training of epoch 5: loss 0.39647 acc 0.84184 auc 0.89236 loss: 0.15548 acc: 0.95833: 100%|█████████████████████████████████████| 9/9 [00:00<00:00, 28.62batch/s] average test of epoch 5: loss 0.30881 acc 0.89623 auc 0.95132 ```
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值