4. 基于数据流分析的平坦化还原

任务目标

本案例针对OLLVM控制流平坦化混淆后的函数，通过数据流分析追踪状态变量（switch变量）的赋值路径，识别各基本块之间的真实跳转关系，从而重构原始分支逻辑，输出还原后的控制流图。

前置环境

需要LLVM/Clang（版本≥14），并安装OLLVM（或使用clang的-mllvm -fla开关模拟平坦化）。本案例使用LLVM Pass框架实现分析，不依赖外部二进制分析工具。

核心思路

平坦化后函数呈现为：一个分发块（dispatch）+ 多个状态块（state），每个状态块末尾对状态变量赋值并跳回分发块。还原的关键在于：对每个状态块，分析其赋值指令（如store i32 5, i32* %state），确定该状态块在何种条件下会执行，以及它赋给状态变量的具体值。通过构建状态转移图，将br指令从“跳回分发块”改写为“跳转到对应状态块”，即可还原原始分支。

完整实现代码

// FlatteningRestore.cpp
#include "llvm/IR/Function.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/CFG.h"
#include "llvm/Pass.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/Transforms/IPO/PassManagerBuilder.h"
#include "llvm/IR/Constants.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/IR/IRBuilder.h"
#include <map>
#include <set>

using namespace llvm;

namespace {
struct FlatteningRestore : public FunctionPass {
    static char ID;
    FlatteningRestore() : FunctionPass(ID) {}

    bool runOnFunction(Function &F) override {
        // 1. 识别状态变量：被store指令频繁写入的整型指针
        Value *stateVar = nullptr;
        std::map<BasicBlock*, int> stateValue; // 每个状态块写入的常量值
        std::set<BasicBlock*> stateBlocks;

        for (auto &BB : F) {
            for (auto &I : BB) {
                if (auto *SI = dyn_cast<StoreInst>(&I)) {
                    Value *ptr = SI->getPointerOperand();
                    if (ptr->getType()->isPointerTy() &&
                        ptr->getType()->getPointerElementType()->isIntegerTy(32)) {
                        if (auto *CI = dyn_cast<ConstantInt>(SI->getValueOperand())) {
                            if (!stateVar) stateVar = ptr;
                            if (ptr == stateVar) {
                                stateValue[&BB] = CI->getZExtValue();
                                stateBlocks.insert(&BB);
                            }
                        }
                    }
                }
            }
        }
        if (!stateVar || stateBlocks.empty()) return false;

        // 2. 找到分发块：包含对stateVar的load和switch指令
        BasicBlock *dispatchBB = nullptr;
        for (auto &BB : F) {
            for (auto &I : BB) {
                if (auto *SI = dyn_cast<SwitchInst>(&I)) {
                    dispatchBB = &BB;
                    break;
                }
            }
            if (dispatchBB) break;
        }
        if (!dispatchBB) return false;

        // 3. 构建状态转移图：从每个state块出发，模拟执行到下一个store
        std::map<int, BasicBlock*> valueToBlock; // 状态值 -> 状态块
        for (auto *BB : stateBlocks) {
            valueToBlock[stateValue[BB]] = BB;
        }

        // 4. 重写分支：将state块末尾的br跳转到dispatch，改为直接跳转到目标state块
        bool changed = false;
        for (auto *BB : stateBlocks) {
            // 找到BB的终止指令（通常是br dispatch）
            Instruction *term = BB->getTerminator();
            if (auto *BI = dyn_cast<BranchInst>(term)) {
                if (BI->isUnconditional()) {
                    BasicBlock *succ = BI->getSuccessor(0);
                    if (succ == dispatchBB) {
                        // 分析该state块的条件分支来源：检查其前驱
                        // 简化：若该state块有多个前驱，则根据前驱的state值确定转移
                        // 此处采用保守策略：若只有一个前驱，直接连接
                        if (BB->hasNPredecessors(1)) {
                            BasicBlock *pred = *pred_begin(BB);
                            if (stateBlocks.count(pred)) {
                                // 从pred的switch case中获取目标值
                                // 实际实现需遍历pred的switch指令
                                // 为演示，直接使用stateValue映射
                                int targetVal = stateValue[BB];
                                // 找到写入targetVal的块
                                if (valueToBlock.count(targetVal)) {
                                    BasicBlock *target = valueToBlock[targetVal];
                                    if (target != BB) {
                                        // 改写跳转
                                        BI->setSuccessor(0, target);
                                        changed = true;
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
        return changed;
    }
};
char FlatteningRestore::ID = 0;
static RegisterPass<FlatteningRestore> X("flattening-restore", "Restore flattened CFG via dataflow", false, false);
}

编译与运行

# 编译Pass
clang++ -shared -fPIC -o libFlatteningRestore.so FlatteningRestore.cpp \
    $(llvm-config --cxxflags --ldflags --libs)

# 生成测试程序（含平坦化）
clang -mllvm -fla -c test.c -o test_flat.bc

# 运行还原Pass
opt -load ./libFlatteningRestore.so -flattening-restore test_flat.bc -o test_restored.bc

# 生成可读CFG对比
llvm-dis test_flat.bc -o - | grep "define" 
llvm-dis test_restored.bc -o - | grep "define"

验证方法

准备一个包含if-else或switch的简单函数（如int foo(int x){ if(x>0) return 1; else return 0; }），编译时加-mllvm -fla。运行还原Pass后，通过llvm-dis查看IR：

• 通过条件：还原后的函数中，switch指令消失，状态块之间的跳转直接对应原始条件分支（如br i1 %cmp, label %then, label %else）。
• 若原始函数有多个条件分支，还原后br指令的条件操作数应能对应原始比较结果。

注意：本实现为教学演示，仅处理单前驱状态块的简单情形；复杂嵌套循环或switch多分支需扩展数据流分析（如使用LLVM的LazyValueInfo或ValueTracking追踪状态变量的到达定值）。