AFLGO源码分析与教学文档

1. AFLGO概述

AFLGO是基于AFL的定向模糊测试工具，通过LLVM的pass机制实现插桩逻辑，并利用模拟退火算法进行目标导向的模糊测试。核心功能包括：

• 预处理阶段：分析目标程序的控制流和调用关系
• 距离计算阶段：计算基本块到目标位置的距离
• 插桩阶段：在目标程序中插入距离测量和覆盖率收集代码

2. AFLGO-PASS.SO.CC分析

2.1 参数设置

AFLGO pass需要以下参数：

• -targets：包含目标位置行号的文件
• -distance：包含基本块到目标位置距离的文件
• -outdir：输出文件路径

2.2 DOT模板特化

AFLGO通过修改LLVM的DefaultDOTGraphTraits基类来生成DOT文件：

• getGraphName：设置DOT文件的图名（如"CFG for 'main' function"）
• getNodeLabel：设置基本块标签（使用基本块名或操作数格式名）

2.3 AFLCoverage类

模块级Pass，继承自ModulePass：

class AFLCoverage : public ModulePass {
    static char ID;  // Pass的唯一标识符
    AFLCoverage() : ModulePass(ID) {}
    bool runOnModule(Module &M) override;
    // 其他成员函数...
};

2.4 getDebugLoc函数

获取IR指令对应的源码位置：

static bool getDebugLoc(const Instruction *I, std::string &Filename, unsigned &Line) {
    if (DILocation *Loc = I->getDebugLoc()) {
        Line = Loc->getLine();
        Filename = Loc->getFilename().str();
        // 处理内联函数情况
        if (Filename.empty() && Loc->getInlinedAt())
            return getDebugLoc(Loc->getInlinedAt(), Filename, Line);
        return true;
    }
    return false;
}

2.5 runOnModule主逻辑

2.5.1 参数校验与准备

检查必需参数并初始化数据结构：

if (!TargetsFile.empty()) {
    // 读取目标位置到targets列表
    is_aflgo_preprocessing = true;
}
if (!DistanceFile.empty()) {
    // 读取距离信息到distance_map哈希表
    // 距离值x100转为int
    is_aflgo = true;
}

2.5.2 预处理阶段

1. 打开输出文件（BBnames.txt, BBcalls.txt等）
2. 遍历模块中的所有函数，跳过黑名单函数
3. 对每个基本块：
   • 获取IR指令的源码位置
   • 跳过系统库代码（/usr/开头的文件）
   • 为无名基本块生成名称（文件名+行号）
   • 检查是否为目标基本块
   • 处理函数调用指令，记录调用关系
   • 写入基本块名称到BBnames.txt
4. 在tracing模式下，向基本块末尾插入桩代码：

IRBuilder<> Builder(BB->getTerminator());
Type *Args[] = {Builder.getInt8PtrTy()};
FunctionType *FTy = FunctionType::get(
    Type::getVoidTy(M.getContext()), Args, false);
Function *instrumented = M.getOrInsertFunction(
    "llvm_profiling_call", FTy);
Builder.CreateCall(instrumented, {bbnameVal});

5. 输出控制流图（.dot文件）
6. 记录目标基本块和函数信息

2.5.3 距离插桩阶段

1. 定义IR数据类型：

   • Int8Ty：1字节（bitmap）
   • Int32Ty：用于cur_loc、prev_loc
   • Int64Ty：距离统计（x86_64平台）

2. 初始化全局变量：

GlobalVariable *AFLMapPtr = new GlobalVariable(
    M, Int8Ty, false, GlobalValue::ExternalLinkage,
    0, "__afl_area_ptr");
GlobalVariable *AFLPrevLoc = new GlobalVariable(
    M, Int32Ty, false, GlobalValue::ExternalLinkage,
    0, "__afl_prev_loc");

3. 遍历函数和基本块：

   • 获取基本块名称
   • 跳过不在BBnames.txt中的基本块（选择性插桩时）
   • 从distance_map获取距离值

4. AFL传统插桩逻辑：

// 设置当前基本块ID
unsigned cur_loc = R(MAP_SIZE);
ConstantInt *CurLoc = ConstantInt::get(Int32Ty, cur_loc);

// 获取上一个基本块ID
LoadInst *PrevLoc = Builder.CreateLoad(AFLPrevLoc);
PrevLoc->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
Value *PrevLocCasted = Builder.CreateZExt(PrevLoc, Int32Ty);

// 加载共享内存指针
LoadInst *MapPtr = Builder.CreateLoad(AFLMapPtr);
MapPtr->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
Value *MapPtrIdx = Builder.CreateGEP(
    MapPtr, Builder.CreateXor(PrevLocCasted, CurLoc));

// 增加bitmap计数器
Value *Counter = Builder.CreateLoad(MapPtrIdx);
Counter->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));
Value *Incr = Builder.CreateAdd(Counter, ConstantInt::get(Int8Ty, 1));
Builder.CreateStore(Incr, MapPtrIdx)
    ->setMetadata(M.getMDKindID("nosanitize"), MDNode::get(C, None));

// 更新prev_loc
Builder.CreateStore(ConstantInt::get(Int32Ty, cur_loc >> 1), AFLPrevLoc);

5. AFLGO距离插桩：

if (distance != -1) {
    // 将距离转为IR常量
    ConstantInt *Distance = ConstantInt::get(LargestType, (unsigned)distance);
    
    // 获取并更新总距离
    LoadInst *MapDist = Builder.CreateLoad(MapDistPtr);
    Value *IncrDist = Builder.CreateAdd(MapDist, Distance);
    Builder.CreateStore(IncrDist, MapDistPtr);
    
    // 更新计数器
    LoadInst *MapCnt = Builder.CreateLoad(MapCntPtr);
    Value *IncrCnt = Builder.CreateAdd(MapCnt, One);
    Builder.CreateStore(IncrCnt, MapCntPtr);
}

2.6 Pass注册

static void registerAFLPass(const PassManagerBuilder &,
                           legacy::PassManagerBase &PM) {
    PM.add(new AFLCoverage());
}

static RegisterStandardPasses RegisterAFLPass(
    PassManagerBuilder::EP_OptimizerLast, registerAFLPass);
static RegisterStandardPasses RegisterAFLPass0(
    PassManagerBuilder::EP_EnabledOnOptLevel0, registerAFLPass);

3. DISTANCE.PY分析

3.1 参数设置

parser.add_argument('-d', '--dot', required=True, help="输入dot文件路径")
parser.add_argument('-t', '--targets', required=True, help="目标位置文件")
parser.add_argument('-o', '--output', required=True, help="输出距离文件")
parser.add_argument('-n', '--names', required=True, help="节点名称文件")
parser.add_argument('-c', '--cg-distance', help="CG距离文件")
parser.add_argument('-s', '--cg-callsites', help="CG调用点文件")

3.2 处理流程

1. 读取dot文件并判断模式（CG或CFG）

2. CG模式处理：

   • 读取目标位置
   • 计算每个节点到目标的距离
   • 使用迪杰斯特拉算法找最短路径
   • 距离公式：1 / (1 + dist)的加权平均倒数

3. CFG模式处理：

   • 检查必需参数（cg_distance和cg_callsites）
   • 建立函数名到距离的映射
   • 处理基本块调用关系
   • 距离计算：

     dist = shortest + 10 * bb_d  # bb_d是函数间距离
     

   • 取所有路径中的最小值作为最终距离

3.3 距离计算函数

def distance(G, name, targets, bbdistance, mode):
    if mode == 'cg' and name in bbdistance:
        return bbdistance[name] * 10
    
    if mode == 'cg':
        # 计算CG距离
        dist = 0.0
        for target in targets:
            try:
                d = nx.shortest_path_length(G, name, target)
                dist += 1.0 / (1.0 + d)
            except:
                continue
        return 1.0 / dist if dist > 0 else sys.maxsize
    
    else:  # CFG模式
        min_dist = sys.maxsize
        for bb, bb_d in bbdistance.items():
            try:
                shortest = nx.shortest_path_length(G, name, bb)
                current_dist = shortest + 10 * bb_d
                if current_dist < min_dist:
                    min_dist = current_dist
            except:
                continue
        return min_dist

4. 关键知识点总结

1. LLVM Pass机制：AFLGO利用LLVM的Pass系统在IR级别进行插桩，比汇编级插桩更高效可靠。

2. 控制流分析：

   • 通过DOT文件生成CFG和CG
   • 使用networkx库进行图分析
   • 基本块命名规则：文件名+行号

3. 距离计算：

   • CG距离：函数调用图级别的距离
   • CFG距离：基本块级别的距离
   • 跨函数调用惩罚：距离×10
   • 最终距离取最小值

4. 插桩策略：

   • 传统AFL边覆盖率插桩（异或计数）
   • AFLGO特有距离插桩（累加距离和计数器）
   • 选择性插桩支持

5. 共享内存布局：

   +---------------------+
   | AFL bitmap (MAP_SIZE)|
   +---------------------+
   | 总距离值 (8字节)     |
   +---------------------+
   | 计数器 (8字节)       |
   +---------------------+
   

5. 教学建议

1. 实践步骤：

   • 编译AFLGO并生成aflgo-pass.so
   • 使用afl-clang-fast编译目标程序
   • 运行预处理阶段生成控制流信息
   • 运行distance.py计算距离
   • 进行定向模糊测试

2. 调试技巧：

   • 检查生成的.dot文件验证控制流
   • 检查BBnames.txt等中间文件
   • 使用LLVM IR查看插桩结果

3. 扩展方向：

   • 修改距离计算公式
   • 添加新的目标选择策略
   • 优化共享内存布局

4. 常见问题：

   • 确保目标程序包含调试信息
   • 检查Pass参数是否正确传递
   • 验证距离计算结果的合理性

通过本教学文档，读者可以全面理解AFLGO的工作原理和实现细节，为进一步研究和定制定向模糊测试工具奠定基础。