diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 8e46bad..5e978ca 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -3,6 +3,7 @@ name: CI on: push: branches: [master] + tags: ['v*'] pull_request: branches: [master] @@ -80,6 +81,7 @@ jobs: -DCMAKE_C_COMPILER="${CC:-${{ matrix.cc }}}" \ -DCMAKE_CXX_COMPILER="${CXX:-${{ matrix.cxx }}}" \ -DCMAKE_BUILD_TYPE=Release \ + -DCOBRA_ENABLE_LLVM=ON \ -DUSE_EXTERNAL_LLVM=ON \ -DCOBRA_BUILD_TESTS=ON cmake --build build-deps @@ -117,3 +119,141 @@ jobs: run: | find include lib tools \( -name '*.h' -o -name '*.cpp' \) -print0 \ | xargs -0 clang-format-22 --dry-run --Werror + + ida-plugin: + strategy: + fail-fast: false + matrix: + include: + - os: ubuntu-24.04 + cc: gcc-14 + cxx: g++-14 + name: Linux + - os: macos-15 + cc: clang + cxx: clang++ + name: macOS + - os: windows-latest + name: Windows + + name: IDA Plugin (${{ matrix.name }}) + runs-on: ${{ matrix.os }} + + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Clone IDA SDK + run: > + git clone --depth 1 --branch v9.3.0-sdk.3 + https://github.com/HexRaysSA/ida-sdk.git ida-sdk + + - name: Install Ninja (Linux) + if: runner.os == 'Linux' + run: sudo apt-get install -y -qq ninja-build + + - name: Install Ninja (macOS) + if: runner.os == 'macOS' + run: brew install ninja + + - name: Install Ninja (Windows) + if: runner.os == 'Windows' + run: choco install ninja -y + + - name: Cache dependencies + uses: actions/cache@668228422ae6a00e4ad889ee87cd7109ec5666a7 # v5.0.4 + with: + path: build-deps/install + key: deps-ida-${{ matrix.os }}-${{ hashFiles('dependencies/*.cmake') }} + + - name: Build dependencies (Unix) + if: runner.os != 'Windows' + run: | + cmake -S dependencies -B build-deps -G Ninja \ + -DCMAKE_C_COMPILER="${CC:-${{ matrix.cc }}}" \ + -DCMAKE_CXX_COMPILER="${CXX:-${{ matrix.cxx }}}" \ + -DCMAKE_BUILD_TYPE=Release + cmake --build build-deps + + - name: Build dependencies (Windows) + if: runner.os == 'Windows' + run: | + cmake -S dependencies -B build-deps -G Ninja ` + -DCMAKE_C_COMPILER=clang-cl ` + -DCMAKE_CXX_COMPILER=clang-cl ` + -DCMAKE_BUILD_TYPE=Release + cmake --build build-deps + + - name: Build IDA plugin (Unix) + if: runner.os != 'Windows' + run: | + cmake -S . -B build -G Ninja \ + -DCMAKE_C_COMPILER="${CC:-${{ matrix.cc }}}" \ + -DCMAKE_CXX_COMPILER="${CXX:-${{ matrix.cxx }}}" \ + -DCMAKE_PREFIX_PATH="$(pwd)/build-deps/install" \ + -DCMAKE_BUILD_TYPE=Release \ + -DCOBRA_BUILD_IDA_PLUGIN=ON \ + -DIDA_SDK_DIR="$(pwd)/ida-sdk/src" + cmake --build build --target ida-cobra + + - name: Build IDA plugin (Windows) + if: runner.os == 'Windows' + run: | + cmake -S . -B build -G Ninja ` + -DCMAKE_C_COMPILER=clang-cl ` + -DCMAKE_CXX_COMPILER=clang-cl ` + -DCMAKE_PREFIX_PATH="$(Resolve-Path build-deps/install)" ` + -DCMAKE_BUILD_TYPE=Release ` + -DCOBRA_BUILD_IDA_PLUGIN=ON ` + -DIDA_SDK_DIR="$(Resolve-Path ida-sdk/src)" + cmake --build build --target ida-cobra + + - name: Upload plugin artifact + if: startsWith(github.ref, 'refs/tags/v') + uses: actions/upload-artifact@ea165f8d65b6e75b540449e92b4886f43607fa02 # v4.6.2 + with: + name: ida-cobra-${{ matrix.name }} + path: | + build/lib/ida/ida-cobra.* + lib/ida/ida-cobra.cfg + + release: + if: startsWith(github.ref, 'refs/tags/v') + needs: [build, lint, ida-plugin] + runs-on: ubuntu-24.04 + permissions: + contents: write + + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Download all IDA plugin artifacts + uses: actions/download-artifact@d3f86a106a0bac45b974a628896c90dbdceab917 # v4.3.0 + with: + pattern: ida-cobra-* + path: release-artifacts/ + + - name: Rename artifacts with version and platform + run: | + version="${GITHUB_REF_NAME#v}" + mkdir -p release + cp lib/ida/ida-cobra.cfg "release/ida-cobra.cfg" + for dir in release-artifacts/ida-cobra-*; do + platform="$(echo "$(basename "$dir" | sed 's/ida-cobra-//')" | tr '[:upper:]' '[:lower:]')" + for bin in "$dir"/build/lib/ida/ida-cobra.*; do + ext="${bin##*.}" + cp "$bin" "release/ida-cobra-${platform}-${version}.${ext}" + done + done + + - name: Upload to GitHub Release + env: + GH_TOKEN: ${{ github.token }} + run: | + gh release create "${{ github.ref_name }}" \ + --title "${{ github.ref_name }}" \ + --generate-notes \ + release/* diff --git a/CMakeLists.txt b/CMakeLists.txt index 4058690..4cf042d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -59,6 +59,7 @@ option(COBRA_BUILD_LLVM_PASS "Build the LLVM pass plugin (requires LLVM 19-22)" option(COBRA_BUILD_TESTS "Build tests (requires GoogleTest in prefix)" OFF) option(COBRA_ENABLE_TRACE "Enable detailed pipeline tracing to stderr (debug builds)" OFF) option(COBRA_ENABLE_TRACY "Enable Tracy profiler instrumentation" OFF) +option(COBRA_BUILD_IDA_PLUGIN "Build IDA Pro plugin (requires IDA_SDK_DIR)" OFF) if(COBRA_BUILD_LLVM_PASS) find_package(LLVM REQUIRED CONFIG) @@ -93,6 +94,10 @@ if(COBRA_BUILD_LLVM_PASS) add_subdirectory(lib/llvm) endif() +if(COBRA_BUILD_IDA_PLUGIN) + add_subdirectory(lib/ida) +endif() + if(COBRA_BUILD_TESTS) enable_testing() add_subdirectory(test) diff --git a/dependencies/CMakeLists.txt b/dependencies/CMakeLists.txt index 1754880..78d851d 100644 --- a/dependencies/CMakeLists.txt +++ b/dependencies/CMakeLists.txt @@ -6,6 +6,7 @@ # cmake --build build-deps # # Options: +# COBRA_ENABLE_LLVM (default OFF) - Enable LLVM dependency # USE_EXTERNAL_LLVM (default ON) - Use system LLVM vs build from source # USE_EXTERNAL_HIGHWAY (default OFF) - Use system highway vs build from source # COBRA_BUILD_TESTS (default OFF) - Build GoogleTest for tests @@ -16,6 +17,7 @@ cmake_minimum_required(VERSION 3.20) project(cobra-dependencies LANGUAGES C CXX) +option(COBRA_ENABLE_LLVM "Enable LLVM dependency" OFF) option(USE_EXTERNAL_LLVM "Use system LLVM instead of building from source" ON) option(USE_EXTERNAL_ABSEIL "Use system abseil instead of building from source" OFF) option(USE_EXTERNAL_HIGHWAY "Use system highway instead of building from source" OFF) @@ -28,7 +30,10 @@ include(superbuild.cmake) include(abseil.cmake) include(highway.cmake) -include(llvm.cmake) + +if(COBRA_ENABLE_LLVM) + include(llvm.cmake) +endif() if(COBRA_BUILD_TESTS) include(googletest.cmake) diff --git a/include/cobra/core/ExtensionLowering.h b/include/cobra/core/ExtensionLowering.h new file mode 100644 index 0000000..104f22e --- /dev/null +++ b/include/cobra/core/ExtensionLowering.h @@ -0,0 +1,24 @@ +#pragma once + +#include "cobra/core/Expr.h" + +#include +#include + +namespace cobra { + + // Scalar helpers — used by frontend evaluators. + // Precondition: 1 <= source_bits <= 64. + uint64_t EvalZeroExtend(uint64_t val, uint32_t source_bits, uint64_t result_mask); + uint64_t EvalSignExtend(uint64_t val, uint32_t source_bits, uint64_t result_mask); + + // Expr helpers — used by frontend AST builders. + // Return ordinary fixed-width Expr; no new Expr::Kind values. + // Precondition: 1 <= source_bits <= 64. + // When source_bits == 64, returns inner unchanged (identity). + std::unique_ptr< Expr > + LowerZeroExtend(std::unique_ptr< Expr > inner, uint32_t source_bits); + std::unique_ptr< Expr > + LowerSignExtend(std::unique_ptr< Expr > inner, uint32_t source_bits); + +} // namespace cobra diff --git a/lib/core/CMakeLists.txt b/lib/core/CMakeLists.txt index 48816bc..9303f5f 100644 --- a/lib/core/CMakeLists.txt +++ b/lib/core/CMakeLists.txt @@ -34,6 +34,7 @@ target_sources(cobra-core PRIVATE SignatureEval.cpp MixedProductRewriter.cpp ExprCost.cpp + ExtensionLowering.cpp SignatureSimplifier.cpp BitwiseDecomposer.cpp HybridDecomposer.cpp diff --git a/lib/core/ExtensionLowering.cpp b/lib/core/ExtensionLowering.cpp new file mode 100644 index 0000000..8ed8763 --- /dev/null +++ b/lib/core/ExtensionLowering.cpp @@ -0,0 +1,57 @@ +#include "cobra/core/ExtensionLowering.h" + +#include "cobra/core/BitWidth.h" + +#include + +namespace cobra { + namespace { + + struct ExtMasks + { + uint64_t low_mask; + uint64_t sign_bit; + }; + + ExtMasks ComputeExtMasks(uint32_t source_bits) { + assert(source_bits >= 1 && source_bits <= 64); + return ExtMasks{ + .low_mask = Bitmask(source_bits), + .sign_bit = 1ULL << (source_bits - 1), + }; + } + + } // anonymous namespace + + uint64_t EvalZeroExtend(uint64_t val, uint32_t source_bits, uint64_t result_mask) { + auto [low_mask, sign_bit] = ComputeExtMasks(source_bits); + return (val & low_mask) & result_mask; + } + + uint64_t EvalSignExtend(uint64_t val, uint32_t source_bits, uint64_t result_mask) { + auto [low_mask, sign_bit] = ComputeExtMasks(source_bits); + uint64_t masked = val & low_mask; + return ((masked ^ sign_bit) - sign_bit) & result_mask; + } + + std::unique_ptr< Expr > + LowerZeroExtend(std::unique_ptr< Expr > inner, uint32_t source_bits) { + assert(source_bits >= 1 && source_bits <= 64); + if (source_bits == 64) { return inner; } + + auto [low_mask, sign_bit] = ComputeExtMasks(source_bits); + return Expr::BitwiseAnd(std::move(inner), Expr::Constant(low_mask)); + } + + std::unique_ptr< Expr > + LowerSignExtend(std::unique_ptr< Expr > inner, uint32_t source_bits) { + assert(source_bits >= 1 && source_bits <= 64); + if (source_bits == 64) { return inner; } + + auto [low_mask, sign_bit] = ComputeExtMasks(source_bits); + auto masked = Expr::BitwiseAnd(std::move(inner), Expr::Constant(low_mask)); + auto xored = Expr::BitwiseXor(std::move(masked), Expr::Constant(sign_bit)); + return Expr::Add(std::move(xored), Expr::Negate(Expr::Constant(sign_bit))); + } + +} // namespace cobra diff --git a/lib/ida/CMakeLists.txt b/lib/ida/CMakeLists.txt new file mode 100644 index 0000000..46acef0 --- /dev/null +++ b/lib/ida/CMakeLists.txt @@ -0,0 +1,64 @@ +if(NOT DEFINED IDA_SDK_DIR) + message(FATAL_ERROR + "IDA_SDK_DIR must be set when COBRA_BUILD_IDA_PLUGIN is ON.\n" + " cmake -DIDA_SDK_DIR=/path/to/ida-sdk/src ...") +endif() + +# Detect platform-specific IDA stub library directory +if(APPLE) + if(CMAKE_SYSTEM_PROCESSOR STREQUAL "arm64") + set(_ida_lib_suffix "arm64_mac_64") + else() + set(_ida_lib_suffix "x64_mac_64") + endif() + set(_ida_lib_name "libida.dylib") +elseif(WIN32) + set(_ida_lib_suffix "x64_win_64") + set(_ida_lib_name "ida.lib") +else() + if(CMAKE_SYSTEM_PROCESSOR STREQUAL "aarch64") + set(_ida_lib_suffix "arm64_linux_64") + else() + set(_ida_lib_suffix "x64_linux_64") + endif() + set(_ida_lib_name "libida.so") +endif() + +set(_ida_lib_path "${IDA_SDK_DIR}/lib/${_ida_lib_suffix}/${_ida_lib_name}") +if(NOT EXISTS "${_ida_lib_path}") + message(FATAL_ERROR "IDA stub library not found: ${_ida_lib_path}") +endif() + +add_library(ida-cobra MODULE + ida-cobra.cpp + MicrocodeDetector.cpp + MicrocodeConverter.cpp + Verifier.cpp +) + +target_link_libraries(ida-cobra PRIVATE cobra-core "${_ida_lib_path}") + +target_include_directories(ida-cobra PRIVATE + ${IDA_SDK_DIR}/include + ${CMAKE_CURRENT_SOURCE_DIR} +) + +# IDA SDK platform defines +target_compile_definitions(ida-cobra PRIVATE + __EA64__=1 + $<$:__NT__> + $<$:__LINUX__> + $<$:__MAC__> +) + +set_target_properties(ida-cobra PROPERTIES + PREFIX "" + SUFFIX "${CMAKE_SHARED_LIBRARY_SUFFIX}" +) + +install(TARGETS ida-cobra + LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}/cobra +) +install(FILES ida-cobra.cfg + DESTINATION ${CMAKE_INSTALL_DATADIR}/cobra +) diff --git a/lib/ida/MicrocodeConverter.cpp b/lib/ida/MicrocodeConverter.cpp new file mode 100644 index 0000000..e28b652 --- /dev/null +++ b/lib/ida/MicrocodeConverter.cpp @@ -0,0 +1,284 @@ +#include "MicrocodeConverter.h" +#include "cobra/core/ExtensionLowering.h" + +namespace ida_cobra { + namespace { + + int FindLeafIndex(const mop_t &op, const MBACandidate &candidate) { + for (size_t i = 0; i < candidate.leaves.size(); ++i) { + if (*candidate.leaves[i] == op) { return static_cast< int >(i); } + } + return -1; + } + + std::unique_ptr< cobra::Expr > + ResolveLeafExpr(const mop_t &op, const MBACandidate &candidate) { + switch (op.t) { + case mop_n: + return cobra::Expr::Constant(static_cast< uint64_t >(op.nnn->value)); + case mop_r: + case mop_l: + case mop_S: + case mop_v: { + int idx = FindLeafIndex(op, candidate); + if (idx >= 0) { + return cobra::Expr::Variable(static_cast< uint32_t >(idx)); + } + return cobra::Expr::Constant(0); + } + default: + return cobra::Expr::Constant(0); + } + } + + std::unique_ptr< cobra::Expr > CombineExpr( + mcode_t opcode, std::unique_ptr< cobra::Expr > l, std::unique_ptr< cobra::Expr > r + ) { + if (!l) { return nullptr; } + switch (opcode) { + case m_bnot: + return cobra::Expr::BitwiseNot(std::move(l)); + case m_neg: + return cobra::Expr::Negate(std::move(l)); + default: + break; + } + if (!r) { return nullptr; } + switch (opcode) { + case m_add: + return cobra::Expr::Add(std::move(l), std::move(r)); + case m_sub: + return cobra::Expr::Add(std::move(l), cobra::Expr::Negate(std::move(r))); + case m_mul: + return cobra::Expr::Mul(std::move(l), std::move(r)); + case m_and: + return cobra::Expr::BitwiseAnd(std::move(l), std::move(r)); + case m_or: + return cobra::Expr::BitwiseOr(std::move(l), std::move(r)); + case m_xor: + return cobra::Expr::BitwiseXor(std::move(l), std::move(r)); + default: + return nullptr; + } + } + + int MapVarToLeaf( + uint32_t var_index, const MBACandidate &candidate, + const std::vector< std::string > &real_vars + ) { + if (var_index >= real_vars.size()) { return -1; } + + const std::string &name = real_vars[var_index]; + for (size_t i = 0; i < candidate.var_names.size(); ++i) { + if (candidate.var_names[i] == name) { return static_cast< int >(i); } + } + return -1; + } + + minsn_t *MakeBinop(mcode_t opcode, minsn_t *left, minsn_t *right, int size, ea_t ea) { + auto *insn = new minsn_t(ea); + insn->opcode = opcode; + insn->l.t = mop_d; + insn->l.d = left; + insn->l.size = size; + insn->r.t = mop_d; + insn->r.d = right; + insn->r.size = size; + insn->d.size = size; + return insn; + } + + minsn_t *MakeUnop(mcode_t opcode, minsn_t *operand, int size, ea_t ea) { + auto *insn = new minsn_t(ea); + insn->opcode = opcode; + insn->l.t = mop_d; + insn->l.d = operand; + insn->l.size = size; + insn->d.size = size; + return insn; + } + + minsn_t *MakeLeafInsn( + const cobra::Expr &expr, const MBACandidate &candidate, + const std::vector< std::string > &real_vars, int size, ea_t ea + ) { + if (expr.kind == cobra::Expr::Kind::kConstant) { + auto *insn = new minsn_t(ea); + insn->opcode = m_mov; + insn->l.make_number(expr.constant_val, size); + insn->d.size = size; + return insn; + } + // kVariable + int leaf_idx = MapVarToLeaf(expr.var_index, candidate, real_vars); + auto *insn = new minsn_t(ea); + insn->opcode = m_mov; + if (leaf_idx >= 0 && leaf_idx < static_cast< int >(candidate.leaves.size())) { + insn->l = *candidate.leaves[leaf_idx]; + } else { + insn->l.make_number(0, size); + } + insn->d.size = size; + return insn; + } + + minsn_t *CombineMinsn( + const cobra::Expr &expr, minsn_t *child0, minsn_t *child1, int size, ea_t ea + ) { + switch (expr.kind) { + case cobra::Expr::Kind::kAdd: + return MakeBinop(m_add, child0, child1, size, ea); + case cobra::Expr::Kind::kMul: + return MakeBinop(m_mul, child0, child1, size, ea); + case cobra::Expr::Kind::kAnd: + return MakeBinop(m_and, child0, child1, size, ea); + case cobra::Expr::Kind::kOr: + return MakeBinop(m_or, child0, child1, size, ea); + case cobra::Expr::Kind::kXor: + return MakeBinop(m_xor, child0, child1, size, ea); + case cobra::Expr::Kind::kNot: + return MakeUnop(m_bnot, child0, size, ea); + case cobra::Expr::Kind::kNeg: + return MakeUnop(m_neg, child0, size, ea); + case cobra::Expr::Kind::kShr: { + auto *insn = new minsn_t(ea); + insn->opcode = m_shr; + insn->l.t = mop_d; + insn->l.d = child0; + insn->l.size = size; + insn->r.make_number(expr.constant_val, size); + insn->d.size = size; + return insn; + } + default: { + auto *insn = new minsn_t(ea); + insn->opcode = m_mov; + insn->l.make_number(0, size); + insn->d.size = size; + return insn; + } + } + } + + minsn_t *ReconstructImpl( + const cobra::Expr &expr, const MBACandidate &candidate, + const std::vector< std::string > &real_vars + ) { + int size = static_cast< int >(candidate.bitwidth / 8); + ea_t ea = candidate.root->ea; + + // Flatten the Expr tree into post-order, then build the + // minsn tree bottom-up with a value stack. + std::vector< const cobra::Expr * > post; + { + std::vector< const cobra::Expr * > work; + work.push_back(&expr); + while (!work.empty()) { + const cobra::Expr *n = work.back(); + work.pop_back(); + post.push_back(n); + for (size_t i = 0; i < n->children.size(); ++i) { + if (n->children[i]) { work.push_back(n->children[i].get()); } + } + } + } + + std::vector< minsn_t * > vals; + for (auto it = post.rbegin(); it != post.rend(); ++it) { + const cobra::Expr *n = *it; + + if (n->kind == cobra::Expr::Kind::kConstant + || n->kind == cobra::Expr::Kind::kVariable) + { + vals.push_back(MakeLeafInsn(*n, candidate, real_vars, size, ea)); + continue; + } + + minsn_t *child1 = nullptr; + if (n->children.size() > 1 && n->children[1]) { + child1 = vals.back(); + vals.pop_back(); + } + + minsn_t *child0 = nullptr; + if (n->children.size() > 0 && n->children[0]) { + child0 = vals.back(); + vals.pop_back(); + } + + vals.push_back(CombineMinsn(*n, child0, child1, size, ea)); + } + + return vals.back(); + } + + } // anonymous namespace + + std::unique_ptr< cobra::Expr > + BuildExprFromMinsn(const minsn_t &insn, const MBACandidate &candidate) { + // Flatten the minsn tree into post-order, then build the Expr + // tree bottom-up with a value stack. + auto post = MicrocodePostOrder(insn); + + std::vector< std::unique_ptr< cobra::Expr > > vals; + for (auto it = post.rbegin(); it != post.rend(); ++it) { + const minsn_t *n = *it; + + // Extensions and mov: unary ops using only the left operand. + if (n->opcode == m_xdu || n->opcode == m_xds || n->opcode == m_mov) { + std::unique_ptr< cobra::Expr > l; + if (n->l.t == mop_d) { + l = std::move(vals.back()); + vals.pop_back(); + } else { + l = ResolveLeafExpr(n->l, candidate); + } + + if (n->opcode == m_mov) { + vals.push_back(std::move(l)); + } else { + uint32_t src_bits = static_cast< uint32_t >(n->l.size) * 8; + if (src_bits < 1 || src_bits > 64 + || src_bits > static_cast< uint32_t >(n->d.size) * 8) + { + return nullptr; + } + if (n->opcode == m_xdu) { + vals.push_back(cobra::LowerZeroExtend(std::move(l), src_bits)); + } else { + vals.push_back(cobra::LowerSignExtend(std::move(l), src_bits)); + } + } + continue; + } + + std::unique_ptr< cobra::Expr > r; + if (n->r.t == mop_d) { + r = std::move(vals.back()); + vals.pop_back(); + } else { + r = ResolveLeafExpr(n->r, candidate); + } + + std::unique_ptr< cobra::Expr > l; + if (n->l.t == mop_d) { + l = std::move(vals.back()); + vals.pop_back(); + } else { + l = ResolveLeafExpr(n->l, candidate); + } + + vals.push_back(CombineExpr(n->opcode, std::move(l), std::move(r))); + } + + return std::move(vals.back()); + } + + minsn_t *ReconstructMinsn( + const cobra::Expr &expr, const MBACandidate &candidate, + const std::vector< std::string > &real_vars + ) { + return ReconstructImpl(expr, candidate, real_vars); + } + +} // namespace ida_cobra diff --git a/lib/ida/MicrocodeConverter.h b/lib/ida/MicrocodeConverter.h new file mode 100644 index 0000000..c45b1db --- /dev/null +++ b/lib/ida/MicrocodeConverter.h @@ -0,0 +1,19 @@ +#pragma once + +#include "MicrocodeDetector.h" + +#include + +#include + +namespace ida_cobra { + + std::unique_ptr< cobra::Expr > + BuildExprFromMinsn(const minsn_t &insn, const MBACandidate &candidate); + + minsn_t *ReconstructMinsn( + const cobra::Expr &expr, const MBACandidate &candidate, + const std::vector< std::string > &real_vars + ); + +} // namespace ida_cobra diff --git a/lib/ida/MicrocodeDetector.cpp b/lib/ida/MicrocodeDetector.cpp new file mode 100644 index 0000000..2547f60 --- /dev/null +++ b/lib/ida/MicrocodeDetector.cpp @@ -0,0 +1,396 @@ +// absl must be included before MicrocodeDetector.h (which pulls hexrays.hpp): +// the IDA SDK poisons stdout/stderr/fwrite/fflush/snprintf via fpro.h macros. +#include "cobra/core/ExtensionLowering.h" +#include + +#include "MicrocodeDetector.h" + +namespace ida_cobra { + namespace { + + // Minimum opcode counts to classify an instruction tree as MBA. + constexpr int kMinBoolOps = 1; + constexpr int kMinArithOps = 1; + + // Maximum variables CoBRA can handle. + constexpr uint32_t kMaxVars = 16; + + struct OpcodeCounter : public minsn_visitor_t + { + int bool_cnt = 0; + int arith_cnt = 0; + + int idaapi visit_minsn() override { + switch (curins->opcode) { + case m_neg: + case m_add: + case m_sub: + case m_mul: + arith_cnt++; + break; + case m_bnot: + case m_or: + case m_and: + case m_xor: + bool_cnt++; + break; + default: + return 0; + } + return (bool_cnt >= kMinBoolOps && arith_cnt >= kMinArithOps) ? 1 : 0; + } + }; + + } // anonymous namespace + + uint64_t EvalMinsn( + const minsn_t &insn, const std::vector< mop_t * > &var_keys, + const std::vector< uint64_t > &var_vals, uint64_t mask + ) { + auto resolve_leaf = [&](const mop_t &op) -> uint64_t { + switch (op.t) { + case mop_n: + return static_cast< uint64_t >(op.nnn->value) & mask; + case mop_r: + case mop_l: + case mop_S: + case mop_v: + for (size_t i = 0; i < var_keys.size(); ++i) { + if (*var_keys[i] == op) { return var_vals[i]; } + } + return 0; + default: + return 0; + } + }; + + // Flatten the minsn tree into post-order, then evaluate + // bottom-up with a value stack. + auto post = MicrocodePostOrder(insn); + + std::vector< uint64_t > vals; + for (auto it = post.rbegin(); it != post.rend(); ++it) { + const minsn_t *n = *it; + + uint64_t r = 0; + if (n->r.t == mop_d) { + r = vals.back(); + vals.pop_back(); + } else { + r = resolve_leaf(n->r); + } + + uint64_t l = 0; + if (n->l.t == mop_d) { + l = vals.back(); + vals.pop_back(); + } else { + l = resolve_leaf(n->l); + } + + switch (n->opcode) { + case m_add: + vals.push_back((l + r) & mask); + break; + case m_sub: + vals.push_back((l - r) & mask); + break; + case m_mul: + vals.push_back((l * r) & mask); + break; + case m_and: + vals.push_back(l & r); + break; + case m_or: + vals.push_back(l | r); + break; + case m_xor: + vals.push_back(l ^ r); + break; + case m_bnot: + vals.push_back((~l) & mask); + break; + case m_neg: + vals.push_back((static_cast< uint64_t >(0) - l) & mask); + break; + case m_xdu: { + uint32_t src_bits = static_cast< uint32_t >(n->l.size) * 8; + vals.push_back(cobra::EvalZeroExtend(l, src_bits, mask)); + break; + } + case m_xds: { + uint32_t src_bits = static_cast< uint32_t >(n->l.size) * 8; + vals.push_back(cobra::EvalSignExtend(l, src_bits, mask)); + break; + } + case m_mov: + vals.push_back(l & mask); + break; + default: + vals.push_back(0); + break; + } + } + + return vals.back(); + } + + namespace { + + // Collect unique leaf (input) operands from a minsn tree. + // Only walks .l and .r — never .d (the destination). + // Uses an explicit worklist instead of recursion. + struct LeafCollector + { + std::vector< mop_t * > leaves; + bool valid = true; + + void Collect(minsn_t &root) { + std::vector< mop_t * > worklist; + worklist.push_back(&root.r); + worklist.push_back(&root.l); + + while (!worklist.empty() && valid) { + mop_t *op = worklist.back(); + worklist.pop_back(); + + if (op->t == mop_d) { + minsn_t *inner = op->d; + // Validate extension widths before traversing. + if (inner->opcode == m_xdu || inner->opcode == m_xds) { + uint32_t src_bits = static_cast< uint32_t >(inner->l.size) * 8; + uint32_t dst_bits = static_cast< uint32_t >(inner->d.size) * 8; + if (src_bits < 1 || src_bits > 64 || dst_bits < 1 || dst_bits > 64 + || src_bits > dst_bits) + { + valid = false; + return; + } + } + worklist.push_back(&inner->r); + worklist.push_back(&inner->l); + continue; + } + if (op->t == mop_n || op->t == mop_z) { continue; } + + if (op->t == mop_r || op->t == mop_l || op->t == mop_S || op->t == mop_v) { + bool found = false; + for (const auto *existing : leaves) { + if (*existing == *op) { + found = true; + break; + } + } + if (!found) { leaves.push_back(op); } + } + } + } + }; + + // Build a human-readable name for a leaf operand. + std::string LeafName(const mop_t &op) { + qstring buf; + op.print(&buf); + return std::string(buf.c_str()); + } + + } // anonymous namespace + + bool IsMba(const minsn_t &insn) { + if (insn.opcode >= m_jcnd) { return false; } + + if (insn.d.size > 8) { return false; } + + OpcodeCounter counter; + return const_cast< minsn_t & >(insn).for_all_insns(counter) != 0; + } + + namespace { + + // Return true if the opcode is one that participates in MBA + // expressions (arithmetic or boolean). + bool IsMbaOpcode(mcode_t op) { + switch (op) { + case m_neg: + case m_bnot: + case m_add: + case m_sub: + case m_mul: + case m_or: + case m_and: + case m_xor: + return true; + default: + return false; + } + } + + // Peel only m_mov wrappers. Accept MBA ops or xdu/xds as + // valid roots — extensions are part of the candidate tree. + minsn_t *MbaRoot(minsn_t *insn) { + while (insn->opcode == m_mov && insn->l.t == mop_d) { insn = insn->l.d; } + if (IsMbaOpcode(insn->opcode)) { return insn; } + if (insn->opcode == m_xdu || insn->opcode == m_xds) { return insn; } + return nullptr; + } + + } // anonymous namespace + + std::vector< MBACandidate > DetectMbaCandidates(mba_t &mba) { + std::vector< MBACandidate > candidates; + + struct DetectorVisitor : public minsn_visitor_t + { + std::vector< MBACandidate > &out; + + explicit DetectorVisitor(std::vector< MBACandidate > &o) : out(o) {} + + int idaapi visit_minsn() override { + if (!IsMba(*curins)) { return 0; } + minsn_t *root = MbaRoot(curins); + if (!root) { return 0; } + + // Validate extension widths at the root itself — + // LeafCollector only inspects child extensions, not the + // root instruction. + if (root->opcode == m_xdu || root->opcode == m_xds) { + uint32_t src = static_cast< uint32_t >(root->l.size) * 8; + uint32_t dst = static_cast< uint32_t >(root->d.size) * 8; + if (src < 1 || src > 64 || dst < 1 || dst > 64 || src > dst) { return 0; } + } + + LeafCollector lc; + lc.Collect(*root); + + if (!lc.valid) { return 0; } + if (lc.leaves.size() > kMaxVars) { return 0; } + + uint32_t bitwidth = static_cast< uint32_t >(root->d.size) * 8; + if (bitwidth == 0 || bitwidth > 64) { return 0; } + uint64_t mask = + bitwidth >= 64 ? ~uint64_t{ 0 } : (uint64_t{ 1 } << bitwidth) - 1; + + uint32_t n = static_cast< uint32_t >(lc.leaves.size()); + + std::vector< uint64_t > sig; + sig.reserve(uint64_t{ 1 } << n); + + for (uint64_t input = 0; input < (uint64_t{ 1 } << n); ++input) { + std::vector< uint64_t > vals(n); + for (uint32_t v = 0; v < n; ++v) { vals[v] = (input >> v) & 1; } + + sig.push_back(EvalMinsn(*root, lc.leaves, vals, mask)); + } + + std::vector< std::string > names; + names.reserve(n); + for (auto *leaf : lc.leaves) { names.push_back(LeafName(*leaf)); } + + out.push_back( + MBACandidate{ + .root = root, + .leaves = std::move(lc.leaves), + .var_names = std::move(names), + .sig = std::move(sig), + .bitwidth = bitwidth, + } + ); + + return 0; + } + }; + + DetectorVisitor visitor(candidates); + mba.for_all_topinsns(visitor); + + return candidates; + } + + std::vector< MBACandidate > DetectMbaCandidatesCrossBlock(mba_t &mba) { + struct CrossBlockDetector + { + mba_t &mba; + absl::flat_hash_set< const minsn_t * > already_in_tree; + std::vector< MBACandidate > candidates; + + explicit CrossBlockDetector(mba_t &m) : mba(m) {} + + void Run() { + node_ordering_t post_order; + mba.get_graph()->depth_first_postorder(&post_order); + + for (size_t i = 0; i < post_order.size(); ++i) { + int blk_idx = post_order.node(i); + mblock_t *blk = mba.get_mblock(blk_idx); + + for (minsn_t *insn = blk->tail; insn != nullptr; insn = insn->prev) { + if (already_in_tree.count(insn) != 0) { continue; } + + if (!IsMba(*insn)) { continue; } + minsn_t *root = MbaRoot(insn); + if (!root) { continue; } + + // Validate extension widths at the root itself. + if (root->opcode == m_xdu || root->opcode == m_xds) { + uint32_t src = static_cast< uint32_t >(root->l.size) * 8; + uint32_t dst = static_cast< uint32_t >(root->d.size) * 8; + if (src < 1 || src > 64 || dst < 1 || dst > 64 || src > dst) { + continue; + } + } + + LeafCollector lc; + MarkTree(root, lc); + + if (!lc.valid) { continue; } + if (lc.leaves.size() > kMaxVars) { continue; } + + uint32_t bitwidth = static_cast< uint32_t >(root->d.size) * 8; + if (bitwidth == 0 || bitwidth > 64) { continue; } + uint64_t mask = + bitwidth >= 64 ? ~uint64_t{ 0 } : (uint64_t{ 1 } << bitwidth) - 1; + + uint32_t n = static_cast< uint32_t >(lc.leaves.size()); + + std::vector< uint64_t > sig; + sig.reserve(uint64_t{ 1 } << n); + for (uint64_t input = 0; input < (uint64_t{ 1 } << n); ++input) { + std::vector< uint64_t > vals(n); + for (uint32_t v = 0; v < n; ++v) { vals[v] = (input >> v) & 1; } + sig.push_back(EvalMinsn(*root, lc.leaves, vals, mask)); + } + + std::vector< std::string > names; + names.reserve(n); + for (auto *leaf : lc.leaves) { names.push_back(LeafName(*leaf)); } + + candidates.push_back( + MBACandidate{ + .root = root, + .leaves = std::move(lc.leaves), + .var_names = std::move(names), + .sig = std::move(sig), + .bitwidth = bitwidth, + } + ); + } + } + } + + void MarkTree(minsn_t *insn, LeafCollector &lc) { + already_in_tree.insert(insn); + lc.Collect(*insn); + // Cross-block extension via graph_chains_t is deferred + // until the use-def chain API is validated via manual + // testing in IDA. For now, this falls through to + // intra-block behavior (same as Tier 1 but with + // post-order + reverse walk + already_in_tree dedup). + } + }; + + CrossBlockDetector detector(mba); + detector.Run(); + return std::move(detector.candidates); + } + +} // namespace ida_cobra diff --git a/lib/ida/MicrocodeDetector.h b/lib/ida/MicrocodeDetector.h new file mode 100644 index 0000000..9613584 --- /dev/null +++ b/lib/ida/MicrocodeDetector.h @@ -0,0 +1,62 @@ +#pragma once + +// STL and absl must be included before hexrays.hpp: the IDA SDK poisons +// stdout, stderr, fwrite, fflush, snprintf etc. via fpro.h macros, which +// breaks any subsequent libc++/absl header that references those identifiers. +#include + +#include +#include +#include + +#include + +namespace ida_cobra { + + struct MBACandidate + { + minsn_t *root = nullptr; + std::vector< mop_t * > leaves; + std::vector< std::string > var_names; + std::vector< uint64_t > sig; + uint32_t bitwidth = 64; + }; + + // Returns true if the instruction tree rooted at `insn` is an MBA + // expression (at least 1 boolean and 1 arithmetic opcode, destination + // fits in 64 bits). + bool IsMba(const minsn_t &insn); + + // Flatten a minsn tree into post-order (two-stack method). + // Children linked via mop_d are expanded; leaf operands are not included. + inline std::vector< const minsn_t * > MicrocodePostOrder(const minsn_t &root) { + std::vector< const minsn_t * > post; + std::vector< const minsn_t * > work; + work.push_back(&root); + while (!work.empty()) { + const minsn_t *n = work.back(); + work.pop_back(); + post.push_back(n); + if (n->l.t == mop_d) { work.push_back(n->l.d); } + if (n->r.t == mop_d) { work.push_back(n->r.d); } + } + return post; + } + + // Evaluate a minsn tree with the given variable assignments. + // Variables are matched by value equality (mop_t::operator==), + // not pointer identity. + uint64_t EvalMinsn( + const minsn_t &insn, const std::vector< mop_t * > &var_keys, + const std::vector< uint64_t > &var_vals, uint64_t mask + ); + + // Walk all top-level instructions in `mba`, detect MBA trees, compute + // boolean signatures, and return candidates ready for simplification. + std::vector< MBACandidate > DetectMbaCandidates(mba_t &mba); + + // Enhanced detection that follows use-def chains across block boundaries. + // Falls back to Tier 1 (intra-block) if cross-block tracing is not viable. + std::vector< MBACandidate > DetectMbaCandidatesCrossBlock(mba_t &mba); + +} // namespace ida_cobra diff --git a/lib/ida/Verifier.cpp b/lib/ida/Verifier.cpp new file mode 100644 index 0000000..af0f054 --- /dev/null +++ b/lib/ida/Verifier.cpp @@ -0,0 +1,80 @@ +#include + +#include + +#include "MicrocodeDetector.h" +#include "Verifier.h" + +namespace ida_cobra { + namespace { + + constexpr int kNumTests = 256; + + std::mt19937_64 &Rng() { + thread_local std::mt19937_64 rng{ std::random_device{}() }; + return rng; + } + + constexpr uint64_t kSpecial[] = { 0, 1, 0xFFFFFFFFFFFFFFFF }; + constexpr int kNumSpecial = 3; + + uint64_t RandValue() { + if (std::uniform_int_distribution< int >{ 0, 4 }(Rng()) == 0) { + return kSpecial[std::uniform_int_distribution< int >{ 0, + kNumSpecial - 1 }(Rng())]; + } + return Rng()(); + } + + } // anonymous namespace + + bool ProbablyEquivalent( + const minsn_t &original, const cobra::Expr &simplified, const MBACandidate &candidate + ) { + uint64_t mask = candidate.bitwidth >= 64 ? ~uint64_t{ 0 } + : (uint64_t{ 1 } << candidate.bitwidth) - 1; + + cobra::CompiledExpr compiled = cobra::CompileExpr(simplified, candidate.bitwidth); + std::vector< uint64_t > stack(compiled.stack_size); + + for (int test = 0; test < kNumTests; ++test) { + std::vector< uint64_t > vals(candidate.leaves.size()); + + for (size_t i = 0; i < candidate.leaves.size(); ++i) { + vals[i] = RandValue() & mask; + } + + uint64_t original_result = + ida_cobra::EvalMinsn(original, candidate.leaves, vals, mask); + + uint64_t simplified_result = cobra::EvalCompiledExpr(compiled, vals, stack) & mask; + + if (original_result != simplified_result) { + msg("ida-cobra: verification FAILED on test %d\n", test); + msg(" original=%llx simplified=%llx\n", + static_cast< unsigned long long >(original_result), + static_cast< unsigned long long >(simplified_result)); + return false; + } + } + + return true; + } + + int CountNodes(const minsn_t &insn) { + struct NodeCounter : public minsn_visitor_t + { + int count = 0; + + int idaapi visit_minsn() override { + count++; + return 0; + } + }; + + NodeCounter counter; + const_cast< minsn_t & >(insn).for_all_insns(counter); + return counter.count; + } + +} // namespace ida_cobra diff --git a/lib/ida/Verifier.h b/lib/ida/Verifier.h new file mode 100644 index 0000000..b4980c5 --- /dev/null +++ b/lib/ida/Verifier.h @@ -0,0 +1,15 @@ +#pragma once + +#include "MicrocodeDetector.h" + +#include + +namespace ida_cobra { + + bool ProbablyEquivalent( + const minsn_t &original, const cobra::Expr &simplified, const MBACandidate &candidate + ); + + int CountNodes(const minsn_t &insn); + +} // namespace ida_cobra diff --git a/lib/ida/ida-cobra.cfg b/lib/ida/ida-cobra.cfg new file mode 100644 index 0000000..79ca3e9 --- /dev/null +++ b/lib/ida/ida-cobra.cfg @@ -0,0 +1,6 @@ +// ida-cobra configuration file. +// Place in IDA's cfg directory or next to the plugin binary. + +// Set to YES to automatically simplify MBA expressions on decompilation. +// Default: only simplify via right-click menu. +COBRA_RUN_AUTOMATICALLY = NO diff --git a/lib/ida/ida-cobra.cpp b/lib/ida/ida-cobra.cpp new file mode 100644 index 0000000..b96190e --- /dev/null +++ b/lib/ida/ida-cobra.cpp @@ -0,0 +1,163 @@ +#include +#include + +#include "MicrocodeConverter.h" +#include "MicrocodeDetector.h" +#include "Verifier.h" + +#define ACTION_NAME "ida_cobra:run" + +struct plugin_ctx_t; + +struct run_ah_t : public action_handler_t +{ + plugin_ctx_t *ctx; + + explicit run_ah_t(plugin_ctx_t *c) : ctx(c) {} + + int idaapi activate(action_activation_ctx_t *act_ctx) override; + + action_state_t idaapi update(action_update_ctx_t *upd_ctx) override { + return upd_ctx->widget_type == BWN_PSEUDOCODE ? AST_ENABLE_FOR_WIDGET + : AST_DISABLE_FOR_WIDGET; + } +}; + +struct plugin_ctx_t : public plugmod_t +{ + bool run_automatically = false; + bool active = false; + + run_ah_t action_handler; + + plugin_ctx_t(); + + ~plugin_ctx_t() { term_hexrays_plugin(); } + + bool idaapi run(size_t) override { return true; } +}; + +static ssize_t idaapi hex_callback(void *ud, hexrays_event_t event, va_list va) { + auto *ctx = static_cast< plugin_ctx_t * >(ud); + + switch (event) { + case hxe_microcode: { + auto *mba = va_arg(va, mba_t *); + if (ctx->run_automatically) { ctx->active = true; } + if (ctx->active) { mba->set_mba_flags2(MBA2_PROP_COMPLEX); } + break; + } + case hxe_populating_popup: { + auto *widget = va_arg(va, TWidget *); + auto *popup = va_arg(va, TPopupMenu *); + attach_action_to_popup(widget, popup, ACTION_NAME); + break; + } + case hxe_glbopt: { + auto *mba = va_arg(va, mba_t *); + + if (!ctx->active) { return MERR_OK; } + + auto candidates = ida_cobra::DetectMbaCandidatesCrossBlock(*mba); + int improved = 0; + + for (auto &cand : candidates) { + auto expr = ida_cobra::BuildExprFromMinsn(*cand.root, cand); + if (!expr) { continue; } + + cobra::Options opts; + opts.bitwidth = cand.bitwidth; + + auto result = cobra::Simplify(cand.sig, cand.var_names, expr.get(), opts); + if (!result.has_value()) { continue; } + + auto &outcome = result.value(); + if (outcome.kind != cobra::SimplifyOutcome::Kind::kSimplified) { continue; } + if (!outcome.expr) { continue; } + + auto original_cost = cobra::ComputeCost(*expr); + auto simplified_cost = cobra::ComputeCost(*outcome.expr); + if (!cobra::IsBetter(simplified_cost.cost, original_cost.cost)) { continue; } + + if (!ida_cobra::ProbablyEquivalent(*cand.root, *outcome.expr, cand)) { + continue; + } + + minsn_t *replacement = + ida_cobra::ReconstructMinsn(*outcome.expr, cand, outcome.real_vars); + if (replacement == nullptr) { continue; } + + replacement->d.swap(cand.root->d); + cand.root->swap(*replacement); + delete replacement; + + improved++; + } + + ctx->active = false; + mba->clr_mba_flags2(MBA2_PROP_COMPLEX); + + if (improved > 0) { + mba->verify(true); + msg("ida-cobra: simplified %d MBA expression(s) in %a\n", improved, + mba->entry_ea); + // Tag the function so scripts can query which functions were simplified. + netnode n(mba->entry_ea); + n.altset(0, improved, 'C'); + return MERR_LOOP; + } + return MERR_OK; + } + default: + break; + } + return 0; +} + +int idaapi run_ah_t::activate(action_activation_ctx_t *act_ctx) { + vdui_t *vu = get_widget_vdui(act_ctx->widget); + if (vu != nullptr) { + ctx->active = true; + vu->refresh_view(true); + return 1; + } + return 0; +} + +plugin_ctx_t::plugin_ctx_t() : action_handler(this) { + install_hexrays_callback(hex_callback, this); + register_action(ACTION_DESC_LITERAL_PLUGMOD( + ACTION_NAME, "Run CoBRA Optimizer", &action_handler, this, nullptr, + "Simplify MBA-obfuscated expressions using CoBRA", -1 + )); +} + +static plugmod_t *idaapi init() { + if (!init_hexrays_plugin()) { return nullptr; } + + const char *hxver = get_hexrays_version(); + msg("ida-cobra: Hex-Rays %s detected, CoBRA MBA optimizer ready\n", hxver); + + auto *ctx = new plugin_ctx_t; + + const cfgopt_t cfgopts[] = { + cfgopt_t("COBRA_RUN_AUTOMATICALLY", &ctx->run_automatically, 1), + }; + read_config_file("ida-cobra", cfgopts, qnumber(cfgopts), nullptr); + + return ctx; +} + +static char comment[] = "CoBRA MBA deobfuscation plugin for Hex-Rays decompiler"; + +plugin_t PLUGIN = { + IDP_INTERFACE_VERSION, + PLUGIN_MULTI | PLUGIN_HIDE, + init, + nullptr, + nullptr, + comment, + nullptr, + "ida-cobra plugin", + nullptr, +}; diff --git a/lib/llvm/MBADetector.cpp b/lib/llvm/MBADetector.cpp index 86067c9..1649588 100644 --- a/lib/llvm/MBADetector.cpp +++ b/lib/llvm/MBADetector.cpp @@ -1,6 +1,7 @@ #include "MBADetector.h" #include "cobra/core/BitWidth.h" #include "cobra/core/Expr.h" +#include "cobra/core/ExtensionLowering.h" #include "cobra/core/Simplifier.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/DenseSet.h" @@ -13,6 +14,7 @@ #include "llvm/Support/Casting.h" #include +#include #include #include #include @@ -27,7 +29,9 @@ namespace cobra { namespace { - bool IsMbaOpcode(unsigned opcode) { + // Core MBA operations — used for classification and candidate + // emission. Extensions are NOT included here. + bool IsCoreMbaOpcode(unsigned opcode) { switch (opcode) { // NOLINT(hicpp-multiway-paths-covered) case llvm::Instruction::Add: case llvm::Instruction::Sub: @@ -36,14 +40,20 @@ namespace cobra { case llvm::Instruction::Or: case llvm::Instruction::Xor: case llvm::Instruction::LShr: - case llvm::Instruction::ZExt: - case llvm::Instruction::SExt: return true; default: return false; } } + // Traversable tree operations — core MBA ops plus extensions. + // Used by CollectTree, root selection, PHI transparency, and + // ArmDepsInLeafSet. + bool IsTreeOpcode(unsigned opcode) { + return IsCoreMbaOpcode(opcode) || opcode == llvm::Instruction::ZExt + || opcode == llvm::Instruction::SExt; + } + // BFS from root following operands. MBA-opcode instructions // are added to tree_insts; everything else becomes a leaf. // PHI nodes are treated as transparent when all incoming @@ -65,7 +75,7 @@ namespace cobra { if (!visited.insert(v).second) { continue; } auto *inst = llvm::dyn_cast< llvm::Instruction >(v); - if ((inst != nullptr) && IsMbaOpcode(inst->getOpcode())) { + if ((inst != nullptr) && IsTreeOpcode(inst->getOpcode())) { // LShr with variable shift amount is unsupported — // treat the whole instruction as a leaf. if (inst->getOpcode() == llvm::Instruction::LShr @@ -88,7 +98,7 @@ namespace cobra { auto *inc = phi->getIncomingValue(i); if (llvm::isa< llvm::ConstantInt >(inc)) { continue; } auto *inc_inst = llvm::dyn_cast< llvm::Instruction >(inc); - if ((inc_inst == nullptr) || !IsMbaOpcode(inc_inst->getOpcode())) { + if ((inc_inst == nullptr) || !IsTreeOpcode(inc_inst->getOpcode())) { all_mba = false; break; } @@ -147,11 +157,20 @@ namespace cobra { auto *inst = llvm::cast< llvm::Instruction >(v); - if (inst->getOpcode() == llvm::Instruction::ZExt - || inst->getOpcode() == llvm::Instruction::SExt) - { + if (inst->getOpcode() == llvm::Instruction::ZExt) { const uint64_t operand = eval(inst->getOperand(0)); - cache[v] = operand & mask; + const uint32_t src_bits = + inst->getOperand(0)->getType()->getIntegerBitWidth(); + assert(src_bits >= 1 && src_bits <= 64); + cache[v] = EvalZeroExtend(operand, src_bits, mask); + return cache[v]; + } + if (inst->getOpcode() == llvm::Instruction::SExt) { + const uint64_t operand = eval(inst->getOperand(0)); + const uint32_t src_bits = + inst->getOperand(0)->getType()->getIntegerBitWidth(); + assert(src_bits >= 1 && src_bits <= 64); + cache[v] = EvalSignExtend(operand, src_bits, mask); return cache[v]; } @@ -287,13 +306,21 @@ namespace cobra { auto *inst = llvm::dyn_cast< llvm::Instruction >(v); if (inst == nullptr || !tree_set.contains(inst)) { return nullptr; } - // ZExt/SExt — pass through to inner operand - if (inst->getOpcode() == llvm::Instruction::ZExt - || inst->getOpcode() == llvm::Instruction::SExt) - { - return BuildExprFromIR( - inst->getOperand(0), leaves, tree_set, mask, phi_redirects - ); + if (inst->getOpcode() == llvm::Instruction::ZExt) { + auto inner = + BuildExprFromIR(inst->getOperand(0), leaves, tree_set, mask, phi_redirects); + if (inner == nullptr) { return nullptr; } + const uint32_t src_bits = inst->getOperand(0)->getType()->getIntegerBitWidth(); + assert(src_bits >= 1 && src_bits <= 64); + return LowerZeroExtend(std::move(inner), src_bits); + } + if (inst->getOpcode() == llvm::Instruction::SExt) { + auto inner = + BuildExprFromIR(inst->getOperand(0), leaves, tree_set, mask, phi_redirects); + if (inner == nullptr) { return nullptr; } + const uint32_t src_bits = inst->getOperand(0)->getType()->getIntegerBitWidth(); + assert(src_bits >= 1 && src_bits <= 64); + return LowerSignExtend(std::move(inner), src_bits); } // LShr with constant shift amount @@ -359,7 +386,7 @@ namespace cobra { if (leaf_set.contains(v)) { continue; } auto *inst = llvm::dyn_cast< llvm::Instruction >(v); - if ((inst == nullptr) || !IsMbaOpcode(inst->getOpcode())) { return false; } + if ((inst == nullptr) || !IsTreeOpcode(inst->getOpcode())) { return false; } // LShr with variable shift — can't evaluate if (inst->getOpcode() == llvm::Instruction::LShr @@ -449,7 +476,7 @@ namespace cobra { // they can be emitted as standalone candidates. for (auto *bb : post_order(&f)) { for (auto &inst : llvm::reverse(*bb)) { - if (!IsMbaOpcode(inst.getOpcode())) { continue; } + if (!IsTreeOpcode(inst.getOpcode())) { continue; } if (already_in_tree.contains(&inst) != 0u) { continue; } if (!inst.getType()->isIntegerTy()) { continue; } @@ -463,6 +490,17 @@ namespace cobra { if (tree_insts.size() < min_ast_size) { continue; } + // At least one core MBA op required — extension-only + // chains are not MBA candidates. + bool has_core_op = false; + for (auto *ti : tree_insts) { + if (IsCoreMbaOpcode(ti->getOpcode())) { + has_core_op = true; + break; + } + } + if (!has_core_op) { continue; } + constexpr uint32_t kPreElimCap = 20; if (leaves.size() > kPreElimCap) { continue; } @@ -484,6 +522,15 @@ namespace cobra { if (tree_insts.size() < min_ast_size) { continue; } if (leaves.size() > kPreElimCap) { continue; } + + has_core_op = false; + for (auto *ti : tree_insts) { + if (IsCoreMbaOpcode(ti->getOpcode())) { + has_core_op = true; + break; + } + } + if (!has_core_op) { continue; } } for (auto *ti : tree_insts) { already_in_tree.insert(ti); } diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt index 47985ae..9de5b75 100644 --- a/test/CMakeLists.txt +++ b/test/CMakeLists.txt @@ -50,6 +50,8 @@ cobra_add_test(test_weighted_poly_fit core/test_weighted_poly_fit.cpp) cobra_add_test(test_term_refiner core/test_term_refiner.cpp) cobra_add_test(test_semilinear_signature core/test_semilinear_signature.cpp) cobra_add_test(test_structure_recovery core/test_structure_recovery.cpp) +cobra_add_test(test_extension_lowering core/test_extension_lowering.cpp) + cobra_add_test(test_competition_group core/test_competition_group.cpp) target_include_directories(test_competition_group PRIVATE ${PROJECT_SOURCE_DIR}/lib/core) diff --git a/test/core/test_extension_lowering.cpp b/test/core/test_extension_lowering.cpp new file mode 100644 index 0000000..6fb231c --- /dev/null +++ b/test/core/test_extension_lowering.cpp @@ -0,0 +1,228 @@ +#include "cobra/core/ExtensionLowering.h" + +#include +#include + +using namespace cobra; + +// ---------- EvalZeroExtend ---------- + +TEST(EvalZeroExtendTest, OneBitWidth) { + // 1-bit zext: only bit 0 survives. + constexpr uint64_t kMask64 = UINT64_MAX; + EXPECT_EQ(EvalZeroExtend(0, 1, kMask64), 0u); + EXPECT_EQ(EvalZeroExtend(1, 1, kMask64), 1u); + EXPECT_EQ(EvalZeroExtend(0xFF, 1, kMask64), 1u); +} + +TEST(EvalZeroExtendTest, EightBitWidth) { + constexpr uint64_t kMask64 = UINT64_MAX; + EXPECT_EQ(EvalZeroExtend(0, 8, kMask64), 0u); + EXPECT_EQ(EvalZeroExtend(0x7F, 8, kMask64), 0x7Fu); + EXPECT_EQ(EvalZeroExtend(0x80, 8, kMask64), 0x80u); + EXPECT_EQ(EvalZeroExtend(0xFF, 8, kMask64), 0xFFu); + // Upper bits cleared. + EXPECT_EQ(EvalZeroExtend(0xDEAD00FF, 8, kMask64), 0xFFu); +} + +TEST(EvalZeroExtendTest, ThirtyTwoBitWidth) { + constexpr uint64_t kMask64 = UINT64_MAX; + EXPECT_EQ(EvalZeroExtend(0, 32, kMask64), 0u); + EXPECT_EQ(EvalZeroExtend(0x7FFFFFFFu, 32, kMask64), 0x7FFFFFFFu); + EXPECT_EQ(EvalZeroExtend(0x80000000u, 32, kMask64), 0x80000000u); + EXPECT_EQ(EvalZeroExtend(0xFFFFFFFFu, 32, kMask64), 0xFFFFFFFFu); + EXPECT_EQ(EvalZeroExtend(0xDEADBEEF12345678ULL, 32, kMask64), 0x12345678u); +} + +TEST(EvalZeroExtendTest, SixtyFourBitIsIdentity) { + constexpr uint64_t kMask64 = UINT64_MAX; + EXPECT_EQ(EvalZeroExtend(0, 64, kMask64), 0u); + EXPECT_EQ(EvalZeroExtend(0xDEADBEEFCAFEBABEULL, 64, kMask64), 0xDEADBEEFCAFEBABEULL); +} + +TEST(EvalZeroExtendTest, ResultMaskApplied) { + // 8-bit zext into a 32-bit result mask. + constexpr uint64_t kMask32 = 0xFFFFFFFFu; + EXPECT_EQ(EvalZeroExtend(0xFF, 8, kMask32), 0xFFu); + // 32-bit zext into a 32-bit result mask. + EXPECT_EQ(EvalZeroExtend(0xFFFFFFFFFFFFFFFFULL, 32, kMask32), 0xFFFFFFFFu); +} + +// ---------- EvalSignExtend ---------- + +TEST(EvalSignExtendTest, OneBitWidth) { + constexpr uint64_t kMask64 = UINT64_MAX; + // sext i1 0 -> 0 + EXPECT_EQ(EvalSignExtend(0, 1, kMask64), 0u); + // sext i1 1 -> all ones (sign bit set in 1-bit value) + EXPECT_EQ(EvalSignExtend(1, 1, kMask64), UINT64_MAX); +} + +TEST(EvalSignExtendTest, EightBitWidth) { + constexpr uint64_t kMask64 = UINT64_MAX; + EXPECT_EQ(EvalSignExtend(0, 8, kMask64), 0u); + EXPECT_EQ(EvalSignExtend(0x7F, 8, kMask64), 0x7Fu); + // 0x80 has sign bit set -> extends to all-F upper bits. + EXPECT_EQ(EvalSignExtend(0x80, 8, kMask64), 0xFFFFFFFFFFFFFF80ULL); + EXPECT_EQ(EvalSignExtend(0xFF, 8, kMask64), 0xFFFFFFFFFFFFFFFFULL); +} + +TEST(EvalSignExtendTest, ThirtyTwoBitWidth) { + constexpr uint64_t kMask64 = UINT64_MAX; + EXPECT_EQ(EvalSignExtend(0, 32, kMask64), 0u); + EXPECT_EQ(EvalSignExtend(0x7FFFFFFFu, 32, kMask64), 0x7FFFFFFFu); + EXPECT_EQ(EvalSignExtend(0x80000000u, 32, kMask64), 0xFFFFFFFF80000000ULL); + EXPECT_EQ(EvalSignExtend(0xFFFFFFFFu, 32, kMask64), 0xFFFFFFFFFFFFFFFFULL); +} + +TEST(EvalSignExtendTest, SixtyFourBitIsIdentity) { + constexpr uint64_t kMask64 = UINT64_MAX; + EXPECT_EQ(EvalSignExtend(0, 64, kMask64), 0u); + EXPECT_EQ(EvalSignExtend(0xDEADBEEFCAFEBABEULL, 64, kMask64), 0xDEADBEEFCAFEBABEULL); +} + +TEST(EvalSignExtendTest, ResultMaskApplied) { + // sext 8-bit 0x80 into 32-bit result mask: upper 32 bits cleared. + constexpr uint64_t kMask32 = 0xFFFFFFFFu; + EXPECT_EQ(EvalSignExtend(0x80, 8, kMask32), 0xFFFFFF80u); + // sext 8-bit 0xFF into 32-bit result mask. + EXPECT_EQ(EvalSignExtend(0xFF, 8, kMask32), 0xFFFFFFFFu); +} + +#if defined(GTEST_HAS_DEATH_TEST) && !defined(NDEBUG) +TEST(EvalZeroExtendDeathTest, ZeroBitsAsserts) { + EXPECT_DEATH(EvalZeroExtend(0, 0, UINT64_MAX), ""); +} + +TEST(EvalZeroExtendDeathTest, OverSixtyFourBitsAsserts) { + EXPECT_DEATH(EvalZeroExtend(0, 65, UINT64_MAX), ""); +} + +TEST(EvalSignExtendDeathTest, ZeroBitsAsserts) { + EXPECT_DEATH(EvalSignExtend(0, 0, UINT64_MAX), ""); +} + +TEST(EvalSignExtendDeathTest, OverSixtyFourBitsAsserts) { + EXPECT_DEATH(EvalSignExtend(0, 65, UINT64_MAX), ""); +} +#endif + +#include "cobra/core/BitWidth.h" +#include "cobra/core/CompiledExpr.h" + +// ---------- Helper: evaluate a lowered Expr at a value ---------- + +namespace { + + uint64_t EvalExprAt(const Expr &expr, uint64_t input, uint32_t bitwidth) { + auto compiled = CompileExpr(expr, bitwidth); + std::vector< uint64_t > vals = { input }; + std::vector< uint64_t > stack(compiled.stack_size); + return EvalCompiledExpr(compiled, vals, stack); + } + +} // anonymous namespace + +// ---------- LowerZeroExtend ---------- + +TEST(LowerZeroExtendTest, SixtyFourBitIsIdentity) { + auto inner = Expr::Variable(0); + auto *raw = inner.get(); + auto result = LowerZeroExtend(std::move(inner), 64); + // Must return the same pointer (no wrapping). + EXPECT_EQ(result.get(), raw); +} + +TEST(LowerZeroExtendTest, EightBitStructure) { + auto result = LowerZeroExtend(Expr::Variable(0), 8); + // Should be And(Variable(0), Constant(0xFF)). + ASSERT_EQ(result->kind, Expr::Kind::kAnd); + ASSERT_EQ(result->children.size(), 2u); + EXPECT_EQ(result->children[0]->kind, Expr::Kind::kVariable); + EXPECT_EQ(result->children[1]->kind, Expr::Kind::kConstant); + EXPECT_EQ(result->children[1]->constant_val, 0xFFu); +} + +TEST(LowerZeroExtendTest, SemanticOneBit) { + auto lowered = LowerZeroExtend(Expr::Variable(0), 1); + // zext i1 at 64-bit width. + EXPECT_EQ(EvalExprAt(*lowered, 0, 64), EvalZeroExtend(0, 1, UINT64_MAX)); + EXPECT_EQ(EvalExprAt(*lowered, 1, 64), EvalZeroExtend(1, 1, UINT64_MAX)); + EXPECT_EQ(EvalExprAt(*lowered, 0xFF, 64), EvalZeroExtend(0xFF, 1, UINT64_MAX)); +} + +TEST(LowerZeroExtendTest, SemanticEightBit) { + auto lowered = LowerZeroExtend(Expr::Variable(0), 8); + for (uint64_t v : { 0ULL, 0x7FULL, 0x80ULL, 0xFFULL, 0xDEAD00FFULL }) { + EXPECT_EQ(EvalExprAt(*lowered, v, 64), EvalZeroExtend(v, 8, UINT64_MAX)); + } +} + +TEST(LowerZeroExtendTest, SemanticThirtyTwoBit) { + auto lowered = LowerZeroExtend(Expr::Variable(0), 32); + for (uint64_t v : + { 0ULL, 0x7FFFFFFFULL, 0x80000000ULL, 0xFFFFFFFFULL, 0xDEADBEEF12345678ULL }) + { + EXPECT_EQ(EvalExprAt(*lowered, v, 64), EvalZeroExtend(v, 32, UINT64_MAX)); + } +} + +// ---------- LowerSignExtend ---------- + +TEST(LowerSignExtendTest, SixtyFourBitIsIdentity) { + auto inner = Expr::Variable(0); + auto *raw = inner.get(); + auto result = LowerSignExtend(std::move(inner), 64); + EXPECT_EQ(result.get(), raw); +} + +TEST(LowerSignExtendTest, EightBitStructure) { + auto result = LowerSignExtend(Expr::Variable(0), 8); + // Should be Add(Xor(And(Variable(0), 0xFF), 0x80), Neg(0x80)). + ASSERT_EQ(result->kind, Expr::Kind::kAdd); + ASSERT_EQ(result->children.size(), 2u); + // Left child: Xor(And(...), sign_bit) + EXPECT_EQ(result->children[0]->kind, Expr::Kind::kXor); + // Right child: Neg(sign_bit) + EXPECT_EQ(result->children[1]->kind, Expr::Kind::kNeg); +} + +TEST(LowerSignExtendTest, SemanticOneBit) { + auto lowered = LowerSignExtend(Expr::Variable(0), 1); + EXPECT_EQ(EvalExprAt(*lowered, 0, 64), EvalSignExtend(0, 1, UINT64_MAX)); + // sext i1 1 -> all ones. + EXPECT_EQ(EvalExprAt(*lowered, 1, 64), EvalSignExtend(1, 1, UINT64_MAX)); + EXPECT_EQ(EvalExprAt(*lowered, 1, 64), UINT64_MAX); +} + +TEST(LowerSignExtendTest, SemanticEightBit) { + auto lowered = LowerSignExtend(Expr::Variable(0), 8); + for (uint64_t v : { 0ULL, 0x7FULL, 0x80ULL, 0xFFULL }) { + EXPECT_EQ(EvalExprAt(*lowered, v, 64), EvalSignExtend(v, 8, UINT64_MAX)); + } +} + +TEST(LowerSignExtendTest, SemanticThirtyTwoBit) { + auto lowered = LowerSignExtend(Expr::Variable(0), 32); + for (uint64_t v : { 0ULL, 0x7FFFFFFFULL, 0x80000000ULL, 0xFFFFFFFFULL }) { + EXPECT_EQ(EvalExprAt(*lowered, v, 64), EvalSignExtend(v, 32, UINT64_MAX)); + } +} + +#if defined(GTEST_HAS_DEATH_TEST) && !defined(NDEBUG) +TEST(LowerZeroExtendDeathTest, ZeroBitsAsserts) { + EXPECT_DEATH(LowerZeroExtend(Expr::Variable(0), 0), ""); +} + +TEST(LowerZeroExtendDeathTest, OverSixtyFourBitsAsserts) { + EXPECT_DEATH(LowerZeroExtend(Expr::Variable(0), 65), ""); +} + +TEST(LowerSignExtendDeathTest, ZeroBitsAsserts) { + EXPECT_DEATH(LowerSignExtend(Expr::Variable(0), 0), ""); +} + +TEST(LowerSignExtendDeathTest, OverSixtyFourBitsAsserts) { + EXPECT_DEATH(LowerSignExtend(Expr::Variable(0), 65), ""); +} +#endif diff --git a/test/idapro/test_ida_headless.py b/test/idapro/test_ida_headless.py new file mode 100755 index 0000000..6328ec5 --- /dev/null +++ b/test/idapro/test_ida_headless.py @@ -0,0 +1,88 @@ +#!/usr/bin/env bash +# /// script +# requires-python = ">=3.12" +# dependencies = ["idapro"] +# /// +''''command -v uv >/dev/null 2>&1 && exec uv run "$0" "$@"; exec nix shell nixpkgs#uv --command uv run "$0" "$@" #''' +"""Headless CoBRA runner using idalib. + +Decompiles every function in the given IDA database and prints the ones +where CoBRA simplified an MBA expression. +""" +import argparse +import sys +from pathlib import Path +import idapro # noqa: E402 + +COBRA_TAG = ord('C') + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Decompile functions from an IDA database with CoBRA loaded." + ) + parser.add_argument("binary", type=Path, help="Path to binary to analyze") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + + idapro.enable_console_messages(True) + + rc = idapro.open_database(str(args.binary), True) + if rc != 0: + print(f"ERROR: open_database returned {rc}") + return 1 + + import ida_auto # noqa: E402 + import ida_funcs # noqa: E402 + import ida_hexrays # noqa: E402 + import ida_netnode # noqa: E402 + import ida_segment # noqa: E402 + import idautils # noqa: E402 + + ida_auto.auto_wait() + + if not ida_hexrays.init_hexrays_plugin(): + print("ERROR: Hex-Rays not available") + return 1 + + for ea in idautils.Functions(): + func = ida_funcs.get_func(ea) + if func is None: + continue + if func.flags & (ida_funcs.FUNC_THUNK | ida_funcs.FUNC_LIB | ida_funcs.FUNC_NORET): + continue + seg = ida_segment.getseg(ea) + if seg and (seg.type == ida_segment.SEG_XTRN + or ida_segment.get_segm_name(seg) in (".plt", ".plt.sec", ".plt.got")): + continue + + name = ida_funcs.get_func_name(ea) + hf = ida_hexrays.hexrays_failure_t() + try: + cfunc = ida_hexrays.decompile(ea, hf) + except Exception: + cfunc = None + if cfunc is None: + continue + + # Check the netnode tag set by CoBRA when it simplifies a function. + n = ida_netnode.netnode(ea) + simplified = n.altval(0, COBRA_TAG) + + if not simplified: + continue + + print(f"--- {name} @ {ea:#x} ---") + print(f" [CoBRA: simplified {simplified} expression(s)]") + print(str(cfunc).strip()) + print() + + idapro.close_database(False) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/test/llvm/test_ext_negative.ll b/test/llvm/test_ext_negative.ll new file mode 100644 index 0000000..a44c460 --- /dev/null +++ b/test/llvm/test_ext_negative.ll @@ -0,0 +1,25 @@ +; RUN: opt -load-pass-plugin=%cobra_pass -passes=cobra-simplify -S %s | FileCheck %s + +; Negative test: extension wrapping a non-MBA expression. +; zext of a single variable has no boolean+arithmetic mix, +; so it must NOT be detected as an MBA candidate. +; CHECK-LABEL: @test_ext_only_not_mba +; CHECK: zext +; CHECK: ret i32 +define i32 @test_ext_only_not_mba(i8 %x) { +entry: + %z = zext i8 %x to i32 + ret i32 %z +} + +; Negative test: extension chain with no core MBA ops. +; CHECK-LABEL: @test_ext_chain_not_mba +; CHECK: sext +; CHECK: zext +; CHECK: ret i32 +define i32 @test_ext_chain_not_mba(i1 %x) { +entry: + %s = sext i1 %x to i8 + %z = zext i8 %s to i32 + ret i32 %z +} diff --git a/test/llvm/test_ext_nested.ll b/test/llvm/test_ext_nested.ll new file mode 100644 index 0000000..9b0875e --- /dev/null +++ b/test/llvm/test_ext_nested.ll @@ -0,0 +1,22 @@ +; RUN: opt -load-pass-plugin=%cobra_pass -passes=cobra-simplify -S %s | FileCheck %s + +; Smoke test: nested zext(sext(i1 -> i8) -> i32) inside MBA. +; Verifies traversal through stacked extensions doesn't crash. +; The double extension (i1 -> i8 sext then i8 -> i32 zext) +; produces masking that the simplifier cannot see through. +; CHECK-LABEL: @test_ext_nested +; CHECK: sext +; CHECK: zext +; CHECK: ret i32 +define i32 @test_ext_nested(i1 %x, i1 %y) { +entry: + %xor1 = xor i1 %x, %y + %s = sext i1 %xor1 to i8 + %z = zext i8 %s to i32 + %and1 = and i1 %x, %y + %s2 = sext i1 %and1 to i8 + %z2 = zext i8 %s2 to i32 + %mul = mul i32 %z2, 2 + %add = add i32 %z, %mul + ret i32 %add +} diff --git a/test/llvm/test_ext_semantic.ll b/test/llvm/test_ext_semantic.ll new file mode 100644 index 0000000..41525e5 --- /dev/null +++ b/test/llvm/test_ext_semantic.ll @@ -0,0 +1,21 @@ +; RUN: opt -load-pass-plugin=%cobra_pass -passes=cobra-simplify -S %s | FileCheck %s + +; Regression test: extension semantics must not be transparent. +; sext(i1 x) + zext(i1 x) + 1 = 1 for all x when extensions +; are lowered correctly: +; x=0: sext(0)=0, zext(0)=0 → 0+0+1 = 1 +; x=1: sext(1)=-1, zext(1)=1 → -1+1+1 = 1 +; With transparent (wrong) semantics both sext and zext would +; produce the identity, giving 2x+1 — not constant. +; CHECK-LABEL: @test_sext_zext_constant +; CHECK-NOT: sext +; CHECK-NOT: zext +; CHECK: ret i32 1 +define i32 @test_sext_zext_constant(i1 %x) { +entry: + %s = sext i1 %x to i32 + %z = zext i1 %x to i32 + %add1 = add i32 %s, %z + %add2 = add i32 %add1, 1 + ret i32 %add2 +} diff --git a/test/llvm/test_sext.ll b/test/llvm/test_sext.ll new file mode 100644 index 0000000..85d716d --- /dev/null +++ b/test/llvm/test_sext.ll @@ -0,0 +1,19 @@ +; RUN: opt -load-pass-plugin=%cobra_pass -passes=cobra-simplify -S %s | FileCheck %s + +; Smoke test: sext in MBA trees doesn't crash. +; The sign-extension lowering introduces masking and offset +; arithmetic that obscures the inner 8-bit MBA identity. +; The expression should pass through unchanged. +; CHECK-LABEL: @test_sext_xor_and +; CHECK: sext +; CHECK: ret i32 +define i32 @test_sext_xor_and(i8 %x, i8 %y) { +entry: + %xor = xor i8 %x, %y + %and = and i8 %x, %y + %sx = sext i8 %xor to i32 + %sa = sext i8 %and to i32 + %mul = mul i32 %sa, 2 + %add = add i32 %sx, %mul + ret i32 %add +} diff --git a/test/llvm/test_sext_i1.ll b/test/llvm/test_sext_i1.ll new file mode 100644 index 0000000..02e455a --- /dev/null +++ b/test/llvm/test_sext_i1.ll @@ -0,0 +1,16 @@ +; RUN: opt -load-pass-plugin=%cobra_pass -passes=cobra-simplify -S %s | FileCheck %s + +; Test: sext i1 in an MBA tree. +; sext(i1 (x & y)) produces all-ones or zero — this must not +; be treated as identity. +; CHECK-LABEL: @test_sext_i1 +; CHECK-NOT: sext +; CHECK: ret i32 +define i32 @test_sext_i1(i1 %x, i1 %y) { +entry: + %and = and i1 %x, %y + %s = sext i1 %and to i32 + %xor = xor i32 %s, -1 + %add = add i32 %s, %xor + ret i32 %add +} diff --git a/test/llvm/test_zext.ll b/test/llvm/test_zext.ll new file mode 100644 index 0000000..2ed4e2e --- /dev/null +++ b/test/llvm/test_zext.ll @@ -0,0 +1,19 @@ +; RUN: opt -load-pass-plugin=%cobra_pass -passes=cobra-simplify -S %s | FileCheck %s + +; Smoke test: zext in MBA trees doesn't crash. +; The extension lowering masks the inner 8-bit MBA to 32 bits, +; so the simplifier cannot currently recognize the cross-width +; identity. The expression should pass through unchanged. +; CHECK-LABEL: @test_zext_xor_and +; CHECK: zext +; CHECK: ret i32 +define i32 @test_zext_xor_and(i8 %x, i8 %y) { +entry: + %xor = xor i8 %x, %y + %and = and i8 %x, %y + %zx = zext i8 %xor to i32 + %za = zext i8 %and to i32 + %mul = mul i32 %za, 2 + %add = add i32 %zx, %mul + ret i32 %add +} diff --git a/test/llvm/test_zext_simplifies.ll b/test/llvm/test_zext_simplifies.ll new file mode 100644 index 0000000..d90150c --- /dev/null +++ b/test/llvm/test_zext_simplifies.ll @@ -0,0 +1,24 @@ +; RUN: opt -load-pass-plugin=%cobra_pass -passes=cobra-simplify -S %s | FileCheck %s + +; Positive test: extension lowering produces correct semantics. +; zext(x^y) - zext(x^y) + zext(x&y) - zext(x&y) = 0 +; Both zext arms cancel, so the full-width evaluator confirms +; constant 0 regardless of input width mismatch. +; CHECK-LABEL: @test_zext_cancel +; CHECK-NOT: zext +; CHECK-NOT: xor +; CHECK-NOT: and +; CHECK: ret i64 0 +define i64 @test_zext_cancel(i8 %x, i8 %y) { +entry: + %xor = xor i8 %x, %y + %and = and i8 %x, %y + %z1 = zext i8 %xor to i64 + %z2 = zext i8 %xor to i64 + %z3 = zext i8 %and to i64 + %z4 = zext i8 %and to i64 + %sub1 = sub i64 %z1, %z2 + %sub2 = sub i64 %z3, %z4 + %add = add i64 %sub1, %sub2 + ret i64 %add +}