本文简单梳理机器学习在基于源码的漏洞挖掘中的应用步骤，和一些案例，同时对机器学习在基于源码的漏洞挖掘中面临的挑战进行了一些简单分析。希望对大家有所帮助。

1. 背景

机器学习，在很多领域得到了很好的应用，比如图形、图像、音频处理识别、自然语言处理等方面。因此，探讨机器学习在源码分析领域的应用，就成了一件自然而然的事情。作者主要的研究领域是基于源码的静态代码分析，因此主要分析下在这方面的应用。

1.1 传统基于源码的分析存在的问题

当前，传统的静态代码分析存在的问题：

• 传统的静态代码分析，在实现时，依赖于安全专家定制的规则，就有了一定的主观判断。安全专家的不同理解，一定程度上导致了误报和漏报的产生；

• 随着现在安全漏洞越来越复杂，新的漏洞形式的不断出现，完全依赖安全专家定制规则就有点儿捉襟见肘；

• 随着软件规则越来越大，传统的在编译原理基础上延伸出来的静态代码分析方法，在软件性能上，无法满足实际的需求；

• 传统的静态代码分析，在生成检查的中间表示时，依赖了编译构建环境（C/C++等），在一定程度上，增加了静态代码分析的依赖和复杂性；

• 传统的静态代码分析，在误报和漏报问题上，居高不下，因此，需要一种新的思路来尝试基于源码的漏洞挖掘。

1.2 机器学习在源码分析中的应用

当前，结合机器学习的实际应用，和源码分析的一些具体场景，针对源码分析，主要有下面的三大类应用：

• 代码抄袭检查(Code Clone Detection)

主要可能存在下面问题：

(1) 代码侵权检查，例如在自研代码中，抄袭 license 并不是很友好的第三方代码；

(2) 代码复制粘贴检查：在同一个项目或者相关项目中，代码没有进行合理的抽象，而是到处复制代码，导致代码的开发和维护效率低下。另外，如果复制的代码存在漏洞，则给项目的稳定性代码很大挑战。

• 代码聚类(Code Classification)

基于一定的特征，对代码进行聚类及打标签，有很多实际的用途。实际上，很多源码，都是一种聚类的效果。例如：漏洞检查，可以针对具有特定漏洞的代码进行聚类后，检查特定的代码是否可以和某类漏洞聚在一起。

• 代码查找(Code Search)

如果已经知道一段代码片段，然后需要在特定的项目中，找到和这段代码片段非常相似的代码，就是代码查找，也有很多的实际用途。举个简单例子，如果已经知道特定的代码片段的写法有问题，就可以拿这段儿代码，到项目里面去找，找出来相似度比较高的代码片段，很大可能也是有问题的代码。

当前，机器学习在上面的三大类应用上，已经取得了一定的效果，从而在一定程度上，我们认为基于基于学习的面向源码的漏洞挖掘成为现实。

2. 基于机器学习的源码漏洞挖掘关键技术

结合我们通常对机器学习的理解，这里大概给出一些自己的理解，有不对的点，请不吝支出。

2.1 机器学习的目标

我理解的机器学习，有两类目标设置：

(1) 我现在有一大堆数据，但是没有明确的目标。我想看看这堆数据之间，有些什么关系，满足什么样的条件，比如非常著名的“啤酒和尿布”的故事；

(2) 我有特定的业务需求，这样我去找特定的数据，来做成这样的一件事情。

面向基于机器学习的源码的分析处理，应该都是适用于第 2 种场景。因为现在有的数据只有代码，并没有做任何抽象，需要有一些特定的业务目标。

比如结合最近看的一些论文，里面提到的一些学习目标：

(1) NL2Type：学习一个模型，用来预测一个 JavaScript 语言的函数的每个入参类型和返回值的类型；

(2) VulDeePecker：学习一个模型，用来检测时否存在特定的针对函数 API 调用的漏洞；

(3) code2vec：提出了一个模型，对 AST 进行 embedding，可以用来预测函数名。

2.2 数据选择

2.2.1 源码数据信息

从源码，可以提取的数据包含下面的一些内容：

• 源码的自然语言信息，包括代码及注释内容

• Token

• AST

• CFG

• PDG/SDG

• 切片数据

• 数据流信息

• 控制流信息

• 其他分析过程中可能产生的中间表示数据

直接拿全量源码训练，效果不可能好，必然是结合特定的业务目标，抽取特定的数据及特征来进行训练。不是有句老话：数据和特征决定了机器学习的上限，而模型和算法只是逼近这个上限而已。

2.2.2 基于业务目标的数据选择

不同的业务目标，或者是训练的目标不同，选择的数据也各不相同，可能产生的效果也会有所不同。所以数据的选择非常重要。

以 2.1 节中选择的三篇论文举例：

(1) NL2Type：源码的自然语言信息

为了对函数的入参及返回值的类型进行预测，使用的数据包含了 函数的 自然语言信息（注释、函数名、参数名 等内容），因为 作者认为，代码也属于某一种自然语言信息，函数名 和 参数名 一般都是定义为具有特定含义的单词组合，具有一定的自解释性，而且注释内容，也可以很好地辅助解释函数作用，基于这些内容进行训练，可以比较好地实现训练目的。

(2) VulDeePecker: 程序切片数据

作者关心的，是特定函数调用可能产生的缺陷信息，在训练的时候，抽取该函数调用及函数所有的参数的程序切片数据，作为训练数据。其实也很好理解，因为这样的一些缺陷检查，大部分都是基于 数据流分析，而通过切片，拿到 切片数据，实际上，就是拿到了和缺陷密切相关的 数据流传递信息。所以，训练使用的 切片数据 和 待检测的缺陷，具有直接关系。

(3) code2vec: AST

在这篇文章中，作者根据 函数体进行训练，想知道 这个函数是做什么事情的。如下图所示（图片内容取自：code2vec）：

如上图所示：函数 f() 的 AST 如右边，对 AST 进行机器学习后，可以识别出来，这个函数 f() 和 sort 函数的相似度达到 98.54%，和其他的 函数 相似度递减。

训练的时候，使用的是函数体 AST 进行 embedding 之后的机器学习训练。

从上面的举例和分析来看，我们可以明确：数据 对 机器学习的训练结果，是具有良好的可解释性的，因此，机器学习要想效果好，训练数据一定要选好，选对。一个基本原则，就是这个数据跟结果，是不是有什么直接或者间接的关系，关系越紧密、越直接，效果应该越好。

2.2.3 选择的数据是确定的吗

对于相同的训练目标，训练使用的数据一定是确定的吗？其实不然！

举上面的论文里面的内容来举例：

(1) NL2Type

作者选择的是 源代码的自然语言信息。在一定程度上，是具有不错的可解释性。但是并不是唯一的数据来源。我们知道，任何一种语言中，特定的数据结构，可以进行的操作不同，因此，基于 源代码的操作数据，也可以用于 变量的类型训练，例如，在 JavaScript 的函数参数和函数返回值的训练中，我们获取到函数参数的所有的切片数据，然后基于切片（其实就是针对函数参数可进行的操作）来进行训练，也可以在一定程度上，也可以训练获取到想要的参数类型。

(2) VulDeePecker

此处，大家可以参考该论文同课题组后来的另一篇论文，实现了一个µVulDeePecker 工具，在训练数据的时候，考虑了更多的情况，不仅仅考虑了 数据依赖，同时考虑了 控制依赖。

(3) 函数名预测

这种情况，论文基于 AST 进行训练，但是是否可以在 AST 基础上，构造出来 CFG 训练呢？

2.3 code embedding

因为机器学习的输入，需要是 数值向量，无法直接支持上面提到的数据，因此，需要将上面的数据转换为 数值向量，才能输入到机器学习模型里面训练。这个过程，就称为 code embedding。

code embedding 非常重要，如果 code embedding 没有做好，那么可能导致部分原数据的信息丢失，从而导致训练模型效果不好。

当前，学术界已经提出了很多 code embedding 方法。甚至对 同一种数据，也有非常多的 embedding 方法。比如 都是 AST，也有各种不同的 embedding 方法。

本文大概介绍一些已有的论文及方向，大家可以自己搜索相关论文了解详细的实现。

2.3.1 部分论文罗列

从现有的 code embedding 论文来看，在 2.2.1 节中，列举了常见的可以从源码中获取到的 数据类型，到目前为止，似乎都可以进行 embedding，而且大部分都有相关的论文支撑。

(1) 源码的自然语言信息

比如已经列出来的 NL2Type 等，对 代码注释、函数名、参数名 等 进行 embedding，主要应用到了 word2vec；

(2) Token

针对 Token 序列的 embedding，很多也还是使用 word2vec，因为直接对 token 使用 id 创建的矩阵比较稀疏，而且无法表达出来相似关系。

相关的一些论文：

Harer J A , Kim L Y , Russell R L , et al. Automated software vulnerability detection with machine learning[J]. 2018.

White M , Tufano M , Martinez M , et al. Sorting and Transforming Program Repair Ingredients via Deep Learning Code Similarities[J]. 2017.

(3) AST

针对 AST 可以进行的 embedding 方法非常多，可以参考的论文：

Alon U , Zilberstein M , Levy O , et al. code2vec: learning distributed representations of code[J]. Proceedings of the ACM on Programming Languages, 2019, 3(POPL):1-29.

Devlin J , Uesato J , Singh R , et al. Semantic Code Repair using Neuro-Symbolic Transformation Networks. 2017.

Buch L , Andrzejak A . Learning-Based Recursive Aggregation of Abstract Syntax Trees for Code Clone Detection[C]// 2019 IEEE 26th International Conference on Software Analysis, Evolution and Reengineering (SANER). IEEE, 2019.（这篇论文同时使用了 AST 和 content 来进行训练的）

(4) CFG

针对 CFG 的 embedding 并不是很多，但是我其实觉得这样的训练蛮有意思的：

Daniel DeFreez, Aditya V Thakur, and Cindy Rubio-González. Path-based function embedding andits application to specification mining. arXiv preprint arXiv:1802.07779, 2018.

我还看到一篇论文，生成 CFG 后，对 BasicBlock 执行 doc2vec，然后对整个 CFG 执行 node2vec 做 embedding 的，忘记论文名字了。

(5) 切片数据

在上面提到的 VulDeePecker，µVulDeePecker 等，都是基于切片数据的，大家可以参考。

(6) PDG/SDG

基于这个层面的论文不多，可以看下面的一篇文章：

Romanov V , Ivanov V , Succi G . Representing Programs with Dependency and Function Call Graphs for Learning Hierarchical Embeddings[C]// 22nd International Conference on Enterprise Information Systems. 2020.

不过，在上面基于切片数据的 embedding 中，也需要综合应用 PDG/SDG 才可以实现的，应该也算是 PDG/SDG 的应用了。

(7) Flow

这种是嫌弃前面在数据依赖上面表示的还不够明显，我只看到一篇文章，参考下面的论文：

Sui Y , Cheng X , Zhang G , et al. Flow2Vec: value-flow-based precise code embedding[J]. Proceedings of the ACM on Programming Languages, 2020, 4(OOPSLA):1-27.

该论文是隋玉磊老师的论文，隋玉磊老师也是 当前非常非常有名的 基于 LLVM 的 静态代码分析工具 SVF 的作者（很多公司搞的自研或商用的静态代码分析工具，可能都是直接使用了 SVF，或者参考了 SVF）。这篇论文是基于 LLVM 的值流关系 进行 embedding 的。

理论上，值流图 距离程序分析最近，效果也是最好的。

2.3.2 部分开源代码

如下的部分链接：

https://github.com/petukhovv/tree2vec

https://github.com/bdqnghi/ast-node-encoding

https://github.com/uohzxela/ast2vec

https://github.com/dazcona/user2code2vec

https://github.com/kth-tcs/3sFix-experiments

https://github.com/tech-srl/code2vec

http://github.com/tech-srl/code2seq

https://github.com/defreez-ucd/func2vec-fse2018-artifact

https://iclr2018anon.github.io/semantic_code_repair/index.html.

https://github.com/jjhenkel/code-vectors-artifact

https://github.com/keowang/dynamic-program-embedding

除了上面的部分直接相关的(基于一些关键字都可以搜索到)，其他还有一些需要关注的通用的 embedding 方法，比如 word2vec（各种自然语言信息、token 等）和 graph2vec（或 node2vec，因为其实除了 自然语言信息 和 token 数据，代码其他所有的表示，包括 AST、CFG、CG、PDG/SDG、VFG 都可以表示成 graph 形式），从而有效利用已有的通用 embedding 方法，快速实现自己的 code embedding。

2.4 机器学习算法选择

上面都完成后，就是到了机器学习算法的选择阶段。

在这方面，我就不过多的描述，因为上面的论文里面，没有一篇对机器学习算法有挑战性的贡献，都是在使用，而且最近的论文，几乎全部都是 RNN、双向 LSTM、双向 LSTM + attention、GRU 等，也不排除其他的算法可以有好的效果。

当然，普通的机器学习算法好像也有人用，但是现在的趋势来看，用普通的机器学习算法，好像就有点儿 low 了，不过适合的就是好的，还是要根据数据特征和业务目标来进行合理选择。

3. 基于机器学习的源码漏洞挖掘挑战

这里大概介绍一部分基于机器学习训练的挑战，个人的观点，可以一起讨论，而且并不系统，也算是抛出自己的一些想法。

(1) 机器学习是否能比传统静态代码分析效果更优有待考验

机器学习，主要应用在逻辑不明确的业务场景上，对于逻辑明确，可以基于配置实现的业务场景，并不适合机器学习。例如垃圾邮件识别，我们无法给出来明确的规则，什么样的邮件内容是垃圾邮件，但是又可以基于邮件内容来判定是否垃圾邮件，此时可以基于机器学习来判定是否垃圾邮件。但是，如果有一个规则，可以判定就是垃圾邮件，那么我们就不需要机器学习，例如来自特定的 *@xx.com 的邮件，全部都是垃圾邮件，此时我们就可以通过配置实现，而不是通过机器学习实现。

到 基于源码的漏洞检测 也是一样的。传统的静态代码分析，就是一种基于配置的检查，有各种不同的检查引擎。因此可以认为基于传统静态代码分析的漏洞检查是一个逻辑性比较强的，可以基于配置来完成的一种业务场景。

只是因为 检查引擎 抽象程度问题（分析精确程度问题） + 配置问题（主观性因素等），导致误报和漏报的增加，但是机器学习是否能比传统静态代码分析效果更好，是个未知数。

例如，涉及到配置问题。我们有个方法 isValid(arg: String)，那么我们调用这个方法，是不是 arg 就安全了？对 SQL 注入安全了，那对 命令注入是不是安全呢？在 A 项目里面，isValid() 的实现 和在 B 项目中的 isValid() 实现一样吗？都作为第三方库的情况下，如果能知道 isValid() 的具体实现？

上面的问题，人工配置有问题，基于机器学习也会有相应的问题的。

(2) 需要的训练漏洞数据集太少

为了能够训练一个好的工具，可以用于缺陷漏洞检查，即使是同一种场景的漏洞，也需要大量的标注数据来进行训练。但是，这样的数据也是非常短缺的，已有的一些开源代码检查用例集也面临漏洞类型短缺、复杂度不够等问题。所以训练数据短缺也是一个非常重要的问题。

(3) 基于机器学习无法给出恰当的缺陷提示

机器学习训练结果，只能告诉我们，这段代码是否有缺陷，最多告诉我们缺陷类型，但是无法告诉我们缺陷是如何产生的，无法正确提示 source、passthrough、sink 等点。这在一定程度上，决定了在很多业务场景中，基于机器学习的漏洞监测，无法替代 传统的基于源码的漏洞检查方法。

(4) 漏洞的不同产生机制决定了训练数据的不同

我们可以看到，当前，为了做 code embedding，产生了大量的 embedding 方法，而且和源码越来越远，和传统静态代码分析的 IR 越来越靠近。因为不同的漏洞，产生的机制、方式千差万别，统一的实现方式显然不大可能，那么针对每一种或几种漏洞训练一套模型，代价会不会太大？

(5) 机器学习在一定程度上依赖了传统静态代码分析的结果

比如 AST、CFG、PDG/SDG、切片、VFG 构造等，都需要依赖传统静态代码分析。需要结合使用。

4. 总结

本文简单总结了机器学习在基于源码的漏洞挖掘中的方法步骤，可以让大家对该方向的工作有一定的认识，同时，也提出了该方法当前面临的一些挑战。总结来说：机器学习肯定有用，但是应该无法替代传统静态代码分析，看如何结合起来，发挥更大作用。