基于LLM大模型的SQL注入检测探索

前言

在SQL注入攻击检测领域，通常分为正则关键词规则检测以及基于机器学习算法检测，随着OpenAI带来的大模型快速发展，使用大模型LLM来检测传统漏洞SQL注入也值得探索研究

LLM检测框架

包含数据预处理、提示工程、指令微调、标签映射和检测能力提升5个部分

数据预处理

选择了公开的Kaggle数据集作为基准数据集（包含30873条样本，分为SQL注入攻击与非SQL注入攻击两类，包含所有常见的SQL注入类型），数据集被划分训练集、测试集和验证集3个部分。

• 训练集被用于大语言模型的微调

• 验证集被用于训练阶段中的模型检测能力提升，以选择最优的推理参数和最佳模型    

• 测试集被用于评估模型的性能和泛化能力

原始数据集可能存在缺失值、异常值、噪声或数据不平衡等问题，通过数据预处理步骤可以提高样本的质量，主要采用下列方式进行处理

• 删除重复样本

• 删除空样本

• URL解码：实际收集的SQL注入样本是经过URL编码的，该编码可能改变SQL注入样本的语义，影响检测效果。因此，本文需要在实施检测之前对样本进行URL解码。例如，解码前是1%27)%20UNION%20ALL%20SELECT%20md5(233)—+，解码后是1') UNION ALL SELECT md5(233)—+。

• 删除冗余空格：过长的输入数据将影响大语言模型的推理速度，并且带来额外噪声

提示工程

提示主要分为任务提示型和答案提示型两种类型。

• 任务提示型的主要功能是基于当前任务为模型建立一个导向。例如，输入一段描述，并添加“翻译”或者“问答”的提示，则大语言模型将生成相应的应答。

• 答案提示型的主要功能是对输出结果进行约束，固定大语言模型的输出格式，便于标签映射

参考了OpenAI的提示设计文档和斯坦福Alpaca微调文档，并考虑微调和推理应该使用同一个提示

• 描述部分引导模型根据指令构建任务，并基于任务和输入生成应答作为输出

• 指令是整个提示的核心，本文分析了大量SQL注入攻击样本，将指令分为任务提示、示例提示以及答案提示，任务提示引导模型优先关注输入部分的URL并检测其中是否包含SQL注入的攻击模式，示例提示展示了一些可能的攻击模式，如嵌入特殊字段‘OR’、‘1=1’，过度使用特殊字符等，答案提示指导模型选择“Yes”或“No”作为回答

指令微调

指令微调仿照人类学习的模式，通过简单的任务定义以及对应示例，能很好地引导大语言模型理解任务要求，进而获得预期的输出([21])。当前，主流的大语言模型微调技术包括低秩自适应（Low-Rank Adaptation，LoRA技术、适配器（Adapter技术、前缀微调（Prefix Tuning）技术等。

指令微调流程如下：

1. 将预期输出拼接至提示的输出部分之后构建微调数据，本文指定SQL注入攻击样本的微调预期输出为Yes，非SQL注入攻击样本的微调预期输出为No

2. 基于微调数据使用LoRA技术对模型进行微调，该技术在微调过程中冻结大语言模型权重参数，通过在旁路加入额外的网络层，并且只改变这些新增的网络层参数，最后将大语言模型的输出与额外网络层的输出求和，得到最终输出    

由此得到核心SQL注入攻击检测的专用大语言模型。

标签映射

由于大语言模型输出具有随机性，即使在模型微调时应用预期应答Yes或No，仍有一定概率输出非预期应答。为在获得检测标签的同时减少非预期应答对实验产生的影响，构建了一系列标签映射，将包含这些关键字的应答映射到相应的标签中

包含Yes、True、1关键字的应答，被映射到SQL注入攻击类；

包含No、False、0关键字的应答被映射到非SQL注入攻击类；

检测能力提升

影响大语言模型SQL注入攻击检测性能的主要因素：提示对齐能力和SQL注入攻击判别能力

定义1 提示对齐能力 设测试样本总数为N，其中测试后未知标签的样本数量为U，则提示条件prompt下的提示对齐能力如公式（1）所示    

定义2 SQL注入攻击判别能力 TP值为模型准确地识别出正例（即SQL注入攻击）的样本数量，TN指模型准确地识别出负例（即非SQL注入攻击）的样本数量，则提示条件prompt下的SQL注入攻击判别能力如公式（2）所示。

AC为介于0到1的值，其大小反映出模型的提示对齐能力，即值越大，提示对齐能力越强；

DC的大小表示模型根据提示分辨SQL注入攻击的准确率，即值越大，模型对SQL注入攻击的判别越准确；

由AC和DC的定义易得，模型准确率如公式（3）所示

综上所述，AC与DC都是影响模型检测准确率的关键因素，很明显提升模型检测识别准确率，要么增大分子、要么减小分母

落地实践

使用公开数据集Kaggle，将其划分为训练集、测试集和验证集，在具体应用时影响正确率的因素有很多，下列会简单介绍几种

迭代轮数扫描

在训练集上对7B和13B版本的模型进行了微调，分别迭代150000轮，在微调阶段，跟踪记录了每轮迭代后模型的损失函数，并采用非线性函数进行拟合，随着迭代轮数的增加，损失函数值迅速下降，当迭代轮次超过25000轮后，模型的损失函数值保持稳定。这表明模型具有良好的收敛性，且参数结构随迭代轮数的增加趋于稳定

在微调过程中，每1280轮评估一次模型，评估标准包括在验证集上的提示对齐能力、SQL注入攻击判别能力、准确率和误报率，可以看到随着轮数增加，准确率有明显得提升；在完成25000轮微调后，这两项能力指标都接近100%    

样本数扫描

探讨微调样本数对模型性能的具体影响，实验首先从训练集中分别随机抽取2的N次方个样本作为子训练集，每个较大的子训练集内包含所有较小的子训练集；然后为确保模型收敛，7B和13B版本的模型在所有子训练集上分别进行了25600轮次的微调；最后在验证集上对模型的性能进行了评估 ，可以很明显看到当微调样本增加至64时，SQL注入攻击判别能力达到95%以上，而在使用全部样本进行微调后，该能力得到了进一步的提升，达到了99.5%    

推理参数扫描

研究推理参数对模型性能的影响。实验选取了不同参数组合，共有66种不同的参数组合；当Temperature =1.0且Top P=0.75时7B版本模型准确率最高，当Temperature =0.4且Top P=0.75时13B版本模型准确率最高；因此，适当的随机性对提升模型性能具有积极影响。由于输出存在一定的随机性，因此减小了模型过度拟合某个输入特征的可能性，进而带来了较好的准确率    

原文始发于微信公众号（暴暴的皮卡丘）：基于LLM大模型的SQL注入检测探索