/*	$OpenBSD: efidev.c,v 1.11 2023/10/26 14:08:48 jsg Exp $	*/

/*
 * Copyright (c) 2015 YASUOKA Masahiko <yasuoka@yasuoka.net>
 * Copyright (c) 2016 Mark Kettenis
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in the
 *    documentation and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
 * IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED.  IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS
 * OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
 * LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 *
 */
#include <sys/param.h>
#include <sys/reboot.h>
#include <sys/disklabel.h>
#include <lib/libz/zlib.h>
#include <isofs/cd9660/iso.h>

#include "libsa.h"

#include <efi.h>

extern EFI_BOOT_SERVICES *BS;

extern int debug;

#include "disk.h"
#include "efidev.h"

#define EFI_BLKSPERSEC(_ed)	((_ed)->blkio->Media->BlockSize / DEV_BSIZE)
#define EFI_SECTOBLK(_ed, _n)	((_n) * EFI_BLKSPERSEC(_ed))

static EFI_STATUS
		 efid_io(int, efi_diskinfo_t, u_int, int, void *);
static int	 efid_diskio(int, struct diskinfo *, u_int, int, void *);
const char *	 efi_getdisklabel(efi_diskinfo_t, struct disklabel *);
static int	 efi_getdisklabel_cd9660(efi_diskinfo_t, struct disklabel *);
static u_int	 findopenbsd(efi_diskinfo_t, const char **);
static u_int	 findopenbsd_gpt(efi_diskinfo_t, const char **);
static int	 gpt_chk_mbr(struct dos_partition *, u_int64_t);

void
efid_init(struct diskinfo *dip, void *handle)
{
	EFI_BLOCK_IO		*blkio = handle;

	memset(dip, 0, sizeof(struct diskinfo));
	dip->ed.blkio = blkio;
	dip->ed.mediaid = blkio->Media->MediaId;
	dip->diskio = efid_diskio;
	dip->strategy = efistrategy;

	if (efi_getdisklabel(&dip->ed, &dip->disklabel) == NULL)
		dip->flags |= DISKINFO_FLAG_GOODLABEL;
}

static EFI_STATUS
efid_io(int rw, efi_diskinfo_t ed, u_int off, int nsect, void *buf)
{
	u_int blks, start, end;
	EFI_PHYSICAL_ADDRESS addr;
	EFI_STATUS status;
	caddr_t data;
	size_t size;

	/* block count of the intrinsic block size in DEV_BSIZE */
	blks = EFI_BLKSPERSEC(ed);
	if (blks == 0)
		/* block size < 512.  HP Stream 13 actually has such a disk. */
		return (EFI_UNSUPPORTED);

	start = off / blks;
	end = (off + nsect + blks - 1) / blks;
	size = (end - start) * ed->blkio->Media->BlockSize;

	status = BS->AllocatePages(AllocateAnyPages, EfiLoaderData,
	    EFI_SIZE_TO_PAGES(size), &addr);
	if (EFI_ERROR(status))
		goto on_eio;
	data = (caddr_t)(uintptr_t)addr;

	switch (rw) {
	case F_READ:
		status = ed->blkio->ReadBlocks(ed->blkio, ed->mediaid, start,
		    size, data);
		if (EFI_ERROR(status))
			goto on_eio;
		memcpy(buf, data + DEV_BSIZE * (off - start * blks),
		    DEV_BSIZE * nsect);
		break;
	case F_WRITE:
		if (ed->blkio->Media->ReadOnly)
			goto on_eio;
		if (off % blks != 0 || nsect % blks != 0) {
			status = ed->blkio->ReadBlocks(ed->blkio, ed->mediaid,
			    start, size, data);
			if (EFI_ERROR(status))
				goto on_eio;
		}
		memcpy(data + DEV_BSIZE * (off - start * blks), buf,
		    DEV_BSIZE * nsect);
		status = ed->blkio->WriteBlocks(ed->blkio, ed->mediaid, start,
		    size, data);
		if (EFI_ERROR(status))
			goto on_eio;
		break;
	}

on_eio:
	BS->FreePages(addr, EFI_SIZE_TO_PAGES(size));

	return (status);
}

static int
efid_diskio(int rw, struct diskinfo *dip, u_int off, int nsect, void *buf)
{
	EFI_STATUS status;

	status = efid_io(rw, &dip->ed, off, nsect, buf);

	return ((EFI_ERROR(status))? -1 : 0);
}

/*
 * Returns 0 if the MBR with the provided partition array is a GPT protective
 * MBR, and returns 1 otherwise. A GPT protective MBR would have one and only
 * one MBR partition, an EFI partition that either covers the whole disk or as
 * much of it as is possible with a 32bit size field.
 *
 * Taken from kern/subr_disk.c.
 *
 * NOTE: MS always uses a size of UINT32_MAX for the EFI partition!**
 */
static int
gpt_chk_mbr(struct dos_partition *dp, u_int64_t dsize)
{
	struct dos_partition *dp2;
	int efi, found, i;
	u_int32_t psize;

	found = efi = 0;
	for (dp2=dp, i=0; i < NDOSPART; i++, dp2++) {
		if (dp2->dp_typ == DOSPTYP_UNUSED)
			continue;
		found++;
		if (dp2->dp_typ != DOSPTYP_EFI)
			continue;
		if (letoh32(dp2->dp_start) != GPTSECTOR)
			continue;
		psize = letoh32(dp2->dp_size);
		if (psize <= (dsize - GPTSECTOR) || psize == UINT32_MAX)
			efi++;
	}
	if (found == 1 && efi == 1)
		return (0);

	return (1);
}

/*
 * Try to find the disk address of the first MBR OpenBSD partition.
 *
 * N.B.: must boot from a partition within first 2^32-1 sectors!
 *
 * Called only if the MBR on sector 0 is *not* a protective MBR
 * and *does* have a valid signature.
 *
 * We don't check the signatures of EBR's, and they cannot be
 * protective MBR's so there is no need to check for that.
 */
static u_int
findopenbsd(efi_diskinfo_t ed, const char **err)
{
	EFI_STATUS status;
	struct dos_mbr mbr;
	struct dos_partition *dp;
	u_int mbroff = DOSBBSECTOR;
	u_int mbr_eoff = DOSBBSECTOR;	/* Offset of MBR extended partition. */
	int i, maxebr = DOS_MAXEBR, nextebr;

again:
	if (!maxebr--) {
		*err = "too many extended partitions";
		return (-1);
	}

	/* Read MBR */
	bzero(&mbr, sizeof(mbr));
	status = efid_io(F_READ, ed, mbroff, 1, &mbr);
	if (EFI_ERROR(status)) {
		*err = "Disk I/O Error";
		return (-1);
	}

	/* Search for OpenBSD partition */
	nextebr = 0;
	for (i = 0; i < NDOSPART; i++) {
		dp = &mbr.dmbr_parts[i];
		if (!dp->dp_size)
			continue;
#ifdef BIOS_DEBUG
		if (debug)
			printf("found partition %u: "
			    "type %u (0x%x) offset %u (0x%x)\n",
			    (int)(dp - mbr.dmbr_parts),
			    dp->dp_typ, dp->dp_typ,
			    dp->dp_start, dp->dp_start);
#endif
		if (dp->dp_typ == DOSPTYP_OPENBSD) {
			if (dp->dp_start > (dp->dp_start + mbroff))
				continue;
			return (dp->dp_start + mbroff);
		}

		/*
		 * Record location of next ebr if and only if this is the first
		 * extended partition in this boot record!
		 */
		if (!nextebr && (dp->dp_typ == DOSPTYP_EXTEND ||
		    dp->dp_typ == DOSPTYP_EXTENDL)) {
			nextebr = dp->dp_start + mbr_eoff;
			if (nextebr < dp->dp_start)
				nextebr = (u_int)-1;
			if (mbr_eoff == DOSBBSECTOR)
				mbr_eoff = dp->dp_start;
		}
	}

	if (nextebr && nextebr != (u_int)-1) {
		mbroff = nextebr;
		goto again;
	}

	return (-1);
}

/*
 * Try to find the disk address of the first GPT OpenBSD partition.
 *
 * N.B.: must boot from a partition within first 2^32-1 sectors!
 *
 * Called only if the MBR on sector 0 *is* a protective MBR
 * with a valid signature and sector 1 is a valid GPT header.
 */
static u_int
findopenbsd_gpt(efi_diskinfo_t ed, const char **err)
{
	EFI_STATUS		 status;
	struct			 gpt_header gh;
	int			 i, part, found;
	uint64_t		 lba;
	uint32_t		 orig_csum, new_csum;
	uint32_t		 ghsize, ghpartsize, ghpartnum, ghpartspersec;
	uint32_t		 gpsectors;
	const char		 openbsd_uuid_code[] = GPT_UUID_OPENBSD;
	struct gpt_partition	 gp;
	static struct uuid	*openbsd_uuid = NULL, openbsd_uuid_space;
	static u_char		 buf[4096];

	/* Prepare OpenBSD UUID */
	if (openbsd_uuid == NULL) {
		/* XXX: should be replaced by uuid_dec_be() */
		memcpy(&openbsd_uuid_space, openbsd_uuid_code,
		    sizeof(openbsd_uuid_space));
		openbsd_uuid_space.time_low =
		    betoh32(openbsd_uuid_space.time_low);
		openbsd_uuid_space.time_mid =
		    betoh16(openbsd_uuid_space.time_mid);
		openbsd_uuid_space.time_hi_and_version =
		    betoh16(openbsd_uuid_space.time_hi_and_version);

		openbsd_uuid = &openbsd_uuid_space;
	}

	if (EFI_BLKSPERSEC(ed) > 8) {
		*err = "disk sector > 4096 bytes\n";
		return (-1);
	}

	/* GPT Header */
	lba = GPTSECTOR;
	status = efid_io(F_READ, ed, EFI_SECTOBLK(ed, lba), EFI_BLKSPERSEC(ed),
	    buf);
	if (EFI_ERROR(status)) {
		*err = "Disk I/O Error";
		return (-1);
	}
	memcpy(&gh, buf, sizeof(gh));

	/* Check signature */
	if (letoh64(gh.gh_sig) != GPTSIGNATURE) {
		*err = "bad GPT signature\n";
		return (-1);
	}

	if (letoh32(gh.gh_rev) != GPTREVISION) {
		*err = "bad GPT revision\n";
		return (-1);
	}

	ghsize = letoh32(gh.gh_size);
	if (ghsize < GPTMINHDRSIZE || ghsize > sizeof(struct gpt_header)) {
		*err = "bad GPT header size\n";
		return (-1);
	}

	/* Check checksum */
	orig_csum = gh.gh_csum;
	gh.gh_csum = 0;
	new_csum = crc32(0, (unsigned char *)&gh, ghsize);
	gh.gh_csum = orig_csum;
	if (letoh32(orig_csum) != new_csum) {
		*err = "bad GPT header checksum\n";
		return (-1);
	}

	lba = letoh64(gh.gh_part_lba);
	ghpartsize = letoh32(gh.gh_part_size);
	ghpartspersec = ed->blkio->Media->BlockSize / ghpartsize;
	ghpartnum = letoh32(gh.gh_part_num);
	gpsectors = (ghpartnum + ghpartspersec - 1) / ghpartspersec;
	new_csum = crc32(0L, Z_NULL, 0);
	found = 0;
	for (i = 0; i < gpsectors; i++, lba++) {
		status = efid_io(F_READ, ed, EFI_SECTOBLK(ed, lba),
		    EFI_BLKSPERSEC(ed), buf);
		if (EFI_ERROR(status)) {
			*err = "Disk I/O Error";
			return (-1);
		}
		for (part = 0; part < ghpartspersec; part++) {
			if (ghpartnum == 0)
				break;
			new_csum = crc32(new_csum, buf + part * sizeof(gp),
			    sizeof(gp));
			ghpartnum--;
			if (found)
				continue;
			memcpy(&gp, buf + part * sizeof(gp), sizeof(gp));
			if (memcmp(&gp.gp_type, openbsd_uuid,
			    sizeof(struct uuid)) == 0)
				found = 1;
		}
	}
	if (new_csum != letoh32(gh.gh_part_csum)) {
		*err = "bad GPT entries checksum\n";
		return (-1);
	}
	if (found) {
		lba = letoh64(gp.gp_lba_start);
		/* Bootloaders do not current handle addresses > UINT_MAX! */
		if (lba > UINT_MAX || EFI_SECTOBLK(ed, lba) > UINT_MAX) {
			*err = "OpenBSD Partition LBA > 2**32 - 1";
			return (-1);
		}
		return (u_int)lba;
	}

	return (-1);
}

const char *
efi_getdisklabel(efi_diskinfo_t ed, struct disklabel *label)
{
	u_int start = 0;
	uint8_t buf[DEV_BSIZE];
	struct dos_partition dosparts[NDOSPART];
	EFI_STATUS status;
	const char *err = NULL;
	int error;

	/*
	 * Read sector 0. Ensure it has a valid MBR signature.
	 *
	 * If it's a protective MBR then try to find the disklabel via
	 * GPT. If it's not a protective MBR, try to find the disklabel
	 * via MBR.
	 */
	memset(buf, 0, sizeof(buf));
	status = efid_io(F_READ, ed, DOSBBSECTOR, 1, buf);
	if (EFI_ERROR(status))
		return ("Disk I/O Error");

	/* Check MBR signature. */
	if (buf[510] != 0x55 || buf[511] != 0xaa) {
		if (efi_getdisklabel_cd9660(ed, label) == 0)
			return (NULL);
		return ("invalid MBR signature");
	}

	memcpy(dosparts, buf+DOSPARTOFF, sizeof(dosparts));

	/* check for GPT protective MBR. */
	if (gpt_chk_mbr(dosparts, ed->blkio->Media->LastBlock + 1) == 0) {
		start = findopenbsd_gpt(ed, &err);
		if (start == (u_int)-1) {
			if (err != NULL)
				return (err);
			return ("no OpenBSD GPT partition");
		}
	} else {
		start = findopenbsd(ed, &err);
		if (start == (u_int)-1) {
			if (err != NULL)
				return (err);
			return "no OpenBSD MBR partition\n";
		}
	}

	/* Load BSD disklabel */
#ifdef BIOS_DEBUG
	if (debug)
		printf("loading disklabel @ %u\n", start + DOS_LABELSECTOR);
#endif
	/* read disklabel */
	error = efid_io(F_READ, ed, EFI_SECTOBLK(ed, start) + DOS_LABELSECTOR,
	    1, buf);

	if (error)
		return "failed to read disklabel";

	/* Fill in disklabel */
	return (getdisklabel(buf, label));
}

static int
efi_getdisklabel_cd9660(efi_diskinfo_t ed, struct disklabel *label)
{
	uint8_t		 buf[DEV_BSIZE];
	EFI_STATUS	 status;

	status = efid_io(F_READ, ed, 64, 1, buf);
	if (EFI_ERROR(status))
		return -1;
	if (buf[0] != ISO_VD_PRIMARY || bcmp(buf + 1, ISO_STANDARD_ID, 5) != 0)
		return -1;

	/* Create an imaginary disk label */
	label->d_secsize = 2048;
	label->d_ntracks = 1;
	label->d_nsectors = 100;
	label->d_ncylinders = 1;
	label->d_secpercyl = label->d_ntracks * label->d_nsectors;

	strncpy(label->d_typename, "ATAPI CD-ROM", sizeof(label->d_typename));
	label->d_type = DTYPE_ATAPI;

	strncpy(label->d_packname, "fictitious", sizeof(label->d_packname));
	DL_SETDSIZE(label, 100);

	/* 'a' partition covering the "whole" disk */
	DL_SETPOFFSET(&label->d_partitions[0], 0);
	DL_SETPSIZE(&label->d_partitions[0], 100);
	label->d_partitions[0].p_fstype = FS_UNUSED;

	/* The raw partition is special */
	DL_SETPOFFSET(&label->d_partitions[RAW_PART], 0);
	DL_SETPSIZE(&label->d_partitions[RAW_PART], 100);
	label->d_partitions[RAW_PART].p_fstype = FS_UNUSED;

	label->d_npartitions = MAXPARTITIONS;

	label->d_magic = DISKMAGIC;
	label->d_magic2 = DISKMAGIC;
	label->d_checksum = dkcksum(label);

	return (0);
}

int
efiopen(struct open_file *f, ...)
{
	struct diskinfo *dip = NULL;
	va_list ap;
	u_int unit, part;
	int i = 0;

	va_start(ap, f);
	unit = va_arg(ap, u_int);
	part = va_arg(ap, u_int);
	va_end(ap);

	if (part >= MAXPARTITIONS)
		return (ENXIO);

	TAILQ_FOREACH(dip, &disklist, list) {
		if (i == unit)
			break;
		i++;
	}

	if (dip == NULL)
		return (ENXIO);

	if ((dip->flags & DISKINFO_FLAG_GOODLABEL) == 0)
		return (ENXIO);

	dip->part = part;
	bootdev_dip = dip;
	f->f_devdata = dip;

	return 0;
}

int
efistrategy(void *devdata, int rw, daddr_t blk, size_t size, void *buf,
    size_t *rsize)
{
	struct diskinfo *dip = (struct diskinfo *)devdata;
	int error = 0;
	size_t nsect;

	nsect = (size + DEV_BSIZE - 1) / DEV_BSIZE;
	blk += DL_SECTOBLK(&dip->disklabel,
	    dip->disklabel.d_partitions[dip->part].p_offset);

	if (blk < 0)
		error = EINVAL;
	else
		error = efid_diskio(rw, dip, blk, nsect, buf);

	if (rsize != NULL)
		*rsize = nsect * DEV_BSIZE;

	return (error);
}

int
eficlose(struct open_file *f)
{
	f->f_devdata = NULL;

	return 0;
}

int
efiioctl(struct open_file *f, u_long cmd, void *data)
{
	return 0;
}