Add report_diff user callback in Python API (#49)

Extend the Python API in a way that allows users to monitor IR-level changes from individual passes. That might be useful for experimentation, validation, testing, etc. The included test shows how to configure it.

Co-authored-by: Antonio Frighetto <[email protected]>

src/core/python/PyConfig.cpp

@@ -298,7 +298,13 @@ void omvll_ctor(py::module_& m) {
          +--------------+-----------------------------------------+
 
          See the :omvll:`anti-hook` documentation.
-         )delim", "module"_a, "function"_a);
+         )delim", "module"_a, "function"_a)
+
+    .def("report_diff",
+         &ObfuscationConfig::report_diff,
+         R"delim(
+         User-callback to monitor IR-level changes from individual obfuscation passes.
+         )delim", "pass_name"_a, "original"_a, "obfuscated"_a);
 
 }
 

src/core/python/PyObfuscationConfig.cpp

@@ -10,6 +10,7 @@
 
 #include <llvm/IR/Function.h>
 #include <llvm/IR/Module.h>
+#include <llvm/Support/Threading.h>
 
 namespace py = pybind11;
 using namespace std::string_literals;
@@ -197,6 +198,31 @@ AntiHookOpt PyObfuscationConfig::anti_hooking(llvm::Module* mod, llvm::Function*
   return false;
 }
 
+bool PyObfuscationConfig::has_report_diff_override() {
+  std::call_once(overrides_report_diff_checked_, [this]() {
+    const auto *base = static_cast<const ObfuscationConfig *>(this);
+    overrides_report_diff_ =
+        static_cast<bool>(py::get_override(base, "report_diff"));
+  });
+  return overrides_report_diff_;
+}
+
+void PyObfuscationConfig::report_diff(const std::string &pass,
+                                      const std::string &original,
+                                      const std::string &obfuscated) {
+  if (has_report_diff_override()) {
+    py::gil_scoped_acquire gil;
+    py::function override = py::get_override(
+        static_cast<const ObfuscationConfig *>(this), "report_diff");
+    assert(override && "Checked once in ctor");
+    try {
+      override(pass, original, obfuscated);
+    } catch (const std::exception &e) {
+      fatalError("Error in 'report_diff': '"s + e.what() + "'");
+    }
+  }
+}
+
 ArithmeticOpt PyObfuscationConfig::obfuscate_arithmetic(llvm::Module* mod, llvm::Function* func)
 {
   py::gil_scoped_acquire gil;

src/core/python/PyObfuscationConfig.hpp

@@ -1,6 +1,7 @@
 #ifndef OMVLL_PY_OBFUSCATION_CONFIG_H
 #define OMVLL_PY_OBFUSCATION_CONFIG_H
 #include "omvll/ObfuscationConfig.hpp"
+#include <mutex>
 
 namespace omvll {
 
@@ -19,6 +20,14 @@ class PyObfuscationConfig : public ObfuscationConfig {
   AntiHookOpt anti_hooking(llvm::Module* mod, llvm::Function* func) override;
   ArithmeticOpt obfuscate_arithmetic(llvm::Module* mod, llvm::Function* func) override;
   OpaqueConstantsOpt obfuscate_constants(llvm::Module* mod, llvm::Function* func) override;
+
+  bool has_report_diff_override() override;
+  void report_diff(const std::string &pass, const std::string &original,
+                   const std::string &obfuscated) override;
+
+private:
+  bool overrides_report_diff_ = false;
+  std::once_flag overrides_report_diff_checked_;
 };
 }
 #endif

src/core/utils.cpp

@@ -1,4 +1,6 @@
 #include "omvll/utils.hpp"
+#include "omvll/ObfuscationConfig.hpp"
+#include "omvll/PyConfig.hpp"
 #include "omvll/log.hpp"
 
 #include <llvm/ADT/Hashing.h>
@@ -442,4 +444,23 @@ generateModule(StringRef Routine, const Triple &Triple, StringRef Extension,
   }
 }
 
+IRChangesMonitor::IRChangesMonitor(const llvm::Module &M,
+                                   llvm::StringRef PassName)
+    : Mod(M), UserConfig(PyConfig::instance().getUserConfig()),
+      PassName(PassName), ChangeReported(false) {
+  if (UserConfig->has_report_diff_override())
+    llvm::raw_string_ostream(OriginalIR) << Mod;
+}
+
+PreservedAnalyses IRChangesMonitor::report() {
+  if (ChangeReported && UserConfig->has_report_diff_override()) {
+    std::string ObfuscatedIR;
+    llvm::raw_string_ostream(ObfuscatedIR) << Mod;
+    if (OriginalIR != ObfuscatedIR)
+      UserConfig->report_diff(PassName, OriginalIR, ObfuscatedIR);
+  }
+
+  return ChangeReported ? PreservedAnalyses::none() : PreservedAnalyses::all();
+}
+
 } // namespace omvll

src/include/omvll/ObfuscationConfig.hpp

@@ -36,6 +36,10 @@ struct ObfuscationConfig {
   virtual ArithmeticOpt obfuscate_arithmetic(llvm::Module* mod, llvm::Function* func) = 0;
 
   virtual AntiHookOpt anti_hooking(llvm::Module* mod, llvm::Function* func) = 0;
+
+  virtual bool has_report_diff_override() { return false; };
+  virtual void report_diff(const std::string &pass, const std::string &original,
+                           const std::string &obfuscated) = 0;
 };
 
 }

src/include/omvll/utils.hpp

@@ -4,6 +4,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Triple.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/Support/Error.h"
 #include <string>
 
@@ -46,5 +47,33 @@ llvm::Expected<std::unique_ptr<llvm::Module>>
 generateModule(llvm::StringRef Routine, const llvm::Triple &Triple,
                llvm::StringRef Extension, llvm::LLVMContext &Ctx,
                llvm::ArrayRef<std::string> ExtraArgs);
+
+struct ObfuscationConfig;
+
+class IRChangesMonitor {
+public:
+  IRChangesMonitor(const llvm::Module &M, llvm::StringRef PassName);
+
+  // Call this function whenever a transformation in the pass reported changes
+  // explicitly. This determines the PreservedAnalyses flag returned from the
+  // report() function.
+  void notify(bool TransformationReportedChange) {
+    ChangeReported |= TransformationReportedChange;
+  }
+
+  llvm::PreservedAnalyses report();
+
+  IRChangesMonitor(IRChangesMonitor &&) = delete;
+  IRChangesMonitor(const IRChangesMonitor &) = delete;
+  IRChangesMonitor &operator=(IRChangesMonitor &&) = delete;
+  IRChangesMonitor &operator=(const IRChangesMonitor &) = delete;
+
+private:
+  const llvm::Module &Mod;
+  ObfuscationConfig *UserConfig;
+  std::string PassName;
+  std::string OriginalIR;
+  bool ChangeReported;
+};
 }
 #endif

src/passes/arithmetic/Arithmetic.cpp

@@ -8,10 +8,11 @@
 #include "llvm/Demangle/Demangle.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/IRBuilder.h"
-#include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/IR/NoFolder.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/Support/RandomNumberGenerator.h"
+#include "llvm/Transforms/Utils/BasicBlockUtils.h"
+#include <llvm/Support/raw_ostream.h>
 
 using namespace llvm;
 using namespace PatternMatch;
@@ -243,7 +244,7 @@ PreservedAnalyses Arithmetic::run(Module &M,
                                   ModuleAnalysisManager &FAM) {
   RNG_ = M.createRNG(name());
   SDEBUG("Running {} on {}", name(), M.getName().str());
-  bool Changed = false;
+  IRChangesMonitor ModuleChanges(M, name());
 
   PyConfig& config = PyConfig::instance();
 
@@ -256,7 +257,6 @@ PreservedAnalyses Arithmetic::run(Module &M,
   std::transform(Fs.begin(), Fs.end(), std::back_inserter(LFs),
                  [] (Function& F) { return &F; });
 
-
   for (Function* F : LFs) {
     ArithmeticOpt opt = config.getUserConfig()->obfuscate_arithmetic(&M, F);
     if (!opt)
@@ -265,13 +265,13 @@ PreservedAnalyses Arithmetic::run(Module &M,
     opts_.insert({F, std::move(opt)});
 
     for (BasicBlock& BB : *F) {
-      Changed |= runOnBasicBlock(BB);
+      bool Changed = runOnBasicBlock(BB);
+      ModuleChanges.notify(Changed);
     }
   }
-  SINFO("[{}] Done!", name());
-  return Changed ? PreservedAnalyses::none() :
-                   PreservedAnalyses::all();
 
+  SINFO("[{}] Done!", name());
+  return ModuleChanges.report();
 }
 }
 

src/test/passes/arithmetic/report_diff-x86_64-linux.c

@@ -0,0 +1,23 @@
+// REQUIRES: x86-registered-target
+
+// RUN: env OMVLL_CONFIG=%S/report_diff.py clang -target x86_64-pc-linux-gnu -fpass-plugin=%libOMVLL -O1 -S %s -o /dev/null -DNEEDS_OBFUSCATION | FileCheck %s
+// RUN: env OMVLL_CONFIG=%S/report_diff.py clang -target x86_64-pc-linux-gnu -fpass-plugin=%libOMVLL -O1 -S %s -o /dev/null | FileCheck --allow-empty --check-prefix=NO_REPORT %s
+
+// Make sure the report_diff function was called if the obfuscation was applied
+// CHECK: omvll::Arithmetic applied obfuscation
+// CHECK: --- original
+// CHECK: +++ obfuscated
+// CHECK: @@
+// CHECK: -
+// CHECK: +
+
+// Make sure the report_diff function was NOT called if the obfuscation was NOT applied
+// NO_REPORT-NOT: omvll::Arithmetic applied obfuscation
+
+void test(char *dst, const char *src, unsigned len) {
+#if defined(NEEDS_OBFUSCATION)
+  *dst = *src ^ 35;
+#else
+  *dst = *src;
+#endif
+}

src/test/passes/arithmetic/report_diff.py

@@ -0,0 +1,20 @@
+import omvll
+from difflib import unified_diff
+from functools import lru_cache
+
+class MyConfig(omvll.ObfuscationConfig):
+    def __init__(self):
+        super().__init__()
+    def obfuscate_arithmetic(self, mod: omvll.Module,
+                                   fun: omvll.Function) -> omvll.ArithmeticOpt:
+        return omvll.ArithmeticOpt(rounds=2)
+    def report_diff(self, pass_name: str, original: str, obfuscated: str):
+        print(pass_name, "applied obfuscation:")
+        diff = unified_diff(original.splitlines(), obfuscated.splitlines(),
+                            'original', 'obfuscated', lineterm='')
+        for line in diff:
+            print(line)
+
+@lru_cache(maxsize=1)
+def omvll_get_config() -> omvll.ObfuscationConfig:
+    return MyConfig()