#include "internal.h"

static khook_stub_t *khook_stub_tbl = NULL;

////////////////////////////////////////////////////////////////////////////////
/*
 * 内核符号表查找函数callback function
 * 通过遍历data数组中的元素，并与name字符串进行比较
 * 如果找到匹配的元素，则将data数组的第二个元素设置为addr，并返回1
 * 如果遍历完整个数组仍然没有找到匹配的元素，则返回0
 */
static int khook_lookup_cb(long data[], const char *name, void *module, long addr)
{
    int i = 0; 
    while (!module && (((const char *)data[0]))[i] == name[i]) {
        if (!name[i++]) return !!(data[1] = addr); // 找到匹配的符号，设置地址并返回1
    } 
    return 0; // 未找到匹配的符号，返回0
}

/*
 * 利用kallsyms_on_each_symbol可以查询符号表，只需要传入查询函数就可以了
 * data[0]表示要查询的地址
 * data[1]表示结果
 * kernel all symbols 内核符号表
 */
static void *khook_lookup_name(const char *name)
{
    long data[2] = { (long)name, 0 };
    kallsyms_on_each_symbol((void *)khook_lookup_cb, data);
    return (void *)data[1]; // 返回找到的符号地址
}

/*
 * 将一个虚拟地址范围[addr, addr + len]映射到 可写 内核内存区域
 */
static void *khook_map_writable(void *addr, size_t len)
{
    struct page *pages[2] = { 0 }; // len << PAGE_SIZE
    long page_offset = offset_in_page(addr);
    int i, nb_pages = DIV_ROUND_UP(page_offset + len, PAGE_SIZE);

    addr = (void *)((long)addr & PAGE_MASK); // 对齐地址到页面边界
    for (i = 0; i < nb_pages; i++, addr += PAGE_SIZE) {
        if ((pages[i] = is_vmalloc_addr(addr) ?
             vmalloc_to_page(addr) : virt_to_page(addr)) == NULL)
            return NULL; // 如果无法获取页面，返回NULL
    }

    addr = vmap(pages, nb_pages, VM_MAP, PAGE_KERNEL); // 映射页面到内核地址空间
    return addr ? addr + page_offset : NULL; // 返回映射后的地址
}

////////////////////////////////////////////////////////////////////////////////

#ifdef CONFIG_X86
# include "x86/hook.c" // 包含x86架构相关的钩子实现
#else
# error Target CPU architecture is NOT supported !!! // 不支持的CPU架构
#endif

////////////////////////////////////////////////////////////////////////////////

/*
 * 唤醒所有进程
 */
static void khook_wakeup(void)
{
    struct task_struct *p;
    rcu_read_lock();
    for_each_process(p) {
        wake_up_process(p); // 唤醒进程
    }
    rcu_read_unlock();
}

/*
 * 初始化所有钩子
 */
static int khook_sm_init_hooks(void *arg)
{
    khook_t *p;
    KHOOK_FOREACH_HOOK(p) {
        if (!p->target.addr_map) continue;
        khook_arch_sm_init_one(p); // 初始化单个钩子
    }
    return 0;
}

/*
 * 清理所有钩子
 */
static int khook_sm_cleanup_hooks(void *arg)
{
    khook_t *p;
    KHOOK_FOREACH_HOOK(p) {
        if (!p->target.addr_map) continue;
        khook_arch_sm_cleanup_one(p); // 清理单个钩子
    }
    return 0;
}

/*
 * 解析所有钩子的目标地址
 */
static void khook_resolve(void)
{
    khook_t *p;
    KHOOK_FOREACH_HOOK(p) {
        p->target.addr = khook_lookup_name(p->target.name); // 查找符号地址
    }
}

/*
 * 映射所有钩子的目标地址到可写内存区域
 */
static void khook_map(void)
{
    khook_t *p;
    KHOOK_FOREACH_HOOK(p) {
        if (!p->target.addr) continue;
        p->target.addr_map = khook_map_writable(p->target.addr, 32); // 映射地址
        khook_debug("target %s@%p -> %p\n", p->target.name, p->target.addr, p->target.addr_map);
    }
}

/*
 * 取消映射所有钩子的目标地址
 */
static void khook_unmap(int wait)
{
    khook_t *p;
    KHOOK_FOREACH_HOOK(p) {
        khook_stub_t *stub = KHOOK_STUB(p);
        if (!p->target.addr_map) continue;
        while (wait && atomic_read(&stub->use_count) > 0) {
            khook_wakeup(); // 唤醒进程
            msleep_interruptible(1000); // 休眠1秒
            khook_debug("waiting for %s...\n", p->target.name);
        }
        vunmap((void *)((long)p->target.addr_map & PAGE_MASK)); // 取消映射
        p->target.addr_map = NULL;
    }
}

////////////////////////////////////////////////////////////////////////////////
/* khook_init()和khook_cleanup()对挂钩引擎进行初始化和注销 */

/*
 * 初始化挂钩引擎
 * 1. malloc分配内存空间用于存储钩子函数的桩（stub）
 * 2. khook_resolve查找并解析内核符号表
 * 3. 映射目标地址到可写的内核内存区域
 * 4. 调用stop_machine函数，将钩子函数的初始化工作交给内核调度器执行
*/
int khook_init(void)
{
    void *(*malloc)(long size) = NULL;
    int   (*set_memory_x)(unsigned long, int) = NULL;

    malloc = khook_lookup_name("module_alloc");
    if (!malloc || KHOOK_ARCH_INIT()) return -EINVAL;

    khook_stub_tbl = malloc(KHOOK_STUB_TBL_SIZE);
    if (!khook_stub_tbl) return -ENOMEM;
    memset(khook_stub_tbl, 0, KHOOK_STUB_TBL_SIZE);

    //
    // Since some point memory allocated by module_alloc() doesn't
    // have eXecutable attributes. That's why we have to mark the
    // region executable explicitly.
    //

    set_memory_x = khook_lookup_name("set_memory_x");
    if (set_memory_x) {
        int numpages = round_up(KHOOK_STUB_TBL_SIZE, PAGE_SIZE) / PAGE_SIZE;
        set_memory_x((unsigned long)khook_stub_tbl, numpages);
    }

    khook_resolve(); // 解析符号表

    khook_map(); // 映射地址
    stop_machine(khook_sm_init_hooks, NULL, NULL); // 初始化钩子
    khook_unmap(0); // 取消映射

    return 0;
}

/*
 * 注销挂钩引擎
 * 1. 映射目标地址到可写的内核内存区域
 * 2. 调用stop_machine函数，将钩子函数的清理工作交给内核调度器执行
 * 3. 取消映射目标地址
 * 4. 释放分配的内存空间
 */
void khook_cleanup(void)
{
    khook_map(); // 映射地址
    stop_machine(khook_sm_cleanup_hooks, NULL, NULL); // 清理钩子
    khook_unmap(1); // 取消映射
    vfree(khook_stub_tbl); // 释放内存
}