/*
 * atom_harness_fixed.c — the SAME amdgpu atom.c expressions as
 * atom_harness.c, but with the DF-4337 central fix ported to the
 * amdgpu twin: an atom_context.bios_size plumbed from the amdgpu_bios.c
 * readers (adev->bios_size already exists — amdgpu.h:860) and every
 * bios-offset dereference gated behind amdgpu_atom_in_bios().
 *
 * Every case that traps in the baseline harness must be cleanly
 * rejected here (ok=0 / ret=-EINVAL / skip) with exit 0.
 *
 * Build:  cc -O1 -g -o atom_harness_fixed atom_harness_fixed.c
 * Run:    ./atom_harness_fixed <case>   (same case names)
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>
#include <sys/mman.h>
#include <unistd.h>
#include <errno.h>

#ifndef EINVAL
#define EINVAL 22
#endif

typedef uint16_t u16;

static inline uint8_t get_u8(void *bios, int ptr)
{
	return ((unsigned char *)bios)[ptr];
}
static inline uint16_t get_u16(void *bios, int ptr)
{
	return get_u8(bios, ptr) | (((uint16_t)get_u8(bios, ptr + 1)) << 8);
}
static inline uint32_t get_u32(void *bios, int ptr)
{
	return get_u16(bios, ptr) | (((uint32_t)get_u16(bios, ptr + 2)) << 16);
}
#define CU8(ptr)  get_u8(ctx->bios, (ptr))
#define CU16(ptr) get_u16(ctx->bios, (ptr))
#define CU32(ptr) get_u32(ctx->bios, (ptr))
#define U16(ptr)  get_u16(ctx->bios, (ptr))
#define U32(ptr)  get_u32(ctx->bios, (ptr))

struct atom_context {
	void *bios;
	int bios_size;			/* <- the fix */
	uint32_t cmd_table, data_table;
	uint16_t *iio;
	uint16_t data_block;
	uint8_t shift;
	char vbios_version[20];
};

/* the fix: every bios offset validated against the allocation length */
static bool amdgpu_atom_in_bios(const struct atom_context *ctx, int off, int len)
{
	return off >= 0 && len >= 0 && off <= ctx->bios_size - len;
}

static void *guarded_alloc(size_t sz)
{
	long ps = 4096;
	size_t guard = (size_t)16 * 1024 * 1024;
	size_t len = (size_t)ps + guard;
	char *base = mmap(NULL, len, PROT_READ | PROT_WRITE,
			  MAP_PRIVATE | MAP_ANON, -1, 0);
	if (base == MAP_FAILED) { perror("mmap"); exit(3); }
	if (mprotect(base + ps, guard, PROT_NONE) != 0) { perror("mprotect"); exit(3); }
	return base + ps - sz;
}

/* ---- fixed amdgpu_atom_parse_data_header ---- */
bool amdgpu_atom_parse_data_header(struct atom_context *ctx, int index,
			    uint16_t * size, uint8_t * frev, uint8_t * crev,
			    uint16_t * data_start)
{
	int offset = index * 2 + 4;
	u16 *mdt = (u16 *)(ctx->bios + ctx->data_table + 4);
	int idx;

	if (index < 0 || !amdgpu_atom_in_bios(ctx, ctx->data_table + offset, 2))
		return false;
	idx = CU16(ctx->data_table + offset);
	if (!mdt[index])
		return false;

	if (!amdgpu_atom_in_bios(ctx, idx, 4) ||
	    !amdgpu_atom_in_bios(ctx, idx, CU16(idx)))
		return false;
	if (size)
		*size = CU16(idx);
	if (frev)
		*frev = CU8(idx + 2);
	if (crev)
		*crev = CU8(idx + 3);
	*data_start = idx;
	return true;
}

/* ---- fixed amdgpu_atom_parse_cmd_header ---- */
bool amdgpu_atom_parse_cmd_header(struct atom_context *ctx, int index,
			   uint8_t * frev, uint8_t * crev)
{
	int offset = index * 2 + 4;
	u16 *mct = (u16 *)(ctx->bios + ctx->cmd_table + 4);
	int idx;

	if (index < 0 || !amdgpu_atom_in_bios(ctx, ctx->cmd_table + offset, 2))
		return false;
	idx = CU16(ctx->cmd_table + offset);
	if (!mct[index])
		return false;

	if (!amdgpu_atom_in_bios(ctx, idx, 4))
		return false;
	if (frev)
		*frev = CU8(idx + 2);
	if (crev)
		*crev = CU8(idx + 3);
	return true;
}

/* ---- fixed ATOM_ARG_ID ---- */
static uint32_t atom_arg_id_read(struct atom_context *ctx, int opnd_off, bool *ok)
{
	uint32_t idx, val = 0;

	idx = get_u16(ctx->bios, opnd_off);
	if (!amdgpu_atom_in_bios(ctx, idx + ctx->data_block, 4)) {
		*ok = false;
		return 0;
	}
	*ok = true;
	val = U32(idx + ctx->data_block);
	return val;
}

/* ---- fixed master-entry reads ---- */
static bool calltable_entry_ok(struct atom_context *ctx, int idx)
{
	return amdgpu_atom_in_bios(ctx, ctx->cmd_table + 4 + 2 * idx, 2);
}
static bool setdatablock_entry_ok(struct atom_context *ctx, int idx)
{
	return amdgpu_atom_in_bios(ctx, ctx->data_table + 4 + 2 * idx, 2);
}

/* ---- fixed atom_index_iio ---- */
#define ATOM_IIO_START	1
#define ATOM_IIO_END	9

static int *atom_iio_len;

static void atom_index_iio(struct atom_context *ctx, int base)
{
	ctx->iio = calloc(256, 2);
	if (!ctx->iio)
		return;
	while (amdgpu_atom_in_bios(ctx, base, 2) && CU8(base) == ATOM_IIO_START) {
		ctx->iio[CU8(base + 1)] = base + 2;
		base += 2;
		while (amdgpu_atom_in_bios(ctx, base, 1) && CU8(base) != ATOM_IIO_END) {
			int op = CU8(base);

			if (op > ATOM_IIO_END)
				return;		/* bogus opcode: stop */
			base += atom_iio_len[op];
		}
		base += 3;
	}
}

/* ---- fixed part-string: require NUL inside the image ---- */
#define ATOM_ROM_PART_NUMBER_PTR	0x6E

static bool amdgpu_atom_parse_partstr(struct atom_context *ctx)
{
	char *str;
	u16 idx;
	int end;

	if (!amdgpu_atom_in_bios(ctx, ATOM_ROM_PART_NUMBER_PTR, 2))
		return false;
	idx = CU16(ATOM_ROM_PART_NUMBER_PTR);
	if (idx == 0)
		idx = 0x80;

	if (!amdgpu_atom_in_bios(ctx, idx, 1))
		return false;
	for (end = idx; end < ctx->bios_size && ((char *)ctx->bios)[end]; end++)
		;
	if (end == ctx->bios_size)
		return false;		/* unterminated: reject */

	str = (char *)ctx->bios + idx;
	if (*str != '\0') {
		printf("ATOM BIOS: %s\n", str);
		strlcpy(ctx->vbios_version, str, sizeof(ctx->vbios_version));
	}
	return true;
}

/* ---- fixed asic_init FWI reads ---- */
#define ATOM_DATA_FWI_PTR	0xC
#define ATOM_FWI_DEFSCLK_PTR	8
#define ATOM_FWI_DEFMCLK_PTR	0xC

static int amdgpu_atom_asic_init_reads(struct atom_context *ctx, uint32_t ps[16])
{
	int hwi;

	if (!amdgpu_atom_in_bios(ctx, ctx->data_table + ATOM_DATA_FWI_PTR, 2))
		return 1;
	hwi = CU16(ctx->data_table + ATOM_DATA_FWI_PTR);
	if (!amdgpu_atom_in_bios(ctx, hwi, 0x28))
		return 1;

	memset(ps, 0, 64);
	ps[0] = CU32(hwi + ATOM_FWI_DEFSCLK_PTR);
	ps[1] = CU32(hwi + ATOM_FWI_DEFMCLK_PTR);
	if (!ps[0] || !ps[1])
		return 1;
	return 0;
}

/* ---- fixed fetch loop ---- */
#define ATOM_OP_CNT	127
#define ATOM_OP_EOT	91

static int fetch_loop(struct atom_context *ctx, int index)
{
	int base, len, ws, ps, ptr;
	unsigned char op;
	int ret = 0;

	if (!amdgpu_atom_in_bios(ctx, ctx->cmd_table + 4 + 2 * index, 2))
		return -EINVAL;
	base = CU16(ctx->cmd_table + 4 + 2 * index);
	if (!base)
		return -EINVAL;
	if (!amdgpu_atom_in_bios(ctx, base, 6))
		return -EINVAL;

	len = CU16(base + 0);
	ws = CU8(base + 4);
	ps = CU8(base + 5) & 0x7F;
	ptr = base + 6;
	(void)len; (void)ws; (void)ps;

	while (1) {
		/* the execution pointer is attacker-controlled VBIOS data
		 * (jumps, switch walks); never let it leave the image */
		if (!amdgpu_atom_in_bios(ctx, ptr, 1)) {
			printf("DRM_ERROR: atombios execution pointer outside BIOS image @ 0x%04X\n", ptr);
			ret = -EINVAL;
			goto out;
		}
		op = CU8(ptr++);
		if (op < ATOM_OP_CNT && op > 0) {
			if (op == ATOM_OP_EOT)
				break;
			continue;
		} else
			break;
	}
out:
	return ret;
}

/* ---- fixed shift helpers (DF-4348) ---- */
static uint32_t atom_shl(uint32_t v, uint32_t s) { return s < 32 ? (v << s) : 0; }
static uint32_t atom_shr(uint32_t v, uint32_t s) { return s < 32 ? (v >> s) : 0; }

/* fixed IIO width mask: widths >= 32 select all bits, width 0 selects none */
static uint32_t atom_iio_mask(uint32_t width)
{
	return width >= 32 ? 0xFFFFFFFFu : (width == 0 ? 0u : (0xFFFFFFFFu >> (32 - width)));
}

static void shift_fixed_demo(void)
{
	printf("  fixed shl(0x%x,40)=0x%08x  shr=0x%08x  IIO width0 mask=0x%08x  WS or(200)=0x%08x\n",
	       0x0000FFFFu, atom_shl(0x0000FFFFu, 40), atom_shr(0x0000FFFFu, 40),
	       0x12345678u & ~(atom_iio_mask(0) << 3),
	       atom_shl(1, 200));
}

int main(int argc, char **argv)
{
	const char *cs = argc > 1 ? argv[1] : "data_header";
	printf("case: %s (fixed logic)\n", cs);

	if (!strcmp(cs, "data_header")) {
		struct atom_context sctx, *ctx = &sctx;
		memset(&sctx, 0, sizeof(sctx));
		ctx->bios = guarded_alloc(512);
		ctx->bios_size = 512;
		ctx->data_table = 0x20;
		((uint8_t *)ctx->bios)[0x24] = 0x00;
		((uint8_t *)ctx->bios)[0x25] = 0x80;
		uint16_t size = 0, data_start = 0;
		bool ok = amdgpu_atom_parse_data_header(ctx, 0, &size, NULL, NULL, &data_start);
		printf("  ok=%d size=0x%04x start=0x%04x (entry rejected vs bios_size=512)\n",
		       ok, size, data_start);
		return 0;
	}
	if (!strcmp(cs, "cmd_header")) {
		struct atom_context sctx, *ctx = &sctx;
		memset(&sctx, 0, sizeof(sctx));
		ctx->bios = guarded_alloc(512);
		ctx->bios_size = 512;
		ctx->cmd_table = 0x20;
		((uint8_t *)ctx->bios)[0x24] = 0x00;
		((uint8_t *)ctx->bios)[0x25] = 0x80;
		uint8_t frev = 0, crev = 0;
		bool ok = amdgpu_atom_parse_cmd_header(ctx, 0, &frev, &crev);
		printf("  ok=%d frev=%u crev=%u (rejected)\n", ok, frev, crev);
		return 0;
	}
	if (!strcmp(cs, "id_operand")) {
		struct atom_context sctx, *ctx = &sctx;
		memset(&sctx, 0, sizeof(sctx));
		ctx->bios = guarded_alloc(512);
		ctx->bios_size = 512;
		ctx->data_block = 0x8000;
		((uint8_t *)ctx->bios)[0x10] = 0xFC;
		((uint8_t *)ctx->bios)[0x11] = 0xFF;
		bool ok = true;
		uint32_t v = atom_arg_id_read(ctx, 0x10, &ok);
		printf("  ok=%d ID value = 0x%08x (rejected -> val=0)\n", ok, v);
		return 0;
	}
	if (!strcmp(cs, "master_entry")) {
		struct atom_context sctx, *ctx = &sctx;
		memset(&sctx, 0, sizeof(sctx));
		ctx->bios = guarded_alloc(512);
		ctx->bios_size = 512;
		ctx->cmd_table = 0x20;
		ctx->data_table = 0x30;
		printf("  calltable entry 250 in_bios=%d (0=skip execution)\n",
		       calltable_entry_ok(ctx, 250));
		printf("  setdatablock entry 250 in_bios=%d (0=data_block=0)\n",
		       setdatablock_entry_ok(ctx, 250));
		return 0;
	}
	if (!strcmp(cs, "iio")) {
		struct atom_context sctx, *ctx = &sctx;
		memset(&sctx, 0, sizeof(sctx));
		ctx->bios = guarded_alloc(512);
		ctx->bios_size = 512;
		atom_iio_len = guarded_alloc(10 * sizeof(int));
		{ static const int lens[10] = {1,2,3,3,3,3,4,4,4,3};
		  memcpy(atom_iio_len, lens, sizeof(lens)); }
		uint8_t *b = ctx->bios;
		b[0x40] = ATOM_IIO_START;
		b[0x41] = 0x05;
		b[0x42] = 200;
		atom_index_iio(ctx, 0x40);
		printf("  iio walk stopped cleanly (opcode > ATOM_IIO_END -> return)\n");
		return 0;
	}
	if (!strcmp(cs, "partstr")) {
		struct atom_context sctx, *ctx = &sctx;
		memset(&sctx, 0, sizeof(sctx));
		size_t sz = 512;
		ctx->bios = guarded_alloc(sz);
		ctx->bios_size = sz;
		((uint8_t *)ctx->bios)[0x6E] = 0x00;
		((uint8_t *)ctx->bios)[0x6F] = 0x01;
		memset((uint8_t *)ctx->bios + 0x100, 'B', sz - 0x100);
		bool ok = amdgpu_atom_parse_partstr(ctx);
		printf("  ok=%d (unterminated message rejected, no %%s walk)\n", ok);
		return 0;
	}
	if (!strcmp(cs, "asic_fwi")) {
		struct atom_context sctx, *ctx = &sctx;
		memset(&sctx, 0, sizeof(sctx));
		ctx->bios = guarded_alloc(512);
		ctx->bios_size = 512;
		ctx->data_table = 0x20;
		((uint8_t *)ctx->bios)[0x2C] = 0xF8;
		((uint8_t *)ctx->bios)[0x2D] = 0xFF;
		uint32_t ps[16];
		int r = amdgpu_atom_asic_init_reads(ctx, ps);
		printf("  ret=%d (FWI pointer 0xFFF8 rejected vs bios_size=512)\n", r);
		return 0;
	}
	if (!strcmp(cs, "fetchloop")) {
		struct atom_context sctx, *ctx = &sctx;
		memset(&sctx, 0, sizeof(sctx));
		size_t sz = 512;
		ctx->bios = guarded_alloc(sz);
		ctx->bios_size = sz;
		ctx->cmd_table = 0x20;
		((uint8_t *)ctx->bios)[0x24] = 0x00;
		((uint8_t *)ctx->bios)[0x25] = 0x02;
		memset((uint8_t *)ctx->bios + 0x200, 90, sz - 0x200);
		int r = fetch_loop(ctx, 0);
		printf("  ret=%d (fetch loop aborted at image end, no trap)\n", r);
		return 0;
	}
	if (!strcmp(cs, "shift")) {
		shift_fixed_demo();
		return 0;
	}
	fprintf(stderr, "unknown case %s\n", cs);
	return 2;
}