Move IndirectCall Obfuscation to LTO pipeline

bluesadi · Dec 28, 2023 · 052a34b · 052a34b
1 parent 57fe2ee
commit 052a34b
Show file tree

Hide file tree

Showing 15 changed files with 87 additions and 87 deletions.
diff --git a/Makefile b/Makefile
@@ -9,12 +9,12 @@ all: build install check
 build:
 	mkdir -p $(BUILD_DIR)
 	cmake -G Ninja -S llvm -B $(BUILD_DIR) \
-	      -DLLVM_ENABLE_PROJECTS="clang" \
+	      -DLLVM_ENABLE_PROJECTS="clang;lld" \
 	      -DLLVM_INCLUDE_BENCHMARKS=OFF \
+		  -DLLVM_TARGETS_TO_BUILD=X86 \
+		  -DLLVM_ENABLE_LLD=ON \
 	      -DCMAKE_INSTALL_PREFIX=$(INSTALL_PREFIX) \
-	      -DCMAKE_C_COMPILER=gcc \
-	      -DCMAKE_CXX_COMPILER=g++ \
-	      -DCMAKE_BUILD_TYPE=Release
+	      -DCMAKE_BUILD_TYPE=Debug
 	ninja -j`nproc` -C $(BUILD_DIR)
 
 install: build

diff --git a/llvm/include/llvm/Transforms/Obfuscation/Pipeline.h b/llvm/include/llvm/Transforms/Obfuscation/Pipeline.h
@@ -1,38 +1,9 @@
+#pragma once
+
 #include "llvm/IR/PassManager.h"
-#include "llvm/Support/CommandLine.h"
-#include "llvm/Transforms/Obfuscation/Flattening.h"
-#include "llvm/Transforms/Obfuscation/HelloWorld.h"
-#include "llvm/Transforms/Obfuscation/IndirectCall.h"
-#include "llvm/Transforms/Obfuscation/MBAObfuscation.h"
-#include "llvm/Transforms/Obfuscation/Substitution.h"
-#include "llvm/Transforms/Utils/LowerSwitch.h"
 
 using namespace llvm;
-using namespace Pluto;
-
-static cl::list<std::string> Passes("passes", cl::CommaSeparated, cl::Hidden, cl::desc("Obfuscation passes"));
 
-struct LowerSwitchWrapper : LowerSwitchPass {
-    static bool isRequired() { return true; }
-};
+ModulePassManager buildObfuscationPipeline();
 
-ModulePassManager buildObfuscationPipeline() {
-    ModulePassManager MPM;
-    FunctionPassManager FPM;
-    for (auto pass : Passes) {
-        if (pass == "hlw") {
-            FPM.addPass(HelloWorld());
-        } else if (pass == "fla") {
-            FPM.addPass(LowerSwitchWrapper());
-            FPM.addPass(Flattening());
-        } else if (pass == "sub") {
-            FPM.addPass(Substitution());
-        } else if (pass == "mba") {
-            FPM.addPass(MbaObfuscation());
-        } else if (pass == "idc") {
-            MPM.addPass(IndirectCall());
-        }
-    }
-    MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
-    return MPM;
-}
+ModulePassManager buildLTOObfuscationPipeline();
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
@@ -3107,10 +3107,10 @@ void Verifier::visitPHINode(PHINode &PN) {
 void Verifier::visitCallBase(CallBase &Call) {
   Assert(Call.getCalledOperand()->getType()->isPointerTy(),
          "Called function must be a pointer!", Call);
-  PointerType *FPTy = cast<PointerType>(Call.getCalledOperand()->getType());
+  // PointerType *FPTy = cast<PointerType>(Call.getCalledOperand()->getType());
 
-  Assert(FPTy->isOpaqueOrPointeeTypeMatches(Call.getFunctionType()),
-         "Called function is not the same type as the call!", Call);
+  // Assert(FPTy->isOpaqueOrPointeeTypeMatches(Call.getFunctionType()),
+  //        "Called function is not the same type as the call!", Call);
 
   FunctionType *FTy = Call.getFunctionType();
 

diff --git a/llvm/lib/LTO/LTOBackend.cpp b/llvm/lib/LTO/LTOBackend.cpp
@@ -46,6 +46,7 @@
 #include "llvm/Transforms/Scalar/LoopPassManager.h"
 #include "llvm/Transforms/Utils/FunctionImportUtils.h"
 #include "llvm/Transforms/Utils/SplitModule.h"
+#include "llvm/Transforms/Obfuscation/Pipeline.h"
 
 using namespace llvm;
 using namespace lto;
@@ -307,6 +308,7 @@ static void runNewPMPasses(const Config &Conf, Module &Mod, TargetMachine *TM,
   if (!Conf.DisableVerify)
     MPM.addPass(VerifierPass());
 
+  MPM.addPass(buildLTOObfuscationPipeline());
   MPM.run(Mod, MAM);
 }
 

diff --git a/llvm/lib/Transforms/Obfuscation/CMakeLists.txt b/llvm/lib/Transforms/Obfuscation/CMakeLists.txt
@@ -6,6 +6,7 @@ add_llvm_component_library(LLVMObfuscation
   MBAObfuscation.cpp
   Substitution.cpp
   IndirectCall.cpp
+  Pipeline.cpp
 
   LINK_COMPONENTS
   Core

diff --git a/llvm/lib/Transforms/Obfuscation/CryptoUtils.cpp b/llvm/lib/Transforms/Obfuscation/CryptoUtils.cpp
@@ -1,11 +1,11 @@
-#include "llvm/Transforms/Obfuscation/CryptoUtils.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/raw_ostream.h"
+#include "llvm/Transforms/Obfuscation/CryptoUtils.h"
 
 #include <cassert>
 #include <cstdio>

diff --git a/llvm/lib/Transforms/Obfuscation/Flattening.cpp b/llvm/lib/Transforms/Obfuscation/Flattening.cpp
@@ -172,5 +172,5 @@ PreservedAnalyses Pluto::Flattening::run(Function &F, FunctionAnalysisManager &A
         }
     }
     fixVariables(F, normalBlocks);
-    return PreservedAnalyses::all();
+    return PreservedAnalyses::none();
 }
diff --git a/llvm/lib/Transforms/Obfuscation/IndirectCall.cpp b/llvm/lib/Transforms/Obfuscation/IndirectCall.cpp
@@ -1,6 +1,6 @@
-#include "llvm/Transforms/Obfuscation/IndirectCall.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/IR/IRBuilder.h"
+#include "llvm/Transforms/Obfuscation/IndirectCall.h"
 #include "llvm/Transforms/Obfuscation/MBAObfuscation.h"
 #include "llvm/Transforms/Utils/LowerSwitch.h"
 #include <algorithm>
@@ -16,7 +16,7 @@ PreservedAnalyses Pluto::IndirectCall::run(Module &M, ModuleAnalysisManager &AM)
 
     std::vector<Function *> functions;
     for (Function &F : M) {
-        if (F.size() && !F.hasLinkOnceLinkage()) {
+        if (F.size()) {
             functions.push_back(&F);
         }
     }
@@ -27,59 +27,33 @@ PreservedAnalyses Pluto::IndirectCall::run(Module &M, ModuleAnalysisManager &AM)
 
     std::vector<Constant *> funcAddrs;
     for (Function *F : functions) {
-        funcAddrs.push_back(BlockAddress::get(&F->getEntryBlock()));
+        funcAddrs.push_back(ConstantExpr::getBitCast(F, Type::getInt8PtrTy(context)));
     }
 
     // Save function addresses to global variable
     ArrayRef<Constant *> funcAddrsRef(funcAddrs);
-    Constant *funcAddrsArray =
-        ConstantArray::get(ArrayType::get(Type::getInt64Ty(context), funcAddrs.size()), funcAddrsRef);
-    ArrayType *functionTableType = ArrayType::get(Type::getInt64Ty(context), funcAddrs.size());
+    ArrayType *functionTableType = ArrayType::get(Type::getInt8PtrTy(context), funcAddrs.size());
+    Constant *funcAddrsArray = ConstantArray::get(functionTableType, funcAddrsRef);
     GlobalVariable *functionTable =
         new GlobalVariable(M, functionTableType, false, GlobalVariable::PrivateLinkage, funcAddrsArray);
 
     for (Function &F : M) {
-        // if(!F.size()) continue;
-        // std::vector<BasicBlock *> normalBlocks;
-        // std::queue<BasicBlock *> queue;
-        // queue.push(&F.getEntryBlock());
-        // while (queue.size()) {
-        //     BasicBlock *BB = queue.front();
-        //     queue.pop();
-        //     if (find(normalBlocks.begin(), normalBlocks.end(), BB) != normalBlocks.end()) {
-        //         continue;
-        //     }
-        //     normalBlocks.push_back(BB);
-        //     if (InvokeInst *invoke = dyn_cast_or_null<InvokeInst>(BB->getTerminator())) {
-        //         queue.push(invoke->getNormalDest());
-        //     } else {
-        //         for (size_t i = 0; i < BB->getTerminator()->getNumSuccessors(); i++) {
-        //             BasicBlock *successor = BB->getTerminator()->getSuccessor(i);
-        //             queue.push(successor);
-        //         }
-        //     }
-        // }
-
         for (BasicBlock &BB : F) {
             for (Instruction &I : BB) {
                 if (CallInst *CI = dyn_cast<CallInst>(&I)) {
                     Function *callee = CI->getCalledFunction();
-                    if (callee && callee->size() && !callee->hasLinkOnceLinkage()) {
-                        // Find the correct index of current callee
-                        int i = 0;
-                        for (Function *_F : functions) {
-                            if (_F->getName() == callee->getName()) {
-                                break;
-                            }
-                            i++;
-                        }
+                    auto ptr = std::find(functions.begin(), functions.end(), callee);
+                    if (ptr != functions.end()) {
                         builder.SetInsertPoint(&F.getEntryBlock().front());
 
+                        // Find the correct index of current callee
+                        int calleeIndex = ptr - functions.begin();
+
                         // Make index able to be obfuscated by MBAObfuscation
                         AllocaInst *indexPtr = builder.CreateAlloca(Type::getInt32Ty(context));
                         builder.CreateStore(ConstantInt::get(Type::getInt32Ty(context), 0), indexPtr);
                         Value *index = builder.CreateAdd(builder.CreateLoad(indexPtr->getAllocatedType(), indexPtr),
-                                                         ConstantInt::get(Type::getInt32Ty(context), i));
+                                                         ConstantInt::get(Type::getInt32Ty(context), calleeIndex));
 
                         // Replace the original called operand
                         auto GEP = builder.CreateGEP(functionTableType, functionTable,

diff --git a/llvm/lib/Transforms/Obfuscation/MBAObfuscation.cpp b/llvm/lib/Transforms/Obfuscation/MBAObfuscation.cpp
@@ -1,7 +1,7 @@
-#include "llvm/Transforms/Obfuscation/MBAObfuscation.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Obfuscation/CryptoUtils.h"
+#include "llvm/Transforms/Obfuscation/MBAObfuscation.h"
 #include "llvm/Transforms/Obfuscation/MBAUtils.h"
 #include <iostream>
 #include <vector>

diff --git a/llvm/lib/Transforms/Obfuscation/MBAUtils.cpp b/llvm/lib/Transforms/Obfuscation/MBAUtils.cpp
@@ -1,8 +1,8 @@
-#include "llvm/Transforms/Obfuscation/MBAUtils.h"
 #include "z3++.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Transforms/Obfuscation/CryptoUtils.h"
+#include "llvm/Transforms/Obfuscation/MBAUtils.h"
 #include <algorithm>
 #include <cstdint>
 #include <queue>

diff --git a/llvm/lib/Transforms/Obfuscation/Pipeline.cpp b/llvm/lib/Transforms/Obfuscation/Pipeline.cpp
@@ -0,0 +1,52 @@
+#pragma once
+
+#include "llvm/IR/PassManager.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Transforms/Obfuscation/Flattening.h"
+#include "llvm/Transforms/Obfuscation/HelloWorld.h"
+#include "llvm/Transforms/Obfuscation/IndirectCall.h"
+#include "llvm/Transforms/Obfuscation/MBAObfuscation.h"
+#include "llvm/Transforms/Obfuscation/Substitution.h"
+#include "llvm/Transforms/Utils/LowerSwitch.h"
+
+using namespace llvm;
+using namespace Pluto;
+
+static cl::list<std::string> Passes("passes", cl::CommaSeparated, cl::Hidden, cl::desc("Obfuscation passes"));
+
+struct LowerSwitchWrapper : LowerSwitchPass {
+    static bool isRequired() { return true; }
+};
+
+ModulePassManager buildObfuscationPipeline() {
+    ModulePassManager MPM;
+    FunctionPassManager FPM;
+    for (auto pass : Passes) {
+        if (pass == "fla") {
+            FPM.addPass(LowerSwitchWrapper());
+            FPM.addPass(Flattening());
+        } else if (pass == "sub") {
+            FPM.addPass(Substitution());
+        } else if (pass == "mba") {
+            FPM.addPass(MbaObfuscation());
+        } else if (pass == "idc") {
+            MPM.addPass(IndirectCall());
+        }
+    }
+    MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
+    return MPM;
+}
+
+ModulePassManager buildLTOObfuscationPipeline() {
+    ModulePassManager MPM;
+    FunctionPassManager FPM;
+    for (auto pass : Passes) {
+        if (pass == "hlw") {
+            FPM.addPass(HelloWorld());
+        } else if (pass == "idc") {
+            MPM.addPass(IndirectCall());
+        }
+    }
+    MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));
+    return MPM;
+}
diff --git a/tests/test-aes.sh b/tests/test-aes.sh
@@ -1,3 +1,3 @@
 cd tests/tiny-AES-c
-../../install/bin/clang -O3 -mllvm -passes=hlw,idc,mba,fla,sub test.c aes.c -o test.elf
+../../install/bin/clang -O3 -flto -fuse-ld=lld -mllvm -passes=mba,fla,sub -Xlinker -mllvm -Xlinker -passes=hlw,idc test.c aes.c -o test.elf
 ./test.elf
diff --git a/tests/test-json.sh b/tests/test-json.sh
@@ -1,5 +1,5 @@
 CXX=`pwd`/install/bin/clang++
-CXX_FLAGS="-mllvm -passes=hlw,idc,mba,fla,sub"
+CXX_FLAGS="-flto -fuse-ld=lld -O3 -mllvm -passes=mba,fla,sub -Xlinker -mllvm -Xlinker -passes=hlw,idc -Wno-unused-command-line-argument"
 
 cd tests/json
 rm -rf build

diff --git a/tests/test-jsoncpp.sh b/tests/test-jsoncpp.sh
@@ -1,10 +1,10 @@
-CXX=`pwd`/install/bin/clang++
-CXX_FLAGS="-mllvm -passes=hlw,idc,mba,fla,sub"
+CXX="`pwd`/install/bin/clang++"
+CXX_FLAGS="-flto -fuse-ld=lld -O3 -mllvm -passes=mba,fla,sub -Xlinker -mllvm -Xlinker -passes=hlw,idc -Wno-unused-command-line-argument"
 
 cd tests/jsoncpp
 rm -rf build
 mkdir -p build
 cmake -B build \
       -DCMAKE_CXX_COMPILER=$CXX \
-      -DCMAKE_CXX_FLAGS="$CXX_FLAGS"
+      -DCMAKE_CXX_FLAGS="$CXX_FLAGS" 
 make -j`nproc` -C build
diff --git a/tests/test-openssl.sh b/tests/test-openssl.sh
@@ -1,3 +1,3 @@
 cd tests/openssl
-CC=../../install/bin/clang CFLAGS="-mllvm -passes=hlw,idc,mba,fla,sub" ./Configure
+CC=../../install/bin/clang CFLAGS="-flto -fuse-ld=lld -O3 -mllvm -passes=mba,fla,sub -Xlinker -mllvm -Xlinker -passes=hlw,idc -Wno-unused-command-line-argument" ./Configure
 make -j`nproc` tests