/*
 * DF-4772 PoC — exclusive-fence reference leak in amdgpu_pasid_free_delayed()
 * via the DragonFly drm-port's reservation_object_get_fences_rcu() shim.
 *
 * amdgpu_ids.c:111-126 (amdgpu_pasid_free_delayed) calls
 *
 *	r = reservation_object_get_fences_rcu(resv, NULL, &count, &fences);
 *	if (r)
 *		goto fallback;
 *	if (count == 0) {
 *		amdgpu_pasid_free(pasid);
 *		return;
 *	}
 *
 * The DFly shim (sys/dev/drm/linux_reservation.c) unconditionally takes a
 * reference on obj->fence_excl (lines 356-357) but only hands that reference
 * to the caller when pfence_excl != NULL (lines 410-411):
 *
 *	if (pfence_excl)
 *		*pfence_excl = fence_excl;
 *
 * Upstream Linux ends this function with
 *
 *	if (pfence_excl)
 *		*pfence_excl = fence_excl;
 *	else if (fence_excl)
 *		dma_fence_put(fence_excl);
 *
 * — the port dropped the else-put.  Every call with pfence_excl == NULL
 * therefore leaks one dma_fence reference whenever the reservation object
 * carries an exclusive fence.  amdgpu hits this on EVERY drm fd close
 * (amdgpu_driver_postclose_kms, amdgpu_kms.c:1076-1077, resv = the VM
 * page-directory's reservation, which carries excl fences from every
 * page-directory update, amdgpu_vm.c:1559/1925 via amdgpu_bo_fence(...,1)).
 *
 * This module replicates the exact call sequence using the port's own
 * primitives (the same shim code the real amdgpu driver executes).  No GPU
 * is needed — the leak is pure shim logic.
 *
 * sysctl kern.df4772_run=1  run N iterations of the amdgpu fd-close path
 * sysctl kern.df4772_n=<n>  iteration count (default 1000, max 4096)
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/systm.h>
#include <sys/module.h>
#include <sys/sysctl.h>
#include <sys/malloc.h>
#include <sys/lock.h>

#include <linux/slab.h>
#include <linux/kref.h>
#include <linux/dma-fence.h>
#include <linux/reservation.h>

#define DF4772_MAX_TRACK	4096

struct df4772_fence {
	struct dma_fence	base;
	struct lock		lock;
};

static volatile int	df4772_release_count;	/* ops->release calls */
static int		df4772_n = 1000;
static struct df4772_fence *df4772_tracked[DF4772_MAX_TRACK];
static int		df4772_tracked_n;	/* still-alive (leaked) fences */

/* kzalloc/kfree (linux/slab.h shim) allocate/free with M_DRM — stay on it */

static const char *
df4772_driver_name(struct dma_fence *f)
{
	return ("df4772");
}

static const char *
df4772_timeline_name(struct dma_fence *f)
{
	return ("pasid-pd");
}

static bool
df4772_enable_signaling(struct dma_fence *f)
{
	return (true);
}

static signed long
df4772_wait(struct dma_fence *f, bool intr, signed long timeout)
{
	return (dma_fence_default_wait(f, intr, timeout));
}

/* == amdgpu_fence_release equivalent: count + free the container */
static void
df4772_fence_release(struct dma_fence *f)
{
	struct df4772_fence *fence = container_of(f, struct df4772_fence, base);

	atomic_add_int(&df4772_release_count, 1);
	kfree(fence);
}

static const struct dma_fence_ops df4772_fence_ops = {
	.get_driver_name	= df4772_driver_name,
	.get_timeline_name	= df4772_timeline_name,
	.enable_signaling	= df4772_enable_signaling,
	.wait			= df4772_wait,
	.release		= df4772_fence_release,
};

/*
 * amdgpu_pasid_free_delayed() replica — amdgpu_ids.c:111-126 verbatim for
 * the count==0 path (the common case at fd close: the PD resv carries an
 * exclusive fence and no unsignaled shared fences).  The pasid ida ops are
 * irrelevant to the fence-ref accounting and are elided.
 */
static void
df4772_pasid_free_delayed(struct reservation_object *resv)
{
	struct dma_fence *fence, **fences;
	unsigned count;
	int r;

	r = reservation_object_get_fences_rcu(resv, NULL, &count, &fences);
	if (r)
		goto fallback;

	if (count == 0) {
		/* amdgpu_pasid_free(pasid); */
		return;			/* <-- excl-fence ref leaks HERE */
	}

	if (count == 1) {
		fence = fences[0];
		kfree(fences);
	} else {
		/* fence-array path — not exercised; DF-1861 territory */
		kfree(fences);
		return;
	}
	dma_fence_put(fence);
	return;

fallback:
	return;
}

static void
df4772_run(void)
{
	int i;
	int before;
	int leaked = 0;

	df4772_release_count = 0;
	df4772_tracked_n = 0;
	before = df4772_n;

	for (i = 0; i < before && i < DF4772_MAX_TRACK; i++) {
		struct df4772_fence *df;
		struct reservation_object resv;
		u64 ctx = dma_fence_context_alloc(1);
		int r0, r1, r2, rel_before;

		rel_before = df4772_release_count;
		df = kzalloc(sizeof(*df), GFP_KERNEL);
		lockinit(&df->lock, "df4772l", 0, LK_CANRECURSE);
		dma_fence_init(&df->base, &df4772_fence_ops, &df->lock,
		    ctx, i + 1);
		r0 = kref_read(&df->base.refcount);		/* 1: ours */

		reservation_object_init(&resv);
		ww_mutex_lock(&resv.lock, NULL);
		reservation_object_add_excl_fence(&resv, &df->base);
		ww_mutex_unlock(&resv.lock);
		r1 = kref_read(&df->base.refcount);		/* 2: +resv */

		/* == amdgpu_driver_postclose_kms -> amdgpu_pasid_free_delayed */
		df4772_pasid_free_delayed(&resv);
		r2 = kref_read(&df->base.refcount);		/* 3 with leak */

		dma_fence_signal(&df->base);
		dma_fence_put(&df->base);			/* drop ours */
		reservation_object_fini(&resv);			/* drop resv's */

		if (i == 0) {
			kprintf("df4772 [iter 0] refcounts: init=%d "
			    "add_excl=%d after_get_fences_rcu_NULL=%d "
			    "(correct value: %d)\n", r0, r1, r2, r1);
		}

		if (df4772_release_count == rel_before) {
			/* not released: object leaked — track for cleanup */
			df4772_tracked[df4772_tracked_n++] = df;
			leaked++;
		}
		/* else: release fired and kfree'd the container */
	}

	kprintf("df4772: iterations=%d released=%d leaked=%d "
	    "(upstream-correct behaviour: released==%d, leaked==0)\n",
	    before, df4772_release_count, leaked, before);

	if (df4772_release_count != before)
		kprintf("df4772: LEAK CONFIRMED — %d dma_fence references "
		    "leaked by reservation_object_get_fences_rcu(pfence_excl="
		    "NULL)\n", before - df4772_release_count);
	else
		kprintf("df4772: BALANCED — every fence released\n");
}

static int
df4772_run_sysctl(SYSCTL_HANDLER_ARGS)
{
	int run = 0;
	int error;

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

static int
df4772_n_sysctl(SYSCTL_HANDLER_ARGS)
{
	return (sysctl_handle_int(oidp, &df4772_n, 0, req));
}

static SYSCTL_NODE(, OID_AUTO, df4772, CTLFLAG_RW, 0, "DF-4772 PoC");
SYSCTL_PROC(_df4772, OID_AUTO, run, CTLTYPE_INT | CTLFLAG_RW,
    0, 0, df4772_run_sysctl, "I", "run the pasid_free_delayed replica");
SYSCTL_PROC(_df4772, OID_AUTO, n, CTLTYPE_INT | CTLFLAG_RW,
    0, 0, df4772_n_sysctl, "I", "iterations");

static int
df4772_modevent(module_t mod, int type, void *data)
{
	int cleaned;

	switch (type) {
	case MOD_LOAD:
		kprintf("df4772: loaded (DF-4772 PoC)\n");
		break;
	case MOD_UNLOAD:
		/*
		 * Cleanup: fences still referenced (the leak) are freed
		 * directly — their release op lives in this module and the
		 * refcount is stuck at 1, so nothing else will ever free
		 * them.  On a fixed kernel nothing is tracked (all were
		 * released) and this loop is a no-op.
		 */
		for (cleaned = 0; cleaned < df4772_tracked_n; cleaned++)
			kfree(df4772_tracked[cleaned]);
		cleaned = df4772_tracked_n;
		df4772_tracked_n = 0;
		kprintf("df4772: unloaded (cleaned %d leaked fence(s))\n",
		    cleaned);
		break;
	default:
		return (EOPNOTSUPP);
	}
	return (0);
}

static moduledata_t df4772_mod = { "df4772", df4772_modevent, 0 };
DECLARE_MODULE(df4772, df4772_mod, SI_SUB_DRIVERS, SI_ORDER_MIDDLE);