Fix bugs in GlobalEncryption: Skip global variables without private o…

…r internal linkage
bluesadi · Jan 4, 2024 · 3bda2c3 · 3bda2c3
1 parent 7f65b41
commit 3bda2c3
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 35 deletions.
diff --git a/llvm/lib/Transforms/Obfuscation/GlobalEncryption.cpp b/llvm/lib/Transforms/Obfuscation/GlobalEncryption.cpp
@@ -10,56 +10,78 @@
 
 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.getGlobalList()) {
-        if (!GV.getName().startswith("llvm.")) {
+        if (!shouldSkip(GV)) {
             GVs.push_back(&GV);
         }
     }
     for (auto &GV : GVs) {
-        if (!GV->getValueType()->isIntegerTy() &&
-            (!GV->getValueType()->isArrayTy() ||
-             !cast<ArrayType>(GV->getValueType())->getElementType()->isIntegerTy())) {
-            continue;
-        }
-        if (GV->hasInitializer() && GV->getInitializer() && !GV->hasSection()) {
-            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())) {
+        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();
-                ConstantInt *enc = ConstantInt::get(dataInt->getType(), key ^ dataInt->getZExtValue());
-                GV->setInitializer(enc);
+                // 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);
-                insertIntDecryption(M, GV, key);
+                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();
 }
 
-static uint64_t cnt = 0;
-
 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.{0:d}", cnt++);
+    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);
@@ -99,9 +121,10 @@ void GlobalEncryption::insertArrayDecryption(Module &M, GlobalVariable *GV, uint
 }
 
 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.{0:d}", cnt++);
+    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);
@@ -116,5 +139,4 @@ void GlobalEncryption::insertIntDecryption(Module &M, GlobalVariable *GV, uint64
 
     appendToGlobalCtors(M, func, 0);
 }
-
 }; // namespace Pluto
diff --git a/llvm/lib/Transforms/Obfuscation/IndirectCall.cpp b/llvm/lib/Transforms/Obfuscation/IndirectCall.cpp
@@ -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() && (F.hasInternalLinkage() || F.hasPrivateLinkage())) {
             functions.push_back(&F);
         }
     }

diff --git a/llvm/lib/Transforms/Obfuscation/Pipeline.cpp b/llvm/lib/Transforms/Obfuscation/Pipeline.cpp
@@ -53,6 +53,8 @@ ModulePassManager buildLTOObfuscationPipeline() {
             FPM.addPass(HelloWorld());
         } else if (pass == "idc") {
             MPM.addPass(IndirectCall());
+        } else if (pass == "gle") {
+            MPM.addPass(GlobalEncryption());
         }
     }
     MPM.addPass(createModuleToFunctionPassAdaptor(std::move(FPM)));