大模型漏洞修复插件是腾讯朱雀实验室在安全垂类场景的一个重要实践。我们希望通过AI大模型，实现研发安全场景的漏洞自动修复，给出修复建议并提供修复代码，帮助更多开发人员提高研发效率。在腾讯混元大模型的支持下，漏洞修复插件通过精调后部署的私有化模型，实现了在帐密硬编码、SQL注入、命令注入等漏洞类型的修复建议输出和修复代码生成等功能，实现安全左移，更有效地在编程中使用插件收敛漏洞风险。
图1. IDE插件示意图
一、研发安全场景的现状和挑战
利用传统方法做漏洞修复提效，只适用于比较简单的场景，比如根据版本号判断使用的开源组件是否存在漏洞，更多高危险的如导致数据泄露的注入类漏洞/账密类等，该方案难以通用。主要原因总结如下：
- 规则限制：传统静态分析技术通常基于预定义的规则和模式进行漏洞检测和修复，无法覆盖所有的漏洞类型和场景；
- 上下文和语义理解限制：传统静态分析技术通常难以理解代码的上下文和语义信息，导致无法准确地理解代码的含义和逻辑；
- 创造性限制：传统静态分析技术通常只能分析已有的代码，无法创造新的代码片段来修复漏洞，限制了漏洞修能力。
相比传统程序分析技术，大模型具备强大的推理能力，尤其是在代码生成方面表现突出，可通过训练来学习漏洞修复的模式和规律。
二、为何在研发安全场景引入大模型？
2.1 大模型可解释性
大模型拥有丰富的预训练语料信息，如：书籍文档、代码素材、网页文本等等。基于预训练语料信息，并对基座大模型进行垂类领域的精调，大模型能结合知识信息进行整合和解释，产生有效的回复。对于研发安全场景来说，我们不仅需要生成有效的修复代码，也希望为更多的开发同学提供具体和明确的修复分析指引，从而更好地提效。
2.2 大模型代码生成能力
各类基座大模型和代码类Codex大模型，具有强大的代码生成能力。这些模型可以理解自然语言指令，并根据这些指令生成相应的代码。这种能力使得它们可以用于各种编程任务，包括但不限于编写新的函数、修复代码中的错误、优化现有代码等。这些模型的代码生成能力基于其在大量代码库上的训练。在训练过程中，模型学习了各种编程语言的语法和语义以及如何将自然语言指令转化为代码。因此，只要给出清晰的指令，这些模型就能生成相应的代码。但这些生成的代码中仍存在一些安全问题，即：潜在的不正确、含漏洞的代码等。因此，我们需要利用大模型的代码生成能力，同时希望通过精调等操作，使大模型能有效修复漏洞代码，尽量生成安全无漏洞的代码片段。
2.3 大模型在安全防护领域的应用已成为趋势
2023年，大模型成为了各行各业的热门关注点，其中安全垂类领域的安全大模型也呈现了百花齐放的态势，主要覆盖安全咨询、安全培训、安全监控、安全修复等能力。微软2023年3月份正式发布集成GPT-4的Microsoft Security Copilot，旨在更好地提供安全工具和专家知识，辅助企业识别和检测安全风险。谷歌2023年4月发布“谷歌云AI安全工作台”，利用 Sec-PaLM 来帮助用户查找、总结和应对安全问题。国内公司也推出了自己独有的安全大模型，主要对于安全场景，提供专家级别的咨询建议，提效运营效率。
三、混元一站式如何快速定制司内自研的研发安全大模型
3.1 腾讯混元大模型能力
训练优化
腾讯混元大模型通过大量实验，对预训练数据进行了语料丰富扩充，已覆盖 100 多种自然语言，32 种编程语言；并经过了数据清洗、过滤、去重等流程，对各类数据进行了大量的数据配比实验，保证了较为平稳的训练过程。同时，项目组对于长文能力、位置编码等技术细节进行了改进和优化。
模型效果
腾讯混元项目组参考业界主流做法，一共选用数十个数据集合，分为中文NLP、英文NLP任务、代码、数学、AGIEval、CMMLU、CEval等维度，综合评估模型在各项维度上的能力。其中，腾讯混元十亿级别大模型在代码能力上的提升更加显著，这点与我们使用中的精调体验也非常一致。
3.2 漏洞相关高质量数据的收集及模型结果
通用大模型的崛起，依赖于大规模的数据灌输；大模型能否在垂直领域落地，同样依赖是否有高质量的领域数据。对于研发安全场景，我们也积累了系统且丰富的高质量数据处理经验。
漏洞修复数据的格式
漏洞修复数据，最基础的是漏洞类型、漏洞代码、修复后代码三个属性。
但在实践过程中仅使用这三部分信息来精调大模型，效果并不好。为此，我们扩至漏洞类型、漏洞代码、修复后代码、代码描述、漏洞信息、修复建议、修复过程七个属性，组成“漏洞代码 -> 代码做了什么 -> 存在什么漏洞 -> 应该怎么修复 -> 修复后代码 -> 具体修复过程”的整体逻辑，实际训练效果更好。
漏洞修复数据源及处理方式
我们采用的漏洞修复数据源，包含GitHub数据、公开数据、业务数据等，数据质量依次上升。下面我们大致介绍各数据源的潜在问题和处理过程。
1.GitHub数据
Github是主流开源社区之一，其历史数据中含有大量的漏洞代码及修复的记录，但不同开发者的编程水平参差不齐。我们根据关键词召回与漏洞相关的数据，但实践中发现大部分(>90%)召回数据存在与漏洞修复无关、错误的修复方式、代码质量低等等问题，因此我们需要一些自动化的方法来过滤这些数据。
缺点
- 不可用数据的占比较高，人工审核成本上升；
- 单个样本可能存在代码块分散、缺失上下文信息等问题；
优点
- 数据量级较大；
- 代码多样性丰富；
- 可按需召回指定类型的漏洞；
2.公开数据
公开数据，更多指的是学术界开源的数据。这类数据，已经被大量学术工作借鉴使用，本身具备一定的可信度。目前，我们使用的主要是CVE-Fixes数据集。
缺点
- 数据量级不大；
- 漏洞类型分布与实际业务漏洞类型分布差距大；
优点
- 数据质量有一定保障；
- 代码多样性丰富；
3.业务数据
业务数据，即收集公司业务历史修复数据，清洗后作为训练数据。业务数据优点非常多，最贴近模型使用的场景、含有更符合公司业务场景的修复方式。但业务数据与GitHub数据有相同的缺点，清洗难度较大。
缺点
- 历史误报、错误修复的数据占比较高，人工审核压力大；
- 符合实际训练需求的高质量数据量级有限；
优点
- 代码多样性丰富；
- 与落地场景完全契合；
训练数据构造
基于上文获取的漏洞数据，我们进行了大量的实验。总结得到，在数据总量和质量不变的情况下，“Prompt工程”和“数据配比”对实验结果会产生影响。同时，我们指出：不同垂直领域下得到的实验结论大概率不同，例如我们在数据配比时配比了部分通用数据和代码数据，而其他团队可能只是用垂直领域的数据。所以，我们建议大家在聚焦的垂直领域多实验、再总结。
Prompt 工程
通过调整对模型提问的方式，使得模型的回答效果更好。具体的，我们通过业界的论文以及经验分享，总结了一个相对通用的模板：expert+COT+输出格式+变量。如图5所示。其中，expert是一种增强提示策略，用于指示LLM像专家一样回答问题。COT，即思维链，通过增加一系列中间推理步骤，能显著提高大型语言模型复杂推理的能力。
数据配比
除了漏洞数据外，我们还收集了开源的通用、代码、数学sft数据集，并进行了相应的配比实验。我们发现，数学数据集对漏洞修复准确率的提升基本无影响，而适当比例的代码、通用数据，有助于提升模型的漏洞修复准确率。
经过调优，我们的最佳数据配比如图6所示。当我们代码：通用：漏洞=1:2:4的时候，结果达到最优。
我们将开源的通用(general)、代码(code)、数学(math)、漏洞(vulnerability)数据集，输入到模型中，抽取倒数第二层对应的embedding，使用t-sne算法对其进行可视化，得到的结果如图7所示。可以看到数学数据集距离漏洞数据集最远，这也解释了为什么数学数据集对漏洞修复准确率的提升基本无影响。
实验结果
我们将实验结果与gpt3.5进行比较，结果如表1所示。从表1中可以看出，我们的平均准确率比gpt3.5高1.67%，具体到每个漏洞类型，均超过或持平gpt3.5。
上下文长度对结果的影响
在前文介绍处理业务数据时，我们提出借助于污点传播技术，只摘取漏洞触发点所在行的上下文代码。那么，这个上下文取多少行合适呢？这里我们在推理阶段取不同大小的上下文进行了测试，如图8所示。从结果中可以看出，不同大小的上下文确实对结果有影响，并且对于不同漏洞，影响不太相同。
对于SQL注入以及命令注入漏洞来说，适中大小的上下文有助于提升测试准确率。当上下文过少时，模型抽取的与漏洞相关的语义信息不足，通过这些信息无法判断是否存在漏洞，导致准确率下降；相反，当上下文过多时，模型定位漏洞所在行可能存在困难，导致准确率下降。
对于帐密硬编码漏洞来说，测试准确率随着上下文行数的增多而减少，猜测可能是因为对于该漏洞类型，较少的行数就足以让模型判断存在帐密硬编码，也足以让模型进行修复。相反，当测试上下文过多时，模型定位漏洞所在行可能存在困难，导致准确率下降。
3.3 插件支持漏洞检测和修复
目前我们基于腾讯混元大模型，推出了漏洞检出和修复功能的插件。插件目前重点发现三大类研发过程中对现网产生严重影响的漏洞：SQL注入、命令注入和账密泄漏，并较业内大幅提升发现风险的准确性。目前插件已支持在vs code内使用，仅需在vs code内安装插件，即可实现在提交代码前检测安全风险并及时修复。
插件的漏洞检测+修复示例
与竞品模型回复的对比
我们将自建模型与gpt3.5模型的回复进行对比，其中一个典型案例如下。gpt3.5与自建模型均能判断出代码中存在SQL注入漏洞，但是gpt3.5无法对其进行正确修复，而自建模型可以。
3.4 插件模型版本的自动化测评和迭代
模型评测迭代的过程中，针对漏洞修复，我们完全自建了benchmark以及对应的评判规则。
（1）待评测数据通过自研模型，拿到修复结果。
（2）将修复前后的代码，输入到代码相似度模型中，防止模型吐出与原始代码不相关的代码。
（3）使用算法/工具判断脱敏后的漏洞代码是否修复。
（4）根据实际测试，我们发现以上流程无法对修复前后代码语义是否更改进行判断。如图11所示，修复后代码漏掉了=号，导致语义的改变。因此我们补充设计了一套规则，来判断修复前后代码语义是否有更改。
四、总结与展望
以研发安全场景为切入点，拥有均衡能力的腾讯混元大模型未来在安全领域会有更广泛的应用场景和前景，帮助安全业务提高安全防护的效率，可以降低业务的投入成本，提高经济效益。
- 在数据安全领域，大模型可以通过对大量数据的分析，发现数据中的异常行为，及时预警并采取措施，防止数据泄露、数据篡改等安全问题的发生。
- 在网络安全领域，大模型可以通过对网络流量的实时监控，发现并阻止网络攻击，保障网络的正常运行。
- 在安全运营领域，大模型可以通过对用户数据、设备状态数据的持续监控，发现并预防各种安全风险，提高安全运营的效率。
- 同时，大模型还可以通过对安全审计的支持，帮助企业更好地遵守各种安全法规，避免因违规操作而带来的法律风险。
总的来说，腾讯混元大模型的出现，将会在安全领域开辟出更广阔的应用场景，充分利用好公司内已有的安全数据，做好业务数据处理和沉淀，将大大提高安全水位和业务效率。
作者：腾讯程序员
文章来源：腾讯技术工程
推荐阅读
更多腾讯AI相关技术干货，请关注专栏腾讯技术工程 欢迎添加极术小姐姐微信（id:aijishu20)加入技术交流群，请备注研究方向。
