Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions mlir/mlir-obfuscator/CMakeLists.txt
Original file line number Diff line number Diff line change
@@ -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
13 changes: 13 additions & 0 deletions mlir/mlir-obfuscator/include/Obfuscator/Passes.h
Original file line number Diff line number Diff line change
@@ -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
Empty file.
68 changes: 68 additions & 0 deletions mlir/mlir-obfuscator/lib/Passes.cpp
Original file line number Diff line number Diff line change
@@ -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());
}
87 changes: 87 additions & 0 deletions mlir/mlir-obfuscator/lib/SymbolPass.cpp
Original file line number Diff line number Diff line change
@@ -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());
}
10 changes: 10 additions & 0 deletions mlir/mlir-obfuscator/test/input.mlir
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
// String test
func.func @hello() attributes { msg = "HELLO WORLD" } {
return
}

// Symbol test
func.func @main() {
%0 = func.call @hello() : () -> ()
return
}