global encryption

CMakeLists.txt

@@ -19,6 +19,7 @@ add_library(passes MODULE
   lib/Pluto/BogusControlFlowPass.cpp
   lib/Pluto/CryptoUtils.cpp
   lib/Pluto/Flattening.cpp
+  lib/Pluto/GlobalEncryption.cpp
   lib/PassRegistration.cpp
 )
 

include/Pluto/GlobalEncryption.h

@@ -0,0 +1,19 @@
+#pragma once
+
+#include "llvm/IR/PassManager.h"
+
+using namespace llvm;
+
+namespace Pluto {
+
+struct GlobalEncryption : PassInfoMixin<GlobalEncryption> {
+    void insertArrayDecryption(Module &M, GlobalVariable *GV, uint64_t key, uint64_t eleNum);
+
+    void insertIntDecryption(Module &M, GlobalVariable *GV, uint64_t key);
+
+    PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM);
+
+    static bool isRequired() { return true; }
+};
+
+}; // namespace Pluto

lib/PassRegistration.cpp

@@ -1,6 +1,6 @@
 #include "Pluto/BogusControlFlowPass.h"
 #include "Pluto/Flattening.h"
-
+#include "Pluto/GlobalEncryption.h"
 #include "ExamplePass.h"
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Passes/PassPlugin.h"
@@ -9,6 +9,18 @@ using namespace llvm;
 
 // Combined registration function
 static void registerPasses(PassBuilder &PB) {
+  // Register the module pass (GlobalEncryption)
+  PB.registerPipelineParsingCallback(
+    [](StringRef Name, ModulePassManager &MPM,
+       ArrayRef<PassBuilder::PipelineElement>) {
+      if (Name == "pluto-global-encryption") {
+        MPM.addPass(Pluto::GlobalEncryption());
+        return true;
+      }
+      return false;
+    });
+
+  // Register the function passes
   PB.registerPipelineParsingCallback(
     [](StringRef Name, FunctionPassManager &FPM,
        ArrayRef<PassBuilder::PipelineElement>) {

lib/Pluto/CryptoUtils.cpp

@@ -4,7 +4,7 @@
 
 #include "llvm/ADT/Twine.h"
 #include "llvm/IR/LLVMContext.h"
-#include "llvm/Support/Debug.h"
+//#include "llvm/Support/Debug.h" // fixme: needed for DEBUG_WITH_TYPE
 
 #include <cassert>
 #include <cstdio>

lib/Pluto/GlobalEncryption.cpp

@@ -0,0 +1,141 @@
+#include "Pluto/GlobalEncryption.h"
+
+#include "llvm/Analysis/AliasAnalysis.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "Pluto/CryptoUtils.h"
+#include "llvm/Transforms/Utils/ModuleUtils.h"
+#include <vector>
+
+namespace Pluto {
+
+bool shouldSkip(GlobalVariable &GV) {
+    // Do not encrypt LLVM-generated GV like llvm.global_ctors
+    if (GV.getName().startswith("llvm.")) {
+        return true;
+    }
+    // Only encrypt GV with internal or private linkage
+    // Other linkages may cause problem. For example, if a GV has LinkOnce linkage, two global variables in two modules
+    // with the same name will be merged into one GV at link-time and the merged GV will be decrypted twice (the two
+    // decrypt functions are not merged).
+    // Reference: https://llvm.org/docs/LangRef.html#linkage-types
+    if (!GV.hasInternalLinkage() && !GV.hasPrivateLinkage()) {
+        return true;
+    }
+    // Encrypt the GV only if it's an integer or integer array
+    if (!GV.getValueType()->isIntegerTy() &&
+        (!GV.getValueType()->isArrayTy() || !cast<ArrayType>(GV.getValueType())->getElementType()->isIntegerTy())) {
+        return true;
+    }
+    // Make sure the GV has an initializer
+    if (!GV.hasInitializer() || !GV.getInitializer()) {
+        return true;
+    }
+    // Make sure the GV doesn't belong to any custom section (which means it belongs .data section by default)
+    // We conservatively skip data in custom section to avoid unexpected behaviors after obfuscation
+    if (GV.hasSection()) {
+        return true;
+    }
+    return false;
+}
+
+PreservedAnalyses GlobalEncryption::run(Module &M, ModuleAnalysisManager &AM) {
+    std::vector<GlobalVariable *> GVs;
+    for (auto &GV : M.globals()) {
+        if (!shouldSkip(GV)) {
+            GVs.push_back(&GV);
+        }
+    }
+    for (auto &GV : GVs) {
+        if (ConstantDataArray *dataArray = dyn_cast<ConstantDataArray>(GV->getInitializer())) {
+            uint64_t eleByteSize = dataArray->getElementByteSize();
+            uint64_t eleNum = dataArray->getNumElements();
+            const char *data = dataArray->getRawDataValues().data();
+            uint64_t dataSize = eleByteSize * eleNum;
+            if (data && eleByteSize <= 8) {
+                char *dataCopy = new char[dataSize];
+                memcpy(dataCopy, data, dataSize);
+                uint64_t key = cryptoutils->get_uint64_t();
+                // A simple xor encryption
+                for (uint32_t i = 0; i < dataSize; i++) {
+                    dataCopy[i] ^= ((char *)&key)[i % eleByteSize];
+                }
+                GV->setInitializer(
+                    ConstantDataArray::getRaw(StringRef(dataCopy, dataSize), eleNum, dataArray->getElementType()));
+                GV->setConstant(false);
+                insertArrayDecryption(M, GV, key, eleNum);
+            }
+        } else if (ConstantInt *dataInt = dyn_cast<ConstantInt>(GV->getInitializer())) {
+            uint64_t key = cryptoutils->get_uint64_t();
+            ConstantInt *enc = ConstantInt::get(dataInt->getType(), key ^ dataInt->getZExtValue());
+            GV->setInitializer(enc);
+            GV->setConstant(false);
+            insertIntDecryption(M, GV, key);
+        }
+    }
+    return PreservedAnalyses::all();
+}
+
+void GlobalEncryption::insertArrayDecryption(Module &M, GlobalVariable *GV, uint64_t key, uint64_t eleNum) {
+    static uint64_t cnt = 0;
+    LLVMContext &context = M.getContext();
+    FunctionType *funcType = FunctionType::get(Type::getVoidTy(context), false);
+    std::string funcName = formatv("decrypt.arr.{0:d}", cnt++);
+    FunctionCallee callee = M.getOrInsertFunction(funcName, funcType);
+    Function *func = cast<Function>(callee.getCallee());
+    func->setLinkage(GlobalValue::LinkageTypes::PrivateLinkage);
+    BasicBlock *head = BasicBlock::Create(context, "head", func);
+    BasicBlock *forCond = BasicBlock::Create(context, "for.cond", func);
+    BasicBlock *forBody = BasicBlock::Create(context, "for.body", func);
+    BasicBlock *forInc = BasicBlock::Create(context, "for.inc", func);
+    BasicBlock *forEnd = BasicBlock::Create(context, "for.inc", func);
+
+    IRBuilder<> builder(context);
+
+    builder.SetInsertPoint(head);
+    AllocaInst *indexPtr = builder.CreateAlloca(Type::getInt32Ty(context));
+    builder.CreateStore(ConstantInt::get(Type::getInt32Ty(context), 0), indexPtr);
+    builder.CreateBr(forCond);
+
+    builder.SetInsertPoint(forCond);
+    LoadInst *index = builder.CreateLoad(Type::getInt32Ty(context), indexPtr);
+    Value *cond = builder.CreateICmpSLT(index, ConstantInt::get(Type::getInt32Ty(context), eleNum));
+    builder.CreateCondBr(cond, forBody, forEnd);
+
+    builder.SetInsertPoint(forBody);
+
+    Value *elePtr = builder.CreateGEP(GV->getValueType(), GV, {ConstantInt::get(Type::getInt32Ty(context), 0), index});
+    Type *eleType = cast<ArrayType>(GV->getValueType())->getElementType();
+    builder.CreateStore(builder.CreateXor(builder.CreateLoad(eleType, elePtr), ConstantInt::get(eleType, key)), elePtr);
+    builder.CreateBr(forInc);
+
+    builder.SetInsertPoint(forInc);
+    builder.CreateStore(builder.CreateAdd(index, ConstantInt::get(Type::getInt32Ty(context), 1)), indexPtr);
+    builder.CreateBr(forCond);
+
+    builder.SetInsertPoint(forEnd);
+    builder.CreateRetVoid();
+
+    appendToGlobalCtors(M, func, 0);
+}
+
+void GlobalEncryption::insertIntDecryption(Module &M, GlobalVariable *GV, uint64_t key) {
+    static uint64_t cnt = 0;
+    LLVMContext &context = M.getContext();
+    FunctionType *funcType = FunctionType::get(Type::getVoidTy(context), false);
+    std::string funcName = formatv("decrypt.int.{0:d}", cnt++);
+    FunctionCallee callee = M.getOrInsertFunction(funcName, funcType);
+    Function *func = cast<Function>(callee.getCallee());
+    func->setLinkage(GlobalValue::LinkageTypes::PrivateLinkage);
+
+    BasicBlock *BB = BasicBlock::Create(context, "BB", func);
+
+    IRBuilder<> builder(context);
+    builder.SetInsertPoint(BB);
+    LoadInst *val = builder.CreateLoad(GV->getValueType(), GV);
+    builder.CreateStore(builder.CreateXor(val, ConstantInt::get(GV->getValueType(), key)), GV);
+    builder.CreateRetVoid();
+
+    appendToGlobalCtors(M, func, 0);
+}
+}; // namespace Pluto

tests/tests.cpp

@@ -6,13 +6,16 @@
 #include <iostream>
 #include <memory>
 #include <array>
+#include <vector>
+
+std::string exec(std::string cmd) {
+    std::cout << "Executing command: " << cmd << std::endl;
 
-std::string exec(const char* cmd) {
     std::array<char, 128> buffer;
     std::string result;
     int exit_status;
 
-    std::unique_ptr<FILE, decltype(&pclose)> pipe(popen(cmd, "r"), pclose);
+    std::unique_ptr<FILE, decltype(&pclose)> pipe(popen(cmd.c_str(), "r"), pclose);
     if (!pipe) {
         throw std::runtime_error("popen() failed!");
     }
@@ -32,56 +35,76 @@ std::string exec(const char* cmd) {
     return result;
 }
 
-std::string run_test(const char* pass_name) {
-    // get the input .ll first: clang -c -emit-llvm test.c -o test.ll
+std::string run_passes(const std::vector<std::string>& pass_names) {
+    std::string input_ll = "test.ll";
     std::string output;
-    output = exec("clang -c -S -emit-llvm -O1 ../tests/test.c -o test.ll");
-    //std::cout << "Output: " << output << std::endl;
-
-    std::string command = "opt -load-pass-plugin=./libpasses.so -passes \"";
-    command += pass_name;
-    command += "\" test.ll -S -o test_obfuscated.ll -debug-pass-manager";
-    std::cout << "Executing command: " << command << std::endl;
-    output = exec(command.c_str());
-    //std::cout << "Command output: " << output << std::endl;
-
-    // compile the obfuscated .ll to the binary
-    output = exec("clang test_obfuscated.ll -o test_obfuscated");
-    //std::cout << "Output: " << output << std::endl;
-
-    // run the obfuscated binary
-    output = exec("./test_obfuscated");
-    //std::cout << "Output: " << output << std::endl;
-    return output;
+
+    // Generate initial .ll file
+    output = exec("clang -c -S -emit-llvm -O1 ../tests/test.c -o " + input_ll);
+
+    std::string passes_names = "";
+    for (size_t i = 0; i < pass_names.size(); ++i) {
+        const std::string& pass_name = pass_names[i];
+        if (passes_names != "") {
+            passes_names += ",";
+        }
+        passes_names += pass_name;
+        std::string output_ll = passes_names + ".ll";
+
+        std::string command = "opt -load-pass-plugin=./libpasses.so -passes \"";
+        command += pass_name;
+        command += "\" " + input_ll + " -S -o " + output_ll + " -debug-pass-manager";
+        output = exec(command.c_str());
+
+        // Set the output as input for the next pass
+        input_ll = output_ll;
+    }
+
+    // Compile the final obfuscated .ll to binary
+    std::string binary_name = passes_names.size() ? passes_names : "test_obfuscated";
+    output = exec("clang " + input_ll + " -o " + binary_name);
+
+    // Run the obfuscated binary
+    return exec("./" + binary_name);
 }
 
 std::string original_output;
 
 TEST(ExamplePassTest, OutputMatches) {
-    std::string obfuscated_output = run_test("example-pass");
+    std::string obfuscated_output = run_passes({"example-pass"});
     EXPECT_EQ(original_output, obfuscated_output);
 }
 
 TEST(BogusControlFlowTest, OutputMatches) {
-
-    std::string obfuscated_output = run_test("pluto-bogus-control-flow");
+    std::string obfuscated_output = run_passes({"pluto-bogus-control-flow"});
     EXPECT_EQ(original_output, obfuscated_output);
-
 }
 
 TEST(FlatteningTest, OutputMatches) {
-    std::string obfuscated_output = run_test("pluto-flattening");
+    std::string obfuscated_output = run_passes({"pluto-flattening"});
+    EXPECT_EQ(original_output, obfuscated_output);
+}
+
+TEST(GlobalEncryptionTest, OutputMatches) {
+    std::string obfuscated_output = run_passes({"pluto-global-encryption"});
     EXPECT_EQ(original_output, obfuscated_output);
 }
 
+// Add a new test case for combined passes
 TEST(CombinedTest, OutputMatches) {
-    std::string obfuscated_output = run_test("example-pass,pluto-bogus-control-flow,pluto-flattening");
+    std::string obfuscated_output = run_passes({
+        "example-pass",
+        "pluto-bogus-control-flow",
+        "pluto-flattening",
+        "pluto-global-encryption"
+    });
     EXPECT_EQ(original_output, obfuscated_output);
 }
 
 int main(int argc, char **argv) {
     testing::InitGoogleTest(&argc, argv);
-    original_output = run_test("");
+    original_output = run_passes({});
+
     // print current working directory
     {
         std::cout << "Current working directory: " << getcwd(nullptr, 0) << std::endl;