// https://syzkaller.appspot.com/bug?id=866cbfa062130e633d2e6b32660ba4e0ded405e7
#define _GNU_SOURCE
#include <pthread.h>
#include <fcntl.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <sys/klog.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <arpa/inet.h>
#include <sched.h>
#include <errno.h>

void write_file(const char *path, const char *val) {
    if (access(path, F_OK) != 0) {
        printf("[!] %s does not exist, skipping.\n", path);
        return;
    }
    int fd = open(path, O_WRONLY);
    if (fd < 0) {
        printf("[-] Failed to open %s: %s\n", path, strerror(errno));
        exit(1);
    }
    ssize_t res = write(fd, val, strlen(val));
    if (res < 0) {
        printf("[-] Failed to write to %s: %s\n", path, strerror(errno));
        exit(1);
    }
    close(fd);
    printf("[+] Successfully wrote to %s\n", path);
}

void pin_to_cpu(int cpu) {
    cpu_set_t cpuset;
    CPU_ZERO(&cpuset);
    CPU_SET(cpu, &cpuset);
    int res = pthread_setaffinity_np(pthread_self(), sizeof(cpu_set_t), &cpuset);
    if (res != 0) {
        printf("[-] Failed to pthread_setaffinity_np: %s\n", strerror(res));
        exit(1);
    }
}

void *martian_sender(void *arg) {
    int cpu = (int)(long)arg;
    pin_to_cpu(cpu);
    
    int s = socket(AF_INET, SOCK_RAW, IPPROTO_RAW);
    if (s < 0) {
        printf("[-] Failed to socket: %s\n", strerror(errno));
        exit(1);
    }
    int opt = 1;
    if (setsockopt(s, IPPROTO_IP, IP_HDRINCL, &opt, sizeof(opt)) < 0) {
        printf("[-] Failed to setsockopt: %s\n", strerror(errno));
        exit(1);
    }
    
    char packet[sizeof(struct iphdr)];
    struct iphdr *ip = (struct iphdr *)packet;
    memset(ip, 0, sizeof(*ip));
    ip->ihl = 5; 
    ip->version = 4;
    ip->tot_len = htons(sizeof(struct iphdr));
    ip->ttl = 64; 
    ip->protocol = IPPROTO_UDP;
    ip->saddr = inet_addr("240.0.0.1"); // Class E martian source
    ip->daddr = inet_addr("127.0.0.1");
    
    struct sockaddr_in dest;
    memset(&dest, 0, sizeof(dest));
    dest.sin_family = AF_INET;
    dest.sin_addr.s_addr = ip->daddr;
    
    /* Continuously send martian packets to trigger printk from NET_RX_SOFTIRQ */
    while (1) {
        ssize_t res = sendto(s, packet, sizeof(packet), 0, (struct sockaddr *)&dest, sizeof(dest));
        if (res < 0) {
            if (errno != ENOBUFS && errno != EAGAIN && errno != EINTR && errno != EPERM) {
                printf("[-] Failed to sendto: %s\n", strerror(errno));
                exit(1);
            }
        }
    }
    return NULL;
}

void *writer(void *arg) {
    int cpu = (int)(long)arg;
    pin_to_cpu(cpu);
    
    int fd = open("/dev/kmsg", O_WRONLY);
    if (fd < 0) {
        printf("[-] Failed to open /dev/kmsg for write: %s\n", strerror(errno));
        exit(1);
    }
    
    /* Continuously write to /dev/kmsg to flood the ringbuffer and cause dropped messages */
    while (1) {
        ssize_t res = write(fd, "test\n", 5);
        if (res < 0) {
            if (errno != EAGAIN && errno != EINTR && errno != EPIPE && errno != EBUSY) {
                printf("[-] Failed to write to /dev/kmsg: %s\n", strerror(errno));
                exit(1);
            }
        }
    }
    return NULL;
}

void *reader_kmsg(void *arg) {
    int cpu = (int)(long)arg;
    pin_to_cpu(cpu);
    
    int fd = open("/dev/kmsg", O_RDONLY);
    if (fd < 0) {
        printf("[-] Failed to open /dev/kmsg for read: %s\n", strerror(errno));
        exit(1);
    }
    
    char buf[8192];
    while (1) {
        ssize_t res = read(fd, buf, sizeof(buf));
        if (res < 0) {
            if (errno != EAGAIN && errno != EINTR && errno != EPIPE && errno != EBUSY) {
                printf("[-] Failed to read from /dev/kmsg: %s\n", strerror(errno));
                exit(1);
            }
        }
    }
    return NULL;
}

void *reader_syslog(void *arg) {
    int cpu = (int)(long)arg;
    pin_to_cpu(cpu);
    
    int size = klogctl(10, NULL, 0); // SYSLOG_ACTION_SIZE_BUFFER
    if (size < 0) {
        printf("[-] Failed to klogctl size: %s\n", strerror(errno));
        size = 8192;
    }
    if (size < 8192) {
        size = 8192;
    }
    
    char *buf = malloc(size);
    if (!buf) {
        printf("[-] Failed to malloc: %s\n", strerror(errno));
        exit(1);
    }
    
    while (1) {
        int res = klogctl(3, buf, size);
        if (res < 0) {
            if (errno != EAGAIN && errno != EINTR && errno != EBUSY) {
                printf("[-] Failed to klogctl read: %s\n", strerror(errno));
                exit(1);
            }
        }
    }
    return NULL;
}

int main() {
    setvbuf(stdout, NULL, _IONBF, 0);
    printf("[*] Starting reproducer...\n");

    /* Disable /dev/kmsg ratelimiting */
    write_file("/proc/sys/kernel/printk_devkmsg", "on\n");
    
    /* Disable Network Ratelimiting so every martian packet triggers a printk */
    write_file("/proc/sys/net/core/message_cost", "0\n");
    
    /* Enable Martian Logging */
    write_file("/proc/sys/net/ipv4/conf/all/log_martians", "1\n");
    write_file("/proc/sys/net/ipv4/conf/lo/log_martians", "1\n");

    int num_cpus = sysconf(_SC_NPROCESSORS_ONLN);
    if (num_cpus < 1) num_cpus = 1;
    printf("[+] Number of CPUs: %d\n", num_cpus);

    pthread_t th[16];
    
    /* Spawn 4 martian sender threads to trigger softirq printks */
    for (int i = 0; i < 4; i++) {
        int res = pthread_create(&th[i], NULL, martian_sender, (void *)(long)(i % num_cpus));
        if (res != 0) {
            printf("[-] Failed to create martian_sender thread: %s\n", strerror(res));
            exit(1);
        }
    }
    printf("[+] Created martian sender threads.\n");
    
    /* Spawn 4 writer threads to flood the ringbuffer */
    for (int i = 4; i < 8; i++) {
        int res = pthread_create(&th[i], NULL, writer, (void *)(long)(i % num_cpus));
        if (res != 0) {
            printf("[-] Failed to create writer thread: %s\n", strerror(res));
            exit(1);
        }
    }
    printf("[+] Created writer threads.\n");
    
    /* Spawn 4 reader threads (kmsg) */
    for (int i = 8; i < 12; i++) {
        int res = pthread_create(&th[i], NULL, reader_kmsg, (void *)(long)(i % num_cpus));
        if (res != 0) {
            printf("[-] Failed to create reader_kmsg thread: %s\n", strerror(res));
            exit(1);
        }
    }
    printf("[+] Created reader_kmsg threads.\n");

    /* Spawn 4 reader threads (syslog) */
    for (int i = 12; i < 16; i++) {
        int res = pthread_create(&th[i], NULL, reader_syslog, (void *)(long)(i % num_cpus));
        if (res != 0) {
            printf("[-] Failed to create reader_syslog thread: %s\n", strerror(res));
            exit(1);
        }
    }
    printf("[+] Created reader_syslog threads.\n");
    
    /* Run for 25 seconds to give KCSAN time to catch the race */
    printf("[*] Running for 25 seconds...\n");
    sleep(25);
    printf("[+] Done.\n");
    
    return 0;
}