// https://syzkaller.appspot.com/bug?id=3d01c4ec9b6a9a03fedf287a0a4d48a2edfa4f4a
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/sysmacros.h>
#include <string.h>
#include <pthread.h>
#include <linux/loop.h>
#include <linux/blkpg.h>
#include <stdint.h>
#include <time.h>
#include <errno.h>

#define BLKTRACESETUP _IOWR(0x12, 115, struct blk_user_trace_setup)
#define BLKTRACESTART _IO(0x12, 116)

struct blk_user_trace_setup {
	char name[32];
	uint16_t act_mask;
	uint32_t buf_size;
	uint32_t buf_nr;
	uint64_t start_lba;
	uint64_t end_lba;
	uint32_t pid;
};

int loop_fd = -1;
int part_fd = -1;
int trace_pipe_fd = -1;

void *reader_thread(void *arg) {
    char buf[4096];
    while (1) {
        ssize_t ret = read(trace_pipe_fd, buf, sizeof(buf));
        if (ret < 0) {
            usleep(10000);
        }
    }
    return NULL;
}

void *io_thread(void *arg) {
    void *buf;
    if (posix_memalign(&buf, 512, 512)) return NULL;
    unsigned long offset = 0;
    while (1) {
        ssize_t ret = pread(part_fd, buf, 512, offset);
        if (ret < 0) {
            usleep(10000);
        }
        offset += 512;
        if (offset >= 1024 * 1024) offset = 0;
    }
    return NULL;
}

void *resize_thread(void *arg) {
    int fd = open("/sys/kernel/tracing/buffer_size_kb", O_WRONLY);
    if (fd < 0)
        fd = open("/sys/kernel/debug/tracing/buffer_size_kb", O_WRONLY);
    if (fd < 0) {
        printf("[-] Failed to open buffer_size_kb: %s\n", strerror(errno));
        return NULL;
    }
    
    char buf[32];
    while (1) {
        int size = 1024 + (rand() % 4096);
        int len = snprintf(buf, sizeof(buf), "%d\n", size);
        ssize_t ret = write(fd, buf, len);
        if (ret < 0) {
            // Ignore
        }
        usleep(100000);
    }
    return NULL;
}

int main() {
    srand(time(NULL));

    // 1. Enable blk tracer
    int fd = open("/sys/kernel/tracing/current_tracer", O_WRONLY);
    if (fd < 0)
        fd = open("/sys/kernel/debug/tracing/current_tracer", O_WRONLY);
    if (fd >= 0) {
        ssize_t ret = write(fd, "blk\n", 4);
        if (ret < 0) {
            printf("[-] Failed to write to current_tracer: %s\n", strerror(errno));
        } else {
            printf("[+] Enabled blk tracer.\n");
        }
        close(fd);
    } else {
        printf("[-] Failed to open current_tracer: %s\n", strerror(errno));
        exit(1);
    }

    // 2. Setup loop device
    int ctl_fd = open("/dev/loop-control", O_RDWR);
    if (ctl_fd < 0) {
        printf("[-] Failed to open /dev/loop-control: %s\n", strerror(errno));
        exit(1);
    }
    int loop_idx = ioctl(ctl_fd, LOOP_CTL_GET_FREE);
    close(ctl_fd);
    if (loop_idx < 0) {
        printf("[-] Failed to get free loop device: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] Got free loop device index: %d\n", loop_idx);

    char loop_name[64];
    snprintf(loop_name, sizeof(loop_name), "/dev/loop%d", loop_idx);
    loop_fd = open(loop_name, O_RDWR);
    if (loop_fd < 0) {
        printf("[-] Failed to open %s: %s\n", loop_name, strerror(errno));
        exit(1);
    }
    printf("[+] Opened %s.\n", loop_name);

    int file_fd = open("/tmp/loop_file", O_RDWR | O_CREAT | O_TRUNC, 0644);
    if (file_fd < 0) {
        printf("[-] Failed to open /tmp/loop_file: %s\n", strerror(errno));
        exit(1);
    }
    if (ftruncate(file_fd, 1024 * 1024) < 0) {
        printf("[-] Failed to ftruncate /tmp/loop_file: %s\n", strerror(errno));
        exit(1);
    }
    if (ioctl(loop_fd, LOOP_SET_FD, file_fd) < 0) {
        printf("[-] Failed to LOOP_SET_FD: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] Setup loop device with backing file.\n");

    // 3. Add partition
    struct blkpg_ioctl_arg arg;
    struct blkpg_partition p;
    memset(&p, 0, sizeof(p));
    p.start = 0;
    p.length = 1024 * 1024;
    p.pno = 1;
    memset(&arg, 0, sizeof(arg));
    arg.op = BLKPG_ADD_PARTITION;
    arg.datalen = sizeof(p);
    arg.data = &p;
    if (ioctl(loop_fd, BLKPG, &arg) < 0) {
        printf("[-] Failed to BLKPG_ADD_PARTITION: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] Added partition to loop device.\n");

    // 4. Start blktrace (version 1)
    struct blk_user_trace_setup setup;
    memset(&setup, 0, sizeof(setup));
    setup.act_mask = 0xffff;
    setup.buf_size = 512 * 1024;
    setup.buf_nr = 4;
    if (ioctl(loop_fd, BLKTRACESETUP, &setup) < 0) {
        printf("[-] Failed to BLKTRACESETUP: %s\n", strerror(errno));
        exit(1);
    }
    if (ioctl(loop_fd, BLKTRACESTART) < 0) {
        printf("[-] Failed to BLKTRACESTART: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] Started blktrace.\n");

    // 5. Open partition device
    char part_sysfs[256];
    snprintf(part_sysfs, sizeof(part_sysfs), "/sys/class/block/loop%d/loop%dp1/dev", loop_idx, loop_idx);
    int maj = 0, min = 0;
    for (int i = 0; i < 100; i++) {
        FILE *f = fopen(part_sysfs, "r");
        if (f) {
            if (fscanf(f, "%d:%d", &maj, &min) == 2) {
                fclose(f);
                break;
            }
            fclose(f);
        }
        usleep(10000);
    }
    if (maj == 0) {
        printf("[-] Failed to find partition devt.\n");
        exit(1);
    }
    printf("[+] Partition devt: %d:%d\n", maj, min);

    if (mknod("/tmp/loop_part", S_IFBLK | 0600, makedev(maj, min)) < 0) {
        if (errno != EEXIST) {
            printf("[-] Failed to mknod /tmp/loop_part: %s\n", strerror(errno));
            exit(1);
        }
    }
    part_fd = open("/tmp/loop_part", O_RDWR | O_DIRECT);
    if (part_fd < 0) {
        part_fd = open("/tmp/loop_part", O_RDWR);
    }
    if (part_fd < 0) {
        printf("[-] Failed to open /tmp/loop_part: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] Opened partition device.\n");

    trace_pipe_fd = open("/sys/kernel/tracing/trace_pipe", O_RDONLY);
    if (trace_pipe_fd < 0)
        trace_pipe_fd = open("/sys/kernel/debug/tracing/trace_pipe", O_RDONLY);
    if (trace_pipe_fd < 0) {
        printf("[-] Failed to open trace_pipe: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] Opened trace_pipe.\n");

    // 6. Spawn threads to trigger the bug
    pthread_t t1, t2, t3, t4, t5;
    pthread_create(&t1, NULL, reader_thread, NULL);
    pthread_create(&t2, NULL, io_thread, NULL);
    pthread_create(&t3, NULL, io_thread, NULL);
    pthread_create(&t4, NULL, io_thread, NULL);
    pthread_create(&t5, NULL, resize_thread, NULL);

    printf("[+] Threads spawned, waiting for 15 seconds...\n");
    sleep(15);
    
    printf("[+] Done.\n");
    return 0;
}