// https://syzkaller.appspot.com/bug?id=2fbb1820458b0dd8b4eed4e65992ac56ba66c8b5
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include <sys/syscall.h>
#include <sys/mman.h>
#include <linux/io_uring.h>
#include <linux/capability.h>
#include <errno.h>
#include <stdatomic.h>

#ifndef _LINUX_CAPABILITY_VERSION_3
#define _LINUX_CAPABILITY_VERSION_3 0x20080522
#endif

#ifndef IOSQE_IO_DRAIN
#define IOSQE_IO_DRAIN (1U << 1)
#endif

// io_uring setup
struct io_sq_ring {
    unsigned *head;
    unsigned *tail;
    unsigned *ring_mask;
    unsigned *ring_entries;
    unsigned *flags;
    unsigned *array;
};

struct io_cq_ring {
    unsigned *head;
    unsigned *tail;
    unsigned *ring_mask;
    unsigned *ring_entries;
    struct io_uring_cqe *cqes;
};

struct io_uring_sq {
    unsigned *khead;
    unsigned *ktail;
    unsigned *kring_mask;
    unsigned *kring_entries;
    unsigned *kflags;
    unsigned *kdropped;
    unsigned *array;
    struct io_uring_sqe *sqes;
    unsigned sqe_head;
    unsigned sqe_tail;
    size_t ring_sz;
};

struct io_uring_cq {
    unsigned *khead;
    unsigned *ktail;
    unsigned *kring_mask;
    unsigned *kring_entries;
    unsigned *koverflow;
    struct io_uring_cqe *cqes;
    size_t ring_sz;
};

struct io_uring {
    struct io_uring_sq sq;
    struct io_uring_cq cq;
    int ring_fd;
};

#define smp_store_release(p, v) atomic_store_explicit((_Atomic typeof(*(p)) *)(p), (v), memory_order_release)
#define smp_load_acquire(p) atomic_load_explicit((_Atomic typeof(*(p)) *)(p), memory_order_acquire)

int io_uring_setup_helper(unsigned entries, struct io_uring *ring) {
    struct io_uring_params p;
    void *sq_ptr, *cq_ptr;

    memset(&p, 0, sizeof(p));
    ring->ring_fd = syscall(__NR_io_uring_setup, entries, &p);
    if (ring->ring_fd < 0) {
        return -1;
    }

    int sq_ring_sz = p.sq_off.array + p.sq_entries * sizeof(unsigned);
    int cq_ring_sz = p.cq_off.cqes + p.cq_entries * sizeof(struct io_uring_cqe);

    if (p.features & IORING_FEAT_SINGLE_MMAP) {
        if (cq_ring_sz > sq_ring_sz)
            sq_ring_sz = cq_ring_sz;
        cq_ring_sz = sq_ring_sz;
    }

    sq_ptr = mmap(0, sq_ring_sz, PROT_READ | PROT_WRITE,
                  MAP_SHARED | MAP_POPULATE, ring->ring_fd, IORING_OFF_SQ_RING);
    if (sq_ptr == MAP_FAILED) return -1;

    if (p.features & IORING_FEAT_SINGLE_MMAP) {
        cq_ptr = sq_ptr;
    } else {
        cq_ptr = mmap(0, cq_ring_sz, PROT_READ | PROT_WRITE,
                      MAP_SHARED | MAP_POPULATE, ring->ring_fd, IORING_OFF_CQ_RING);
        if (cq_ptr == MAP_FAILED) return -1;
    }

    ring->sq.khead = sq_ptr + p.sq_off.head;
    ring->sq.ktail = sq_ptr + p.sq_off.tail;
    ring->sq.kring_mask = sq_ptr + p.sq_off.ring_mask;
    ring->sq.kring_entries = sq_ptr + p.sq_off.ring_entries;
    ring->sq.kflags = sq_ptr + p.sq_off.flags;
    ring->sq.kdropped = sq_ptr + p.sq_off.dropped;
    ring->sq.array = sq_ptr + p.sq_off.array;

    ring->sq.sqes = mmap(0, p.sq_entries * sizeof(struct io_uring_sqe),
                         PROT_READ | PROT_WRITE, MAP_SHARED | MAP_POPULATE,
                         ring->ring_fd, IORING_OFF_SQES);
    if (ring->sq.sqes == MAP_FAILED) return -1;

    ring->cq.khead = cq_ptr + p.cq_off.head;
    ring->cq.ktail = cq_ptr + p.cq_off.tail;
    ring->cq.kring_mask = cq_ptr + p.cq_off.ring_mask;
    ring->cq.kring_entries = cq_ptr + p.cq_off.ring_entries;
    ring->cq.koverflow = cq_ptr + p.cq_off.overflow;
    ring->cq.cqes = cq_ptr + p.cq_off.cqes;

    return 0;
}

struct io_uring ring;
volatile int stop = 0;

void *uring_thread(void *arg) {
    while (!stop) {
        unsigned tail = *ring.sq.ktail;
        unsigned index = tail & *ring.sq.kring_mask;
        struct io_uring_sqe *sqe = &ring.sq.sqes[index];

        memset(sqe, 0, sizeof(*sqe));
        sqe->opcode = IORING_OP_NOP;
        sqe->flags = IOSQE_IO_DRAIN;

        ring.sq.array[index] = index;
        tail++;
        smp_store_release(ring.sq.ktail, tail);

        syscall(__NR_io_uring_enter, ring.ring_fd, 1, 0, 0, NULL, 0);

        // Consume CQEs
        unsigned head = smp_load_acquire(ring.cq.khead);
        unsigned cq_tail = smp_load_acquire(ring.cq.ktail);
        if (head != cq_tail) {
            smp_store_release(ring.cq.khead, cq_tail);
        }
    }
    return NULL;
}

int main() {
    if (io_uring_setup_helper(8, &ring) < 0) {
        printf("[-] Failed to io_uring_setup: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] io_uring_setup successful.\n");

    pthread_t th;
    if (pthread_create(&th, NULL, uring_thread, NULL) != 0) {
        printf("[-] Failed to create thread: %s\n", strerror(errno));
        exit(1);
    }
    printf("[+] Thread created successfully.\n");

    struct __user_cap_header_struct header = {
        .version = _LINUX_CAPABILITY_VERSION_3,
        .pid = 0,
    };
    
    struct __user_cap_data_struct data[2] = {
        { .effective = 1, .permitted = 0, .inheritable = 0 },
        { .effective = 0, .permitted = 0, .inheritable = 0 }
    };

    printf("[+] Starting capset loop.\n");
    for (int i = 0; i < 100000; i++) {
        syscall(__NR_capset, &header, data);
    }
    
    stop = 1;
    pthread_join(th, NULL);
    printf("[+] Done.\n");

    return 0;
}