/*
 * Copyright (c) 2012 Andrey V. Elsukov <ae@FreeBSD.org>
 * 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 AUTHORS AND CONTRIBUTORS ``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 AUTHORS 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/cdefs.h>

#include <stand.h>
#include <stddef.h>
#include <sys/param.h>
#include <sys/diskmbr.h>
#include <sys/disklabel.h>
#include <sys/endian.h>
#include <sys/gpt.h>
#include <sys/queue.h>
#include <sys/vtoc.h>

#include <fs/cd9660/iso.h>

#include <zlib.h>
#include <part.h>
#include <uuid.h>

#ifdef PART_DEBUG
#define DPRINTF(fmt, args...) printf("%s: " fmt "\n", __func__, ## args)
#else
#define DPRINTF(fmt, args...)   ((void)0)
#endif

#ifdef LOADER_GPT_SUPPORT
#define MAXTBLSZ        64
static const uuid_t gpt_uuid_unused = GPT_ENT_TYPE_UNUSED;
static const uuid_t gpt_uuid_ms_basic_data = GPT_ENT_TYPE_MS_BASIC_DATA;
static const uuid_t gpt_uuid_freebsd_ufs = GPT_ENT_TYPE_FREEBSD_UFS;
static const uuid_t gpt_uuid_efi = GPT_ENT_TYPE_EFI;
static const uuid_t gpt_uuid_freebsd = GPT_ENT_TYPE_FREEBSD;
static const uuid_t gpt_uuid_freebsd_boot = GPT_ENT_TYPE_FREEBSD_BOOT;
static const uuid_t gpt_uuid_freebsd_swap = GPT_ENT_TYPE_FREEBSD_SWAP;
static const uuid_t gpt_uuid_freebsd_zfs = GPT_ENT_TYPE_FREEBSD_ZFS;
static const uuid_t gpt_uuid_freebsd_vinum = GPT_ENT_TYPE_FREEBSD_VINUM;
static const uuid_t gpt_uuid_illumos_boot = GPT_ENT_TYPE_ILLUMOS_BOOT;
static const uuid_t gpt_uuid_illumos_ufs = GPT_ENT_TYPE_ILLUMOS_UFS;
static const uuid_t gpt_uuid_illumos_zfs = GPT_ENT_TYPE_ILLUMOS_ZFS;
static const uuid_t gpt_uuid_reserved = GPT_ENT_TYPE_RESERVED;
static const uuid_t gpt_uuid_apple_apfs = GPT_ENT_TYPE_APPLE_APFS;
#endif

struct pentry {
        struct ptable_entry     part;
        uint64_t                flags;
        union {
                uint8_t bsd;
                uint8_t mbr;
                uuid_t  gpt;
                uint16_t vtoc8;
                uint16_t vtoc;
        } type;
        STAILQ_ENTRY(pentry)    entry;
};

struct ptable {
        enum ptable_type        type;
        uint16_t                sectorsize;
        uint64_t                sectors;

        STAILQ_HEAD(, pentry)   entries;
};

static struct parttypes {
        enum partition_type     type;
        const char              *desc;
} ptypes[] = {
        { PART_UNKNOWN,         "Unknown" },
        { PART_EFI,             "EFI" },
        { PART_FREEBSD,         "FreeBSD" },
        { PART_FREEBSD_BOOT,    "FreeBSD boot" },
        { PART_FREEBSD_UFS,     "FreeBSD UFS" },
        { PART_FREEBSD_ZFS,     "FreeBSD ZFS" },
        { PART_FREEBSD_SWAP,    "FreeBSD swap" },
        { PART_FREEBSD_VINUM,   "FreeBSD vinum" },
        { PART_LINUX,           "Linux" },
        { PART_LINUX_SWAP,      "Linux swap" },
        { PART_DOS,             "DOS/Windows" },
        { PART_ISO9660,         "ISO9660" },
        { PART_SOLARIS2,        "Solaris 2" },
        { PART_ILLUMOS_UFS,     "illumos UFS" },
        { PART_ILLUMOS_ZFS,     "illumos ZFS" },
        { PART_RESERVED,        "Reserved" },
        { PART_VTOC_BOOT,       "boot" },
        { PART_VTOC_ROOT,       "root" },
        { PART_VTOC_SWAP,       "swap" },
        { PART_VTOC_USR,        "usr" },
        { PART_VTOC_STAND,      "stand" },
        { PART_VTOC_VAR,        "var" },
        { PART_VTOC_HOME,       "home" },
        { PART_APFS,            "APFS" }
};

const char *
parttype2str(enum partition_type type)
{
        size_t i;

        for (i = 0; i < nitems(ptypes); i++)
                if (ptypes[i].type == type)
                        return (ptypes[i].desc);
        return (ptypes[0].desc);
}

#ifdef LOADER_GPT_SUPPORT
static void
uuid_letoh(uuid_t *uuid)
{

        uuid->time_low = le32toh(uuid->time_low);
        uuid->time_mid = le16toh(uuid->time_mid);
        uuid->time_hi_and_version = le16toh(uuid->time_hi_and_version);
}

static enum partition_type
gpt_parttype(uuid_t type)
{

        if (uuid_equal(&type, &gpt_uuid_efi, NULL))
                return (PART_EFI);
        else if (uuid_equal(&type, &gpt_uuid_ms_basic_data, NULL))
                return (PART_DOS);
        else if (uuid_equal(&type, &gpt_uuid_freebsd_boot, NULL))
                return (PART_FREEBSD_BOOT);
        else if (uuid_equal(&type, &gpt_uuid_freebsd_ufs, NULL))
                return (PART_FREEBSD_UFS);
        else if (uuid_equal(&type, &gpt_uuid_freebsd_zfs, NULL))
                return (PART_FREEBSD_ZFS);
        else if (uuid_equal(&type, &gpt_uuid_freebsd_swap, NULL))
                return (PART_FREEBSD_SWAP);
        else if (uuid_equal(&type, &gpt_uuid_freebsd_vinum, NULL))
                return (PART_FREEBSD_VINUM);
        else if (uuid_equal(&type, &gpt_uuid_freebsd, NULL))
                return (PART_FREEBSD);
        else if (uuid_equal(&type, &gpt_uuid_illumos_boot, NULL))
                return (PART_VTOC_BOOT);
        else if (uuid_equal(&type, &gpt_uuid_illumos_ufs, NULL))
                return (PART_ILLUMOS_UFS);
        else if (uuid_equal(&type, &gpt_uuid_illumos_zfs, NULL))
                return (PART_ILLUMOS_ZFS);
        else if (uuid_equal(&type, &gpt_uuid_reserved, NULL))
                return (PART_RESERVED);
        else if (uuid_equal(&type, &gpt_uuid_apple_apfs, NULL))
                return (PART_APFS);
        return (PART_UNKNOWN);
}

static struct gpt_hdr *
gpt_checkhdr(struct gpt_hdr *hdr, uint64_t lba_self,
    uint64_t lba_last __attribute((unused)), uint16_t sectorsize)
{
        uint32_t sz, crc;

        if (memcmp(hdr->hdr_sig, GPT_HDR_SIG, sizeof (hdr->hdr_sig)) != 0) {
                DPRINTF("no GPT signature");
                return (NULL);
        }
        sz = le32toh(hdr->hdr_size);
        if (sz < 92 || sz > sectorsize) {
                DPRINTF("invalid GPT header size: %u", sz);
                return (NULL);
        }
        crc = le32toh(hdr->hdr_crc_self);
        hdr->hdr_crc_self = crc32(0, Z_NULL, 0);
        if (crc32(hdr->hdr_crc_self, (const Bytef *)hdr, sz) != crc) {
                DPRINTF("GPT header's CRC doesn't match");
                return (NULL);
        }
        hdr->hdr_crc_self = crc;
        hdr->hdr_revision = le32toh(hdr->hdr_revision);
        if (hdr->hdr_revision < GPT_HDR_REVISION) {
                DPRINTF("unsupported GPT revision %u", hdr->hdr_revision);
                return (NULL);
        }
        hdr->hdr_lba_self = le64toh(hdr->hdr_lba_self);
        if (hdr->hdr_lba_self != lba_self) {
                DPRINTF("self LBA doesn't match");
                return (NULL);
        }
        hdr->hdr_lba_alt = le64toh(hdr->hdr_lba_alt);
        if (hdr->hdr_lba_alt == hdr->hdr_lba_self) {
                DPRINTF("invalid alternate LBA");
                return (NULL);
        }
        hdr->hdr_entries = le32toh(hdr->hdr_entries);
        hdr->hdr_entsz = le32toh(hdr->hdr_entsz);
        if (hdr->hdr_entries == 0 ||
            hdr->hdr_entsz < sizeof (struct gpt_ent) ||
            sectorsize % hdr->hdr_entsz != 0) {
                DPRINTF("invalid entry size or number of entries");
                return (NULL);
        }
        hdr->hdr_lba_start = le64toh(hdr->hdr_lba_start);
        hdr->hdr_lba_end = le64toh(hdr->hdr_lba_end);
        hdr->hdr_lba_table = le64toh(hdr->hdr_lba_table);
        hdr->hdr_crc_table = le32toh(hdr->hdr_crc_table);
        uuid_letoh(&hdr->hdr_uuid);
        return (hdr);
}

static int
gpt_checktbl(const struct gpt_hdr *hdr, uint8_t *tbl, size_t size,
    uint64_t lba_last __attribute((unused)))
{
        struct gpt_ent *ent;
        uint32_t i, cnt;

        cnt = size / hdr->hdr_entsz;
        if (hdr->hdr_entries <= cnt) {
                cnt = hdr->hdr_entries;
                /* Check CRC only when buffer size is enough for table. */
                if (hdr->hdr_crc_table !=
                    crc32(0, tbl, hdr->hdr_entries * hdr->hdr_entsz)) {
                        DPRINTF("GPT table's CRC doesn't match");
                        return (-1);
                }
        }
        for (i = 0; i < cnt; i++) {
                ent = (struct gpt_ent *)(tbl + i * hdr->hdr_entsz);
                uuid_letoh(&ent->ent_type);
                if (uuid_equal(&ent->ent_type, &gpt_uuid_unused, NULL))
                        continue;
                ent->ent_lba_start = le64toh(ent->ent_lba_start);
                ent->ent_lba_end = le64toh(ent->ent_lba_end);
        }
        return (0);
}

static struct ptable *
ptable_gptread(struct ptable *table, void *dev, diskread_t dread)
{
        struct pentry *entry;
        struct gpt_hdr *phdr, hdr;
        struct gpt_ent *ent;
        uint8_t *buf, *tbl;
        uint64_t offset;
        int pri, sec;
        size_t size, i;

        buf = malloc(table->sectorsize);
        if (buf == NULL)
                return (NULL);
        tbl = malloc(table->sectorsize * MAXTBLSZ);
        if (tbl == NULL) {
                free(buf);
                return (NULL);
        }
        /* Read the primary GPT header. */
        if (dread(dev, buf, 1, 1) != 0) {
                ptable_close(table);
                table = NULL;
                goto out;
        }
        pri = sec = 0;
        /* Check the primary GPT header. */
        phdr = gpt_checkhdr((struct gpt_hdr *)buf, 1, table->sectors - 1,
            table->sectorsize);
        if (phdr != NULL) {
                /* Read the primary GPT table. */
                size = MIN(MAXTBLSZ, (phdr->hdr_entries * phdr->hdr_entsz +
                    table->sectorsize - 1) / table->sectorsize);
                if (dread(dev, tbl, size, phdr->hdr_lba_table) == 0 &&
                    gpt_checktbl(phdr, tbl, size * table->sectorsize,
                    table->sectors - 1) == 0) {
                        memcpy(&hdr, phdr, sizeof (hdr));
                        pri = 1;
                }
        }
        offset = pri ? hdr.hdr_lba_alt: table->sectors - 1;
        /* Read the backup GPT header. */
        if (dread(dev, buf, 1, offset) != 0)
                phdr = NULL;
        else
                phdr = gpt_checkhdr((struct gpt_hdr *)buf, offset,
                    table->sectors - 1, table->sectorsize);
        if (phdr != NULL) {
                /*
                 * Compare primary and backup headers.
                 * If they are equal, then we do not need to read backup
                 * table. If they are different, then prefer backup header
                 * and try to read backup table.
                 */
                if (pri == 0 ||
                    uuid_equal(&hdr.hdr_uuid, &phdr->hdr_uuid, NULL) == 0 ||
                    hdr.hdr_revision != phdr->hdr_revision ||
                    hdr.hdr_size != phdr->hdr_size ||
                    hdr.hdr_lba_start != phdr->hdr_lba_start ||
                    hdr.hdr_lba_end != phdr->hdr_lba_end ||
                    hdr.hdr_entries != phdr->hdr_entries ||
                    hdr.hdr_entsz != phdr->hdr_entsz ||
                    hdr.hdr_crc_table != phdr->hdr_crc_table) {
                        /* Read the backup GPT table. */
                        size = MIN(MAXTBLSZ, (phdr->hdr_entries *
                            phdr->hdr_entsz + table->sectorsize - 1) /
                            table->sectorsize);
                        if (dread(dev, tbl, size, phdr->hdr_lba_table) == 0 &&
                            gpt_checktbl(phdr, tbl, size * table->sectorsize,
                            table->sectors - 1) == 0) {
                                memcpy(&hdr, phdr, sizeof (hdr));
                                sec = 1;
                        }
                }
        }
        if (pri == 0 && sec == 0) {
                /* Both primary and backup tables are invalid. */
                table->type = PTABLE_NONE;
                goto out;
        }
        DPRINTF("GPT detected");
        size = MIN(hdr.hdr_entries * hdr.hdr_entsz,
            MAXTBLSZ * table->sectorsize);

        /*
         * If the disk's sector count is smaller than the sector count recorded
         * in the disk's GPT table header, set the table->sectors to the value
         * recorded in GPT tables. This is done to work around buggy firmware
         * that returns truncated disk sizes.
         *
         * Note, this is still not a foolproof way to get disk's size. For
         * example, an image file can be truncated when copied to smaller media.
         */
        table->sectors = hdr.hdr_lba_alt + 1;

        for (i = 0; i < size / hdr.hdr_entsz; i++) {
                ent = (struct gpt_ent *)(tbl + i * hdr.hdr_entsz);
                if (uuid_equal(&ent->ent_type, &gpt_uuid_unused, NULL))
                        continue;

                /* Simple sanity checks. */
                if (ent->ent_lba_start < hdr.hdr_lba_start ||
                    ent->ent_lba_end > hdr.hdr_lba_end ||
                    ent->ent_lba_start > ent->ent_lba_end)
                        continue;

                entry = malloc(sizeof (*entry));
                if (entry == NULL)
                        break;
                entry->part.start = ent->ent_lba_start;
                entry->part.end = ent->ent_lba_end;
                entry->part.index = i + 1;
                entry->part.type = gpt_parttype(ent->ent_type);
                entry->flags = le64toh(ent->ent_attr);
                memcpy(&entry->type.gpt, &ent->ent_type, sizeof (uuid_t));
                STAILQ_INSERT_TAIL(&table->entries, entry, entry);
                DPRINTF("new GPT partition added");
        }
out:
        free(buf);
        free(tbl);
        return (table);
}
#endif /* LOADER_GPT_SUPPORT */

#ifdef LOADER_MBR_SUPPORT
/* We do not need to support too many EBR partitions in the loader */
#define MAXEBRENTRIES           8
static enum partition_type
mbr_parttype(uint8_t type)
{

        switch (type) {
        case DOSPTYP_386BSD:
                return (PART_FREEBSD);
        case DOSPTYP_LINSWP:
                return (PART_LINUX_SWAP);
        case DOSPTYP_LINUX:
                return (PART_LINUX);
        case DOSPTYP_SUNIXOS2:
                return (PART_SOLARIS2);
        case 0x01:
        case 0x04:
        case 0x06:
        case 0x07:
        case 0x0b:
        case 0x0c:
        case 0x0e:
                return (PART_DOS);
        }
        return (PART_UNKNOWN);
}

static struct ptable *
ptable_ebrread(struct ptable *table, void *dev, diskread_t dread)
{
        struct dos_partition *dp;
        struct pentry *e1, *entry;
        uint32_t start, end, offset;
        uint8_t *buf;
        int i, idx;

        STAILQ_FOREACH(e1, &table->entries, entry) {
                if (e1->type.mbr == DOSPTYP_EXT ||
                    e1->type.mbr == DOSPTYP_EXTLBA)
                        break;
        }
        if (e1 == NULL)
                return (table);
        idx = 5;
        offset = e1->part.start;
        buf = malloc(table->sectorsize);
        if (buf == NULL)
                return (table);
        DPRINTF("EBR detected");
        for (i = 0; i < MAXEBRENTRIES; i++) {
#if 0   /* Some BIOSes return an incorrect number of sectors */
                if (offset >= table->sectors)
                        break;
#endif
                if (dread(dev, buf, 1, offset) != 0)
                        break;
                dp = (struct dos_partition *)(buf + DOSPARTOFF);
                if (dp[0].dp_typ == 0)
                        break;
                start = le32toh(dp[0].dp_start);
                if (dp[0].dp_typ == DOSPTYP_EXT &&
                    dp[1].dp_typ == 0) {
                        offset = e1->part.start + start;
                        continue;
                }
                end = le32toh(dp[0].dp_size);
                entry = malloc(sizeof (*entry));
                if (entry == NULL)
                        break;
                entry->part.start = offset + start;
                entry->part.end = entry->part.start + end - 1;
                entry->part.index = idx++;
                entry->part.type = mbr_parttype(dp[0].dp_typ);
                entry->flags = dp[0].dp_flag;
                entry->type.mbr = dp[0].dp_typ;
                STAILQ_INSERT_TAIL(&table->entries, entry, entry);
                DPRINTF("new EBR partition added");
                if (dp[1].dp_typ == 0)
                        break;
                offset = e1->part.start + le32toh(dp[1].dp_start);
        }
        free(buf);
        return (table);
}
#endif /* LOADER_MBR_SUPPORT */

static enum partition_type
bsd_parttype(uint8_t type)
{

        switch (type) {
        case FS_SWAP:
                return (PART_FREEBSD_SWAP);
        case FS_BSDFFS:
                return (PART_FREEBSD_UFS);
        case FS_VINUM:
                return (PART_FREEBSD_VINUM);
        case FS_ZFS:
                return (PART_FREEBSD_ZFS);
        }
        return (PART_UNKNOWN);
}

static struct ptable *
ptable_bsdread(struct ptable *table, void *dev, diskread_t dread)
{
        struct disklabel *dl;
        struct partition *part;
        struct pentry *entry;
        uint8_t *buf;
        uint32_t raw_offset;
        int i;

        if (table->sectorsize < sizeof (struct disklabel)) {
                DPRINTF("Too small sectorsize");
                return (table);
        }
        buf = malloc(table->sectorsize);
        if (buf == NULL)
                return (table);
        if (dread(dev, buf, 1, 1) != 0) {
                DPRINTF("read failed");
                ptable_close(table);
                table = NULL;
                goto out;
        }
        dl = (struct disklabel *)buf;
        if (le32toh(dl->d_magic) != DISKMAGIC &&
            le32toh(dl->d_magic2) != DISKMAGIC)
                goto out;
        if (le32toh(dl->d_secsize) != table->sectorsize) {
                DPRINTF("unsupported sector size");
                goto out;
        }
        dl->d_npartitions = le16toh(dl->d_npartitions);
        if (dl->d_npartitions > 20 || dl->d_npartitions < 8) {
                DPRINTF("invalid number of partitions");
                goto out;
        }
        DPRINTF("BSD detected");
        part = &dl->d_partitions[0];
        raw_offset = le32toh(part[RAW_PART].p_offset);
        for (i = 0; i < dl->d_npartitions; i++, part++) {
                if (i == RAW_PART)
                        continue;
                if (part->p_size == 0)
                        continue;
                entry = malloc(sizeof (*entry));
                if (entry == NULL)
                        break;
                entry->part.start = le32toh(part->p_offset) - raw_offset;
                entry->part.end = entry->part.start +
                    le32toh(part->p_size) - 1;
                entry->part.type = bsd_parttype(part->p_fstype);
                entry->part.index = i; /* starts from zero */
                entry->type.bsd = part->p_fstype;
                STAILQ_INSERT_TAIL(&table->entries, entry, entry);
                DPRINTF("new BSD partition added");
        }
        table->type = PTABLE_BSD;
out:
        free(buf);
        return (table);
}

#ifdef LOADER_VTOC8_SUPPORT
static enum partition_type
vtoc8_parttype(uint16_t type)
{

        switch (type) {
        case VTOC_TAG_FREEBSD_SWAP:
                return (PART_FREEBSD_SWAP);
        case VTOC_TAG_FREEBSD_UFS:
                return (PART_FREEBSD_UFS);
        case VTOC_TAG_FREEBSD_VINUM:
                return (PART_FREEBSD_VINUM);
        case VTOC_TAG_FREEBSD_ZFS:
                return (PART_FREEBSD_ZFS);
        };
        return (PART_UNKNOWN);
}

static struct ptable *
ptable_vtoc8read(struct ptable *table, void *dev, diskread_t dread)
{
        struct pentry *entry;
        struct vtoc8 *dl;
        uint8_t *buf;
        uint16_t sum, heads, sectors;
        int i;

        if (table->sectorsize != sizeof (struct vtoc8))
                return (table);
        buf = malloc(table->sectorsize);
        if (buf == NULL)
                return (table);
        if (dread(dev, buf, 1, 0) != 0) {
                DPRINTF("read failed");
                ptable_close(table);
                table = NULL;
                goto out;
        }
        dl = (struct vtoc8 *)buf;
        /* Check the sum */
        for (i = sum = 0; i < sizeof (struct vtoc8); i += sizeof (sum))
                sum ^= be16dec(buf + i);
        if (sum != 0) {
                DPRINTF("incorrect checksum");
                goto out;
        }
        if (be16toh(dl->nparts) != VTOC8_NPARTS) {
                DPRINTF("invalid number of entries");
                goto out;
        }
        sectors = be16toh(dl->nsecs);
        heads = be16toh(dl->nheads);
        if (sectors * heads == 0) {
                DPRINTF("invalid geometry");
                goto out;
        }
        DPRINTF("VTOC8 detected");
        for (i = 0; i < VTOC8_NPARTS; i++) {
                dl->part[i].tag = be16toh(dl->part[i].tag);
                if (i == VTOC_RAW_PART ||
                    dl->part[i].tag == VTOC_TAG_UNASSIGNED)
                        continue;
                entry = malloc(sizeof (*entry));
                if (entry == NULL)
                        break;
                entry->part.start = be32toh(dl->map[i].cyl) * heads * sectors;
                entry->part.end = be32toh(dl->map[i].nblks) +
                    entry->part.start - 1;
                entry->part.type = vtoc8_parttype(dl->part[i].tag);
                entry->part.index = i; /* starts from zero */
                entry->type.vtoc8 = dl->part[i].tag;
                STAILQ_INSERT_TAIL(&table->entries, entry, entry);
                DPRINTF("new VTOC8 partition added");
        }
        table->type = PTABLE_VTOC8;
out:
        free(buf);
        return (table);

}
#endif /* LOADER_VTOC8_SUPPORT */

static enum partition_type
vtoc_parttype(uint16_t type)
{
        switch (type) {
        case VTOC_TAG_BOOT:
                return (PART_VTOC_BOOT);
        case VTOC_TAG_ROOT:
                return (PART_VTOC_ROOT);
        case VTOC_TAG_SWAP:
                return (PART_VTOC_SWAP);
        case VTOC_TAG_USR:
                return (PART_VTOC_USR);
        case VTOC_TAG_BACKUP:
                return (PART_VTOC_BACKUP);
        case VTOC_TAG_STAND:
                return (PART_VTOC_STAND);
        case VTOC_TAG_VAR:
                return (PART_VTOC_VAR);
        case VTOC_TAG_HOME:
                return (PART_VTOC_HOME);
        };
        return (PART_UNKNOWN);
}

static struct ptable *
ptable_dklabelread(struct ptable *table, void *dev, diskread_t dread)
{
        struct pentry *entry;
        struct dk_label *dl;
        struct dk_vtoc *dv;
        uint8_t *buf;
        int i;

        if (table->sectorsize < sizeof (struct dk_label)) {
                DPRINTF("Too small sectorsize");
                return (table);
        }
        buf = malloc(table->sectorsize);
        if (buf == NULL)
                return (table);
        if (dread(dev, buf, 1, DK_LABEL_LOC) != 0) {
                DPRINTF("read failed");
                ptable_close(table);
                table = NULL;
                goto out;
        }
        dl = (struct dk_label *)buf;
        dv = (struct dk_vtoc *)&dl->dkl_vtoc;

        if (dl->dkl_magic != VTOC_MAGIC) {
                DPRINTF("dk_label magic error");
                goto out;
        }
        if (dv->v_sanity != VTOC_SANITY) {
                DPRINTF("this vtoc is not sane");
                goto out;
        }
        if (dv->v_nparts != NDKMAP) {
                DPRINTF("invalid number of entries");
                goto out;
        }
        DPRINTF("VTOC detected");
        for (i = 0; i < NDKMAP; i++) {
                if (i == VTOC_RAW_PART ||       /* skip slice 2 and empty */
                    dv->v_part[i].p_size == 0)
                        continue;
                entry = malloc(sizeof (*entry));
                if (entry == NULL)
                        break;
                entry->part.start = dv->v_part[i].p_start;
                entry->part.end = dv->v_part[i].p_size +
                    entry->part.start - 1;
                entry->part.type = vtoc_parttype(dv->v_part[i].p_tag);
                entry->part.index = i; /* starts from zero */
                entry->type.vtoc = dv->v_part[i].p_tag;
                STAILQ_INSERT_TAIL(&table->entries, entry, entry);
                DPRINTF("new VTOC partition added");
        }
        table->type = PTABLE_VTOC;
out:
        free(buf);
        return (table);
}

#define cdb2devb(bno)   ((bno) * ISO_DEFAULT_BLOCK_SIZE / table->sectorsize)

static struct ptable *
ptable_iso9660read(struct ptable *table, void *dev, diskread_t dread)
{
        uint8_t *buf;
        struct iso_primary_descriptor *vd;
        struct pentry *entry;

        buf = malloc(table->sectorsize);
        if (buf == NULL)
                return (table);

        if (dread(dev, buf, 1, cdb2devb(16)) != 0) {
                DPRINTF("read failed");
                ptable_close(table);
                table = NULL;
                goto out;
        }
        vd = (struct iso_primary_descriptor *)buf;
        if (bcmp(vd->id, ISO_STANDARD_ID, sizeof (vd->id)) != 0)
                goto out;

        entry = malloc(sizeof (*entry));
        if (entry == NULL)
                goto out;
        entry->part.start = 0;
        entry->part.end = table->sectors;
        entry->part.type = PART_ISO9660;
        entry->part.index = 0;
        STAILQ_INSERT_TAIL(&table->entries, entry, entry);

        table->type = PTABLE_ISO9660;

out:
        free(buf);
        return (table);
}

struct ptable *
ptable_open(void *dev, uint64_t sectors, uint16_t sectorsize, diskread_t *dread)
{
        struct dos_partition *dp;
        struct ptable *table;
        uint8_t *buf;
        int i;
#ifdef LOADER_MBR_SUPPORT
        struct pentry *entry;
        uint32_t start, end;
        int has_ext;
#endif
        table = NULL;
        dp = NULL;
        buf = malloc(sectorsize);
        if (buf == NULL)
                return (NULL);
        /* First, read the MBR. */
        if (dread(dev, buf, 1, DOSBBSECTOR) != 0) {
                DPRINTF("read failed");
                goto out;
        }

        table = malloc(sizeof (*table));
        if (table == NULL)
                goto out;
        table->sectors = sectors;
        table->sectorsize = sectorsize;
        table->type = PTABLE_NONE;
        STAILQ_INIT(&table->entries);

        if (ptable_iso9660read(table, dev, dread) == NULL) {
                /* Read error. */
                table = NULL;
                goto out;
        } else if (table->type == PTABLE_ISO9660)
                goto out;

        if (ptable_dklabelread(table, dev, dread) == NULL) { /* Read error. */
                table = NULL;
                goto out;
        } else if (table->type == PTABLE_VTOC)
                goto out;

#ifdef LOADER_VTOC8_SUPPORT
        if (be16dec(buf + offsetof(struct vtoc8, magic)) == VTOC_MAGIC) {
                if (ptable_vtoc8read(table, dev, dread) == NULL) {
                        /* Read error. */
                        table = NULL;
                        goto out;
                } else if (table->type == PTABLE_VTOC8)
                        goto out;
        }
#endif
        /* Check the BSD label. */
        if (ptable_bsdread(table, dev, dread) == NULL) { /* Read error. */
                table = NULL;
                goto out;
        } else if (table->type == PTABLE_BSD)
                goto out;

#if defined(LOADER_GPT_SUPPORT) || defined(LOADER_MBR_SUPPORT)
        /* Check the MBR magic. */
        if (buf[DOSMAGICOFFSET] != 0x55 ||
            buf[DOSMAGICOFFSET + 1] != 0xaa) {
                DPRINTF("magic sequence not found");
#if defined(LOADER_GPT_SUPPORT)
                /* There is no PMBR, check that we have backup GPT */
                table->type = PTABLE_GPT;
                table = ptable_gptread(table, dev, dread);
#endif
                goto out;
        }
        /* Check that we have PMBR. Also do some validation. */
        dp = malloc(NDOSPART * sizeof (struct dos_partition));
        if (dp == NULL)
                goto out;
        bcopy(buf + DOSPARTOFF, dp, NDOSPART * sizeof (struct dos_partition));

        /*
         * macOS can create PMBR partition in a hybrid MBR; that is, an MBR
         * partition which has a DOSTYP_PMBR entry defined to start at sector 1.
         * After the DOSTYP_PMBR, there may be other paritions. A UEFI
         * compliant PMBR has no other partitions.
         */
        for (i = 0; i < NDOSPART; i++) {
                if (dp[i].dp_flag != 0 && dp[i].dp_flag != 0x80) {
                        DPRINTF("invalid partition flag %x", dp[i].dp_flag);
                        goto out;
                }
#ifdef LOADER_GPT_SUPPORT
                if (dp[i].dp_typ == DOSPTYP_PMBR && dp[i].dp_start == 1) {
                        table->type = PTABLE_GPT;
                        DPRINTF("PMBR detected");
                }
#endif
        }
#ifdef LOADER_GPT_SUPPORT
        if (table->type == PTABLE_GPT) {
                table = ptable_gptread(table, dev, dread);
                goto out;
        }
#endif
#ifdef LOADER_MBR_SUPPORT
        /* Read MBR. */
        DPRINTF("MBR detected");
        table->type = PTABLE_MBR;
        for (i = has_ext = 0; i < NDOSPART; i++) {
                if (dp[i].dp_typ == 0)
                        continue;
                start = le32dec(&(dp[i].dp_start));
                end = le32dec(&(dp[i].dp_size));
                if (start == 0 || end == 0)
                        continue;
#if 0   /* Some BIOSes return an incorrect number of sectors */
                if (start + end - 1 >= sectors)
                        continue;       /* XXX: ignore */
#endif
                if (dp[i].dp_typ == DOSPTYP_EXT ||
                    dp[i].dp_typ == DOSPTYP_EXTLBA)
                        has_ext = 1;
                entry = malloc(sizeof (*entry));
                if (entry == NULL)
                        break;
                entry->part.start = start;
                entry->part.end = start + end - 1;
                entry->part.index = i + 1;
                entry->part.type = mbr_parttype(dp[i].dp_typ);
                entry->flags = dp[i].dp_flag;
                entry->type.mbr = dp[i].dp_typ;
                STAILQ_INSERT_TAIL(&table->entries, entry, entry);
                DPRINTF("new MBR partition added");
        }
        if (has_ext) {
                table = ptable_ebrread(table, dev, dread);
                /* FALLTHROUGH */
        }
#endif /* LOADER_MBR_SUPPORT */
#endif /* LOADER_MBR_SUPPORT || LOADER_GPT_SUPPORT */
out:
        free(dp);
        free(buf);
        return (table);
}

void
ptable_close(struct ptable *table)
{
        struct pentry *entry;

        if (table == NULL)
                return;

        while (!STAILQ_EMPTY(&table->entries)) {
                entry = STAILQ_FIRST(&table->entries);
                STAILQ_REMOVE_HEAD(&table->entries, entry);
                free(entry);
        }
        free(table);
}

enum ptable_type
ptable_gettype(const struct ptable *table)
{

        return (table->type);
}

int
ptable_getsize(const struct ptable *table, uint64_t *sizep)
{
        uint64_t tmp = table->sectors * table->sectorsize;

        if (tmp < table->sectors)
                return (EOVERFLOW);

        if (sizep != NULL)
                *sizep = tmp;
        return (0);
}

int
ptable_getpart(const struct ptable *table, struct ptable_entry *part, int idx)
{
        struct pentry *entry;

        if (part == NULL || table == NULL)
                return (EINVAL);

        STAILQ_FOREACH(entry, &table->entries, entry) {
                if (entry->part.index != idx)
                        continue;
                memcpy(part, &entry->part, sizeof (*part));
                return (0);
        }
        return (ENOENT);
}

/*
 * Search for a slice with the following preferences:
 *
 * 1: Active illumos slice
 * 2: Non-active illumos slice
 * 3: Active Linux slice
 * 4: non-active Linux slice
 * 5: Active FAT/FAT32 slice
 * 6: non-active FAT/FAT32 slice
 */
#define PREF_RAWDISK    0
#define PREF_ILLUMOS_ACT        1
#define PREF_ILLUMOS    2
#define PREF_LINUX_ACT  3
#define PREF_LINUX      4
#define PREF_DOS_ACT    5
#define PREF_DOS        6
#define PREF_NONE       7
int
ptable_getbestpart(const struct ptable *table, struct ptable_entry *part)
{
        struct pentry *entry, *best;
        int pref, preflevel;

        if (part == NULL || table == NULL)
                return (EINVAL);

        best = NULL;
        preflevel = pref = PREF_NONE;
        STAILQ_FOREACH(entry, &table->entries, entry) {
#ifdef LOADER_MBR_SUPPORT
                if (table->type == PTABLE_MBR) {
                        switch (entry->type.mbr) {
                        case DOSPTYP_SUNIXOS2:
                                pref = entry->flags & 0x80 ? PREF_ILLUMOS_ACT:
                                    PREF_ILLUMOS;
                                break;
                        case DOSPTYP_LINUX:
                                pref = entry->flags & 0x80 ? PREF_LINUX_ACT:
                                    PREF_LINUX;
                                break;
                        case 0x01:              /* DOS/Windows */
                        case 0x04:
                        case 0x06:
                        case 0x0c:
                        case 0x0e:
                        case DOSPTYP_FAT32:
                                pref = entry->flags & 0x80 ? PREF_DOS_ACT:
                                    PREF_DOS;
                                break;
                        default:
                                pref = PREF_NONE;
                        }
                }
#endif /* LOADER_MBR_SUPPORT */
#ifdef LOADER_GPT_SUPPORT
                if (table->type == PTABLE_GPT) {
                        if (entry->part.type == PART_DOS)
                                pref = PREF_DOS;
                        else if (entry->part.type == PART_ILLUMOS_ZFS)
                                pref = PREF_ILLUMOS;
                        else
                                pref = PREF_NONE;
                }
#endif /* LOADER_GPT_SUPPORT */
                if (pref < preflevel) {
                        preflevel = pref;
                        best = entry;
                }
        }
        if (best != NULL) {
                memcpy(part, &best->part, sizeof (*part));
                return (0);
        }
        return (ENOENT);
}

/*
 * iterate will stop if iterator will return non 0.
 */
int
ptable_iterate(const struct ptable *table, void *arg, ptable_iterate_t *iter)
{
        struct pentry *entry;
        char name[32];
        int ret = 0;

        name[0] = '\0';
        STAILQ_FOREACH(entry, &table->entries, entry) {
#ifdef LOADER_MBR_SUPPORT
                if (table->type == PTABLE_MBR)
                        sprintf(name, "s%d", entry->part.index);
                else
#endif
#ifdef LOADER_GPT_SUPPORT
                if (table->type == PTABLE_GPT)
                        sprintf(name, "p%d", entry->part.index);
                else
#endif
#ifdef LOADER_VTOC8_SUPPORT
                if (table->type == PTABLE_VTOC8)
                        sprintf(name, "%c", (uint8_t)'a' +
                            entry->part.index);
                else
#endif
                if (table->type == PTABLE_VTOC)
                        sprintf(name, "%c", (uint8_t)'a' +
                            entry->part.index);
                else
                if (table->type == PTABLE_BSD)
                        sprintf(name, "%c", (uint8_t)'a' +
                            entry->part.index);
                ret = iter(arg, name, &entry->part);
                if (ret != 0)
                        return (ret);
        }
        return (ret);
}