用于漏洞生成的强化学习
几个月前,在一场 CTF 决赛中,我们的团队在一个 pwn 二进制文件上卡了三个小时。栈金丝雀、NX、部分 PIE,在一个不起眼的函数中恰好 72 字节的溢出。我们什么都试过了。毫无进展。凌晨 03:14,我们中的一个人几乎是开玩笑地启动了一个五行脚本,用 AFL 式的位翻转来变异输入,然后不假思索地重发。二十分钟后,这个脚本凭着纯粹的统计运气拿到了 flag。
那天晚上我们带着一个挥之不去的问题回了家:如果我们拥有的不是一个愚蠢的进程,而是一个能从每次失败尝试中学习的进程,会怎么样?
为什么 RL 与控制流劫持契合
漏洞的搜索空间残酷但不均匀。大多数输入产生相同的行为,只有一个狭窄的区域能引导到你想要的机器状态:指令寄存器指向你控制的内存。区分该子空间的是信息梯度:接近破坏 RIP 的输入会表现出可观察的副作用——部分被覆盖的寄存器、与输入字节匹配的栈值、跳转到无效但接近受控字节的地址。
这正是设计良好的奖励函数所能捕捉的。
当前技术水平
Avgerinos 和 Brumley 在 2011 年发表了第一个基于符号执行的端到端系统 AEG(NDSS 2011)。Mayhem(Cha 等,S&P 2012)扩展了该分析。两者都是准确定性的,但要付出符号执行的经典代价:路径爆炸和被噎住的 SMT 求解器。
Böttinger 用 Deep Reinforcement Fuzzing(2018)使用 DQN 将模糊测试形式化为 MDP。PwnGPT(Shao 等,ACL 2025)报告使用 o1-preview 将漏洞生成率从 26.3% 提升至 57.9%。PPO(Schulman 等,2017)+ RLHF(Ouyang 等,2022)已经成熟到足以使在二进制环境上训练智能体成为严肃的话题。
我们的环境
基于 ptrace 的 Gym 风格环境和一个 eBPF 边车。观测是一个密集张量,包含:
- 16 个通用寄存器的状态(x86_64),已归一化。
- RSP 周围的栈窗口(256 字节)作为原始字节。
- 控制标志:
NX、canary_present、RIP_corrupted、register_controlled[]。 - 从 eBPF(在每个基本块上的 uprobe)取得的最后一次覆盖追踪的哈希。
动作空间在输入字节上是离散的。奖励结合了新覆盖率(正向、递减)、可利用性启发式和长度惩罚。
import gym
from gwaihir.tools.rl_env import BinaryPwnEnv
class CTFPwnEnv(BinaryPwnEnv):
def step(self, action):
mutated = self.apply_action(self.current_input, action)
trace = self.run_with_ebpf(mutated)
new_bbs = len(trace.basic_blocks - self.coverage_seen)
cov_reward = np.log1p(new_bbs) * 0.3
self.coverage_seen |= trace.basic_blocks
expl = 0.0
if trace.rip_controlled_bytes > 0:
expl += 5.0 + 0.1 * trace.rip_controlled_bytes
if trace.canary_leaked: expl += 2.5
if trace.tainted_syscall_args: expl += 1.5
reward = cov_reward + expl - 0.001 * len(mutated)
done = trace.rip_fully_controlled or self.steps > 4096
return self.observe(trace), reward, done, {"trace": trace}
PPO vs DQN
我们从 DQN 开始。一旦动作空间变大,DQN 就变得脆弱:每次可利用性奖励触发一个巨大且孤立的值,目标网络就会发散。PPO 通过对策略比率的裁剪,可以吸收这些尖峰而不崩溃。使用 GAE(λ=0.95)、熵系数 0.01 和 4096 的批次大小的 PPO 在一张 RTX 4090 上为我们带来了 24 小时稳定的训练。
一个鲜少有人提及的细节:大部分时间不是花在梯度上,而是花在运行二进制文件上。每一步都需要 fork、ptrace、eBPF 收集、解析。我们用 SubprocVecEnv 在每个 GPU 上并行 32 个环境,即便如此瓶颈仍然是 I/O。
在 10 个 CTF 二进制文件上的结果
测试集:两个 baby、三个需要 canary 泄露、两个格式化字符串、两个需要最小 ROP 链,以及一个堆漏洞(UAF)。该智能体逐个二进制从零训练,解决了 10 个中的 7 个。失败的三个:堆漏洞(在预期内)、一个需要不存在的 gadget 的 ROP,以及一个格式化字符串,在该题中策略陷入了永远泄露地址的局部最小值。
与纯 angr 相比:RL 赢了 5 个,平 2 个,输了 3 个(angr 更早解决了堆和格式化写)。它不是一个更优越的工具,而是一个互补的工具。
残酷的局限
- 每个二进制一个策略。我们不在挑战之间迁移知识。
- 训练时间数小时,而资深人类在几分钟内就能解决的挑战。
- 奖励塑造很脆弱。自定义分配器或非标准缓解措施会破坏这些启发式。
- 当路径可确定性推导时,符号执行仍然胜出。
真正的权衡:廉价的广覆盖(RL)与昂贵的深度推理(symbolic)。我们在 Gwaihir 内部并行运行它们。
下一步
三个方向:(a) 在 exploit-db 和 CTF writeup 上进行预训练,(b) RLHF 风格的奖励学习,(c) 将智能体作为协调者 LLM 的子工具集成。
目前我们继续训练。下次一支团队在凌晨三点卡了三个小时时,我们希望是智能体找到 flag——而不是运气。
参考文献
- Avgerinos, T. et al. (2011). AEG: Automatic Exploit Generation. NDSS.
- Cha, S.K. et al. (2012). Unleashing Mayhem. IEEE S&P.
- Böttinger, K. et al. (2018). Deep Reinforcement Fuzzing. IEEE SPW. arXiv:1801.04589.
- Schulman, J. et al. (2017). PPO. arXiv:1707.06347.
- Ouyang, L. et al. (2022). InstructGPT/RLHF. NeurIPS.
- Shao, Y. et al. (2025). PwnGPT. ACL.
