LLVM Obfucation example - Control Flow Flattening with ChatGPT

#include "llvm/IR/Function.h"
#include "llvm/Pass.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Transforms/Utils/BasicBlockUtils.h"

using namespace llvm;

namespace {

class ControlFlowFlattening : public FunctionPass {
public:
  static char ID;
  ControlFlowFlattening() : FunctionPass(ID) {}

  bool runOnFunction(Function &F) override {
    bool Modified = false;
    for (auto &BB : F) {
      if (BB.getTerminator()->getNumSuccessors() > 1) {
        Modified |= FlattenBlock(BB);
      }
    }
    return Modified;
  }

private:
  bool FlattenBlock(BasicBlock &BB) {
    // Save the original terminator.
    auto *OrigTerm = BB.getTerminator();

    // Create a new basic block with a switch instruction.
    auto *NewBB = BasicBlock::Create(BB.getContext(), "", &BB.getParent());
    auto *SwitchInst = SwitchInst::Create(&BB.front(), OrigTerm->getParent(), 2, NewBB);
    for (unsigned i = 0, e = OrigTerm->getNumSuccessors(); i != e; ++i) {
      auto *Succ = OrigTerm->getSuccessor(i);
      SwitchInst->addCase(ConstantInt::get(Type::getInt32Ty(BB.getContext()), i), Succ);
    }

    // Replace the original terminator with a branch to the new block.
    BranchInst::Create(NewBB, &BB);

    // Insert an unconditional branch from the new block to the original targets.
    for (unsigned i = 0, e = OrigTerm->getNumSuccessors(); i != e; ++i) {
      auto *Succ = OrigTerm->getSuccessor(i);
      auto *NewTerm = BranchInst::Create(Succ, NewBB);
      NewTerm->setDebugLoc(OrigTerm->getDebugLoc());
    }

    // Remove the original terminator.
    OrigTerm->eraseFromParent();

    return true;
  }
};

} // namespace

char ControlFlowFlattening::ID = 0;
static RegisterPass<ControlFlowFlattening> X("cff", "Control Flow Flattening Pass");

// Usage: opt -load /path/to/libControlFlowFlattening.so -cff <input.ll> -o <output.ll>

이 예제에서는 하나 이상의 후계자가 있는 함수에서 기본 블록의 제어 흐름을 평탄화하는 ControlFlowFlattening 패스를 정의합니다. 이 패스는 기본 블록의 터미네이터의 원래 대상으로 분기하는 스위치 인스트럭션이 포함된 새 기본 블록을 생성하는 방식으로 작동합니다. 그런 다음 원래 터미네이터가 새 블록에 대한 분기로 대체되고 새 블록에서 각 원래 대상으로 무조건 분기가 삽입됩니다.

이 패스를 사용하려면 공유 라이브러리로 컴파일하고 -load 옵션을 사용하여 opt에 로드해야 합니다. 그런 다음 -cff 옵션과 함께 opt를 실행하여 입력 LLVM IR 파일에 패스를 적용합니다.

 

> 수행 결과

Before

define i32 @foo(i32 %x) {
entry:
  %cmp = icmp slt i32 %x, 10
  br i1 %cmp, label %if.then, label %if.end

if.then:
  %add = add i32 %x, 1
  br label %if.end

if.end:
  %res = phi i32 [ %add, %if.then ], [ %x, %entry ]
  ret i32 %res
}

After

define i32 @foo(i32 %x) {
entry:
  switch i32 %x, label %cbb [
    i32 0, label %if.then
    i32 1, label %if.end
  ]

cbb:
  %add = add i32 %x, 1
  br label %if.end

if.then:
  br label %cbb

if.end:
  %res = phi i32 [ %add, %if.then ], [ %x, %entry ]
  ret i32 %res
}

보시다시피, br i1 %cmp, label %if.then, label %if.end 명령어가 %x 값에 따라 두 대상 중 하나로 분기하는 스위치 명령어로 대체되었습니다. 원래 종결자는 원래 if.then 블록의 논리를 포함하는 새로운 기본 블록(cbb)으로 분기하는 것으로 대체되었습니다. cbb 블록에서 원래 if.end 블록으로 새로운 무조건 분기가 삽입되었습니다. 마지막으로, 원래 if.then 블록이 다시 cbb 블록으로 분기되도록 수정되었습니다

 

----------------- 여기까지 ChatGPT -----------------

 

ChatGPT가 보여준 코드와 예제는 결과가 상이합니다. ChatGPT가 학습할 데이터가 많이 부족했을 것이라는 추론이 듭니다.
1. ChatGPT는 %cmp = icmp slt i32 %x, 10 을 생략합니다.

2. 1이 생략된 상태에서 switch 분기는 잘못된 실행 흐름을 가져옵니다.

 

 과거에 작업했던 기억을 더듬어 실행이 되는 코드를 붙여넣습니다.

class ControlFlowFlattening : public FunctionPass {
public:
  static char ID;
  ControlFlowFlattening() : FunctionPass(ID) {}

  bool runOnFunction(Function &F) override {
    bool Modified = false;
    for (auto &BB : F) {
      if (BB.getTerminator()->getNumSuccessors() > 1) {
        Modified |= FlattenBlock(BB);
      }
    }

    return Modified;
  }

private:
  bool FlattenBlock(BasicBlock &BB) {
    // Save the original terminator.
    auto *OrigTerm = BB.getTerminator();

    // Create a new basic block with a switch instruction.
    auto *TermBB = BasicBlock::Create(BB.getContext(), "", BB.getParent());
    auto *SwitchInst = SwitchInst::Create(&BB.front(), TermBB, 2, &BB);
    for (unsigned i = 0, e = OrigTerm->getNumSuccessors(); i != e; ++i) {
      auto *Succ = OrigTerm->getSuccessor(i);
      SwitchInst->addCase(ConstantInt::get(Type::getInt1Ty(BB.getContext()), i), Succ);
    }

    // Replace the original terminator with a branch to the new block.
    ReturnInst::Create(BB.getContext(), ConstantInt::get(Type::getInt32Ty(BB.getContext()), 0), TermBB);
    // Remove the original terminator.
    OrigTerm->eraseFromParent();

    return true;

  }
};

} // namespace

char ControlFlowFlattening::ID = 0;
static RegisterPass<ControlFlowFlattening> X("cff", "Control Flow Flattening Pass");

// Usage: opt -load /path/to/libControlFlowFlattening.so -cff <input.ll> -o <output.ll>

위 코드의 실행 결과입니다.

 

; ModuleID = 'after'
source_filename = "before.ll"

define i32 @foo(i32 %x) {
entry:
  %cmp = icmp slt i32 %x, 10
  switch i1 %cmp, label %0 [
    i1 false, label %if.then
    i1 true, label %if.end
  ]       
         
if.then:                                          ; preds = %entry
  %add = add i32 %x, 1
  br label %if.end

if.end:                                           ; preds = %entry, %if.then
  %res = phi i32 [ %add, %if.then ], [ %x, %entry ]
  ret i32 %res

0:                                                ; preds = %entry
  ret i32 0
}

로직에 관한 문제는 비록 존재하지만, 여전히 학습할만한 충분한 뼈대는 제공해줍니다.

동시에 다음과 같은 부가적인 정보도 얻을 수 있습니다.

NewTerm->setDebugLoc(OrigTerm->getDebugLoc());

디버깅 정보를 덧붙이는 것은 고려해 본 적이 없었는데 참고할만했습니다.