// https://syzkaller.appspot.com/bug?id=9704c177dc2f59a00d1daec8051c2ef4b4f62436
#define _GNU_SOURCE
#include <pthread.h>
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/prctl.h>
#include <fcntl.h>
#include <linux/futex.h>
#include <sys/syscall.h>
#include <string.h>
#include <sys/mount.h>
#include <sys/stat.h>
#include <errno.h>

int futex_val = 0;
int cpu0 = -1;
int cpu1 = -1;

void set_affinity(int cpu) {
    cpu_set_t set;
    CPU_ZERO(&set);
    CPU_SET(cpu, &set);
    if (sched_setaffinity(0, sizeof(set), &set) < 0) {
        printf("[-] sched_setaffinity failed: %s\n", strerror(errno));
        exit(1);
    }
}

void *thread0_func(void *arg) {
    set_affinity(cpu0);
    while (1) {
        prctl(PR_SET_NAME, "AAAAAAAAAA");
        prctl(PR_SET_NAME, "BBBBBBBBBB");
    }
    return NULL;
}

void *thread1_func(void *arg) {
    set_affinity(cpu1);
    while (1) {
        syscall(SYS_futex, &futex_val, FUTEX_WAKE, 1, NULL, NULL, 0);
    }
    return NULL;
}

void *thread2_func(void *arg) {
    set_affinity(cpu0);
    while (1) {
        futex_val = 0;
        syscall(SYS_futex, &futex_val, FUTEX_WAIT, 0, NULL, NULL, 0);
    }
    return NULL;
}

void setup_tracefs() {
    struct stat st;
    if (stat("/sys/kernel/tracing", &st) == 0) {
        mount("tracefs", "/sys/kernel/tracing", "tracefs", 0, NULL);
    } else if (stat("/sys/kernel/debug", &st) == 0) {
        mount("debugfs", "/sys/kernel/debug", "debugfs", 0, NULL);
    }
}

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

    cpu_set_t set;
    if (sched_getaffinity(0, sizeof(set), &set) == 0) {
        for (int i = 0; i < CPU_SETSIZE; i++) {
            if (CPU_ISSET(i, &set)) {
                if (cpu0 == -1) cpu0 = i;
                else if (cpu1 == -1) { cpu1 = i; break; }
            }
        }
    }
    if (cpu0 == -1) cpu0 = 0;
    if (cpu1 == -1) cpu1 = cpu0;

    if (cpu0 == cpu1) {
        printf("[-] Need at least 2 CPUs to reproduce cross-CPU wakeup race reliably.\n");
    } else {
        printf("[+] Found CPUs: %d and %d\n", cpu0, cpu1);
    }

    setup_tracefs();

    int fd = open("/sys/kernel/tracing/events/sched/sched_stat_runtime/enable", O_WRONLY);
    if (fd < 0) {
        fd = open("/sys/kernel/debug/tracing/events/sched/sched_stat_runtime/enable", O_WRONLY);
    }
    if (fd >= 0) {
        if (write(fd, "1\n", 2) < 0) {
            printf("[-] Failed to write to tracepoint enable: %s\n", strerror(errno));
        } else {
            printf("[+] Successfully enabled sched_stat_runtime tracepoint.\n");
        }
        close(fd);
    } else {
        printf("[-] Failed to open tracepoint enable file: %s\n", strerror(errno));
        exit(1);
    }

    pthread_t t0, t1, t2;
    if (pthread_create(&t0, NULL, thread0_func, NULL) != 0) {
        printf("[-] Failed to create thread 0\n");
        exit(1);
    }
    if (pthread_create(&t1, NULL, thread1_func, NULL) != 0) {
        printf("[-] Failed to create thread 1\n");
        exit(1);
    }
    if (pthread_create(&t2, NULL, thread2_func, NULL) != 0) {
        printf("[-] Failed to create thread 2\n");
        exit(1);
    }

    printf("[+] Threads started. Running for 10 seconds...\n");
    sleep(10);

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