/*
 * DF-5033 deterministic interleaving proof (concurrency oracle).
 *
 * idr_find() in sys/libkern/linux_idr.c performs, with NO lock:
 *
 *   (L1)  if (id < 0 || id >= idp->idr_count) ...
 *   (L2)  ... idp->idr_nodes[id].allocated ...      pointer load + deref
 *
 * Two unsynchronized interleavings with idr_grow() (called under the token
 * by idr_alloc()) turn this into a use-after-free / OOB read:
 *
 *  V1 (preemption): reader loads idp->idr_nodes (L2, 1st half), is
 *     preempted; grower swaps in the new array and kfree()s the old one
 *     (linux_idr.c:460-463); reader resumes and dereferences the freed
 *     array -> UAF read; the garbage "data" pointer is returned and, in
 *     amdgpu_bo_list.c:170, handed to kref_get_unless_zero().
 *
 *  V2 (load reorder): nothing in idr_find() forbids the compiler from
 *     loading idp->idr_nodes BEFORE idp->idr_count (independent loads, no
 *     barriers, no READ_ONCE).  Reader caches the OLD (255-entry) pointer,
 *     then observes the NEW idr_count (e.g. 1023), passes the bound check
 *     with id=400 and reads nodes[400] -> OOB read 2320 bytes past the
 *     freed 4080-byte chunk, plus UAF.
 *
 * This oracle executes those exact interleavings against the byte-identical
 * core (idr_test_core.c) with the same poison-on-free instrumentation as
 * idr_race.c.  Basline build -> both variants trip the poison.  Fixed build
 * (-DIDR_FIXED_BUILD: idr_find() takes the lwkt token, here a real mutex)
 * -> idr_find() serializes against idr_alloc()/idr_grow() and both variants
 * return clean data.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <pthread.h>
#include <sys/types.h>

#include "idr_test_core.c"

struct mycpu_s mycpu_stub;

struct ablk {
	void		*p;
	size_t		sz;
	struct ablk	*next;
};
static struct ablk *allocs;
static pthread_mutex_t amtx = PTHREAD_MUTEX_INITIALIZER;

void *
xmalloc(size_t bytes)
{
	struct ablk *a = malloc(sizeof(*a) + bytes);
	void *p;

	memset(a, 0, sizeof(*a));
	a->p = a + 1;
	a->sz = bytes;
	pthread_mutex_lock(&amtx);
	a->next = allocs;
	allocs = a;
	pthread_mutex_unlock(&amtx);
	p = a->p;
	memset(p, 0, bytes);
	return p;
}

void
xkfree(void *ptr)
{
	struct ablk *a, **pp;

	if (ptr == NULL)
		return;
	pthread_mutex_lock(&amtx);
	for (pp = &allocs; (a = *pp) != NULL; pp = &a->next) {
		if (a->p == ptr) {
			*pp = a->next;
			break;
		}
	}
	pthread_mutex_unlock(&amtx);
	if (a == NULL) {
		fprintf(stderr, "xkfree: untracked %p\n", ptr);
		abort();
	}
	memset(ptr, 0xDD, a->sz);
}

static struct idr idr;

static void
fill(int n)
{
	int i, r;

	for (i = 0; i < n; i++) {
		r = idr_alloc(&idr, (void *)(0x5100000000000000ULL + i),
		    1, 0, 0);
		if (r < 0)
			abort();
	}
}

static int
poisoned(uintptr_t v)
{
	return v == 0xDDDDDDDDDDDDDDDDULL;
}

int
main(void)
{
	int id, allocated, r = 0;
	void *data;

	setvbuf(stdout, NULL, _IONBF, 0);

#ifdef IDR_FIXED_BUILD
	printf("== FIXED build: idr_find() holds the idr token ==\n");
	idr_init(&idr);
	fill(100);
	/* reader preempted for the whole grow; with the fix, idr_find()
	 * serializes against it internally */
	fill(300);			/* forces idr_grow + kfree(old) */
	id = 4;
	data = idr_find(&idr, id);
	printf("V1 call-after-grow idr_find(4) = %016lx -> %s\n",
	    (unsigned long)(uintptr_t)data,
	    poisoned((uintptr_t)data) ? "POISON (UAF READ)" :
	    "clean (valid tag / NULL)");
	r |= poisoned((uintptr_t)data);

	idr_destroy(&idr);
	idr_init(&idr);
	fill(100);
	{
		struct idr_node *nodes = idr.idr_nodes;

		(void)nodes;
		fill(300);		/* grow: old array freed */
		id = 400;		/* valid only under the NEW count */
		/* with the fix the caller cannot bypass idr_find()'s token;
		 * equivalent fixed-mode operation is the plain call: */
		data = idr_find(&idr, id);
		printf("V2 call idr_find(400)            = %016lx -> %s\n",
		    (unsigned long)(uintptr_t)data,
		    poisoned((uintptr_t)data) ? "POISON (UAF/OOB READ)" :
		    "clean (valid tag / NULL)");
		r |= poisoned((uintptr_t)data);
	}
	printf("FIXED RESULT: %s\n", r ? "FAILED (still racy)" :
	    "no unsynchronized read possible; both variants clean");
	return r;

#else
	printf("== BASELINE build: idr_find() takes NO lock (as in the "
	    "kernel) ==\n");

	/* ---- V1: preemption between pointer load and dereference ---- */
	idr_init(&idr);
	fill(100);
	id = 4;
	if (id >= idr.idr_count)
		abort();	/* bound check (L1): 4 < 255 */
	{
		/* idr_find()'s pointer load (L2, first half) */
		struct idr_node *nodes = idr.idr_nodes;

		/* === reader preempted here: grower runs idr_alloc() === */
		fill(300);	/* idr_grow(): swap + kfree(old array) */

		/* reader resumes; idr_find()'s dereference (L2, second half) */
		allocated = nodes[id].allocated;
		data = allocated ? nodes[id].data : NULL;
	}
	printf("V1 preemption-in-window: allocated=%08x data=%016lx -> %s\n",
	    (unsigned)allocated, (unsigned long)(uintptr_t)data,
	    poisoned((uintptr_t)data) ?
	    "*** USE-AFTER-FREE READ of freed idr_nodes[] (value returned "
	    "to caller; kernel hands it to kref_get_unless_zero) ***" :
	    "clean");
	r |= poisoned((uintptr_t)data) ? 1 : 0;

	/* ---- V2: compiler load reorder (new count, old nodes) ---- */
	idr_destroy(&idr);
	idr_init(&idr);
	fill(100);
	{
		/* hoisted pointer load (legal: independent loads, no
		 * barriers/READ_ONCE in idr_find) */
		struct idr_node *nodes = idr.idr_nodes;

		/* grower runs: idr_nodes = new (1023 entries), old freed */
		fill(300);
		id = 400;		/* valid only under the NEW count */
		if (id >= idr.idr_count)
			abort();	/* 400 < 1023 passes */
		allocated = nodes[id].allocated;   /* OLD, freed, 255 entries */
		data = allocated ? nodes[id].data : NULL;
	}
	printf("V2 load-reorder:          allocated=%08x data=%016lx -> %s\n",
	    (unsigned)allocated, (unsigned long)(uintptr_t)data,
	    poisoned((uintptr_t)data) ?
	    "*** OOB READ past freed 255-entry chunk (id 400 vs 255 "
	    "entries) + UAF ***" : "clean");
	r |= poisoned((uintptr_t)data) ? 1 : 0;

	printf("BASELINE RESULT: %s\n", r ?
	    "both interleavings dereference freed idr_nodes[] memory" :
	    "unexpected: no poison observed");
	return r;
#endif
}