以下内容为阅读论文所做记录和实验过程记录，非最终实验报告内容，实验报告最终版本见PDF文件.
部分CS224W课程学习笔记：
本篇论文的方向是安全邻域的代码漏洞识别问题。
最近的研究聚焦了深度神经网络在更自动化的漏洞识别方式上的应用，但现有的方法在学习综合程序语义以表征高多样性、脆弱性以及真实世界的代码的复杂性方面都有很大的局限。可以将现有方法的缺点总结为以下三点：
1）对源代码结构和逻辑的简化：传统的学习方法将源代码视为平面序列或仅使用部分信息来表示代码，忽略了源代码的结构性和逻辑性。源代码通常具有抽象语法树、数据流、控制流等复杂的表示方式，这些信息对于准确理解代码含义是至关重要的。
2）缺乏综合检测漏洞的能力：源代码中的漏洞有时是微妙的缺陷，需要从语义的多个维度进行综合调查才能发现。然而，现有方法在设计上存在缺陷，限制了其发现各种漏洞的潜力。
3）训练数据的问题：部分数据是由静态分析器标注的，这导致了高比例的误分类，即被标记为漏洞但实际上并不是真正的漏洞。另一部分数据是人工编写的简单代码，相比于真实世界的代码过于简单。
模型包含三个部分：1）复合代码语义的图嵌入层；2）门控图循环层；3）卷积层模块。各自功能如下：
- 将源代码编码成具有综合程序语义的联合图结构；
- 通过消息传递和消息聚合邻居节点的信息以学习每个节点特征；
- 提取有用的节点表征用于图层面的预测。
创新点：1）首次使用图的数据结构，并将不同层次的程序控制和数据依赖关系编码成一个由异质边
[1] 异质边（Heterogeneous edges）是指在图结构中连接不同类型节点之间的边。在异质图（Heterogeneous graph）中，节点可以表示不同的实体或概念，而异质边则表示这些实体或概念之间的关系或连接。异质边的存在使得我们能够在图结构中更好地捕捉不同实体之间的复杂关系和连接模式，从而提供更准确和全面的图分析和挖掘能力。异质图在许多现实世界的应用中都很常见。例如，在社交网络中，可以将用户表示为一个类型的节点，将用户之间的关系（如好友、关注关系）表示为另一类型的节点。这样，连接不同类型的节点的边就是异质边。在推荐系统中，可以将用户、商品和标签表示为不同类型的节点，并通过异质边来表示用户与商品之间的交互行为或商品与标签之间的关联。
转换为一个二元分类问题，给定一个源代码，判定其是否存在函数级别的代码漏洞。
假设代码数据样本定义为：
其中
- 把每个函数
$c_i$ 编码成一个多边图$g_i(V,X,A) \in \mathcal{G}$ ； -
$m$ 是$V$ 的节点个数； -
$X \in\mathbb{R}^{m \times d}$ 是初始节点特征矩阵，每个节点$v_j$ 被表示为一个$d$ 维的向量$x_j \in \mathbb{R}^{d}$ （即每个节点的特征是$1\times d$ ，$V$ 中一共有$m$ 个节点，因此$\mathbb{R}$ 的大小是$m \times d$ ）； - 邻接矩阵
$A \in \{ 0,1\}^{k \times m \times m}$ ，其中$k$ 是边类型的总数，$e^{p} _ {s,t} \in A$ 等于1表示节点$v_s$ 和$v_t$ 通过类型$p$ 的边相连，否则为0（每个节点都有可能与其它节点有$p$ 种类型的边相连，因此大小为$k \times m \times m$ ）
所以模型的目标是学习从
常见的代码程序表示方法有抽象语法树、控制流和数据流图（编译原理学过的那些）。比较优秀的漏洞检测方法需要结合这三种表示方法形成一个联合数据结构，同时本模型还考虑了源代码的自然序列，即一共四种表示方法。具体做法为：
将函数
给定一个嵌入图
在每个时刻
具体，到达
我对以上的公式进行了改造以助于理解：
这里
然后对传递的信息进行聚合，这里
因此当前时刻节点的隐藏状态就可以通过前一时刻的隐藏状态和聚合后的邻居信息表现为：
在经历了$T$时刻这样的更新后，
这里就是实验要求实现的部分
应用一维卷积和全连接层来学习与图级任务相关的特征，以获得更有效的预测。定义
它将输入的特征进行卷积操作，并通过ReLU激活函数进行非线性变换，然后使用最大池化对卷积结果进行下采样。
同时为了利用每个节点本身的信息，devign分别训练了两个卷积模块，得到特征：
最后再将两个特征输入到MLP中，然后对乘法结果进行平均聚合，最后使用Sigmoid激活函数将其映射到[0, 1]的范围内，得到节点的预测结果。
configs.json 配置文件包含了各个配置类所需的参数值，用于配置代码的行为和功能。它定义了不同类别的配置参数，如Devign模型、数据创建、路径、文件、代码嵌入和处理等。configs.py文件就是安装 configs.json 文件中最外层的键创建了一些对应的类供再上层获取参数，调用的地方就在 main.py 。
CREATE_PARAMS = configs.Create()
PROCESS_PARAMS = configs.Process()
PATHS = configs.Paths()
FILES = configs.Files()
DEVICE = PROCESS_PARAMS.device这里的 DEVICE 是根据 torch.cuda.is_available() and self.get_property("use_gpu")得到的。
这里我总结了一下 configs.py 里的类的功能，大致如下：
- Config一个基础配置类，用于加载和访问配置文件中的参数；
- Create类用于访问创建CPG数据相关的配置参数，如过滤条件、切片大小、Joern CLI目录和数据集大小等；
- Data类用于访问数据相关的配置参数，包括CPG数据、原始数据、输入数据、模型数据、标记数据和词向量数据等；
- Paths类用于访问路径相关的配置参数，如CPG数据路径、Joern路径和其他数据路径等；
- Files类用于访问文件相关的配置参数，如原始数据文件、CPG文件、标记文件、词向量文件和模型文件等；
- Embed类用于访问代码嵌入相关的配置参数，如最大节点数、Word2Vec参数和边类型等；
- Process类用于访问训练和评估模型相关的配置参数，如训练轮数、早停策略、批处理大小、设备选择等；
- Devign类用于访问Devign模型相关的配置参数，包括学习率、权重衰减、损失函数参数和模型参数等。
创建 Path() 里定义的文件夹，我们到 configs.json 里可以看到，路径为 data/*。
"paths": {
"cpg": "data/cpg/",
"joern": "data/joern/",
"raw": "data/raw/",
"input": "data/input/",
"model": "data/model/",
"tokens": "data/tokens/",
"w2v": "data/w2v/"
},data/raw/ 存放着我们在第一步里探索的原始代码数据集，data/cpg/ 存放着我们从根目录的压缩包里解压出来的用joern工具生成的CPG图。
执行完第二步的 main.setup() 就会创建其它文件夹以供后续使用。
画了个代码流程图
graph TB
A(开始)-->B[获取文件路径]
B[获取文件路径]--调用data.get_directory_file-->C[获取.pkl文件路径列表]
C-->D{遍历列表完成?}
D--未完成-->E[加载CPG数据集]
E--源代码分词-->F[获取分词结果并写入token文件夹]
F--初始化Word2Vec-->G[创建模型对象w2vmodel]
G-->I[学习每个标记的初始嵌入]
E-->J[将CPG转换为节点表示]
I-->K
J-->K[使用Word2Vec模型对每一个节点的文本进行编码]
K-->L[保存结果到input文件夹]
D--完成-->M[保存Word2Vec模型到为w2v文件夹]
L-->N(结束)
M-->N
前半部分是Devign模型初始化，模型初始化的时候用到了configs.Devign()里用到的参数。
"devign": {
"learning_rate": 1e-4,
"weight_decay": 1.3e-6,
"loss_lambda": 1.3e-6,
"model": {
"gated_graph_conv_args": {
"out_channels": 200,
"num_layers": 6,
"aggr": "add",
"bias": true
},
"emb_size": 101
}
},后半部分是训练的代码。
好久没用，check一下本地的环境
根据torch-geometric官方的教程安装一下这个包（看起来已经迭代到了比较好安装的样子，网上很多以前的安装教程还挺麻烦的）同时根据requirement.txt文件安装一下对应的包。
conda activate pytorch
conda install pyg -c pyg
pip install -r ./requirements.txt解压一下数据集
按指南探索一下代码数据集
执行第二步代码文本编码，完成word2vec模型的训练
这是控制台调用 embed_task() 打印的东西
根据论文里的内容，利用每个节点本身的信息，分别训练了两个卷积模块，最后再将两个特征输入到MLP中，然后对乘法结果进行平均聚合，最后使用Sigmoid激活函数将其映射到[0, 1]的范围内，得到节点的预测结果。按照公式摸出来就行。这里我设计了两个两层的卷积模块和对应的池化层，参数如下：
# 设置卷积层和池化层参数
self.conv1_size = {
"in_channels": self.max_nodes,
"out_channels": 64,
"kernel_size": 3,
"padding": 1
}
self.conv2_size = {
"in_channels": 64,
"out_channels": 16,
"kernel_size": 2,
"padding": 1
}
self.maxp1_size = {
"kernel_size": 3,
"stride": 2
}
self.maxp2_size = {
"kernel_size": 2,
"stride": 2
}
self.feature1 = nn.Conv1d(**self.conv1_size)
self.maxpool1 = nn.MaxPool1d(**self.maxp1_size)
self.feature2 = nn.Conv1d(**self.conv2_size)
self.maxpool2 = nn.MaxPool1d(**self.maxp2_size)对于两个特征 get_conv_mp_out_size() 计算（但是在后面实验的过程中其实遇到了一些小问题）调用传入的第一个参数需要由前面的图神经网络的参数计算，所以需要在初始化函数里添加图神经网络的 out_channels 和 emb_size，根据公式：
graph_out_chs + emb_size，而 graph_out_chs。
# 根据conv和maxpool参数计算mlp尺寸
self.mlp1_size = get_conv_mp_out_size(
graph_out_chs + emb_size,
self.conv2_size,
[self.maxp1_size, self.maxp2_size]
)
self.mlp2_size = get_conv_mp_out_size(
graph_out_chs,
self.conv2_size,
[self.maxp1_size, self.maxp2_size]
)
self.mlp1 = nn.Linear(1200, 1)
self.mlp2 = nn.Linear(self.mlp2_size, 1)再补充forward函数
def forward(self, h, x):
z_feature = torch.cat([h, x], 1)
z_feature = z_feature.view(-1, self.max_nodes, h.shape[1] + x.shape[1])
out_z = self.maxpool1(F.relu(self.feature1(z_feature)))
out_z = self.maxpool2(F.relu(self.feature2(out_z)))
out_z = out_z.view(-1, int(out_z.shape[1] * out_z.shape[-1]))
out_z = self.mlp1(out_z)
y_feature = h.view(-1, self.max_nodes, h.shape[1])
out_y = self.maxpool1(F.relu(self.feature1(y_feature)))
out_y = self.maxpool2(F.relu(self.feature2(out_y)))
out_y = out_y.view(-1, int(out_y.shape[1] * out_y.shape[-1]))
out_y = self.mlp2(out_y)
out = out_z * out_y
out = torch.sigmoid(torch.flatten(out))
return out嗯...好的，把配置文件的 process.use_gpu 改成 false（破电脑是这样的ToT）然后重启ipynb文件
打印显示了 train with cpu
出了bug，调用老师给的函数 get_conv_mp_out_size() 算出来的是1216，但是实际是1200所以对不上，暴力修改Linear层的第一个参数为1200。
第二个倒是没出问题
跑出一个batch_loss了
继续放着跑吧，看到 EarlyStopping=True 应该是会跑十轮然后自己停
跑完了
结果
