// https://syzkaller.appspot.com/bug?id=38bac744801a26576f017cacffa13732ec3a3ea9
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/syscall.h>
#include <sys/prctl.h>
#include <stdint.h>
#include <sched.h>
#include <string.h>
#include <errno.h>

#ifndef PR_RSEQ_SLICE_EXTENSION
# define PR_RSEQ_SLICE_EXTENSION        79
# define PR_RSEQ_SLICE_EXTENSION_SET    2
# define PR_RSEQ_SLICE_EXT_ENABLE       0x01
#endif

#ifndef SYS_rseq
#define SYS_rseq 334
#endif

struct my_rseq_slice_ctrl {
    union {
        uint32_t all;
        struct {
            uint8_t request;
            uint8_t granted;
            uint16_t __reserved;
        };
    };
};

struct my_rseq {
    uint32_t cpu_id_start;
    uint32_t cpu_id;
    uint64_t rseq_cs;
    uint32_t flags;
    uint32_t node_id;
    uint32_t mm_cid;
    struct my_rseq_slice_ctrl slice_ctrl;
    uint8_t __reserved;
    char end[];
} __attribute__((aligned(64)));

int thread_fn(void *arg) {
    struct my_rseq *rseq_area = NULL;
    if (posix_memalign((void **)&rseq_area, 64, sizeof(struct my_rseq)) != 0) {
        printf("[-] posix_memalign failed: %s\n", strerror(errno));
        exit(1);
    }
    memset(rseq_area, 0, sizeof(struct my_rseq));

    /* Register our custom RSEQ area */
    int ret = syscall(SYS_rseq, rseq_area, sizeof(struct my_rseq), 0, 0x53053053);
    if (ret != 0) {
        printf("[-] rseq syscall failed: %s\n", strerror(errno));
        exit(1);
    }

    /* Enable the RSEQ time slice extension feature */
    if (prctl(PR_RSEQ_SLICE_EXTENSION, PR_RSEQ_SLICE_EXTENSION_SET, PR_RSEQ_SLICE_EXT_ENABLE, 0, 0) != 0) {
        printf("[-] prctl PR_RSEQ_SLICE_EXTENSION failed: %s\n", strerror(errno));
        exit(1);
    }

    /* 
     * Spin in user space requesting a slice extension. 
     * A timer interrupt will eventually hit, set _TIF_NEED_RESCHED, 
     * and trigger the lockdep warning on return to user space.
     */
    while (1) {
        rseq_area->slice_ctrl.request = 1;
        /* Prevent compiler from optimizing out the store */
        asm volatile("" ::: "memory");
    }
    return 0;
}

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

    for (int i = 0; i < 2; i++) {
        char *stack = malloc(1024 * 1024);
        if (!stack) {
            printf("[-] malloc failed: %s\n", strerror(errno));
            exit(1);
        }
        void *stack_top = (void *)(((uintptr_t)stack + 1024 * 1024) & ~0xF);

        pid_t pid = clone(thread_fn, stack_top, 
                          CLONE_VM | CLONE_FS | CLONE_FILES | CLONE_SIGHAND | CLONE_THREAD, NULL);
        if (pid == -1) {
            printf("[-] clone failed: %s\n", strerror(errno));
            exit(1);
        }
    }
    printf("[+] Created 2 spinning threads on CPU 0.\n");
    
    /* Wait enough time for the scheduler tick to hit the spinning threads */
    sleep(5);
    printf("[+] Done.\n");
    return 0;
}