Added String and Symbol Obfuscation

Author: Zrahay

mlir/mlir-obfuscator/CMakeLists.txt

@@ -0,0 +1,21 @@
+cmake_minimum_required(VERSION 3.13)
+project(MLIROBF LANGUAGES CXX C)
+
+# Use C++17
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+
+
+find_package(MLIR REQUIRED CONFIG) # What this does is that it tells the code that I want to use MLIR. Go find it and link to it.
+
+
+message(STATUS "Found MLIR: ${MLIR_PACKAGE_PREFIX}")
+
+
+include_directories(${PROJECT_SOURCE_DIR}/include) # THis helps it find include/Obfuscator/Passes.h
+
+
+add_subdirectory(lib) # This tells the CMake to build the code in  mlir-obfuscator/lib/
+
+
+mlir_print_global_pass_registry() # Optional: Thjis is just to print MLIR passes for debugging

mlir/mlir-obfuscator/include/Obfuscator/Passes.h

@@ -0,0 +1,13 @@
+#pragma once
+
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+
+// Creates the String Encryption Pass
+std::unique_ptr<Pass> createStringEncryptPass(llvm::StringRef key = "");
+
+// Creates the Symbol Obfuscation / Renaming Pass
+std::unique_ptr<Pass> createSymbolObfuscatePass(llvm::StringRef key = "");
+
+} // namespace mlir

mlir/mlir-obfuscator/lib/CMakeLists.txt

@@ -0,0 +1,68 @@
+#include "Obfuscator/Passes.h"
+
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/IR/MLIRContext.h"
+
+#include <string>
+
+using namespace mlir;
+
+namespace {
+
+/// Simple XOR encryption for demonstration
+static std::string xorEncrypt(const std::string &input, const std::string &key) {
+  std::string out = input;
+  for (size_t i = 0; i < input.size(); i++) {
+    out[i] = input[i] ^ key[i % key.size()];
+  }
+  return out;
+}
+
+/// String Encryption Pass
+struct StringEncryptPass 
+    : public PassWrapper<StringEncryptPass, OperationPass<ModuleOp>> { // This lines makes a skeleton for the pass
+
+  StringEncryptPass() = default;
+  StringEncryptPass(const std::string &k) : key(k) {}
+
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+    MLIRContext *ctx = module.getContext();
+
+    module.walk([&](Operation *op) { // Walk visits every operation in the IR, including the nested ones
+      bool changed = false;
+      SmallVector<NamedAttribute> newAttrs;
+
+      for (auto &attr : op->getAttrs()) {
+        // Only encrypt string attributes
+        if (auto strAttr = attr.getValue().dyn_cast<StringAttr>()) {
+          std::string original = strAttr.getValue().str();
+          std::string encrypted = xorEncrypt(original, key);
+
+          auto newValue = StringAttr::get(ctx, encrypted);
+          newAttrs.emplace_back(attr.getName(), newValue);
+          changed = true;
+        } else {
+          newAttrs.push_back(attr);
+        }
+      }
+
+      // Replace attribute dictionary if something changed
+      if (changed) {
+        op->setAttrDictionary(DictionaryAttr::get(ctx, newAttrs));
+      }
+    });
+  }
+
+  std::string key = "default_key";
+};
+
+} // namespace
+
+/// Factory function exposed to the outside world
+std::unique_ptr<Pass> mlir::createStringEncryptPass(llvm::StringRef key) {
+  return std::make_unique<StringEncryptPass>(key.str());
+}

mlir/mlir-obfuscator/lib/Passes.cpp

@@ -0,0 +1,87 @@
+#include "Obfuscator/Passes.h"
+
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/SymbolTable.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+
+#include <random>
+
+using namespace mlir;
+
+namespace {
+
+/// Utility: generate random obfuscated names (hex-based)
+static std::string generateObfuscatedName(std::mt19937 &rng) {
+  std::uniform_int_distribution<uint32_t> dist(0, 0xFFFFFFFF);
+  uint32_t num = dist(rng);
+
+  // format as hex string: f_a1b2c3d4
+  char buffer[16];
+  snprintf(buffer, sizeof(buffer), "f_%08x", num);
+  return std::string(buffer);
+}
+
+/// Symbol Obfuscation Pass
+struct SymbolObfuscatePass
+    : public PassWrapper<SymbolObfuscatePass, OperationPass<ModuleOp>> {
+
+  SymbolObfuscatePass() = default;
+  SymbolObfuscatePass(const std::string &k) : key(k) {}
+
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+    MLIRContext *ctx = module.getContext();
+    SymbolTable symbolTable(module);
+
+    // initialize RNG with deterministic seed (from key)
+    std::seed_seq seq(key.begin(), key.end());
+    std::mt19937 rng(seq);
+
+    // Mapping: oldName -> newName
+    llvm::StringMap<std::string> renameMap;
+
+    // Step 1: Rename symbol definitions (functions)
+    module.walk([&](func::FuncOp func) {
+      StringRef oldName = func.getName();
+      std::string newName = generateObfuscatedName(rng);
+
+      renameMap[oldName] = newName;
+      symbolTable.setSymbolName(func, newName);
+    });
+
+    // Step 2: Update symbol references everywhere
+    module.walk([&](Operation *op) {
+      SmallVector<NamedAttribute> updatedAttrs;
+      bool changed = false;
+
+      for (auto &attr : op->getAttrs()) {
+        if (auto symAttr = attr.getValue().dyn_cast<SymbolRefAttr>()) {
+          StringRef old = symAttr.getRootReference();
+          if (renameMap.count(old)) {
+            auto newRef = SymbolRefAttr::get(ctx, renameMap[old]);
+            updatedAttrs.emplace_back(attr.getName(), newRef);
+            changed = true;
+            continue;
+          }
+        }
+        // no change -> keep original
+        updatedAttrs.push_back(attr);
+      }
+
+      if (changed) {
+        op->setAttrDictionary(DictionaryAttr::get(ctx, updatedAttrs));
+      }
+    });
+  }
+
+  std::string key = "seed";
+};
+
+} // namespace
+
+/// Public factory
+std::unique_ptr<Pass> mlir::createSymbolObfuscatePass(llvm::StringRef key) {
+  return std::make_unique<SymbolObfuscatePass>(key.str());
+}

mlir/mlir-obfuscator/lib/SymbolPass.cpp

@@ -0,0 +1,10 @@
+// String test
+func.func @hello() attributes { msg = "HELLO WORLD" } {
+  return
+}
+
+// Symbol test
+func.func @main() {
+  %0 = func.call @hello() : () -> ()
+  return
+}