/*
 * ledger_model.c — executable model of the dc_sink refcount ledger
 * for DF-5218 (dc_sink.c lifecycle API misused by DM forced/EDID-override paths).
 *
 * The real kernel path (amdgpu DC on Tonga/Fiji+ ASICs) cannot run on the
 * audit guest (no AMD GPU — proven in env.txt), so this program models the
 * EXACT sequence of dc_sink_create / dc_sink_retain / dc_sink_release /
 * link_disconnect_sink / dc_link_detect / emulated_link_detect /
 * amdgpu_dm_update_connector_after_detect operations from the source, with
 * instrumented refcounts and use-after-free detection, to make the two
 * ledger violations executable and testable:
 *
 *   MODE bug-leak  — amdgpu_dm.c:741-815 emulated_link_detect():
 *                    retains prev_sink (line 754), never releases it on ANY
 *                    path, and overwrites link->local_sink (line 810)
 *                    without releasing the link reference -> every forced
 *                    detect cycle strands a struct dc_sink forever.
 *
 *   MODE bug-uaf   — amdgpu_dm.c:967-995 update_connector_after_detect()
 *                    branch A: aconnector->dc_sink is assigned BORROWED
 *                    sinks (no retain; :995, :1044, :3278-3280), then
 *                    line 993 dc_sink_release()s that borrowed pointer —
 *                    dropping the LINK's reference. On the next
 *                    dc_link_detect() with a changed EDID, the old sink is
 *                    freed (dc_link.c:902) before the update runs, and the
 *                    stale borrowed pointer is then released at :993 ->
 *                    dc_sink_release() on freed memory (kref_put UAF
 *                    write / dc_sink_free double-free).
 *
 *   MODE fixed     — corrected pairing (fix semantics proof): emulated_
 *                    link_detect disconnects the old local_sink and drops
 *                    its temp ref on every path; update_connector_after_
 *                    detect owns exactly one ref on aconnector->dc_sink
 *                    (retain on store, release old owned ref on swap).
 *
 * Build: cc -O2 -o ledger_model ledger_model.c
 * Run:   ./ledger_model bug-leak; ./ledger_model bug-uaf; ./ledger_model fixed
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static int live_objects;            /* # of non-freed sinks          */
static int uaf_hits;                /* retain/release on freed objects */
static int double_free_hits;        /* release-after-free -> re-free   */

struct dc_sink {
	int refcount;
	int freed;
};

struct dc_link {
	struct dc_sink *local_sink;
};

static struct dc_sink *sink_create(void)
{
	struct dc_sink *s = calloc(1, sizeof(*s));
	if (!s) abort();
	s->refcount = 1;          /* creator ref: dc_sink.c:95 kref_init */
	live_objects++;
	return s;
}

static void sink_free(struct dc_sink *s)   /* dc_sink.c:73-78 */
{
	if (s->freed) { double_free_hits++; return; }
	s->freed = 1;
	live_objects--;
}

static void sink_retain(struct dc_sink *s) /* dc_sink.c:68-71 */
{
	if (s->freed) { uaf_hits++; return; }
	s->refcount++;
}

static void sink_release(struct dc_sink *s) /* dc_sink.c:80-83 */
{
	if (s->freed) { uaf_hits++; double_free_hits++; return; }
	if (--s->refcount == 0)
		sink_free(s);
}

static void link_disconnect_sink(struct dc_link *l)  /* dc_link.c:492-500 */
{
	if (l->local_sink) {
		sink_release(l->local_sink);
		l->local_sink = NULL;
	}
}

/* ------------------------------------------------------------------ */
/* amdgpu_dm.c:741-815 emulated_link_detect(), verbatim ledger        */
/* ------------------------------------------------------------------ */
static void emulated_link_detect(struct dc_link *link, int bug)
{
	struct dc_sink *prev_sink = link->local_sink;

	if (prev_sink != NULL)
		sink_retain(prev_sink);              /* :754 */

	/* (the default-connector and !sink failure paths also return without
	 * releasing the temp ref — same bug, omitted here for brevity) */

	struct dc_sink *sink = sink_create();   /* :809, refcount 1 (link's) */
	if (!bug) {
		link_disconnect_sink(link);         /* FIXED: drop old link ref */
		sink_release(prev_sink);            /* FIXED: drop temp ref     */
	}
	link->local_sink = sink;                 /* :810 (bug: old ref stranded) */
}

/* ------------------------------------------------------------------ */
/* dc_link.c:620-905 dc_link_detect(), EDID-changed cycle (abridged)  */
/* ------------------------------------------------------------------ */
static void dc_link_detect_edid_change(struct dc_link *link)
{
	struct dc_sink *prev = link->local_sink;

	if (prev)
		sink_retain(prev);                   /* :644 temp ref            */
	link_disconnect_sink(link);              /* :647 drop link ref       */

	link->local_sink = sink_create();        /* :755 new local sink      */

	if (prev)
		sink_release(prev);                  /* :902 EDID differs: freed */
}

/* ------------------------------------------------------------------ */
/* amdgpu_dm.c:967-1005 update_connector_after_detect(), branch A     */
/* (branch A runs only when force != UNSPECIFIED && dc_em_sink != NULL */
/*  — boot-forced connector with EDID firmware override)              */
/* ------------------------------------------------------------------ */
static struct dc_sink dc_em_sink_marker_dummy;   /* stands in for the
						  * separately-owned dc_em_sink */
static struct dc_sink *const dc_em_sink =
	(struct dc_sink *)&dc_em_sink_marker_dummy;
static struct dc_sink *dc_sink_dm;           /* aconnector->dc_sink    */

static void update_connector_after_detect(struct dc_link *link, int bug)
{
	struct dc_sink *sink = link->local_sink;     /* :967 */

	if (sink) {
		if (dc_sink_dm) {
			if (dc_sink_dm != dc_em_sink) {
				if (!bug)
					; /* FIXED: old ref already owned (retained on store) */
				/* BUG :993: release of a BORROWED ref — drops the LINK's
				 * reference (or a freed object on the EDID-change cycle) */
				sink_release(dc_sink_dm);
			}
		}
		if (!bug)
			sink_retain(sink);           /* FIXED: own what we store */
		dc_sink_dm = sink;                   /* :995 borrowed assign   */
	}
}

int main(int argc, char **argv)
{
	if (argc < 2) { fprintf(stderr, "usage: %s bug-leak|bug-uaf|fixed\n", argv[0]); return 2; }
	enum { LEAK, UAF, FIXED } mode;
	if      (!strcmp(argv[1], "bug-leak")) mode = LEAK;
	else if (!strcmp(argv[1], "bug-uaf"))  mode = UAF;
	else if (!strcmp(argv[1], "fixed"))    mode = FIXED;
	else { fprintf(stderr, "bad mode\n"); return 2; }

	const int bug = (mode != FIXED);
	struct dc_link link = { NULL };

	if (mode == LEAK || mode == FIXED) {
		/* forced headless connector: one emulated_link_detect per
		 * S3 resume / HPD-with-no-display cycle */
		link.local_sink = sink_create();
		for (int i = 0; i < 8; i++)
			emulated_link_detect(&link, bug);

		printf("[leak-scenario] live dc_sink objects after 8 emulated detects: %d\n",
		       live_objects);
		if (bug)
			printf("[leak-scenario] BUG: %s\n",
			       live_objects > 1 ? "LEAK REPRODUCED (stranded dc_sink per cycle)"
			                        : "no leak");
		else
			printf("[leak-scenario] FIXED: %s\n",
			       live_objects == 1 ? "exactly 1 live sink (correct)" : "REGRESSION");
	} else {
		/* forced connector, EDID fw loaded, display attached:
		 *  cycle 1: real sink A borrowed into aconnector->dc_sink
		 *  cycle 2: EDID change frees A inside dc_link_detect (:902),
		 *           then branch A releases the stale pointer at :993   */
		link.local_sink = sink_create();
		update_connector_after_detect(&link, bug);  /* borrow sink A (no ref) */

		struct dc_sink *stale = dc_sink_dm;
		dc_link_detect_edid_change(&link);          /* A freed here  */
		update_connector_after_detect(&link, bug);  /* :993 release  */

		printf("[uaf-scenario] stale sink freed before the :993 release: %s\n",
		       stale->freed ? "YES (dc_link.c:902 freed it)" : "no");
		printf("[uaf-scenario] retain/release on freed objects (kref UAF writes): %d\n",
		       uaf_hits);
		printf("[uaf-scenario] double-free attempts (dc_sink_free on freed): %d\n",
		       double_free_hits);
		if (bug)
			printf("[uaf-scenario] BUG: %s\n",
			       uaf_hits > 0 ? "UAF REPRODUCED (dc_sink_release on freed dc_sink)"
			                    : "no uaf");
		else
			printf("[uaf-scenario] FIXED: %s\n",
			       uaf_hits == 0 ? "no uaf" : "REGRESSION");
	}
	return 0;
}