#include <linux/module.h>
#include <linux/kernel.h>
#include <linux/init.h>
#include <linux/debugfs.h>
#include <linux/seq_file.h>
#include <linux/fs.h>
#include <linux/kprobes.h>
#include <linux/uaccess.h>
#include <drm/drm_buddy.h>

/* THIS IS THE MAKEFILE
obj-m += vramgaze_buddy.o

all:
	make -C /lib/modules/$(shell uname -r)/build M=$(PWD) modules

clean:
	make -C /lib/modules/$(shell uname -r)/build M=$(PWD) clean
*/

/* This interface reimplements drm_mm-style VRAM dumps for amdgpu buddy allocator (Linux 6.x)
$ sudo bash -c 'cat /sys/kernel/debug/vramgaze_buddy_nodes'
*/

static struct dentry *vramgaze_debugfs_file;
static struct drm_buddy *global_mm = NULL;

/* ---- kprobe on drm_buddy_print to sniff the drm_buddy pointer ---------- */

static struct kprobe kp_buddy_print;

static int buddy_print_pre(struct kprobe *p, struct pt_regs *regs)
{
#ifdef CONFIG_X86_64
    if (!global_mm && regs->di)
        global_mm = (struct drm_buddy *)(uintptr_t)regs->di;
#elif defined(CONFIG_ARM64)
    if (!global_mm && regs->regs[0])
        global_mm = (struct drm_buddy *)(uintptr_t)regs->regs[0];
#endif
    return 0;
}

/*
 * Trigger drm_buddy_print by opening the amdgpu_vram_mm seq_file and
 * directly invoking its show op ? no VFS read, no kernel_read(), no
 * usermodehelper. We just need the seq_show to run once.
 *
 * seq_file internals (include/linux/seq_file.h):
 *   struct seq_file { ... struct seq_operations *op; ... }
 *   op->show(seq_file *, void *) is what we want.
 *
 * filp_open() succeeds for root-context kernel code on debugfs.
 * The blocked path was kernel_read() / vfs_read(), NOT filp_open().
 * We open, grab the seq_file, call show(m, NULL) directly.
 */
static int trigger_via_seq_show(void)
{
    struct file *f;
    struct seq_file *m;
    char path[64];
    int i, ret = -ENODEV;

    for (i = 0; i < 16; i++) {
        snprintf(path, sizeof(path),
                 "/sys/kernel/debug/dri/%d/amdgpu_vram_mm", i);

        f = filp_open(path, O_RDONLY, 0);
        if (IS_ERR(f))
            continue;

        /*
         * For a seq_file, private_data is the struct seq_file *.
         * Cast and call show directly ? this runs the amdgpu seq_show
         * which calls drm_buddy_print, firing our kprobe.
         */
        m = f->private_data;
        if (m && m->op && m->op->show) {
            m->op->show(m, NULL);
            if (global_mm) {
                pr_info("vramgaze_buddy: captured drm_buddy* from dri/%d\n", i);
                ret = 0;
            }
        }

        filp_close(f, NULL);

        if (ret == 0)
            break;
    }

    return ret;
}

/* ---- Tree walker ------------------------------------------------------- */

static void dump_buddy_block(struct seq_file *m, struct drm_buddy_block *block)
{
    struct drm_buddy_block safe;

    if (!block)
        return;
    if (copy_from_kernel_nofault(&safe, block, sizeof(safe)))
        return;

    if (drm_buddy_block_is_split(&safe)) {
        dump_buddy_block(m, safe.left);
        dump_buddy_block(m, safe.right);
    } else {
        u64 offset     = drm_buddy_block_offset(&safe);
        u64 size       = 1ULL << (drm_buddy_block_order(&safe) + PAGE_SHIFT);
        bool allocated = drm_buddy_block_is_allocated(&safe);

        seq_printf(m, "node 0x%012llx: %llu bytes (%llu pages), %s\n",
                   offset, size, size >> PAGE_SHIFT,
                   allocated ? "used" : "free");
    }
}

static int vramgaze_show(struct seq_file *m, void *v)
{
    struct drm_buddy safe_mm;
    unsigned int i;

    if (!global_mm) {
        seq_puts(m, "drm_buddy not yet captured.\n"
                    "Run: sudo cat /sys/kernel/debug/dri/1/amdgpu_vram_mm\n"
                    "Then re-read this file.\n");
        return 0;
    }

    if (copy_from_kernel_nofault(&safe_mm, global_mm, sizeof(safe_mm))) {
        seq_puts(m, "Failed to read drm_buddy structure.\n");
        return 0;
    }

    seq_printf(m, "drm_buddy @ %px  size=%llu  chunk=%llu  n_roots=%u\n",
               global_mm, safe_mm.size, safe_mm.chunk_size, safe_mm.n_roots);

    for (i = 0; i < safe_mm.n_roots; i++) {
        void *root_ptr = NULL;
        if (copy_from_kernel_nofault(&root_ptr,
                                     &safe_mm.roots[i], sizeof(void *)))
            continue;
        seq_printf(m, "--- root[%u] ---\n", i);
        dump_buddy_block(m, root_ptr);
    }

    return 0;
}

static int vramgaze_open(struct inode *inode, struct file *file)
{
    return single_open(file, vramgaze_show, NULL);
}

static const struct file_operations vramgaze_fops = {
    .owner   = THIS_MODULE,
    .open    = vramgaze_open,
    .read    = seq_read,
    .llseek  = seq_lseek,
    .release = single_release,
};

/* ---- Module init / exit ------------------------------------------------ */

static int __init vramgaze_buddy_init(void)
{
    int err;

    /* Register kprobe first so it's live before we trigger show() */
    memset(&kp_buddy_print, 0, sizeof(kp_buddy_print));
    kp_buddy_print.symbol_name = "drm_buddy_print";
    kp_buddy_print.pre_handler = buddy_print_pre;

    err = register_kprobe(&kp_buddy_print);
    if (err) {
        pr_err("vramgaze_buddy: kprobe on drm_buddy_print failed (%d)\n", err);
        return err;
    }

    /* Now trigger the seq_show directly ? no VFS read path involved */
    err = trigger_via_seq_show();
    if (err) {
        pr_warn("vramgaze_buddy: seq_show trigger failed; "
                "pointer will be captured on next natural read of amdgpu_vram_mm\n");
        /* Not fatal ? kprobe remains active */
    }

    vramgaze_debugfs_file = debugfs_create_file(
        "vramgaze_buddy_nodes", 0444, NULL, NULL, &vramgaze_fops);
    if (IS_ERR_OR_NULL(vramgaze_debugfs_file)) {
        pr_err("vramgaze_buddy: Failed to create debugfs file\n");
        unregister_kprobe(&kp_buddy_print);
        return -ENOMEM;
    }

    if (global_mm)
        pr_info("vramgaze_buddy: loaded, drm_buddy captured successfully.\n");
    else
        pr_info("vramgaze_buddy: loaded, waiting for first amdgpu_vram_mm read.\n");

    return 0;
}

static void __exit vramgaze_buddy_exit(void)
{
    unregister_kprobe(&kp_buddy_print);
    debugfs_remove(vramgaze_debugfs_file);
    pr_info("vramgaze_buddy: unloaded.\n");
}

module_init(vramgaze_buddy_init);
module_exit(vramgaze_buddy_exit);

MODULE_LICENSE("GPL");
MODULE_AUTHOR("superkuh");
MODULE_DESCRIPTION("Reimplements drm_mm-style VRAM dumps for amdgpu buddy allocator (Linux 6.x)");