From 24fe099872e9e1f93cb1d783fef29c0fabb4ce6e Mon Sep 17 00:00:00 2001
From: Shane Slattery <Slattz@users.noreply.github.com>
Date: Tue, 11 Apr 2023 17:06:33 +0100
Subject: [PATCH] Fix CFF, add Test and improve TestAll

---
 CMakeLists.txt         |  2 +-
 source/CFF/CffPass.cpp |  9 ++++-----
 tests/TestAll.cpp      | 15 +++++++++++++++
 tests/cffTest.cpp      | 30 ++++++++++++++++++++++++++++++
 4 files changed, 50 insertions(+), 6 deletions(-)
 create mode 100644 tests/cffTest.cpp

diff --git a/CMakeLists.txt b/CMakeLists.txt
index c8b0c25..883499d 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -28,7 +28,7 @@ execute_process(
   OUTPUT_STRIP_TRAILING_WHITESPACE
 )
   
-add_definitions(-DOBFUSC_VERSION_MAJOR=0 -DOBFUSC_VERSION_MINOR=1 -DOBFUSC_VERSION_MICRO=0 -DOBFUSC_GIT_REV=${GIT_COMMIT_HASH})
+add_definitions(-DOBFUSC_VERSION_MAJOR=0 -DOBFUSC_VERSION_MINOR=2 -DOBFUSC_VERSION_MICRO=0 -DOBFUSC_GIT_REV=${GIT_COMMIT_HASH})
   
 
   # Add Target and source dir 
diff --git a/source/CFF/CffPass.cpp b/source/CFF/CffPass.cpp
index 20a6310..412085b 100644
--- a/source/CFF/CffPass.cpp
+++ b/source/CFF/CffPass.cpp
@@ -10,9 +10,7 @@ namespace obfusc {
     CffPass::~CffPass() {}
 
     bool CffPass::obfuscate(llvm::Module& mod, llvm::Function& func) {
-        // Lower switch
-        llvm::createLowerSwitchPass()->runOnFunction(func);
-
+        // Copy original blocks
         std::vector<llvm::BasicBlock *> origBB;
         for (auto& block : func) {
             if (isa<llvm::InvokeInst>(block.getTerminator())) {
@@ -56,7 +54,7 @@ namespace obfusc {
 
         // Create switch variable and set as it
         llvm::AllocaInst* switchVar = new llvm::AllocaInst(llvm::Type::getInt32Ty(func.getContext()), 0, "switchVar", insert);
-        new llvm::StoreInst(llvm::ConstantInt::get(llvm::Type::getInt32Ty(func.getContext()), m_randGen64()), switchVar, insert);
+        new llvm::StoreInst(llvm::ConstantInt::get(llvm::Type::getInt32Ty(func.getContext()), 0), switchVar, insert);
 
         // Create main loop
         llvm::BasicBlock* loopEntry = llvm::BasicBlock::Create(func.getContext(), "loopEntry", &func, insert);
@@ -147,6 +145,7 @@ namespace obfusc {
         }
 
         fixStack(func);
+        //func.viewCFG();
 
         return true;
     }
@@ -195,7 +194,7 @@ namespace obfusc {
                 llvm::DemotePHIToStack(phi, func.begin()->getTerminator());
             }
 
-            if (tmpReg.size() != 0 || tmpPhi.size() != 0) {
+            if (tmpReg.size() == 0 || tmpPhi.size() == 0) {
                 break;
             }
         }
diff --git a/tests/TestAll.cpp b/tests/TestAll.cpp
index 045a0f9..721407d 100644
--- a/tests/TestAll.cpp
+++ b/tests/TestAll.cpp
@@ -5,10 +5,25 @@ namespace MbaTest {
     void MbaTestAll();
 }
 
+namespace CffTest {
+    void CffTestAll();
+}
+
+namespace iSubTest {
+    void iSubTestAll();
+}
+
+namespace BcfTest {
+    void BcfTestAll();
+}
+
 #ifdef OBFUSC_TEST_BUILD_ALL
 
 int main(int argc, char *argv[]) {
     MbaTest::MbaTestAll();
+    CffTest::CffTestAll();
+    iSubTest::iSubTestAll();
+    BcfTest::BcfTestAll();
     return 0;
 }
 
diff --git a/tests/cffTest.cpp b/tests/cffTest.cpp
new file mode 100644
index 0000000..528638d
--- /dev/null
+++ b/tests/cffTest.cpp
@@ -0,0 +1,30 @@
+#include <cstdio>
+#include <cstdint>
+#include <cassert>
+
+namespace CffTest {
+    
+    [[obfusc::cff]] int CffTestVal(int numLoops) {
+        int ret = 0;
+        for (int i = 0; i < numLoops; i++) {
+            ret++;
+        }
+
+        return ret;
+    }
+    
+    
+    void CffTestAll() {
+        int ret = CffTestVal(100);
+        printf("Cff Ret: %d\n", ret);
+    }
+}
+
+#ifndef OBFUSC_TEST_BUILD_ALL
+
+int main(int argc, char *argv[]) {
+    CffTest::CffTestAll();
+    return 0;
+}
+
+#endif
\ No newline at end of file