ucc.h

/*
    By JayWaves <jay.waves@qq.com>, 2025
    Class for Uc Engine
*/

#pragma once

#include "unicorn/unicorn.h"
#include "unicorn/sparc.h"
#include "uc_pub.h"

#include <thread>
#include <stop_token>
#include <atomic>

#include "include/prelude.h"


class ucc_error : public std::runtime_error {
public:
    ucc_error(const char* msg, uc_err err)
        : std::runtime_error(uc_strerror(err)), err_(err) {

        LOG_ERROR("uc_error: {}: {}", msg, uc_strerror(err));
    }

    ucc_error(uc_err err)
        : std::runtime_error(uc_strerror(err)), err_(err) {

        LOG_ERROR("uc_error: {}", uc_strerror(err));
    }

    uc_err code() const noexcept { return err_; }

private:
    uc_err err_;
};


/*
    unicorn_hook 由外部管理生命周期, 不能晚于 Unicorn 释放.
    Unicorn 直接用 lambda 从本地捕获即可.

    unicorn uc;
    uc_hook_intr h1([&](uint32_t intno) {
        uc.get_pc();
        ...
    });
    uc.hook_install(h1);
*/
class ucc_hook {
public:
    ucc_hook() = default;
    virtual ~ucc_hook() = default;

private:
    friend class ucc; // 只允许 unicorn 进行安装或卸载
    virtual void install(ucc* uc) = 0;
    virtual void uninstall(ucc* uc);

protected:
    // 默认是全地址空间 (begin > end) 插入钩子
    void hook_add(ucc* uc, int type, void* cb, void *self, 
                 uint64_t begin = 1, uint64_t end = 0);

    uc_hook hh_{0};
};


// tracing interrupts
class ucc_hook_intr : public ucc_hook {
public:
    using cb_t = std::function<void(uint32_t)>;

    explicit ucc_hook_intr(cb_t cb) : cb_(std::move(cb)) {}

    void install(ucc* uc) {
        ucc_hook::hook_add(uc, UC_HOOK_INTR, (void*)&cb_c, this);
    }

private:
    static void cb_c(uc_engine* /* uc */, uint32_t intno, void *ud) {
        auto* self = (ucc_hook_intr*)ud;
        if (self && self->cb_) self->cb_(intno);
    }

    cb_t cb_;
};

// new edges between tbs
class ucc_hook_egde_gen : public ucc_hook {
public:
    using cb_t = std::function<void(uc_tb*, uc_tb*)>;
    explicit ucc_hook_egde_gen(cb_t cb) : cb_(std::move(cb)) {}
    void install(ucc* uc) {
        ucc_hook::hook_add(uc, UC_HOOK_EDGE_GENERATED, (void*)&cb_c, this);
    }
private:
    static void cb_c(uc_engine* /* uc */, uc_tb* cur_tb, uc_tb* prev_tb, void *ud) {
        ucc_hook_egde_gen* self = (ucc_hook_egde_gen*)ud;
        if (self && self->cb_) 
            self->cb_(cur_tb, prev_tb);
    }
    cb_t cb_;
};

// 这个是 Unicorn 推荐的 hookmem 用法. 但需要手动获取所有的 uc_mem_regions
class ucc_hook_tlbevent : public ucc_hook {
public:
    using cb_t = std::function<bool(uint64_t, uc_mem_type, uc_tlb_entry*)>;
    explicit ucc_hook_tlbevent(cb_t cb) : cb_(std::move(cb)) {}
    /*
        vtlb 类型, 如果不指定 hook, 默认 vaddr == paddr. 
        详见 unicorn/qemu/softmmu/unicorn_vtlb.c

        使用该 hook, 需要开启 UC_TLB_VIRTUAL 模式. 
    */
    void install(ucc* uc) {
        ucc_hook::hook_add(uc, UC_HOOK_TLB_FILL, (void*)&cb_c, this);
    }
private:
    cb_t cb_;

    static void cb_c(uc_engine* /*uc*/, uint64_t vaddr, uc_mem_type type,
                         uc_tlb_entry* result, void *ud) {
        ucc_hook_tlbevent* self = (ucc_hook_tlbevent*)ud;
        if (self && self->cb_) 
            self->cb_(vaddr, type, result);
    }
};

/*
    回调可能多次重复 (按字节), 不如 tlbevent, 但符合我们需求.
*/
class ucc_hook_mem_invalid : public ucc_hook {
public:
    using cb_t = std::function<bool(uc_mem_type, uint64_t, int, int64_t)>;
    explicit ucc_hook_mem_invalid(cb_t cb) : cb_(std::move(cb)) {}
    /*
        type: UC_HOOK_MEM_UNMAPPED + UC_HOOK_MEM_PROT
    */
    void install(ucc* uc) {
        ucc_hook::hook_add(uc, UC_HOOK_MEM_INVALID, (void*)&cb_c, this);
    }
private:
    cb_t cb_;

    static void cb_c(uc_engine* /*uc*/, uc_mem_type type, uint64_t address,
                     int size, int64_t value, void *ud) {
        ucc_hook_mem_invalid* self = (ucc_hook_mem_invalid*)ud;
        if (self && self->cb_) 
            self->cb_(type, address, size, value);
    }
};

/*
    MMIO 映射 + 回调. 这是对一段 MMIO 内存的完整建模
*/
class ucc_mmio {
public:
    using cb_t = std::function<void(uint64_t offset, unsigned sz, uint64_t* val)>;
    ucc_mmio(uint64_t addr, uint64_t sz, cb_t read_cb, cb_t write_cb)
        : addr_(addr), sz_(sz), 
          r_cb_(std::move(read_cb)), w_cb_(std::move(write_cb)) {}
private:
    cb_t r_cb_, w_cb_;
    uint64_t addr_, sz_;

     // 提供 MMIO 回调注册机制. 
    static uint64_t mmio_r_cb_c(uc_engine *uc, uint64_t o, unsigned s, void *ud); 
    static void mmio_w_cb_c(uc_engine *uc, uint64_t o, unsigned s, uint64_t v, void *ud);     

    friend class ucc;
    void install(ucc* uc); 
    void uninstall(ucc* uc); 
};


enum class ucc_mode : int { thumb, sparc32 };

/*
    不会封装 Unicorn 的所有类型, 要不会累死...
    这是一个 bare-metal MCU, 没有 TLB, 没有外设. 
*/
class ucc {
public:
    /*
        uc_open: 硬件复位. 
    */
    ucc(ucc_mode mode) {
        uc_err err;
        switch (mode) {
        case ucc_mode::thumb:
            err = uc_open(UC_ARCH_ARM, UC_MODE_THUMB, &uc_);
        case ucc_mode::sparc32:
            err = uc_open(UC_ARCH_SPARC, 
                static_cast<uc_mode>(UC_MODE_SPARC32 | UC_MODE_BIG_ENDIAN), &uc_);
        }
        if (err != UC_ERR_OK)  throw ucc_error(err);

        // 启用多点退出机制
        uc_ctl_exits_enable(uc_);

        // vaddr == paddr
        uc_ctl_tlb_mode(uc_, UC_TLB_VIRTUAL);

        // map_memory();
        // init_registers(); // entry point & sp
        LOG_INFO("Unicorn engine created.");
    }

    ~ucc() {
        emu_stop_async();
        if (uc_) uc_close(uc_);
    }

    void mem_write(uint64_t address, const void *data, uint64_t size);
    void mem_read(uint64_t address, void *data, uint64_t size); 
    void reg_write(int regid, const void *value);
    void reg_read(int regid, void *value); 
    // 阻塞执行
    void emu_start(uint64_t begin, uint64_t timeout = 0, size_t count = 0);
    // 后台持续运行
    void emu_start_async(uint64_t begin);
    
    /*  
        在 CPU TB 块结束后, 异步触发中断请求. 
        通常外部中断触发路径为: 外设 -> IRQ Line -> CPU -> cpu_interrupt() -> ...
        这里直接强制切换. 也没有 raise/lower 电平逻辑.
    */
    virtual void raise_irq(int irq_num); 

    void mem_map(uint64_t addr, uint64_t size, uint32_t perms); 

    void emu_stop_async(); 

    /*
        支持动态设置退出点
    */
    void set_emu_exits(uint64_t* exits, size_t len);

    uc_engine* c_ptr() { return uc_; }
    void hook_install(ucc_hook& hook) { hook.install(this); }
    void hook_uninstall(ucc_hook& hook) { hook.uninstall(this); }

    void mmio_install(ucc_mmio& mmio) { mmio.install(this); }

    /*
        打印出所有 CPU 内部状态. 供调试使用.
    */
    void coredump();

protected:
    uc_engine *uc_ = nullptr;

private:
    std::jthread runner_;

    // 用来防止 emu_start / emu_start_async 重入
    std::atomic<bool> running_{false};
    bool is_running() const { return running_.load(); }
    bool cas_running() {
        bool expected = false;
        return running_.compare_exchange_strong(expected, true);
    }
};