#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/version.h>
#include <linux/kallsyms.h>
#include <linux/kprobes.h>
#include <linux/mm.h>
#include <linux/slab.h>

#ifdef CONFIG_XEN
#include <xen/xen.h>
#include <xen/interface/xen.h>
#include <asm/xen/hypercall.h>
#endif

MODULE_LICENSE("GPL");
MODULE_AUTHOR("turnstile");
MODULE_DESCRIPTION("turnstile guest module for livepatch support");

#define TURNSTILE_HYPERCALL_NR    50
#define TURNSTILE_OP_REQUEST_WRITE 3

struct turnstile_op_request_write {
    uint64_t gpa_start;
    uint64_t length;
    uint32_t timeout_ms;
    uint32_t _pad;
};

static unsigned long kernel_text_start;
static unsigned long kernel_text_end;
static unsigned long kernel_rodata_start;
static unsigned long kernel_rodata_end;

static bool turnstile_enabled;

#ifdef CONFIG_KPROBES
static unsigned long kprobe_get_symbol(const char *name)
{
    struct kprobe kp = { .symbol_name = name };
    unsigned long addr = 0;

    if (register_kprobe(&kp) == 0) {
        addr = (unsigned long)kp.addr;
        unregister_kprobe(&kp);
    }
    return addr;
}
#endif

static unsigned long resolve_symbol(const char *name)
{
#ifdef CONFIG_KALLSYMS
#if LINUX_VERSION_CODE >= KERNEL_VERSION(5, 7, 0)
#ifdef CONFIG_KPROBES
    return kprobe_get_symbol(name);
#else
    return 0;
#endif
#else
    return kallsyms_lookup_name(name);
#endif
#else
    return 0;
#endif
}

static bool addr_in_protected_range(unsigned long addr)
{
    return (addr >= kernel_text_start && addr < kernel_text_end) ||
           (addr >= kernel_rodata_start && addr < kernel_rodata_end);
}

static unsigned long virt_to_gpa(unsigned long vaddr)
{
    if (vaddr < PAGE_OFFSET)
        return 0;

#ifdef CONFIG_X86_64
    return __pa(vaddr);
#else
    return vaddr - PAGE_OFFSET;
#endif
}

#ifdef CONFIG_XEN
static int turnstile_hypercall(unsigned int op, void *arg)
{
    return _hypercall3(long, TURNSTILE_HYPERCALL_NR, op, DOMID_SELF, arg);
}
#else
static int turnstile_hypercall(unsigned int op, void *arg)
{
    return -ENOSYS;
}
#endif

static int request_write_access(unsigned long addr, unsigned long len)
{
    struct turnstile_op_request_write op;
    unsigned long gpa;

    if (!turnstile_enabled)
        return 0;

    gpa = virt_to_gpa(addr);
    if (!gpa) {
        pr_warn("turnstile: invalid address %pK\n", (void *)addr);
        return -EINVAL;
    }

    op.gpa_start = gpa;
    op.length = len;
    op.timeout_ms = 100;
    op._pad = 0;

    return turnstile_hypercall(TURNSTILE_OP_REQUEST_WRITE, &op);
}

int turnstile_request_patch_access(unsigned long addr, unsigned long len)
{
    int ret;

    if (!addr_in_protected_range(addr))
        return 0;

    pr_info("turnstile: requesting write access for %pK len %lu\n", (void *)addr, len);

    ret = request_write_access(addr, len);
    if (ret < 0)
        pr_warn("turnstile: write access request failed: %d\n", ret);

    return ret;
}
EXPORT_SYMBOL_GPL(turnstile_request_patch_access);

int turnstile_request_patch_access_range(unsigned long start, unsigned long end)
{
    if (end <= start)
        return -EINVAL;

    return turnstile_request_patch_access(start, end - start);
}
EXPORT_SYMBOL_GPL(turnstile_request_patch_access_range);

static int resolve_kernel_symbols(void)
{
    unsigned long addr;

    addr = resolve_symbol("_stext");
    if (!addr) {
        pr_warn("turnstile: failed to resolve _stext\n");
        return -ENOENT;
    }
    kernel_text_start = addr;

    addr = resolve_symbol("_etext");
    if (!addr) {
        pr_warn("turnstile: failed to resolve _etext\n");
        return -ENOENT;
    }
    kernel_text_end = addr;

    addr = resolve_symbol("__start_rodata");
    if (!addr) {
        pr_warn("turnstile: __start_rodata not found, using _etext\n");
        kernel_rodata_start = kernel_text_end;
    } else {
        kernel_rodata_start = addr;
    }

    addr = resolve_symbol("__end_rodata");
    if (!addr) {
        addr = resolve_symbol("__init_begin");
        if (!addr) {
            pr_warn("turnstile: failed to resolve rodata end\n");
            kernel_rodata_end = kernel_rodata_start;
        } else {
            kernel_rodata_end = addr;
        }
    } else {
        kernel_rodata_end = addr;
    }

    return 0;
}

static int __init turnstile_guest_init(void)
{
    int ret;

#ifdef CONFIG_XEN
    if (!xen_domain()) {
        pr_info("turnstile: not running on xen, module inactive\n");
        turnstile_enabled = false;
        return 0;
    }

    if (!xen_pv_domain() && !xen_hvm_domain()) {
        pr_info("turnstile: unsupported xen domain type\n");
        turnstile_enabled = false;
        return 0;
    }
#else
    pr_info("turnstile: compiled without xen support, module inactive\n");
    turnstile_enabled = false;
    return 0;
#endif

    ret = resolve_kernel_symbols();
    if (ret) {
        pr_warn("turnstile: symbol resolution failed, protection ranges unknown\n");
        kernel_text_start = 0;
        kernel_text_end = 0;
        kernel_rodata_start = 0;
        kernel_rodata_end = 0;
    }

    turnstile_enabled = true;

    pr_info("turnstile: guest module loaded\n");
    pr_info("turnstile: text range: %pK - %pK\n",
            (void *)kernel_text_start, (void *)kernel_text_end);
    pr_info("turnstile: rodata range: %pK - %pK\n",
            (void *)kernel_rodata_start, (void *)kernel_rodata_end);

    return 0;
}

static void __exit turnstile_guest_exit(void)
{
    turnstile_enabled = false;
    pr_info("turnstile: guest module unloaded\n");
}

module_init(turnstile_guest_init);
module_exit(turnstile_guest_exit);