// https://syzkaller.appspot.com/bug?id=21238395430da9a8e6464175db069e28f80d4599
#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <sys/ioctl.h>
#include <sys/mount.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <linux/types.h>
#include <linux/ioctl.h>

#define IOMMUFD_TYPE (';')
#define IOMMUFD_CMD_BASE 0x80

#define IOMMU_IOAS_ALLOC _IO(IOMMUFD_TYPE, IOMMUFD_CMD_BASE + 1)
#define IOMMU_TEST_CMD _IO(IOMMUFD_TYPE, IOMMUFD_CMD_BASE + 32)
#define IOMMU_TEST_OP_MOCK_DOMAIN 2

struct iommu_ioas_alloc {
	__u32 size;
	__u32 flags;
	__u32 out_ioas_id;
};

struct iommu_test_cmd {
	__u32 size;
	__u32 op;
	__u32 id;
	__u32 __reserved;
	union {
		struct {
			__aligned_u64 start;
			__aligned_u64 length;
		} add_reserved;
		struct {
			__u32 out_stdev_id;
			__u32 out_hwpt_id;
			__u32 out_idev_id;
		} mock_domain;
		__u8 payload[48];
	};
	__u32 last;
};

static void write_file(const char *path, const char *val)
{
	int fd = open(path, O_WRONLY);
	if (fd < 0) {
		printf("[-] Notice: unable to open %s: %s\n", path, strerror(errno));
		return;
	}
	ssize_t len = strlen(val);
	ssize_t written = write(fd, val, len);
	if (written < 0) {
		printf("[-] Failed to write to %s: %s\n", path, strerror(errno));
	} else {
		printf("[+] Configured %s successfully\n", path);
	}
	int res = close(fd);
	if (res < 0) {
		printf("[-] Failed to close %s: %s\n", path, strerror(errno));
	}
}

int main(void)
{
	setvbuf(stdout, NULL, _IONBF, 0);
	printf("[*] Starting lean reproducer for memory leak in iommu_fwspec_init...\n");

	int res = mount("debugfs", "/sys/kernel/debug", "debugfs", 0, NULL);
	if (res < 0 && errno != EBUSY) {
		printf("[-] Failed to mount debugfs: %s\n", strerror(errno));
		exit(1);
	}
	printf("[+] debugfs ready at /sys/kernel/debug\n");

	/* Suppress stack traces on fault injection to avoid serial console flooding */
	write_file("/sys/kernel/debug/failslab/verbose", "0\n");

	/* Disable printk ratelimiting for /dev/kmsg */
	write_file("/proc/sys/kernel/printk_devkmsg", "on\n");

	/* Promote unreferenced objects on the first scan */
	write_file("/sys/module/kmemleak/parameters/min_unref_scans", "1\n");

	/* Clear previous kmemleak records */
	write_file("/sys/kernel/debug/kmemleak", "clear\n");

	int iommu_fd = open("/dev/iommu", O_RDWR);
	if (iommu_fd < 0) {
		printf("[-] Failed to open /dev/iommu: %s\n", strerror(errno));
		exit(1);
	}
	printf("[+] Opened /dev/iommu (fd=%d)\n", iommu_fd);

	int fail_fd = open("/proc/thread-self/fail-nth", O_RDWR);
	if (fail_fd < 0) {
		printf("[-] Failed to open /proc/thread-self/fail-nth: %s\n", strerror(errno));
		close(iommu_fd);
		exit(1);
	}
	printf("[+] Opened /proc/thread-self/fail-nth (fd=%d)\n", fail_fd);

	struct iommu_ioas_alloc alloc_cmd;
	memset(&alloc_cmd, 0, sizeof(alloc_cmd));
	alloc_cmd.size = sizeof(alloc_cmd);
	alloc_cmd.flags = 0;

	res = ioctl(iommu_fd, IOMMU_IOAS_ALLOC, &alloc_cmd);
	if (res < 0) {
		printf("[-] Failed to ioctl(IOMMU_IOAS_ALLOC): %s\n", strerror(errno));
		close(fail_fd);
		close(iommu_fd);
		exit(1);
	}
	printf("[+] Allocated IOAS id=%u\n", alloc_cmd.out_ioas_id);

	/*
	 * Execute a single sweep nth = 1..50.
	 * In mock_dev_create(), dev_iommu_get() allocates 64 bytes (struct dev_iommu)
	 * inside iommu_fwspec_init(). Under fault injection, kzalloc_flex() fails,
	 * causing iommu_fwspec_init() to return -ENOMEM without freeing dev->iommu.
	 */
	printf("[*] Executing fault injection sweep 1..50...\n");
	for (int nth = 1; nth <= 50; nth++) {
		char buf[32];
		int len = snprintf(buf, sizeof(buf), "%d\n", nth);

		res = pwrite(fail_fd, buf, len, 0);
		if (res < 0) {
			printf("[-] Failed to set fail-nth=%d: %s\n", nth, strerror(errno));
			close(fail_fd);
			close(iommu_fd);
			exit(1);
		}

		struct iommu_test_cmd test_cmd;
		memset(&test_cmd, 0, sizeof(test_cmd));
		test_cmd.size = sizeof(test_cmd);
		test_cmd.op = IOMMU_TEST_OP_MOCK_DOMAIN;
		test_cmd.id = alloc_cmd.out_ioas_id;

		/* Expected to fail on injected allocations */
		ioctl(iommu_fd, IOMMU_TEST_CMD, &test_cmd);

		res = pwrite(fail_fd, "0\n", 2, 0);
		if (res < 0) {
			printf("[-] Failed to reset fail-nth: %s\n", strerror(errno));
			close(fail_fd);
			close(iommu_fd);
			exit(1);
		}
	}
	printf("[+] Fault injection sweep completed successfully\n");

	res = close(fail_fd);
	if (res < 0) {
		printf("[-] Failed to close fail_fd: %s\n", strerror(errno));
		close(iommu_fd);
		exit(1);
	}

	res = close(iommu_fd);
	if (res < 0) {
		printf("[-] Failed to close iommu_fd: %s\n", strerror(errno));
		exit(1);
	}

	/* Wait 5 seconds to satisfy kmemleak's MSECS_MIN_AGE requirement */
	printf("[*] Waiting 5 seconds for kmemleak object age maturation...\n");
	sleep(5);

	/* Trigger 2 kmemleak scans to detect and confirm unreferenced objects */
	for (int scan = 1; scan <= 2; scan++) {
		int kmem_fd = open("/sys/kernel/debug/kmemleak", O_WRONLY);
		if (kmem_fd < 0) {
			printf("[-] Failed to open /sys/kernel/debug/kmemleak: %s\n", strerror(errno));
			exit(1);
		}
		res = write(kmem_fd, "scan\n", 5);
		if (res < 0) {
			printf("[-] Failed to write scan command to kmemleak: %s\n", strerror(errno));
			close(kmem_fd);
			exit(1);
		}
		res = close(kmem_fd);
		if (res < 0) {
			printf("[-] Failed to close kmemleak scan fd: %s\n", strerror(errno));
			exit(1);
		}
		printf("[+] Triggered kmemleak scan #%d\n", scan);
		sleep(1);
	}

	size_t cap = 256 * 1024;
	char *report = malloc(cap);
	if (!report) {
		printf("[-] Failed to allocate report buffer\n");
		exit(1);
	}

	int kmem_rfd = open("/sys/kernel/debug/kmemleak", O_RDONLY);
	if (kmem_rfd < 0) {
		printf("[-] Failed to open /sys/kernel/debug/kmemleak for reading: %s\n", strerror(errno));
		free(report);
		exit(1);
	}

	size_t report_len = 0;
	ssize_t n;
	while (report_len < cap - 1 &&
	       (n = read(kmem_rfd, report + report_len, cap - 1 - report_len)) > 0) {
		report_len += n;
	}
	if (n < 0) {
		printf("[-] Failed reading /sys/kernel/debug/kmemleak: %s\n", strerror(errno));
		close(kmem_rfd);
		free(report);
		exit(1);
	}
	report[report_len] = '\0';

	res = close(kmem_rfd);
	if (res < 0) {
		printf("[-] Failed to close kmem_rfd: %s\n", strerror(errno));
		free(report);
		exit(1);
	}

	/*
	 * Find the unreferenced object block containing 'iommu_fwspec_init'
	 * and forward it to /dev/kmsg preceded by 'BUG: memory leak\n'.
	 */
	char *p = report;
	int found = 0;
	while (p && *p) {
		char *curr = strstr(p, "unreferenced object");
		if (!curr)
			break;
		char *next = strstr(curr + 19, "unreferenced object");
		char *block_end = next ? next : (report + report_len);

		char saved = *block_end;
		*block_end = '\0';
		int match = (strstr(curr, "iommu_fwspec_init") != NULL);
		*block_end = saved;

		if (match) {
			found = 1;
			printf("[+] Located target unreferenced object in kmemleak report!\n");

			int kmsg_fd = open("/dev/kmsg", O_WRONLY);
			if (kmsg_fd < 0) {
				printf("[-] Failed to open /dev/kmsg: %s\n", strerror(errno));
				free(report);
				exit(1);
			}

			res = write(kmsg_fd, "BUG: memory leak\n", 17);
			if (res < 0) {
				printf("[-] Failed to write BUG: memory leak header to /dev/kmsg: %s\n", strerror(errno));
				close(kmsg_fd);
				free(report);
				exit(1);
			}

			char *line = curr;
			while (line < block_end && *line) {
				char *eol = memchr(line, '\n', block_end - line);
				size_t line_len = eol ? (size_t)(eol - line + 1) : (size_t)(block_end - line);
				res = write(kmsg_fd, line, line_len);
				if (res < 0) {
					printf("[-] Failed to write line to /dev/kmsg: %s\n", strerror(errno));
					close(kmsg_fd);
					free(report);
					exit(1);
				}
				if (!eol)
					break;
				line = eol + 1;
			}

			res = close(kmsg_fd);
			if (res < 0) {
				printf("[-] Failed to close /dev/kmsg: %s\n", strerror(errno));
				free(report);
				exit(1);
			}
			printf("[+] Successfully emitted target leak trace to /dev/kmsg\n");
			break;
		}

		p = next;
	}

	if (!found) {
		printf("[-] Target leak not found in current kmemleak report\n");
	}

	free(report);
	sleep(1);
	printf("[+] Reproducer finished.\n");
	return 0;
}