// https://syzkaller.appspot.com/bug?id=dbb019a91e6ad82255d58aedf052975f4717c971
#define _GNU_SOURCE
#include <errno.h>
#include <fcntl.h>
#include <pthread.h>
#include <sched.h>
#include <signal.h>
#include <stdatomic.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/eventfd.h>
#include <sys/mman.h>
#include <sys/resource.h>
#include <sys/stat.h>
#include <sys/sysmacros.h>
#include <sys/syscall.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

#define MSR_MAJOR               202
#define MSR_IA32_APICBASE       0x0000001b
#define MSR_APIC_ICR            0x00000830
#define MSR_APIC_SELF_IPI       0x0000083f
#define MSR_KVM_ASYNC_PF_EN     0x4b564d02
#define MSR_KVM_ASYNC_PF_INT    0x4b564d06
#define HYPERVISOR_CALLBACK_VEC 0xf3

#define MMAP_SZ                 (64 * 1024 * 1024)
#define PAGE_SZ                 4096

static atomic_int keep_running = 1;
static char *shared_buf = NULL;

static void sig_handler(int sig)
{
	(void)sig;
	atomic_store_explicit(&keep_running, 0, memory_order_relaxed);
}

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

static void enable_msr_writes(void)
{
	const char *paths[] = {
		"/sys/module/msr/parameters/allow_writes",
		"/sys/module/msr_core/parameters/allow_writes"
	};
	int enabled = 0;

	for (size_t i = 0; i < sizeof(paths) / sizeof(paths[0]); i++) {
		int fd = open(paths[i], O_WRONLY);
		if (fd >= 0) {
			ssize_t ret = write(fd, "on\n", 3);
			if (ret == 3) {
				printf("[+] Successfully wrote 'on' to %s.\n", paths[i]);
				enabled = 1;
			} else {
				printf("[-] Failed to write to %s: %s\n", paths[i], strerror(errno));
			}
			close(fd);
		} else {
			printf("[*] Path %s not available: %s\n", paths[i], strerror(errno));
		}
	}

	if (!enabled) {
		printf("[-] Warning: Failed to enable allow_writes on any known sysfs path!\n");
	} else {
		printf("[+] MSR writes successfully enabled.\n");
		int rfd = open("/sys/module/msr/parameters/allow_writes", O_RDONLY);
		if (rfd >= 0) {
			char buf[32] = {0};
			ssize_t r = read(rfd, buf, sizeof(buf) - 1);
			if (r > 0)
				printf("[+] Verified allow_writes is now: %s", buf);
			close(rfd);
		}
	}
}

static void setup_msr_nodes(void)
{
	int res = mkdir("/dev/cpu", 0755);
	if (res < 0 && errno != EEXIST) {
		printf("[-] Failed to mkdir /dev/cpu: %s\n", strerror(errno));
		exit(1);
	}
	printf("[+] mkdir /dev/cpu successful or already exists.\n");

	for (int cpu = 0; cpu < 2; cpu++) {
		char dir[64];
		snprintf(dir, sizeof(dir), "/dev/cpu/%d", cpu);
		res = mkdir(dir, 0755);
		if (res < 0 && errno != EEXIST) {
			printf("[-] Failed to mkdir %s: %s\n", dir, strerror(errno));
			exit(1);
		}

		char msr_path[64];
		snprintf(msr_path, sizeof(msr_path), "/dev/cpu/%d/msr", cpu);
		res = mknod(msr_path, S_IFCHR | 0666, makedev(MSR_MAJOR, cpu));
		if (res < 0 && errno != EEXIST) {
			printf("[-] Failed to mknod %s: %s\n", msr_path, strerror(errno));
			exit(1);
		}
		chmod(msr_path, 0666);

		char alt_path[64];
		snprintf(alt_path, sizeof(alt_path), "/dev/msr%d", cpu);
		res = mknod(alt_path, S_IFCHR | 0666, makedev(MSR_MAJOR, cpu));
		if (res < 0 && errno != EEXIST) {
			printf("[-] Failed to mknod %s: %s\n", alt_path, strerror(errno));
			exit(1);
		}
		chmod(alt_path, 0666);
	}
	printf("[+] MSR device nodes configured successfully.\n");

	enable_msr_writes();
}

static int open_msr(int cpu)
{
	char path[64];
	snprintf(path, sizeof(path), "/dev/cpu/%d/msr", cpu);
	int fd = open(path, O_RDWR);
	if (fd < 0) {
		snprintf(path, sizeof(path), "/dev/msr%d", cpu);
		fd = open(path, O_RDWR);
	}
	if (fd < 0) {
		printf("[-] Failed to open MSR device for CPU %d: %s\n", cpu, strerror(errno));
		exit(1);
	}
	printf("[+] Successfully opened MSR device for CPU %d (fd=%d).\n", cpu, fd);
	return fd;
}

/*
 * Worker 1 & 2 (pinned to CPU 0 and CPU 1):
 * Ensure x2APIC mode is enabled in APICBASE MSR (bits 10 & 11) and flood vector
 * 0xf3 (HYPERVISOR_CALLBACK_VECTOR) interrupts via MSR_APIC_SELF_IPI (0x83f) and
 * MSR_APIC_ICR (0x830).
 * In sysvec_kvm_asyncpf_interrupt(), kvm_async_pf_task_wake(0) repeatedly allocates
 * dummy entries (104 bytes -> kmalloc-128) under hardirq context (LD_WAIT_SPIN).
 * When kmalloc-128 per-CPU sheaves are depleted, refill_sheaf() invokes
 * __refill_objects_node() -> get_partial_node_bulk(), which acquires
 * spin_lock_irqsave(&n->list_lock) (LD_WAIT_CONFIG), triggering lockdep:
 * "WARNING: locking bug in __refill_objects_node" / "[ BUG: Invalid wait context ]".
 */
static void *msr_interrupt_worker(void *arg)
{
	int cpu = (int)(intptr_t)arg;
	set_cpu_affinity(cpu);
	printf("[+] msr_interrupt_worker started on CPU %d.\n", cpu);

	int fd = open_msr(cpu);

	uint64_t apicbase = 0;
	if (pread(fd, &apicbase, 8, MSR_IA32_APICBASE) == 8) {
		printf("[+] CPU %d APICBASE: 0x%llx\n", cpu, (unsigned long long)apicbase);
		if (!(apicbase & (1ULL << 10))) {
			apicbase |= (1ULL << 10) | (1ULL << 11);
			pwrite(fd, &apicbase, 8, MSR_IA32_APICBASE);
			printf("[+] CPU %d enabled x2APIC via APICBASE.\n", cpu);
		}
	}

	uint64_t apf_int = HYPERVISOR_CALLBACK_VEC;
	pwrite(fd, &apf_int, 8, MSR_KVM_ASYNC_PF_INT);

	uint64_t self_ipi = HYPERVISOR_CALLBACK_VEC;
	uint64_t icr_self = (1ULL << 18) | HYPERVISOR_CALLBACK_VEC;
	uint64_t icr_all  = (2ULL << 18) | HYPERVISOR_CALLBACK_VEC;
	uint64_t icr_dest = (3ULL << 18) | HYPERVISOR_CALLBACK_VEC;
	uint64_t icr_other = ((uint64_t)(1 - cpu) << 32) | HYPERVISOR_CALLBACK_VEC;

	ssize_t wr = pwrite(fd, &self_ipi, 8, MSR_APIC_SELF_IPI);
	printf("[*] CPU %d initial write MSR_APIC_SELF_IPI: %zd (%s)\n",
	       cpu, wr, wr < 0 ? strerror(errno) : "success");

	wr = pwrite(fd, &icr_self, 8, MSR_APIC_ICR);
	printf("[*] CPU %d initial write MSR_APIC_ICR (self): %zd (%s)\n",
	       cpu, wr, wr < 0 ? strerror(errno) : "success");

	wr = pwrite(fd, &icr_all, 8, MSR_APIC_ICR);
	printf("[*] CPU %d initial write MSR_APIC_ICR (all): %zd (%s)\n",
	       cpu, wr, wr < 0 ? strerror(errno) : "success");

	wr = pwrite(fd, &icr_other, 8, MSR_APIC_ICR);
	printf("[*] CPU %d initial write MSR_APIC_ICR (other): %zd (%s)\n",
	       cpu, wr, wr < 0 ? strerror(errno) : "success");

	unsigned long iters = 0;
	while (atomic_load_explicit(&keep_running, memory_order_relaxed)) {
		pwrite(fd, &self_ipi, 8, MSR_APIC_SELF_IPI);
		pwrite(fd, &icr_self, 8, MSR_APIC_ICR);
		pwrite(fd, &icr_all, 8, MSR_APIC_ICR);
		pwrite(fd, &icr_dest, 8, MSR_APIC_ICR);
		pwrite(fd, &icr_other, 8, MSR_APIC_ICR);

		iters++;
		if (iters % 25000 == 0) {
			printf("[+] CPU %d msr_interrupt_worker reached %lu iterations.\n", cpu, iters);
		}
	}

	int res = close(fd);
	if (res < 0) {
		printf("[-] Failed to close MSR fd on CPU %d: %s\n", cpu, strerror(errno));
		exit(1);
	}
	printf("[+] Closed MSR fd on CPU %d.\n", cpu);
	return NULL;
}

/*
 * Worker 3 (CPU 0): Touches shared memory buffer across both CPUs, keeping mm
 * referenced and active on CPU 0 so dup_mmap() on CPU 1 issues remote TLB flushes
 * via kvm_flush_tlb_multi() -> smp_call_function_many_cond() -> csd_lock_wait().
 */
static void *cpu0_background_worker(void *arg)
{
	(void)arg;
	set_cpu_affinity(0);
	printf("[+] cpu0_background_worker started on CPU 0.\n");

	volatile char val = 0;
	unsigned long iters = 0;

	while (atomic_load_explicit(&keep_running, memory_order_relaxed)) {
		if (shared_buf) {
			for (size_t i = 0; i < MMAP_SZ; i += PAGE_SZ * 16) {
				val += shared_buf[i];
				shared_buf[i] = val;
			}
		}

		int fds[16];
		for (int i = 0; i < 16; i++)
			fds[i] = eventfd(0, EFD_NONBLOCK | EFD_CLOEXEC);
		for (int i = 0; i < 16; i++) {
			if (fds[i] >= 0)
				close(fds[i]);
		}

		iters++;
		if (iters % 2000 == 0) {
			printf("[+] CPU 0 background_worker reached %lu iterations.\n", iters);
		}
	}
	return NULL;
}

/*
 * Worker 4 (CPU 1): Continual fork storm with flags 0x01200011
 * (CLONE_CHILD_SETTID | CLONE_CHILD_CLEARTID | SIGCHLD).
 * Replicates syzkaller's task state holding dup_mmap_sem and mm->mmap_lock
 * while spinning in csd_lock_wait() with interrupts enabled.
 */
static void *fork_worker(void *arg)
{
	(void)arg;
	set_cpu_affinity(1);
	printf("[+] fork_worker started on CPU 1.\n");

	pid_t ctid = 0;
	unsigned long forks = 0;

	while (atomic_load_explicit(&keep_running, memory_order_relaxed)) {
		pid_t pid = (pid_t)syscall(SYS_clone, 0x01200011, NULL, NULL, &ctid, NULL);
		if (pid < 0) {
			if (errno == EAGAIN || errno == ENOMEM) {
				usleep(50);
				continue;
			}
			printf("[-] Failed to clone on CPU 1: %s\n", strerror(errno));
			exit(1);
		}
		if (pid == 0) {
			_exit(0);
		}
		int status = 0;
		pid_t w = waitpid(pid, &status, __WALL);
		if (w < 0 && errno != ECHILD && errno != EINTR) {
			printf("[-] Failed to waitpid on CPU 1: %s\n", strerror(errno));
			exit(1);
		}
		forks++;
		if (forks % 1000 == 0) {
			printf("[+] CPU 1 fork_worker reached %lu iterations.\n", forks);
		}
	}
	return NULL;
}

int main(void)
{
	setvbuf(stdout, NULL, _IONBF, 0);
	printf("[*] Initializing reproducer for locking bug in __refill_objects_node...\n");

	long nprocs = sysconf(_SC_NPROCESSORS_ONLN);
	if (nprocs < 1) {
		printf("[-] Failed to sysconf(_SC_NPROCESSORS_ONLN): %s\n", strerror(errno));
		exit(1);
	}
	printf("[+] sysconf(_SC_NPROCESSORS_ONLN) = %ld online CPU(s).\n", nprocs);

	int cpu0 = 0;
	int cpu1 = (nprocs > 1) ? 1 : 0;

	struct sigaction sa;
	memset(&sa, 0, sizeof(sa));
	sa.sa_handler = sig_handler;
	int res = sigaction(SIGINT, &sa, NULL);
	if (res < 0) {
		printf("[-] Failed to sigaction(SIGINT): %s\n", strerror(errno));
		exit(1);
	}
	res = sigaction(SIGTERM, &sa, NULL);
	if (res < 0) {
		printf("[-] Failed to sigaction(SIGTERM): %s\n", strerror(errno));
		exit(1);
	}
	res = sigaction(SIGALRM, &sa, NULL);
	if (res < 0) {
		printf("[-] Failed to sigaction(SIGALRM): %s\n", strerror(errno));
		exit(1);
	}
	printf("[+] Signal handlers installed successfully.\n");

	alarm(30);
	printf("[+] Alarm set to 30 seconds.\n");

	struct rlimit rlim;
	rlim.rlim_cur = 65536;
	rlim.rlim_max = 65536;
	res = setrlimit(RLIMIT_NOFILE, &rlim);
	if (res < 0) {
		printf("[*] setrlimit(RLIMIT_NOFILE) failed (%s), continuing with default limit.\n", strerror(errno));
	} else {
		printf("[+] setrlimit(RLIMIT_NOFILE, 65536) successful.\n");
	}

	setup_msr_nodes();

	shared_buf = mmap(NULL, MMAP_SZ, PROT_READ | PROT_WRITE,
	                  MAP_SHARED | MAP_ANONYMOUS, -1, 0);
	if (shared_buf == MAP_FAILED) {
		printf("[-] Failed to mmap shared_buf: %s\n", strerror(errno));
		exit(1);
	}
	printf("[+] mmap shared_buf (%zu bytes) successful.\n", (size_t)MMAP_SZ);

	for (size_t i = 0; i < MMAP_SZ; i += PAGE_SZ)
		shared_buf[i] = (char)i;
	printf("[+] Pre-faulted shared_buf pages successfully.\n");

	pthread_t th_msr0, th_msr1, th_bg, th_fork;

	int s = pthread_create(&th_msr0, NULL, msr_interrupt_worker, (void *)(intptr_t)cpu0);
	if (s != 0) {
		printf("[-] Failed to pthread_create th_msr0: %s\n", strerror(s));
		exit(1);
	}
	printf("[+] pthread_create th_msr0 successful.\n");

	s = pthread_create(&th_msr1, NULL, msr_interrupt_worker, (void *)(intptr_t)cpu1);
	if (s != 0) {
		printf("[-] Failed to pthread_create th_msr1: %s\n", strerror(s));
		exit(1);
	}
	printf("[+] pthread_create th_msr1 successful.\n");

	s = pthread_create(&th_bg, NULL, cpu0_background_worker, NULL);
	if (s != 0) {
		printf("[-] Failed to pthread_create th_bg: %s\n", strerror(s));
		exit(1);
	}
	printf("[+] pthread_create th_bg successful.\n");

	s = pthread_create(&th_fork, NULL, fork_worker, (void *)(intptr_t)cpu1);
	if (s != 0) {
		printf("[-] Failed to pthread_create th_fork: %s\n", strerror(s));
		exit(1);
	}
	printf("[+] pthread_create th_fork successful.\n");

	printf("[*] All workers active across CPU %d and CPU %d. Running stress race...\n", cpu0, cpu1);

	while (atomic_load_explicit(&keep_running, memory_order_relaxed)) {
		usleep(50000);
	}

	printf("[*] Shutting down workers...\n");

	s = pthread_join(th_msr0, NULL);
	if (s != 0) {
		printf("[-] Failed to pthread_join th_msr0: %s\n", strerror(s));
		exit(1);
	}
	s = pthread_join(th_msr1, NULL);
	if (s != 0) {
		printf("[-] Failed to pthread_join th_msr1: %s\n", strerror(s));
		exit(1);
	}
	s = pthread_join(th_bg, NULL);
	if (s != 0) {
		printf("[-] Failed to pthread_join th_bg: %s\n", strerror(s));
		exit(1);
	}
	s = pthread_join(th_fork, NULL);
	if (s != 0) {
		printf("[-] Failed to pthread_join th_fork: %s\n", strerror(s));
		exit(1);
	}
	printf("[+] All worker threads joined successfully.\n");

	res = munmap(shared_buf, MMAP_SZ);
	if (res < 0) {
		printf("[-] Failed to munmap shared_buf: %s\n", strerror(errno));
		exit(1);
	}
	printf("[+] munmap shared_buf successful.\n");

	printf("[+] Reproducer run completed successfully.\n");
	return 0;
}