/*
 * DF-1442 trigger harness.
 *
 * Replicates the EXACT vulnerable code path from
 * sys/dev/disk/ahci/ahci_cam.c::ahci_xpt_scsi_disk_io() INQUIRY case, and the
 * EXACT allocation that CAM Domain Validation (DV1) makes for the inquiry
 * buffer in sys/bus/cam/cam_xpt.c::probescsi() (PROBE_INQUIRY_BASIC_DV1/DV2).
 *
 * The kernel structs (struct scsi_inquiry_data, union scsi_data) are copied
 * verbatim from sys/bus/cam/scsi/scsi_all.h so that the field offsets are
 * bit-identical to the kernel. Every field is u_int8_t/char, so there is no
 * padding: the offsets are exact.
 *
 * WHAT THIS PROVES:
 *   1. CAM DV1 allocates a 38-byte inquiry buffer (kmalloc(38) -> 64-byte slab
 *      chunk) for an AHCI disk (additional_length=32 => SID_ADDITIONAL_LENGTH=37
 *      => roundup2(,2)=38).
 *   2. The AHCI INQUIRY handler writes vendor_specific1[0]/[1] at byte offsets
 *      96/97 of that buffer, i.e. 58/59 bytes PAST the 38-byte allocation,
 *      overflowing into the adjacent slab object.
 *
 * This is a userspace stand-in for the kernel write because the audit guest has
 * NO AHCI SATA disk (only virtio-blk + a legacy-ata DVD-ROM), so the in-kernel
 * path cannot be exercised at runtime on this guest. The offset arithmetic is
 * identical in-kernel.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>
#include <stddef.h>

/* ---- verbatim from sys/bus/cam/scsi/scsi_all.h ---- */
#define SHORT_INQUIRY_LENGTH	36
#define SID_VENDOR_SIZE		8
#define SID_PRODUCT_SIZE	16
#define SID_REVISION_SIZE	4
#define SID_VENDOR_SPECIFIC_0_SIZE	20
#define SID_VENDOR_SPECIFIC_1_SIZE	160

struct scsi_inquiry_data {
	uint8_t device;
	uint8_t dev_qual2;
	uint8_t version;
	uint8_t response_format;
	uint8_t additional_length;
#define SID_ADDITIONAL_LENGTH(iqd) \
	((iqd)->additional_length + \
	 offsetof(struct scsi_inquiry_data, additional_length) + 1)
	uint8_t reserved;
	uint8_t spc2_flags;
	uint8_t flags;
	char    vendor[SID_VENDOR_SIZE];
	char    product[SID_PRODUCT_SIZE];
	char    revision[SID_REVISION_SIZE];
	uint8_t vendor_specific0[SID_VENDOR_SPECIFIC_0_SIZE];
	uint8_t spi3data;
	uint8_t reserved2;
	uint8_t version1[2];
	uint8_t version2[2];
	uint8_t version3[2];
	uint8_t version4[2];
	uint8_t version5[2];
	uint8_t version6[2];
	uint8_t version7[2];
	uint8_t version8[2];
	uint8_t reserved3[22];
	uint8_t vendor_specific1[SID_VENDOR_SPECIFIC_1_SIZE];
};

/* ---- verbatim union scsi_data (only inquiry_data matters) ---- */
typedef union scsi_data {
	struct scsi_inquiry_data inquiry_data;
	/* other members elided; inquiry_data is the largest relevant member */
} scsi_data;

#define ATA_SUPPORT_DSM_TRIM	0x0001

/* ---- mimic ata_port.at_identify subset that the bug reads ---- */
struct ata_identify_dsm {
	uint16_t support_dsm;	/* word 169 (nominal) */
	uint16_t max_dsm_blocks;	/* word 105 */
};

/*
 * Reproduce the vulnerable write exactly as ahci_cam.c:1134-1139 does it,
 * but with a canary guard so we can DETECT the out-of-bounds write rather
 * than silently corrupt memory.
 *
 * mode = 0 : UNPATCHED logic (no bounds check)  -> expect OVERFLOW
 * mode = 1 : PATCHED   logic (rdata_len>=offsetof(vs1)+2 guard) -> expect clean
 */
static int run_case(int mode)
{
	printf("sizeof(struct scsi_inquiry_data) = %zu\n",
	       sizeof(struct scsi_inquiry_data));
	printf("offsetof(vendor_specific1)       = %zu\n",
	       offsetof(struct scsi_inquiry_data, vendor_specific1));
	printf("vendor_specific1[0] byte offset  = %zu\n",
	       offsetof(struct scsi_inquiry_data, vendor_specific1) + 0);
	printf("vendor_specific1[1] byte offset  = %zu\n",
	       offsetof(struct scsi_inquiry_data, vendor_specific1) + 1);
	printf("SHORT_INQUIRY_LENGTH             = %d\n", SHORT_INQUIRY_LENGTH);

	/* ---- the exact CAM DV1 allocation ---- */
	/* AHCI sets additional_length=32 (ahci_cam.c:1119), so:               */
	/*   SID_ADDITIONAL_LENGTH = 32 + offsetof(additional_length)+1 = 37   */
	/*   roundup2(37,2) = 38                                               */
	uint8_t additional_length = 32;
	int inquiry_len = (additional_length +
			   offsetof(struct scsi_inquiry_data, additional_length)
			   + 1);
	inquiry_len = (inquiry_len + 1) & ~1;	/* roundup2(,2) */
	printf("\nCAM DV1 inquiry_len = %d (additional_length=%u)\n",
	       inquiry_len, additional_length);

	/*
	 * Allocate a 4096-byte arena so we can place a canary AFTER the
	 * inquiry buffer and detect the overflow precisely. In the kernel the
	 * kmalloc(38) lands in the 64-byte slab bucket; the write at offset 96
	 * hits byte 32-33 of the NEXT slab chunk. We emulate that adjacency.
	 */
#define CANARY 0xCD
#define ARENA_SZ 256
	uint8_t *arena = malloc(ARENA_SZ);
	memset(arena, CANARY, ARENA_SZ);

	scsi_data *rdata = (scsi_data *)arena;	/* data_ptr = arena base */
	int rdata_len = inquiry_len;		/* = 38 */

	/* emulate the normal INQUIRY body (ahci_cam.c:1109-1128) */
	memset(rdata, 0, rdata_len);		/* bzero(rdata, rdata_len) */
	rdata->inquiry_data.additional_length = additional_length;

	/* emulate a TRIM-capable AHCI disk's at_identify */
	struct ata_identify_dsm id;
	id.support_dsm    = ATA_SUPPORT_DSM_TRIM;	/* TRIM supported */
	id.max_dsm_blocks = 0x0102;			/* arbitrary >0 */

	int over_off0 = offsetof(struct scsi_inquiry_data, vendor_specific1) + 0;
	int over_off1 = offsetof(struct scsi_inquiry_data, vendor_specific1) + 1;

	printf("\n[%s] mode=%d, buffer=%d bytes, write at offset %d/%d\n",
	       mode == 0 ? "UNPATCHED" : "PATCHED",
	       mode, inquiry_len, over_off0, over_off1);

	/* === THE (VULNERABLE / FIXED) WRITE (ahci_cam.c:1134-1139) === */
	if (mode == 0) {
		/* UNPATCHED: no bounds check -> writes at offset 96/97 always */
		if (id.support_dsm) {
			rdata->inquiry_data.vendor_specific1[0] =
				id.support_dsm & ATA_SUPPORT_DSM_TRIM;	/* 96 */
			rdata->inquiry_data.vendor_specific1[1] =
				id.max_dsm_blocks;			/* 97 */
		}
	} else {
		/* PATCHED: guard rdata_len >= offsetof(vs1)+2 (==98) */
		if (id.support_dsm &&
		    rdata_len >= (int)(offsetof(struct scsi_inquiry_data,
						vendor_specific1) + 2)) {
			rdata->inquiry_data.vendor_specific1[0] =
				id.support_dsm & ATA_SUPPORT_DSM_TRIM;
			rdata->inquiry_data.vendor_specific1[1] =
				id.max_dsm_blocks;
		}
	}

	/* ---- detect the overflow ---- */
	int last_inbuf = inquiry_len - 1;	/* 37 */

	/*
	 * Detection: the BUG is that the write actually lands past the buffer.
	 * We detect a real overflow iff the canary bytes at offset 96/97 were
	 * actually clobbered (write happened) AND the offset is past the
	 * allocation. In PATCHED mode the guard skips the write, so the canary
	 * stays 0xCD even though the offset would be out-of-range.
	 */
	int wrote0 = (arena[over_off0] != CANARY);
	int wrote1 = (arena[over_off1] != CANARY);
	int overflowed = (over_off0 > last_inbuf) && (wrote0 || wrote1);

	printf("\n--- result (mode=%d %s) ---\n", mode,
	       mode == 0 ? "UNPATCHED" : "PATCHED");
	printf("buffer length (inquiry_len) = %d bytes ([0..%d])\n",
	       inquiry_len, last_inbuf);
	printf("write target offsets = %d / %d (%s / %s past end)\n",
	       over_off0, over_off1,
	       over_off0 > last_inbuf ? "OUT-OF-RANGE" : "in-bounds",
	       over_off1 > last_inbuf ? "OUT-OF-RANGE" : "in-bounds");
	printf("canary clobbered at %d/%d: %s / %s (0x%02x/0x%02x vs canary 0x%02x)\n",
	       over_off0, over_off1,
	       wrote0 ? "YES" : "no", wrote1 ? "YES" : "no",
	       arena[over_off0], arena[over_off1], CANARY);

	if (overflowed) {
		printf("\n*** HEAP OVERFLOW CONFIRMED (mode %d) ***\n", mode);
		printf("wrote 0x%02x at arena[%d] and 0x%02x at arena[%d]\n",
		       arena[over_off0], over_off0,
		       arena[over_off1], over_off1);
		printf("(canary was 0x%02x; these bytes are %d/%d past the "
		       "%d-byte allocation)\n",
		       CANARY, over_off0 - inquiry_len, over_off1 - inquiry_len,
		       inquiry_len);
		printf("In-kernel: kmalloc(%d) -> slab chunk; offset %d/%d "
		       "lands in the ADJACENT slab object.\n",
		       inquiry_len, over_off0, over_off1);
		free(arena);
		return 1;	/* overflow = BUG PRESENT */
	}
	printf("\n*** NO OVERFLOW (mode %d): write correctly skipped/guarded ***\n",
	       mode);
	printf("arena[%d]=0x%02x arena[%d]=0x%02x (canary intact=0x%02x)\n",
	       over_off0, arena[over_off0],
	       over_off1, arena[over_off1], CANARY);
	free(arena);
	return 0;	/* clean = FIX HOLDS */
}

int main(void)
{
	printf("sizeof(struct scsi_inquiry_data) = %zu\n",
	       sizeof(struct scsi_inquiry_data));
	printf("offsetof(vendor_specific1)       = %zu\n",
	       offsetof(struct scsi_inquiry_data, vendor_specific1));
	printf("vendor_specific1[0] byte offset  = %zu\n",
	       offsetof(struct scsi_inquiry_data, vendor_specific1) + 0);
	printf("vendor_specific1[1] byte offset  = %zu\n",
	       offsetof(struct scsi_inquiry_data, vendor_specific1) + 1);
	printf("SHORT_INQUIRY_LENGTH             = %d\n", SHORT_INQUIRY_LENGTH);

	int r0 = run_case(0);	/* UNPATCHED: expect overflow */
	int r1 = run_case(1);	/* PATCHED:   expect clean     */

	printf("\n=== SUMMARY ===\n");
	printf("UNPATCHED logic: %s\n", r0 ? "OVERFLOW (bug present)" : "clean (unexpected)");
	printf("PATCHED   logic: %s\n", r1 ? "OVERFLOW (fix failed)"  : "clean (fix holds)");
	if (r0 == 1 && r1 == 0) {
		printf("VERDICT: bug reproduced on unpatched logic, fixed on "
		       "patched logic.\n");
		return 0;
	}
	return 1;
}