/*
 * DF-5689 — lockless-reader vs idr_grow() use-after-free race rig for
 * sys/libkern/linux_idr.c (kernel branch, compiled UNMODIFIED).
 *
 * Defect under test: idr_find() (linux_idr.c:515-528) and idr_for_each()
 * (linux_idr.c:530-546) read idp->idr_count / idp->idr_nodes / the node
 * array with NO idr_token held, while idr_grow() — called under the
 * token from idr_pre_get()/idr_alloc() (linux_idr.c:275, 375) — swaps
 * idp->idr_nodes and kfree()s the OLD array (linux_idr.c:449-459).
 * A concurrent lockless reader that has loaded the old idr_nodes pointer
 * dereferences freed memory (UAF READ): idr_find() returns whatever the
 * freed slab chunk contains as a kernel object pointer.
 *
 * In-tree unguarded consumer: i915_gem_object_lookup_rcu()
 * (sys/dev/drm/i915/i915_gem_object.h:307-314) calls idr_find() under
 * rcu_read_lock() only (i915_gem.c:4532-4533, 4971-4972); DragonFly's
 * kfree() has no RCU grace period, so the Linux RCU argument does not
 * transfer.  Any other consumer that calls idr_find()/idr_for_each()
 * without the external lock that serializes its idr_alloc()s is equally
 * exposed.
 *
 * Detection: the shim allocator poisons every kfree()d chunk with
 * 0xDD (see shim/sys/malloc.h).  Any read of a freed idr_nodes array
 * observes allocated=0xDDDDDDDD (passes the !=0 guard) and
 * data=0xDDDDDDDDDDDDDDDD — the rigs flag that pointer value on sight.
 * (On a host with ASAN the same run reports heap-use-after-free
 * directly; the guest has no ASAN, hence the poison canary.)
 *
 * usage: df5689_race <seconds> [grow_target]
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <pthread.h>
#include <unistd.h>
#include <time.h>

#include "linux_idr.c"			/* file under audit, kernel branch */

static struct idr	R;
static volatile int	stop;
static volatile int	pause_flag;	/* readers quiesce while set so
					 * idr_destroy()/idr_init() (a
					 * caller-teardown operation, out
					 * of scope) cannot masquerade as
					 * the idr_grow() race */
static volatile int	in_readers;	/* readers between enter-inc and
					 * exit-dec (two-sided barrier) */
static volatile int	hit;		/* poison pointer observed */
static volatile unsigned long long find_iters, foreach_iters, grows;

static int
reader_enter(void)
{
	for (;;) {
		__sync_fetch_and_add(&in_readers, 1);
		__sync_synchronize();
		if (!pause_flag)
			return (1);
		__sync_fetch_and_add(&in_readers, -1);
		while (pause_flag && !stop && !hit)
			usleep(50);
		if (stop || hit)
			return (0);
	}
}

static void
reader_exit(void)
{
	__sync_synchronize();
	__sync_fetch_and_add(&in_readers, -1);
}

unsigned long rig_live_chunks, rig_live_bytes, rig_free_poisons;
struct rig_mycpu rig_mycpu;

#define	TARGET_COUNT	2047		/* 255->511->1023->2047: 3 grows */

static void *
find_thread(void *arg)
{
	unsigned int seed = (unsigned int)(uintptr_t)arg;
	unsigned long long n = 0;

	while (!stop && !hit) {
		if (!reader_enter())
			break;
		{
			int id = (seed = seed * 1103515245 + 12345) %
			    TARGET_COUNT;
			void *p = idr_find(&R, id);

			reader_exit();
			n++;
			if (p == RIG_POISON_PTR) {
				hit = 1;
				fprintf(stderr,
				    "UAF-READ: idr_find(%d) returned data "
				    "from a FREED idr_nodes array (%p)\n",
				    id, p);
			}
		}
	}
	__sync_fetch_and_add(&find_iters, n);
	return (NULL);
}

struct fectx {
	volatile int	dummy;
};

static int
fe_cb(int id, void *p, void *data __attribute__((unused)))
{
	if (p == RIG_POISON_PTR) {
		if (!hit)
			fprintf(stderr,
			    "UAF-READ: idr_for_each() node %d returned data "
			    "from a FREED idr_nodes array (%p)\n", id, p);
		hit = 1;
	}
	return (0);
}

static void *
foreach_thread(void *arg)
{
	unsigned long long n = 0;

	while (!stop && !hit) {
		if (!reader_enter())
			break;
		idr_for_each(&R, fe_cb, NULL);
		reader_exit();
		n++;
	}
	__sync_fetch_and_add(&foreach_iters, n);
	return (NULL);
}

int
main(int ac, char **av)
{
	long seconds = 20;
	int run_foreach = 0;
	pthread_t t1, t2, t3;
	unsigned long long cycles = 0;
	time_t t0;

	if (ac >= 2)
		seconds = strtol(av[1], NULL, 0);
	if (ac >= 3)
		run_foreach = 1;	/* argv[2]: enable the idr_for_each
					 * reader — NOTE: against the
					 * stock source this variant
					 * usually dies with SIGSEGV
					 * (OOB read past the freed old
					 * array: for_each caches the
					 * nodes pointer but re-reads
					 * idr_count every iteration) */

	t0 = time(NULL);

	idr_init(&R);

	pthread_create(&t1, NULL, find_thread, (void *)0x1234abcdUL);
	pthread_create(&t2, NULL, find_thread, (void *)0x9e3779b9UL);
	if (run_foreach)
		pthread_create(&t3, NULL, foreach_thread, NULL);

	idr_init(&R);

	while (!hit && time(NULL) - t0 < seconds) {
		int id = 0;

		while (R.idr_count < TARGET_COUNT && !hit) {
			idr_alloc(&R,
			    (void *)(0x515000000000UL + (unsigned)id++),
			    0, 0, 0);
		}
		grows = R.idr_nexpands;
		/* quiesce readers before teardown (destroy/init is a
		 * caller-contract operation, not the race under test);
		 * the idr_grow() kfrees above happened with readers
		 * running freely */
		pause_flag = 1;
		__sync_synchronize();
		while (in_readers > 0)
			usleep(50);
		idr_destroy(&R);
		idr_init(&R);
		pause_flag = 0;
		cycles++;
	}
	stop = 1;
	pthread_join(t1, NULL);
	pthread_join(t2, NULL);
	if (run_foreach)
		pthread_join(t3, NULL);

	printf("cycles=%llu grows/cycle=%lu finds=%llu for_eachs=%llu "
	    "kfrees_poisoned=%lu -> %s\n", (unsigned long long)cycles,
	    (unsigned long)grows, find_iters, foreach_iters,
	    rig_free_poisons,
	    hit ? "USE-AFTER-FREE READ OBSERVED" : "no hit");
	return (hit ? 2 : 0);
}