// Unit tests for scanGadgetsCapstone — the exhaustive overlapping gadget scan.
// Feeds hand-assembled x86-64 byte blobs (no r2/binary) and checks that clean
// ret/syscall-terminated gadgets are found, control-flow-terminated sequences
// are rejected, and the relevance filter works.
#include <catch2/catch_test_macros.hpp>
#include <cstdint>
#include <string>
#include <vector>

#include "fold/binary/gadget_source.hpp"

namespace {
using namespace fold::binary;

BinaryArchInfo x64() { return BinaryArchInfo{.arch = "x86", .bits = 64, .little = true}; }

bool hasText(const std::vector<RopGadgetRaw>& gs, const std::string& needle) {
    for (const auto& g : gs)
        if (g.text.find(needle) != std::string::npos)
            return true;
    return false;
}
const RopGadgetRaw* at(const std::vector<RopGadgetRaw>& gs, std::uint64_t addr) {
    for (const auto& g : gs)
        if (g.addr == addr)
            return &g;
    return nullptr;
}

// pop rdi;ret | pop rsi;ret | pop rax;ret | xor rsi,rsi;ret | mov [rdi],rsi;ret |
// pop rbx;ret | syscall   — laid out at known offsets from base 0x400000.
std::vector<std::uint8_t> sampleCode() {
    return {0x5f, 0xc3,             // 0:  pop rdi ; ret
            0x5e, 0xc3,             // 2:  pop rsi ; ret
            0x58, 0xc3,             // 4:  pop rax ; ret
            0x48, 0x31, 0xf6, 0xc3, // 6:  xor rsi, rsi ; ret
            0x48, 0x89, 0x37, 0xc3, // 10: mov [rdi], rsi ; ret
            0x5b, 0xc3,             // 14: pop rbx ; ret
            0x0f, 0x05};            // 16: syscall
}
} // namespace

TEST_CASE("capstone: finds clean pop;ret gadgets at exact addresses", "[capstone]") {
    auto code = sampleCode();
    RopGadgetSearchOptions opts; // goal_regs empty → keep all
    auto gs = scanGadgetsCapstone(x64(), code.data(), code.size(), 0x400000, opts);

    const RopGadgetRaw* pop_rdi = at(gs, 0x400000);
    REQUIRE(pop_rdi != nullptr);
    REQUIRE(pop_rdi->text.find("pop rdi") != std::string::npos);
    REQUIRE(pop_rdi->text.find("ret") != std::string::npos);
    REQUIRE(pop_rdi->size == 2);
    REQUIRE(at(gs, 0x400002) != nullptr); // pop rsi ; ret
    REQUIRE(at(gs, 0x400004) != nullptr); // pop rax ; ret
    REQUIRE(hasText(gs, "xor rsi, rsi")); // zeroing gadget
    REQUIRE(hasText(gs, "syscall"));      // syscall gadget
}

TEST_CASE("capstone: a write-what-where gadget (mov [reg], reg ; ret) is found", "[capstone]") {
    auto code = sampleCode();
    RopGadgetSearchOptions opts;
    auto gs = scanGadgetsCapstone(x64(), code.data(), code.size(), 0x400000, opts);
    const RopGadgetRaw* w = at(gs, 0x40000a);
    REQUIRE(w != nullptr);
    REQUIRE(w->text.find('[') != std::string::npos); // memory write
    REQUIRE(w->text.find("rdi") != std::string::npos);
}

TEST_CASE("capstone: a control-flow-terminated sequence is NOT a clean gadget", "[capstone]") {
    // pop rdi ; jmp rax  (5f ff e0) — the jmp aborts the gadget; no clean
    // ret-terminated gadget should start here.
    std::vector<std::uint8_t> code = {0x5f, 0xff, 0xe0};
    RopGadgetSearchOptions opts;
    auto gs = scanGadgetsCapstone(x64(), code.data(), code.size(), 0x400000, opts);
    REQUIRE(at(gs, 0x400000) == nullptr); // pop rdi ; jmp rax is rejected
}

TEST_CASE("capstone: relevance filter keeps only goal-register / mem / syscall gadgets",
          "[capstone]") {
    auto code = sampleCode();
    RopGadgetSearchOptions opts;
    opts.goal_regs = {"rdi"}; // only rdi is a goal
    auto gs = scanGadgetsCapstone(x64(), code.data(), code.size(), 0x400000, opts);
    REQUIRE(at(gs, 0x400000) != nullptr);  // pop rdi ; ret kept
    REQUIRE(hasText(gs, "[rdi]"));         // write gadget kept (memory access)
    REQUIRE(hasText(gs, "syscall"));       // syscall kept (want_syscall default)
    REQUIRE_FALSE(hasText(gs, "pop rbx")); // rbx not a goal → filtered out
}