/* gpt.c - Read/Verify/Write GUID Partition Tables (GPT).  */
/*
 *  VasEBoot  --  GRand Unified Bootloader
 *  Copyright (C) 2002,2005,2006,2007,2008  Free Software Foundation, Inc.
 *  Copyright (C) 2014 CoreOS, Inc.
 *
 *  VasEBoot is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  VasEBoot is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with VasEBoot.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <VasEBoot/charset.h>
#include <VasEBoot/crypto.h>
#include <VasEBoot/device.h>
#include <VasEBoot/disk.h>
#include <VasEBoot/misc.h>
#include <VasEBoot/mm.h>
#include <VasEBoot/dl.h>
#include <VasEBoot/msdos_partition.h>
#include <VasEBoot/gpt_partition.h>

VasEBoot_MOD_LICENSE ("GPLv3+");

static VasEBoot_uint8_t VasEBoot_gpt_magic[] = VasEBoot_GPT_HEADER_MAGIC;

static VasEBoot_err_t
VasEBoot_gpt_read_entries (VasEBoot_disk_t disk, VasEBoot_gpt_t gpt,
		       struct VasEBoot_gpt_header *header,
		       void **ret_entries,
		       VasEBoot_size_t *ret_entries_size);

char *
VasEBoot_gpt_guid_to_str (VasEBoot_gpt_guid_t *guid)
{
  return VasEBoot_xasprintf ("%08x-%04x-%04x-%02x%02x-%02x%02x%02x%02x%02x%02x",
			 VasEBoot_le_to_cpu32 (guid->data1),
			 VasEBoot_le_to_cpu16 (guid->data2),
			 VasEBoot_le_to_cpu16 (guid->data3),
			 guid->data4[0], guid->data4[1],
			 guid->data4[2], guid->data4[3],
			 guid->data4[4], guid->data4[5],
			 guid->data4[6], guid->data4[7]);
}

static VasEBoot_err_t
VasEBoot_gpt_device_partentry (VasEBoot_device_t device,
			   struct VasEBoot_gpt_partentry *entry)
{
  VasEBoot_disk_t disk = device->disk;
  VasEBoot_partition_t p;
  VasEBoot_err_t err;

  if (!disk || !disk->partition)
    return VasEBoot_error (VasEBoot_ERR_BUG, "not a partition");

  if (VasEBoot_strcmp (disk->partition->partmap->name, "gpt"))
    return VasEBoot_error (VasEBoot_ERR_BAD_ARGUMENT, "not a GPT partition");

  p = disk->partition;
  disk->partition = p->parent;
  err = VasEBoot_disk_read (disk, p->offset, p->index, sizeof (*entry), entry);
  disk->partition = p;

  return err;
}

VasEBoot_err_t
VasEBoot_gpt_part_label (VasEBoot_device_t device, char **label)
{
  struct VasEBoot_gpt_partentry entry;
  const VasEBoot_size_t name_len = ARRAY_SIZE (entry.name);
  const VasEBoot_size_t label_len = name_len * VasEBoot_MAX_UTF8_PER_UTF16 + 1;
  VasEBoot_size_t i;
  VasEBoot_uint8_t *end;

  if (VasEBoot_gpt_device_partentry (device, &entry))
    return VasEBoot_errno;

  *label = VasEBoot_malloc (label_len);
  if (!*label)
    return VasEBoot_errno;

  for (i = 0; i < name_len; i++)
    entry.name[i] = VasEBoot_le_to_cpu16 (entry.name[i]);

  end = VasEBoot_utf16_to_utf8 ((VasEBoot_uint8_t *) *label, entry.name, name_len);
  *end = '\0';

  return VasEBoot_ERR_NONE;
}

VasEBoot_err_t
VasEBoot_gpt_part_uuid (VasEBoot_device_t device, char **uuid)
{
  struct VasEBoot_gpt_partentry entry;

  if (VasEBoot_gpt_device_partentry (device, &entry))
    return VasEBoot_errno;

  *uuid = VasEBoot_gpt_guid_to_str (&entry.guid);
  if (!*uuid)
    return VasEBoot_errno;

  return VasEBoot_ERR_NONE;
}

static struct VasEBoot_gpt_header *
VasEBoot_gpt_get_header (VasEBoot_gpt_t gpt)
{
  if (gpt->status & VasEBoot_GPT_PRIMARY_HEADER_VALID)
    return &gpt->primary;
  else if (gpt->status & VasEBoot_GPT_BACKUP_HEADER_VALID)
    return &gpt->backup;

  VasEBoot_error (VasEBoot_ERR_BUG, "No valid GPT header");
  return NULL;
}

VasEBoot_err_t
VasEBoot_gpt_disk_uuid (VasEBoot_device_t device, char **uuid)
{
  struct VasEBoot_gpt_header *header;

  VasEBoot_gpt_t gpt = VasEBoot_gpt_read (device->disk);
  if (!gpt)
    goto done;

  header = VasEBoot_gpt_get_header (gpt);
  if (!header)
    goto done;

  *uuid = VasEBoot_gpt_guid_to_str (&header->guid);

done:
  VasEBoot_gpt_free (gpt);
  return VasEBoot_errno;
}

static VasEBoot_uint64_t
VasEBoot_gpt_size_to_sectors (VasEBoot_gpt_t gpt, VasEBoot_size_t size)
{
  unsigned int sector_size;
  VasEBoot_uint64_t sectors;

  sector_size = 1U << gpt->log_sector_size;
  sectors = size / sector_size;
  if (size % sector_size)
    sectors++;

  return sectors;
}

/* Copied from VasEBoot-core/kern/disk_common.c VasEBoot_disk_adjust_range so we can
 * avoid attempting to use disk->total_sectors when VasEBoot won't let us.
 * TODO: Why is disk->total_sectors not set to VasEBoot_DISK_SIZE_UNKNOWN?  */
static int
VasEBoot_gpt_disk_size_valid (VasEBoot_disk_t disk)
{
  VasEBoot_disk_addr_t total_sectors;

  /* Transform total_sectors to number of 512B blocks.  */
  total_sectors = disk->total_sectors << (disk->log_sector_size - VasEBoot_DISK_SECTOR_BITS);

  /* Some drivers have problems with disks above reasonable.
     Treat unknown as 1EiB disk. While on it, clamp the size to 1EiB.
     Just one condition is enough since VasEBoot_DISK_UNKNOWN_SIZE << ls is always
     above 9EiB.
  */
  if (total_sectors > (1ULL << 51))
    return 0;

  return 1;
}

static void
VasEBoot_gpt_lecrc32 (VasEBoot_uint32_t *crc, const void *data, VasEBoot_size_t len)
{
  VasEBoot_uint32_t crc32_val;

  VasEBoot_crypto_hash (VasEBoot_MD_CRC32, &crc32_val, data, len);

  /* VasEBoot_MD_CRC32 always uses big endian, gpt is always little.  */
  *crc = VasEBoot_swap_bytes32 (crc32_val);
}

static void
VasEBoot_gpt_header_lecrc32 (VasEBoot_uint32_t *crc, struct VasEBoot_gpt_header *header)
{
  VasEBoot_uint32_t old, new;

  /* crc32 must be computed with the field cleared.  */
  old = header->crc32;
  header->crc32 = 0;
  VasEBoot_gpt_lecrc32 (&new, header, sizeof (*header));
  header->crc32 = old;

  *crc = new;
}

/* Make sure the MBR is a protective MBR and not a normal MBR.  */
VasEBoot_err_t
VasEBoot_gpt_pmbr_check (struct VasEBoot_msdos_partition_mbr *mbr)
{
  unsigned int i;

  if (mbr->signature !=
      VasEBoot_cpu_to_le16_compile_time (VasEBoot_PC_PARTITION_SIGNATURE))
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid MBR signature");

  for (i = 0; i < sizeof (mbr->entries); i++)
    if (mbr->entries[i].type == VasEBoot_PC_PARTITION_TYPE_GPT_DISK)
      return VasEBoot_ERR_NONE;

  return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid protective MBR");
}

static VasEBoot_uint64_t
VasEBoot_gpt_entries_size (struct VasEBoot_gpt_header *gpt)
{
  return (VasEBoot_uint64_t) VasEBoot_le_to_cpu32 (gpt->maxpart) *
         (VasEBoot_uint64_t) VasEBoot_le_to_cpu32 (gpt->partentry_size);
}

static VasEBoot_uint64_t
VasEBoot_gpt_entries_sectors (struct VasEBoot_gpt_header *gpt,
			  unsigned int log_sector_size)
{
  VasEBoot_uint64_t sector_bytes, entries_bytes;

  sector_bytes = 1ULL << log_sector_size;
  entries_bytes = VasEBoot_gpt_entries_size (gpt);
  return VasEBoot_divmod64(entries_bytes + sector_bytes - 1, sector_bytes, NULL);
}

static int
is_pow2 (VasEBoot_uint32_t n)
{
  return (n & (n - 1)) == 0;
}

VasEBoot_err_t
VasEBoot_gpt_header_check (struct VasEBoot_gpt_header *gpt,
		       unsigned int log_sector_size)
{
  VasEBoot_uint32_t crc = 0, size;
  VasEBoot_uint64_t start, end;

  if (VasEBoot_memcmp (gpt->magic, VasEBoot_gpt_magic, sizeof (VasEBoot_gpt_magic)) != 0)
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid GPT signature");

  if (gpt->version != VasEBoot_GPT_HEADER_VERSION)
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "unknown GPT version");

  VasEBoot_gpt_header_lecrc32 (&crc, gpt);
  if (gpt->crc32 != crc)
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid GPT header crc32");

  /* The header size "must be greater than or equal to 92 and must be less
   * than or equal to the logical block size."  */
  size = VasEBoot_le_to_cpu32 (gpt->headersize);
  if (size < 92U || size > (1U << log_sector_size))
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid GPT header size");

  /* The partition entry size must be "a value of 128*(2^n) where n is an
   * integer greater than or equal to zero (e.g., 128, 256, 512, etc.)."  */
  size = VasEBoot_le_to_cpu32 (gpt->partentry_size);
  if (size < 128U || size % 128U || !is_pow2 (size / 128U))
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid GPT entry size");

  /* The minimum entries table size is specified in terms of bytes,
   * regardless of how large the individual entry size is.  */
  if (VasEBoot_gpt_entries_size (gpt) < VasEBoot_GPT_DEFAULT_ENTRIES_SIZE)
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid GPT entry table size");

  /* And of course there better be some space for partitions!  */
  start = VasEBoot_le_to_cpu64 (gpt->start);
  end = VasEBoot_le_to_cpu64 (gpt->end);
  if (start > end)
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid usable sectors");

  return VasEBoot_ERR_NONE;
}

static int
VasEBoot_gpt_headers_equal (VasEBoot_gpt_t gpt)
{
  /* Assume headers passed VasEBoot_gpt_header_check so skip magic and version.
   * Individual fields must be checked instead of just using memcmp because
   * crc32, header, alternate, and partitions will all normally differ.  */

  if (gpt->primary.headersize != gpt->backup.headersize ||
      gpt->primary.header_lba != gpt->backup.alternate_lba ||
      gpt->primary.alternate_lba != gpt->backup.header_lba ||
      gpt->primary.start != gpt->backup.start ||
      gpt->primary.end != gpt->backup.end ||
      gpt->primary.maxpart != gpt->backup.maxpart ||
      gpt->primary.partentry_size != gpt->backup.partentry_size ||
      gpt->primary.partentry_crc32 != gpt->backup.partentry_crc32)
    return 0;

  return VasEBoot_memcmp(&gpt->primary.guid, &gpt->backup.guid,
                     sizeof(VasEBoot_gpt_guid_t)) == 0;
}

static VasEBoot_err_t
VasEBoot_gpt_check_primary (VasEBoot_gpt_t gpt)
{
  VasEBoot_uint64_t backup, primary, entries, entries_len, start, end;

  primary = VasEBoot_le_to_cpu64 (gpt->primary.header_lba);
  backup = VasEBoot_le_to_cpu64 (gpt->primary.alternate_lba);
  entries = VasEBoot_le_to_cpu64 (gpt->primary.partitions);
  entries_len = VasEBoot_gpt_entries_sectors(&gpt->primary, gpt->log_sector_size);
  start = VasEBoot_le_to_cpu64 (gpt->primary.start);
  end = VasEBoot_le_to_cpu64 (gpt->primary.end);

  VasEBoot_dprintf ("gpt", "Primary GPT layout:\n"
		"primary header = 0x%llx backup header = 0x%llx\n"
		"entries location = 0x%llx length = 0x%llx\n"
		"first usable = 0x%llx last usable = 0x%llx\n",
		(unsigned long long) primary,
		(unsigned long long) backup,
		(unsigned long long) entries,
		(unsigned long long) entries_len,
		(unsigned long long) start,
		(unsigned long long) end);

  if (VasEBoot_gpt_header_check (&gpt->primary, gpt->log_sector_size))
    return VasEBoot_errno;
  if (primary != 1)
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid primary GPT LBA");
  if (entries <= 1 || entries+entries_len > start)
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid entries location");
  if (backup <= end)
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid backup GPT LBA");

  return VasEBoot_ERR_NONE;
}

static VasEBoot_err_t
VasEBoot_gpt_check_backup (VasEBoot_gpt_t gpt)
{
  VasEBoot_uint64_t backup, primary, entries, entries_len, start, end;

  backup = VasEBoot_le_to_cpu64 (gpt->backup.header_lba);
  primary = VasEBoot_le_to_cpu64 (gpt->backup.alternate_lba);
  entries = VasEBoot_le_to_cpu64 (gpt->backup.partitions);
  entries_len = VasEBoot_gpt_entries_sectors(&gpt->backup, gpt->log_sector_size);
  start = VasEBoot_le_to_cpu64 (gpt->backup.start);
  end = VasEBoot_le_to_cpu64 (gpt->backup.end);

  VasEBoot_dprintf ("gpt", "Backup GPT layout:\n"
		"primary header = 0x%llx backup header = 0x%llx\n"
		"entries location = 0x%llx length = 0x%llx\n"
		"first usable = 0x%llx last usable = 0x%llx\n",
		(unsigned long long) primary,
		(unsigned long long) backup,
		(unsigned long long) entries,
		(unsigned long long) entries_len,
		(unsigned long long) start,
		(unsigned long long) end);

  if (VasEBoot_gpt_header_check (&gpt->backup, gpt->log_sector_size))
    return VasEBoot_errno;
  if (primary != 1)
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid primary GPT LBA");
  if (entries <= end || entries+entries_len > backup)
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid entries location");
  if (backup <= end)
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid backup GPT LBA");

  /* If both primary and backup are valid but differ prefer the primary.  */
  if ((gpt->status & VasEBoot_GPT_PRIMARY_HEADER_VALID) &&
      !VasEBoot_gpt_headers_equal (gpt))
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "backup GPT out of sync");

  return VasEBoot_ERR_NONE;
}

static VasEBoot_err_t
VasEBoot_gpt_read_primary (VasEBoot_disk_t disk, VasEBoot_gpt_t gpt)
{
  VasEBoot_disk_addr_t addr;

  /* TODO: The gpt partmap module searches for the primary header instead
   * of relying on the disk's sector size. For now trust the disk driver
   * but eventually this code should match the existing behavior.  */
  gpt->log_sector_size = disk->log_sector_size;

  VasEBoot_dprintf ("gpt", "reading primary GPT from sector 0x1\n");

  addr = VasEBoot_gpt_sector_to_addr (gpt, 1);
  if (VasEBoot_disk_read (disk, addr, 0, sizeof (gpt->primary), &gpt->primary))
    return VasEBoot_errno;

  if (VasEBoot_gpt_check_primary (gpt))
    return VasEBoot_errno;

  gpt->status |= VasEBoot_GPT_PRIMARY_HEADER_VALID;

  if (VasEBoot_gpt_read_entries (disk, gpt, &gpt->primary,
			     &gpt->entries, &gpt->entries_size))
    return VasEBoot_errno;

  gpt->status |= VasEBoot_GPT_PRIMARY_ENTRIES_VALID;

  return VasEBoot_ERR_NONE;
}

static VasEBoot_err_t
VasEBoot_gpt_read_backup (VasEBoot_disk_t disk, VasEBoot_gpt_t gpt)
{
  void *entries = NULL;
  VasEBoot_size_t entries_size;
  VasEBoot_uint64_t sector;
  VasEBoot_disk_addr_t addr;

  /* Assumes gpt->log_sector_size == disk->log_sector_size  */
  if (gpt->status & VasEBoot_GPT_PRIMARY_HEADER_VALID)
    {
      sector = VasEBoot_le_to_cpu64 (gpt->primary.alternate_lba);
      if (VasEBoot_gpt_disk_size_valid (disk) && sector >= disk->total_sectors)
	return VasEBoot_error (VasEBoot_ERR_OUT_OF_RANGE,
			   "backup GPT located at 0x%llx, "
			   "beyond last disk sector at 0x%llx",
			   (unsigned long long) sector,
			   (unsigned long long) disk->total_sectors - 1);
    }
  else if (VasEBoot_gpt_disk_size_valid (disk))
    sector = disk->total_sectors - 1;
  else
    return VasEBoot_error (VasEBoot_ERR_OUT_OF_RANGE,
		       "size of disk unknown, cannot locate backup GPT");

  VasEBoot_dprintf ("gpt", "reading backup GPT from sector 0x%llx\n",
		(unsigned long long) sector);

  addr = VasEBoot_gpt_sector_to_addr (gpt, sector);
  if (VasEBoot_disk_read (disk, addr, 0, sizeof (gpt->backup), &gpt->backup))
    return VasEBoot_errno;

  if (VasEBoot_gpt_check_backup (gpt))
    return VasEBoot_errno;

  /* Ensure the backup header thinks it is located where we found it.  */
  if (VasEBoot_le_to_cpu64 (gpt->backup.header_lba) != sector)
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid backup GPT LBA");

  gpt->status |= VasEBoot_GPT_BACKUP_HEADER_VALID;

  if (VasEBoot_gpt_read_entries (disk, gpt, &gpt->backup,
			     &entries, &entries_size))
    return VasEBoot_errno;

  if (gpt->status & VasEBoot_GPT_PRIMARY_ENTRIES_VALID)
    {
      if (entries_size != gpt->entries_size ||
	  VasEBoot_memcmp (entries, gpt->entries, entries_size) != 0)
	return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "backup GPT out of sync");

      VasEBoot_free (entries);
    }
  else
    {
      gpt->entries = entries;
      gpt->entries_size = entries_size;
    }

  gpt->status |= VasEBoot_GPT_BACKUP_ENTRIES_VALID;

  return VasEBoot_ERR_NONE;
}

static VasEBoot_err_t
VasEBoot_gpt_read_entries (VasEBoot_disk_t disk, VasEBoot_gpt_t gpt,
		       struct VasEBoot_gpt_header *header,
		       void **ret_entries,
		       VasEBoot_size_t *ret_entries_size)
{
  void *entries = NULL;
  VasEBoot_uint32_t count, size, crc;
  VasEBoot_uint64_t sector;
  VasEBoot_disk_addr_t addr;
  VasEBoot_size_t entries_size;

  /* VasEBoot doesn't include calloc, hence the manual overflow check.  */
  count = VasEBoot_le_to_cpu32 (header->maxpart);
  size = VasEBoot_le_to_cpu32 (header->partentry_size);
  entries_size = count *size;
  if (size && entries_size / size != count)
    {
      VasEBoot_error (VasEBoot_ERR_OUT_OF_MEMORY, N_("out of memory"));
      goto fail;
    }

  /* Double check that the header was validated properly.  */
  if (entries_size < VasEBoot_GPT_DEFAULT_ENTRIES_SIZE)
    return VasEBoot_error (VasEBoot_ERR_BUG, "invalid GPT entries table size");

  entries = VasEBoot_malloc (entries_size);
  if (!entries)
    goto fail;

  sector = VasEBoot_le_to_cpu64 (header->partitions);
  VasEBoot_dprintf ("gpt", "reading GPT %lu entries from sector 0x%llx\n",
		(unsigned long) count,
		(unsigned long long) sector);

  addr = VasEBoot_gpt_sector_to_addr (gpt, sector);
  if (VasEBoot_disk_read (disk, addr, 0, entries_size, entries))
    goto fail;

  VasEBoot_gpt_lecrc32 (&crc, entries, entries_size);
  if (crc != header->partentry_crc32)
    {
      VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "invalid GPT entry crc32");
      goto fail;
    }

  *ret_entries = entries;
  *ret_entries_size = entries_size;
  return VasEBoot_ERR_NONE;

fail:
  VasEBoot_free (entries);
  return VasEBoot_errno;
}

VasEBoot_gpt_t
VasEBoot_gpt_read (VasEBoot_disk_t disk)
{
  VasEBoot_gpt_t gpt;

  VasEBoot_dprintf ("gpt", "reading GPT from %s\n", disk->name);

  gpt = VasEBoot_zalloc (sizeof (*gpt));
  if (!gpt)
    goto fail;

  if (VasEBoot_disk_read (disk, 0, 0, sizeof (gpt->mbr), &gpt->mbr))
    goto fail;

  /* Check the MBR but errors aren't reported beyond the status bit.  */
  if (VasEBoot_gpt_pmbr_check (&gpt->mbr))
    VasEBoot_errno = VasEBoot_ERR_NONE;
  else
    gpt->status |= VasEBoot_GPT_PROTECTIVE_MBR;

  /* If both the primary and backup fail report the primary's error.  */
  if (VasEBoot_gpt_read_primary (disk, gpt))
    {
      VasEBoot_error_push ();
      VasEBoot_gpt_read_backup (disk, gpt);
      VasEBoot_error_pop ();
    }
  else
    VasEBoot_gpt_read_backup (disk, gpt);

  /* If either succeeded clear any possible error from the other.  */
  if (VasEBoot_gpt_primary_valid (gpt) || VasEBoot_gpt_backup_valid (gpt))
    VasEBoot_errno = VasEBoot_ERR_NONE;
  else
    goto fail;

  return gpt;

fail:
  VasEBoot_gpt_free (gpt);
  return NULL;
}

struct VasEBoot_gpt_partentry *
VasEBoot_gpt_get_partentry (VasEBoot_gpt_t gpt, VasEBoot_uint32_t n)
{
  struct VasEBoot_gpt_header *header;
  VasEBoot_size_t offset;

  header = VasEBoot_gpt_get_header (gpt);
  if (!header)
    return NULL;

  if (n >= VasEBoot_le_to_cpu32 (header->maxpart))
    return NULL;

  offset = (VasEBoot_size_t) VasEBoot_le_to_cpu32 (header->partentry_size) * n;
  return (struct VasEBoot_gpt_partentry *) ((char *) gpt->entries + offset);
}

VasEBoot_err_t
VasEBoot_gpt_repair (VasEBoot_disk_t disk, VasEBoot_gpt_t gpt)
{
  /* Skip if there is nothing to do.  */
  if (VasEBoot_gpt_both_valid (gpt))
    return VasEBoot_ERR_NONE;

  VasEBoot_dprintf ("gpt", "repairing GPT for %s\n", disk->name);

  if (disk->log_sector_size != gpt->log_sector_size)
    return VasEBoot_error (VasEBoot_ERR_NOT_IMPLEMENTED_YET,
		       "GPT sector size must match disk sector size");

  if (VasEBoot_gpt_primary_valid (gpt))
    {
      VasEBoot_uint64_t backup_header;

      VasEBoot_dprintf ("gpt", "primary GPT is valid\n");

      /* Relocate backup to end if disk if the disk has grown.  */
      backup_header = VasEBoot_le_to_cpu64 (gpt->primary.alternate_lba);
      if (VasEBoot_gpt_disk_size_valid (disk) &&
	  disk->total_sectors - 1 > backup_header)
	{
	  backup_header = disk->total_sectors - 1;
	  VasEBoot_dprintf ("gpt", "backup GPT header relocated to 0x%llx\n",
			(unsigned long long) backup_header);

	  gpt->primary.alternate_lba = VasEBoot_cpu_to_le64 (backup_header);
	}

      VasEBoot_memcpy (&gpt->backup, &gpt->primary, sizeof (gpt->backup));
      gpt->backup.header_lba = gpt->primary.alternate_lba;
      gpt->backup.alternate_lba = gpt->primary.header_lba;
      gpt->backup.partitions = VasEBoot_cpu_to_le64 (backup_header -
	  VasEBoot_gpt_size_to_sectors (gpt, gpt->entries_size));
    }
  else if (VasEBoot_gpt_backup_valid (gpt))
    {
      VasEBoot_dprintf ("gpt", "backup GPT is valid\n");

      VasEBoot_memcpy (&gpt->primary, &gpt->backup, sizeof (gpt->primary));
      gpt->primary.header_lba = gpt->backup.alternate_lba;
      gpt->primary.alternate_lba = gpt->backup.header_lba;
      gpt->primary.partitions = VasEBoot_cpu_to_le64_compile_time (2);
    }
  else
    return VasEBoot_error (VasEBoot_ERR_BUG, "No valid GPT");

  if (VasEBoot_gpt_update (gpt))
    return VasEBoot_errno;

  VasEBoot_dprintf ("gpt", "repairing GPT for %s successful\n", disk->name);

  return VasEBoot_ERR_NONE;
}

VasEBoot_err_t
VasEBoot_gpt_update (VasEBoot_gpt_t gpt)
{
  VasEBoot_uint32_t crc;

  /* Clear status bits, require revalidation of everything.  */
  gpt->status &= ~(VasEBoot_GPT_PRIMARY_HEADER_VALID |
		   VasEBoot_GPT_PRIMARY_ENTRIES_VALID |
		   VasEBoot_GPT_BACKUP_HEADER_VALID |
		   VasEBoot_GPT_BACKUP_ENTRIES_VALID);

  /* Writing headers larger than our header structure are unsupported.  */
  gpt->primary.headersize =
    VasEBoot_cpu_to_le32_compile_time (sizeof (gpt->primary));
  gpt->backup.headersize =
    VasEBoot_cpu_to_le32_compile_time (sizeof (gpt->backup));

  VasEBoot_gpt_lecrc32 (&crc, gpt->entries, gpt->entries_size);
  gpt->primary.partentry_crc32 = crc;
  gpt->backup.partentry_crc32 = crc;

  VasEBoot_gpt_header_lecrc32 (&gpt->primary.crc32, &gpt->primary);
  VasEBoot_gpt_header_lecrc32 (&gpt->backup.crc32, &gpt->backup);

  if (VasEBoot_gpt_check_primary (gpt))
    {
      VasEBoot_error_push ();
      return VasEBoot_error (VasEBoot_ERR_BUG, "Generated invalid GPT primary header");
    }

  gpt->status |= (VasEBoot_GPT_PRIMARY_HEADER_VALID |
		  VasEBoot_GPT_PRIMARY_ENTRIES_VALID);

  if (VasEBoot_gpt_check_backup (gpt))
    {
      VasEBoot_error_push ();
      return VasEBoot_error (VasEBoot_ERR_BUG, "Generated invalid GPT backup header");
    }

  gpt->status |= (VasEBoot_GPT_BACKUP_HEADER_VALID |
		  VasEBoot_GPT_BACKUP_ENTRIES_VALID);

  return VasEBoot_ERR_NONE;
}

static VasEBoot_err_t
VasEBoot_gpt_write_table (VasEBoot_disk_t disk, VasEBoot_gpt_t gpt,
		      struct VasEBoot_gpt_header *header)
{
  VasEBoot_disk_addr_t addr;

  if (VasEBoot_le_to_cpu32 (header->headersize) != sizeof (*header))
    return VasEBoot_error (VasEBoot_ERR_NOT_IMPLEMENTED_YET,
		       "Header size is %u, must be %u",
		       VasEBoot_le_to_cpu32 (header->headersize),
		       sizeof (*header));

  addr = VasEBoot_gpt_sector_to_addr (gpt, VasEBoot_le_to_cpu64 (header->header_lba));
  if (addr == 0)
    return VasEBoot_error (VasEBoot_ERR_BUG,
		       "Refusing to write GPT header to address 0x0");
  if (VasEBoot_disk_write (disk, addr, 0, sizeof (*header), header))
    return VasEBoot_errno;

  addr = VasEBoot_gpt_sector_to_addr (gpt, VasEBoot_le_to_cpu64 (header->partitions));
  if (addr < 2)
    return VasEBoot_error (VasEBoot_ERR_BUG,
		       "Refusing to write GPT entries to address 0x%llx",
		       (unsigned long long) addr);
  if (VasEBoot_disk_write (disk, addr, 0, gpt->entries_size, gpt->entries))
    return VasEBoot_errno;

  return VasEBoot_ERR_NONE;
}

VasEBoot_err_t
VasEBoot_gpt_write (VasEBoot_disk_t disk, VasEBoot_gpt_t gpt)
{
  VasEBoot_uint64_t backup_header;

  /* TODO: update/repair protective MBRs too.  */

  if (!VasEBoot_gpt_both_valid (gpt))
    return VasEBoot_error (VasEBoot_ERR_BAD_PART_TABLE, "Invalid GPT data");

  /* Write the backup GPT first so if writing fails the update is aborted
   * and the primary is left intact.  However if the backup location is
   * inaccessible we have to just skip and hope for the best, the backup
   * will need to be repaired in the OS.  */
  backup_header = VasEBoot_le_to_cpu64 (gpt->backup.header_lba);
  if (VasEBoot_gpt_disk_size_valid (disk) &&
      backup_header >= disk->total_sectors)
    {
      VasEBoot_printf ("warning: backup GPT located at 0x%llx, "
		   "beyond last disk sector at 0x%llx\n",
		   (unsigned long long) backup_header,
		   (unsigned long long) disk->total_sectors - 1);
      VasEBoot_printf ("warning: only writing primary GPT, "
	           "the backup GPT must be repaired from the OS\n");
    }
  else
    {
      VasEBoot_dprintf ("gpt", "writing backup GPT to %s\n", disk->name);
      if (VasEBoot_gpt_write_table (disk, gpt, &gpt->backup))
	return VasEBoot_errno;
    }

  VasEBoot_dprintf ("gpt", "writing primary GPT to %s\n", disk->name);
  if (VasEBoot_gpt_write_table (disk, gpt, &gpt->primary))
    return VasEBoot_errno;

  return VasEBoot_ERR_NONE;
}

void
VasEBoot_gpt_free (VasEBoot_gpt_t gpt)
{
  if (!gpt)
    return;

  VasEBoot_free (gpt->entries);
  VasEBoot_free (gpt);
}