Implement cgptlib for MTD devices.

Defines MTD on-disk structures & API in mtdlib.h/c that closely mirrors
the existing cgpt implementation. Currently, the disk structures do not
contain guids or labels, and the number of available partition types and
quantities are limited, exactly what we want to support should be decided
before we ship this.

Adds appropriate test coverage to the unit test library - either by modifying
existing tests, or copying them and changing them accordingly.

BUG=chromium:221745
TEST=added appropriate tests to the unittests
BRANCH=cros/embedded

Change-Id: I031eca69d6c8e825b02bd0522d57e92b05eb191a
Reviewed-on: https://gerrit.chromium.org/gerrit/46082
Tested-by: Albert Chaulk <achaulk@chromium.org>
Reviewed-by: Bill Richardson <wfrichar@chromium.org>
Reviewed-by: Randall Spangler <rspangler@chromium.org>
Commit-Queue: Albert Chaulk <achaulk@chromium.org>

8 files changed, 913 insertions, 25 deletions

diff --git a/Makefile b/Makefile
index 2fba7b37..6618d158 100644
--- a/Makefile
+++ b/Makefile
@@ -249,6 +249,7 @@ VBSLK_SRCS = \
 	firmware/lib/cgptlib/cgptlib.c \
 	firmware/lib/cgptlib/cgptlib_internal.c \
 	firmware/lib/cgptlib/crc32.c \
+	firmware/lib/cgptlib/mtdlib.c \
 	firmware/lib/utility_string.c \
 	firmware/lib/vboot_api_kernel.c \
 	firmware/lib/vboot_audio.c \
diff --git a/firmware/lib/cgptlib/cgptlib.c b/firmware/lib/cgptlib/cgptlib.c
index 370530f1..49fa322b 100644
--- a/firmware/lib/cgptlib/cgptlib.c
+++ b/firmware/lib/cgptlib/cgptlib.c
@@ -112,10 +112,9 @@ int GptNextKernelEntry(GptData *gpt, uint64_t *start_sector, uint64_t *size)
 
 int GptUpdateKernelEntry(GptData *gpt, uint32_t update_type)
 {
-	GptHeader *header = (GptHeader *)gpt->primary_header;
 	GptEntry *entries = (GptEntry *)gpt->primary_entries;
 	GptEntry *e = entries + gpt->current_kernel;
-	uint16_t previous_attr = e->attrs.fields.gpt_att;
+	int modified = 0;
 
 	if (gpt->current_kernel == CGPT_KERNEL_ENTRY_NOT_FOUND)
 		return GPT_ERROR_INVALID_UPDATE_TYPE;
@@ -136,6 +135,7 @@ int GptUpdateKernelEntry(GptData *gpt, uint32_t update_type)
 		tries = GetEntryTries(e);
 		if (tries > 1) {
 			/* Still have tries left */
+			modified = 1;
 			SetEntryTries(e, tries - 1);
 			break;
 		}
@@ -148,9 +148,9 @@ int GptUpdateKernelEntry(GptData *gpt, uint32_t update_type)
 			 * Only clear tries and priority if the successful bit
 			 * is not set.
 			 */
-			e->attrs.fields.gpt_att = previous_attr &
-				~(CGPT_ATTRIBUTE_TRIES_MASK |
-				  CGPT_ATTRIBUTE_PRIORITY_MASK);
+			modified = 1;
+			SetEntryTries(e, 0);
+			SetEntryPriority(e, 0);
 		}
 		break;
 	}
@@ -158,25 +158,9 @@ int GptUpdateKernelEntry(GptData *gpt, uint32_t update_type)
 		return GPT_ERROR_INVALID_UPDATE_TYPE;
 	}
 
-	/* If no change to attributes, we're done */
-	if (e->attrs.fields.gpt_att == previous_attr)
-		return GPT_SUCCESS;
-
-	/* Update the CRCs */
-	header->entries_crc32 = Crc32((const uint8_t *)entries,
-				      header->size_of_entry *
-				      header->number_of_entries);
-	header->header_crc32 = HeaderCrc(header);
-	gpt->modified |= GPT_MODIFIED_HEADER1 | GPT_MODIFIED_ENTRIES1;
-
-	/*
-	 * Use the repair function to update the other copy of the GPT.  This
-	 * is a tad inefficient, but is much faster than the disk I/O to update
-	 * the GPT on disk so it doesn't matter.
-	 */
-	gpt->valid_headers = MASK_PRIMARY;
-	gpt->valid_entries = MASK_PRIMARY;
-	GptRepair(gpt);
+	if (modified) {
+		GptModified(gpt);
+	}
 
 	return GPT_SUCCESS;
 }
diff --git a/firmware/lib/cgptlib/cgptlib_internal.c b/firmware/lib/cgptlib/cgptlib_internal.c
index d51ce33b..dca1ce25 100644
--- a/firmware/lib/cgptlib/cgptlib_internal.c
+++ b/firmware/lib/cgptlib/cgptlib_internal.c
@@ -373,6 +373,27 @@ void GetCurrentKernelUniqueGuid(GptData *gpt, void *dest)
 	Memcpy(dest, &e->unique, sizeof(Guid));
 }
 
+void GptModified(GptData *gpt) {
+	GptHeader *header = (GptHeader *)gpt->primary_header;
+
+	/* Update the CRCs */
+	header->entries_crc32 = Crc32(gpt->primary_entries,
+				      header->size_of_entry *
+				      header->number_of_entries);
+	header->header_crc32 = HeaderCrc(header);
+	gpt->modified |= GPT_MODIFIED_HEADER1 | GPT_MODIFIED_ENTRIES1;
+
+	/*
+	 * Use the repair function to update the other copy of the GPT.  This
+	 * is a tad inefficient, but is much faster than the disk I/O to update
+	 * the GPT on disk so it doesn't matter.
+	 */
+	gpt->valid_headers = MASK_PRIMARY;
+	gpt->valid_entries = MASK_PRIMARY;
+	GptRepair(gpt);
+}
+
+
 const char *GptErrorText(int error_code)
 {
 	switch(error_code) {
@@ -412,6 +433,12 @@ const char *GptErrorText(int error_code)
 	case GPT_ERROR_DUP_GUID:
 		return "Duplicated GUID";
 
+	case GPT_ERROR_INVALID_FLASH_GEOMETRY:
+		return "Invalid flash geometry";
+
+	case GPT_ERROR_NO_SUCH_ENTRY:
+		return "No entry found";
+
 	default:
 		break;
 	};
diff --git a/firmware/lib/cgptlib/include/cgptlib.h b/firmware/lib/cgptlib/include/cgptlib.h
index ccaa2beb..ef329783 100644
--- a/firmware/lib/cgptlib/include/cgptlib.h
+++ b/firmware/lib/cgptlib/include/cgptlib.h
@@ -21,6 +21,8 @@ enum {
 	GPT_ERROR_START_LBA_OVERLAP,
 	GPT_ERROR_END_LBA_OVERLAP,
 	GPT_ERROR_DUP_GUID,
+	GPT_ERROR_INVALID_FLASH_GEOMETRY,
+	GPT_ERROR_NO_SUCH_ENTRY,
 	/* Number of errors */
 	GPT_ERROR_COUNT
 };
diff --git a/firmware/lib/cgptlib/include/cgptlib_internal.h b/firmware/lib/cgptlib/include/cgptlib_internal.h
index c7606287..b2691419 100644
--- a/firmware/lib/cgptlib/include/cgptlib_internal.h
+++ b/firmware/lib/cgptlib/include/cgptlib_internal.h
@@ -131,6 +131,12 @@ int GptSanityCheck(GptData *gpt);
  */
 void GptRepair(GptData *gpt);
 
+/**
+ * Called when the primary entries are modified and the CRCs need to be
+ * recalculated and propagated to the secondary entries
+ */
+void GptModified(GptData *gpt);
+
 /* Getters and setters for partition attribute fields. */
 
 int GetEntrySuccessful(const GptEntry *e);
diff --git a/firmware/lib/cgptlib/include/mtdlib.h b/firmware/lib/cgptlib/include/mtdlib.h
new file mode 100644
index 00000000..135a530b
--- /dev/null
+++ b/firmware/lib/cgptlib/include/mtdlib.h
@@ -0,0 +1,148 @@
+/* Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+#ifndef VBOOT_REFERENCE_MTDLIB_H_
+#define VBOOT_REFERENCE_MTDLIB_H_
+
+#include "cgptlib.h"
+#include "sysincludes.h"
+
+
+#define MTD_DRIVE_SIGNATURE "CrOSPart" /* This must be exactly 8 chars */
+
+/*
+ * Bit definitions and masks for MTD attributes.
+ *
+ * 13-16 -- partition number
+ *  9-12 -- partition type
+ *    8  -- success
+ *  7-4  -- tries
+ *  3-0  -- priority
+ */
+#define MTD_ATTRIBUTE_PRIORITY_OFFSET (0)
+#define MTD_ATTRIBUTE_MAX_PRIORITY (15UL)
+#define MTD_ATTRIBUTE_PRIORITY_MASK (MTD_ATTRIBUTE_MAX_PRIORITY << \
+                                     MTD_ATTRIBUTE_PRIORITY_OFFSET)
+
+#define MTD_ATTRIBUTE_TRIES_OFFSET (4)
+#define MTD_ATTRIBUTE_MAX_TRIES (15UL)
+#define MTD_ATTRIBUTE_TRIES_MASK (MTD_ATTRIBUTE_MAX_TRIES << \
+                                  MTD_ATTRIBUTE_TRIES_OFFSET)
+
+#define MTD_ATTRIBUTE_SUCCESSFUL_OFFSET (8)
+#define MTD_ATTRIBUTE_MAX_SUCCESSFUL (1UL)
+#define MTD_ATTRIBUTE_SUCCESSFUL_MASK (MTD_ATTRIBUTE_MAX_SUCCESSFUL << \
+                                       MTD_ATTRIBUTE_SUCCESSFUL_OFFSET)
+
+#define MTD_ATTRIBUTE_TYPE_OFFSET (9)
+#define MTD_ATTRIBUTE_MAX_TYPE (15UL)
+#define MTD_ATTRIBUTE_TYPE_MASK (MTD_ATTRIBUTE_MAX_TYPE << \
+                                 MTD_ATTRIBUTE_TYPE_OFFSET)
+
+#define MTD_ATTRIBUTE_NUMBER_OFFSET (13)
+#define MTD_ATTRIBUTE_MAX_NUMBER (15UL)
+#define MTD_ATTRIBUTE_NUMBER_MASK (MTD_ATTRIBUTE_MAX_NUMBER << \
+                                   MTD_ATTRIBUTE_NUMBER_OFFSET)
+
+
+#define MTD_PARTITION_TYPE_UNUSED             0
+#define MTD_PARTITION_TYPE_CHROMEOS_KERNEL    1
+#define MTD_PARTITION_TYPE_CHROMEOS_FIRMWARE  2
+#define MTD_PARTITION_TYPE_CHROMEOS_ROOTFS    3
+#define MTD_PARTITION_TYPE_CHROMEOS_RESERVED  4
+#define MTD_PARTITION_TYPE_CHROMEOS_FLAGSTORE 5
+#define MTD_PARTITION_TYPE_LINUX_DATA         6
+#define MTD_PARTITION_TYPE_EFI                7
+
+/* This is mostly arbitrary at the moment, but gives a little room to expand. */
+#define MTD_MAX_PARTITIONS 16
+
+
+
+typedef struct {
+  uint32_t starting_lba;
+  uint32_t ending_lba;
+  uint32_t flags;
+  uint32_t reserved;
+} __attribute__((packed)) MtdDiskPartition;
+
+typedef struct {
+  unsigned char signature[8];
+  /* For compatibility, this is only ever the CRC of the first
+   * MTD_DRIVE_V1_SIZE bytes. Further extensions must include their own CRCs,
+   * so older FW can boot newer layouts if we expand in the future.
+   */
+  uint32_t crc32;
+  uint32_t size;
+  uint32_t first_lba;
+  uint32_t last_lba;
+  MtdDiskPartition partitions[MTD_MAX_PARTITIONS];
+} __attribute__((packed)) MtdDiskLayout;
+
+#define MTD_DRIVE_V1_SIZE (24 + 16*16)
+
+#define MTDENTRY_EXPECTED_SIZE (16)
+#define MTDLAYOUT_EXPECTED_SIZE (24 + 16 * MTDENTRY_EXPECTED_SIZE)
+
+
+typedef struct {
+  /* Specifies the flash geometry, in erase blocks & write pages */
+  uint32_t flash_block_bytes;
+  uint32_t flash_page_bytes;
+
+  /* Location, in blocks, of FTS partition */
+  uint32_t fts_block_offset;
+  /* Size, in blocks, of FTS partition */
+  uint32_t fts_block_size;
+
+  /* Size of a LBA sector, in bytes */
+  uint32_t sector_bytes;
+  /* Size of drive in LBA sectors, in sectors */
+  uint64_t drive_sectors;
+
+  /*
+   * The current chromeos kernel index in partition table.  -1 means not
+   * found on drive.
+   */
+  int current_kernel;
+  int current_priority;
+
+  /* If set, the flags partition has been modified and needs to be flushed */
+  int modified;
+
+  /* Internal variables */
+  MtdDiskLayout primary;
+} MtdData;
+
+
+/* APIs are documented in cgptlib.h & cgptlib_internal.h */
+int MtdInit(MtdData *mtd);
+int MtdCheckParameters(MtdData *mtd);
+
+int MtdNextKernelEntry(MtdData *gpt, uint64_t *start_sector, uint64_t *size);
+int MtdUpdateKernelEntry(MtdData *gpt, uint32_t update_type);
+
+int MtdGetEntryPriority(const MtdDiskPartition *e);
+int MtdGetEntryTries(const MtdDiskPartition *e);
+int MtdGetEntrySuccessful(const MtdDiskPartition *e);
+int MtdGetEntryType(const MtdDiskPartition *e);
+int MtdIsKernelEntry(const MtdDiskPartition *e);
+void MtdSetEntrySuccessful(MtdDiskPartition *e, int successful) ;
+void MtdSetEntryPriority(MtdDiskPartition *e, int priority);
+void MtdSetEntryTries(MtdDiskPartition *e, int tries);
+void MtdSetEntryType(MtdDiskPartition *e, int type);
+
+void MtdModified(MtdData *mtd);
+int MtdGptInit(MtdData *mtd);
+int MtdIsPartitionValid(const MtdDiskPartition *part);
+int MtdCheckEntries(MtdDiskPartition *entries, MtdDiskLayout *h);
+int MtdSanityCheck(MtdData *disk);
+void MtdRepair(MtdData *gpt);
+void MtdGetCurrentKernelUniqueGuid(MtdData *gpt, void *dest);
+uint32_t MtdHeaderCrc(MtdDiskLayout *h);
+
+
+#endif
+
diff --git a/firmware/lib/cgptlib/mtdlib.c b/firmware/lib/cgptlib/mtdlib.c
new file mode 100644
index 00000000..bb5955b5
--- /dev/null
+++ b/firmware/lib/cgptlib/mtdlib.c
@@ -0,0 +1,337 @@
+/* Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ */
+
+#include "mtdlib.h"
+
+#include "cgptlib.h"
+#include "cgptlib_internal.h"
+#include "crc32.h"
+#include "utility.h"
+
+
+int MtdInit(MtdData *mtd) {
+  int ret;
+
+  mtd->modified = 0;
+  mtd->current_kernel = CGPT_KERNEL_ENTRY_NOT_FOUND;
+  mtd->current_priority = 999;
+
+  ret = MtdSanityCheck(mtd);
+  if (GPT_SUCCESS != ret) {
+    VBDEBUG(("MtdInit() failed sanity check\n"));
+    return ret;
+  }
+
+  return GPT_SUCCESS;
+}
+
+int MtdCheckParameters(MtdData *disk) {
+  if (disk->sector_bytes != 512) {
+    return GPT_ERROR_INVALID_SECTOR_SIZE;
+  }
+
+  /* At minimum, the disk must consist of at least one erase block */
+  if (disk->drive_sectors < disk->flash_block_bytes / disk->sector_bytes) {
+    return GPT_ERROR_INVALID_SECTOR_NUMBER;
+  }
+
+  /* Write pages must be an integer multiple of sector size */
+  if (disk->flash_page_bytes == 0 ||
+      disk->flash_page_bytes % disk->sector_bytes != 0) {
+    return GPT_ERROR_INVALID_FLASH_GEOMETRY;
+  }
+
+  /* Erase blocks must be an integer multiple of write pages */
+  if (disk->flash_block_bytes == 0 ||
+      disk->flash_block_bytes % disk->flash_page_bytes != 0) {
+    return GPT_ERROR_INVALID_FLASH_GEOMETRY;
+  }
+
+  /* Without a FTS region, why are you using MTD? */
+  if (disk->fts_block_size == 0) {
+    return GPT_ERROR_INVALID_FLASH_GEOMETRY;
+  }
+  return GPT_SUCCESS;
+}
+
+int MtdGetEntryPriority(const MtdDiskPartition *e) {
+  return ((e->flags & MTD_ATTRIBUTE_PRIORITY_MASK) >>
+          MTD_ATTRIBUTE_PRIORITY_OFFSET);
+}
+
+int MtdGetEntryTries(const MtdDiskPartition *e) {
+  return ((e->flags & MTD_ATTRIBUTE_TRIES_MASK) >>
+          MTD_ATTRIBUTE_TRIES_OFFSET);
+}
+
+int MtdGetEntrySuccessful(const MtdDiskPartition *e) {
+  return ((e->flags & MTD_ATTRIBUTE_SUCCESSFUL_MASK) >>
+          MTD_ATTRIBUTE_SUCCESSFUL_OFFSET);
+}
+
+int MtdGetEntryType(const MtdDiskPartition *e) {
+  return ((e->flags & MTD_ATTRIBUTE_TYPE_MASK) >> MTD_ATTRIBUTE_TYPE_OFFSET);
+}
+
+int MtdIsKernelEntry(const MtdDiskPartition *e) {
+  return MtdGetEntryType(e) == MTD_PARTITION_TYPE_CHROMEOS_KERNEL;
+}
+
+
+static void SetBitfield(MtdDiskPartition *e,
+                        uint32_t offset, uint32_t mask, uint32_t v) {
+  e->flags = (e->flags & ~mask) | ((v << offset) & mask);
+}
+void MtdSetEntrySuccessful(MtdDiskPartition *e, int successful) {
+  SetBitfield(e, MTD_ATTRIBUTE_SUCCESSFUL_OFFSET, MTD_ATTRIBUTE_SUCCESSFUL_MASK,
+              successful);
+}
+void MtdSetEntryPriority(MtdDiskPartition *e, int priority) {
+  SetBitfield(e, MTD_ATTRIBUTE_PRIORITY_OFFSET, MTD_ATTRIBUTE_PRIORITY_MASK,
+              priority);
+}
+void MtdSetEntryTries(MtdDiskPartition *e, int tries) {
+  SetBitfield(e, MTD_ATTRIBUTE_TRIES_OFFSET, MTD_ATTRIBUTE_TRIES_MASK, tries);
+}
+void MtdSetEntryType(MtdDiskPartition *e, int type) {
+  SetBitfield(e, MTD_ATTRIBUTE_TYPE_OFFSET, MTD_ATTRIBUTE_TYPE_MASK, type);
+}
+
+void MtdModified(MtdData *mtd) {
+  mtd->modified = 1;
+  mtd->primary.crc32 = MtdHeaderCrc(&mtd->primary);
+}
+
+int MtdIsPartitionValid(const MtdDiskPartition *part) {
+  return MtdGetEntryType(part) != 0;
+}
+
+int MtdCheckEntries(MtdDiskPartition *entries, MtdDiskLayout *h) {
+  uint32_t i, j;
+
+  for (i = 0; i < MTD_MAX_PARTITIONS; i++) {
+    for (j = 0; j < MTD_MAX_PARTITIONS; j++) {
+      if (i != j) {
+        MtdDiskPartition *entry = entries + i;
+        MtdDiskPartition *e2 = entries + j;
+
+        if (!MtdIsPartitionValid(entry) || !MtdIsPartitionValid(e2))
+          continue;
+
+        if((entry->starting_lba == 0 && entry->ending_lba == 0) ||
+           (e2->starting_lba == 0 && e2->ending_lba == 0)) {
+          continue;
+        }
+
+        if (entry->ending_lba > h->last_lba) {
+          return GPT_ERROR_OUT_OF_REGION;
+        }
+        if (entry->starting_lba < h->first_lba) {
+          return GPT_ERROR_OUT_OF_REGION;
+        }
+        if (entry->starting_lba > entry->ending_lba) {
+          return GPT_ERROR_OUT_OF_REGION;
+        }
+
+        if ((entry->starting_lba >= e2->starting_lba) &&
+            (entry->starting_lba <= e2->ending_lba)) {
+          return GPT_ERROR_START_LBA_OVERLAP;
+        }
+        if ((entry->ending_lba >= e2->starting_lba) &&
+            (entry->ending_lba <= e2->ending_lba)) {
+          return GPT_ERROR_END_LBA_OVERLAP;
+        }
+      }
+    }
+  }
+  return GPT_SUCCESS;
+}
+
+int MtdSanityCheck(MtdData *disk) {
+  MtdDiskLayout *h = &disk->primary;
+  int ret;
+
+  ret = MtdCheckParameters(disk);
+  if(GPT_SUCCESS != ret)
+    return ret;
+
+  if (Memcmp(disk->primary.signature, MTD_DRIVE_SIGNATURE,
+              sizeof(disk->primary.signature))) {
+    return GPT_ERROR_INVALID_HEADERS;
+  }
+
+  if (disk->primary.first_lba > disk->primary.last_lba ||
+      disk->primary.last_lba > disk->drive_sectors) {
+    return GPT_ERROR_INVALID_SECTOR_NUMBER;
+  }
+
+  if (h->crc32 != MtdHeaderCrc(h)) {
+    return GPT_ERROR_CRC_CORRUPTED;
+  }
+  if (h->size < MTD_DRIVE_V1_SIZE) {
+    return GPT_ERROR_INVALID_HEADERS;
+  }
+  return MtdCheckEntries(h->partitions, h);
+}
+
+void MtdRepair(MtdData *gpt) {
+
+}
+
+void MtdGetCurrentKernelUniqueGuid(MtdData *gpt, void *dest) {
+  Memset(dest, 0, 16);
+}
+
+uint32_t MtdHeaderCrc(MtdDiskLayout *h) {
+  uint32_t crc32, original_crc32;
+
+  /* Original CRC is calculated with the CRC field 0. */
+  original_crc32 = h->crc32;
+  h->crc32 = 0;
+  crc32 = Crc32((const uint8_t *)h, h->size);
+  h->crc32 = original_crc32;
+
+  return crc32;
+}
+
+
+
+int MtdNextKernelEntry(MtdData *mtd, uint64_t *start_sector, uint64_t *size)
+{
+  MtdDiskLayout *header = &mtd->primary;
+  MtdDiskPartition *entries = header->partitions;
+  MtdDiskPartition *e;
+  int new_kernel = CGPT_KERNEL_ENTRY_NOT_FOUND;
+  int new_prio = 0;
+  uint32_t i;
+
+  /*
+   * If we already found a kernel, continue the scan at the current
+   * kernel's priority, in case there is another kernel with the same
+   * priority.
+   */
+  if (mtd->current_kernel != CGPT_KERNEL_ENTRY_NOT_FOUND) {
+    for (i = mtd->current_kernel + 1;
+         i < MTD_MAX_PARTITIONS; i++) {
+      e = entries + i;
+      if (!MtdIsKernelEntry(e))
+        continue;
+      VBDEBUG(("GptNextKernelEntry looking at same prio "
+         "partition %d\n", i+1));
+      VBDEBUG(("GptNextKernelEntry s%d t%d p%d\n",
+         MtdGetEntrySuccessful(e), MtdGetEntryTries(e),
+         MtdGetEntryPriority(e)));
+      if (!(MtdGetEntrySuccessful(e) || MtdGetEntryTries(e)))
+        continue;
+      if (MtdGetEntryPriority(e) == mtd->current_priority) {
+        mtd->current_kernel = i;
+        *start_sector = e->starting_lba;
+        *size = e->ending_lba - e->starting_lba + 1;
+        VBDEBUG(("GptNextKernelEntry likes it\n"));
+        return GPT_SUCCESS;
+      }
+    }
+  }
+
+  /*
+   * We're still here, so scan for the remaining kernel with the highest
+   * priority less than the previous attempt.
+   */
+  for (i = 0, e = entries; i < MTD_MAX_PARTITIONS; i++, e++) {
+    int current_prio = MtdGetEntryPriority(e);
+    if (!MtdIsKernelEntry(e))
+      continue;
+    VBDEBUG(("GptNextKernelEntry looking at new prio "
+       "partition %d\n", i+1));
+    VBDEBUG(("GptNextKernelEntry s%d t%d p%d\n",
+       MtdGetEntrySuccessful(e), MtdGetEntryTries(e),
+       MtdGetEntryPriority(e)));
+    if (!(MtdGetEntrySuccessful(e) || MtdGetEntryTries(e)))
+      continue;
+    if (current_prio >= mtd->current_priority) {
+      /* Already returned this kernel in a previous call */
+      continue;
+    }
+    if (current_prio > new_prio) {
+      new_kernel = i;
+      new_prio = current_prio;
+    }
+  }
+
+  /*
+   * Save what we found.  Note that if we didn't find a new kernel,
+   * new_prio will still be -1, so future calls to this function will
+   * also fail.
+   */
+  mtd->current_kernel = new_kernel;
+  mtd->current_priority = new_prio;
+
+  if (CGPT_KERNEL_ENTRY_NOT_FOUND == new_kernel) {
+    VBDEBUG(("GptNextKernelEntry no more kernels\n"));
+    return GPT_ERROR_NO_VALID_KERNEL;
+  }
+
+  VBDEBUG(("GptNextKernelEntry likes partition %d\n", new_kernel + 1));
+  e = entries + new_kernel;
+  *start_sector = e->starting_lba;
+  *size = e->ending_lba - e->starting_lba + 1;
+  return GPT_SUCCESS;
+}
+
+int MtdUpdateKernelEntry(MtdData *mtd, uint32_t update_type)
+{
+  MtdDiskLayout *header = &mtd->primary;
+  MtdDiskPartition *entries = header->partitions;
+  MtdDiskPartition *e = entries + mtd->current_kernel;
+  int modified = 0;
+
+  if (mtd->current_kernel == CGPT_KERNEL_ENTRY_NOT_FOUND)
+    return GPT_ERROR_INVALID_UPDATE_TYPE;
+  if (!MtdIsKernelEntry(e))
+    return GPT_ERROR_INVALID_UPDATE_TYPE;
+
+  switch (update_type) {
+  case GPT_UPDATE_ENTRY_TRY: {
+    /* Used up a try */
+    int tries;
+    if (MtdGetEntrySuccessful(e)) {
+      /*
+       * Successfully booted this partition, so tries field
+       * is ignored.
+       */
+      return GPT_SUCCESS;
+    }
+    tries = MtdGetEntryTries(e);
+    if (tries > 1) {
+      /* Still have tries left */
+      modified = 1;
+      MtdSetEntryTries(e, tries - 1);
+      break;
+    }
+    /* Out of tries, so drop through and mark partition bad. */
+  }
+  case GPT_UPDATE_ENTRY_BAD: {
+    /* Giving up on this partition entirely. */
+    if (!MtdGetEntrySuccessful(e)) {
+      /*
+       * Only clear tries and priority if the successful bit
+       * is not set.
+       */
+      modified = 1;
+      MtdSetEntryTries(e, 0);
+      MtdSetEntryPriority(e, 0);
+    }
+    break;
+  }
+  default:
+    return GPT_ERROR_INVALID_UPDATE_TYPE;
+  }
+
+  if (modified) {
+    MtdModified(mtd);
+  }
+
+  return GPT_SUCCESS;
+}
diff --git a/tests/cgptlib_test.c b/tests/cgptlib_test.c
index 415b1996..66f8a27f 100644
--- a/tests/cgptlib_test.c
+++ b/tests/cgptlib_test.c
@@ -13,6 +13,7 @@
 #include "errno.h"
 #include "flash_ts.h"
 #include "gpt.h"
+#include "mtdlib.h"
 #include "test_common.h"
 #define _STUB_IMPLEMENTATION_
 #include "utility.h"
@@ -140,6 +141,13 @@ static GptData *GetEmptyGptData(void)
 	return &gpt;
 }
 
+static MtdData *GetEmptyMtdData() {
+	static MtdData mtd;
+	Memset(&mtd, 0, sizeof(mtd));
+	mtd.current_kernel = CGPT_KERNEL_ENTRY_NOT_FOUND;
+	return &mtd;
+}
+
 /*
  * Fill in most of fields and creates the layout described in the top of this
  * file. Before calling this function, primary/secondary header/entries must
@@ -206,6 +214,49 @@ static void BuildTestGptData(GptData *gpt)
 	RefreshCrc32(gpt);
 }
 
+static void BuildTestMtdData(MtdData *mtd) {
+	MtdDiskPartition *partitions;
+
+	mtd->sector_bytes = DEFAULT_SECTOR_SIZE;
+	mtd->drive_sectors = DEFAULT_DRIVE_SECTORS;
+	mtd->current_kernel = CGPT_KERNEL_ENTRY_NOT_FOUND;
+	mtd->modified = 0;
+	Memset(&mtd->primary, 0, sizeof(mtd->primary));
+
+	Memcpy(mtd->primary.signature, MTD_DRIVE_SIGNATURE,
+		sizeof(mtd->primary.signature));
+	mtd->primary.first_lba = 32;
+	mtd->primary.last_lba = DEFAULT_DRIVE_SECTORS - 1;
+	mtd->primary.size = MTD_DRIVE_V1_SIZE;
+
+	/* These values are not used directly by the library, but they are checked */
+	mtd->flash_page_bytes = mtd->sector_bytes * 8;
+	mtd->flash_block_bytes = mtd->flash_page_bytes * 8;
+	mtd->fts_block_offset = 1;
+	mtd->fts_block_size = 1;
+
+	partitions = &mtd->primary.partitions[0];
+	partitions[0].starting_lba = 34;
+	partitions[0].ending_lba = 133;
+	partitions[0].flags =
+		MTD_PARTITION_TYPE_CHROMEOS_KERNEL << MTD_ATTRIBUTE_TYPE_OFFSET;
+	partitions[1].starting_lba = 134;
+	partitions[1].ending_lba = 232;
+	partitions[1].flags =
+		MTD_PARTITION_TYPE_CHROMEOS_ROOTFS << MTD_ATTRIBUTE_TYPE_OFFSET;
+	partitions[2].starting_lba = 234;
+	partitions[2].ending_lba = 331;
+	partitions[2].flags =
+		MTD_PARTITION_TYPE_CHROMEOS_KERNEL << MTD_ATTRIBUTE_TYPE_OFFSET;
+	partitions[3].starting_lba = 334;
+	partitions[3].ending_lba = 430;
+	partitions[3].flags =
+		MTD_PARTITION_TYPE_CHROMEOS_ROOTFS << MTD_ATTRIBUTE_TYPE_OFFSET;
+
+	mtd->primary.crc32 = 0;
+	mtd->primary.crc32 = Crc32(&mtd->primary, MTD_DRIVE_V1_SIZE);
+}
+
 
 /*
  * Test if the structures are the expected size; if this fails, struct packing
@@ -217,6 +268,8 @@ static int StructSizeTest(void)
 	EXPECT(GUID_EXPECTED_SIZE == sizeof(Guid));
 	EXPECT(GPTHEADER_EXPECTED_SIZE == sizeof(GptHeader));
 	EXPECT(GPTENTRY_EXPECTED_SIZE == sizeof(GptEntry));
+	EXPECT(MTDENTRY_EXPECTED_SIZE == sizeof(MtdDiskPartition));
+	EXPECT(MTDLAYOUT_EXPECTED_SIZE == sizeof(MtdDiskLayout));
 	return TEST_OK;
 }
 
@@ -234,15 +287,26 @@ static int TestBuildTestGptData(void)
 	return TEST_OK;
 }
 
+static int TestBuildTestMtdData() {
+	MtdData *mtd = GetEmptyMtdData();
+
+	BuildTestMtdData(mtd);
+	EXPECT(GPT_SUCCESS == MtdInit(mtd));
+	return TEST_OK;
+}
+
 /*
  * Test if wrong sector_bytes or drive_sectors is detected by GptInit().
  * Currently we only support 512 bytes per sector.  In the future, we may
  * support other sizes.  A too small drive_sectors should be rejected by
  * GptInit().
+ * For MtdInit(), additionally test various flash geometries to verify
+ * that only valid ones are accepted.
  */
 static int ParameterTests(void)
 {
 	GptData *gpt;
+	MtdData *mtd;
 	struct {
 		uint32_t sector_bytes;
 		uint64_t drive_sectors;
@@ -256,6 +320,32 @@ static int ParameterTests(void)
 		 GPT_ENTRIES_SECTORS * 2, GPT_SUCCESS},
 		{4096, DEFAULT_DRIVE_SECTORS, GPT_ERROR_INVALID_SECTOR_SIZE},
 	};
+	struct {
+		uint32_t sector_bytes;
+		uint32_t drive_sectors;
+		uint32_t flash_page_bytes;
+		uint32_t flash_block_bytes;
+		int expected_retval;
+	} mtdcases[] = {
+		{512, DEFAULT_DRIVE_SECTORS, 8*512,
+			8*512, GPT_SUCCESS},
+		{510, DEFAULT_DRIVE_SECTORS, 8*512,
+			8*512, GPT_ERROR_INVALID_SECTOR_SIZE},
+		{512, DEFAULT_DRIVE_SECTORS, 8*512,
+			8*512, GPT_SUCCESS},
+		{512, DEFAULT_DRIVE_SECTORS, 512,
+			8*512, GPT_SUCCESS},
+		{512, DEFAULT_DRIVE_SECTORS, 8*512,
+			10*512, GPT_ERROR_INVALID_FLASH_GEOMETRY},
+		{512, DEFAULT_DRIVE_SECTORS, 3*512,
+			9*512, GPT_SUCCESS},
+		{512, DEFAULT_DRIVE_SECTORS, 8*512,
+			6*512, GPT_ERROR_INVALID_FLASH_GEOMETRY},
+		{512, DEFAULT_DRIVE_SECTORS, 256,
+			6*512, GPT_ERROR_INVALID_FLASH_GEOMETRY},
+		{512, DEFAULT_DRIVE_SECTORS, 512,
+			6*512 + 256, GPT_ERROR_INVALID_FLASH_GEOMETRY},
+	};
 	int i;
 
 	gpt = GetEmptyGptData();
@@ -266,6 +356,19 @@ static int ParameterTests(void)
 		EXPECT(cases[i].expected_retval == CheckParameters(gpt));
 	}
 
+	mtd = GetEmptyMtdData();
+	for (i = 0; i < ARRAY_SIZE(cases); ++i) {
+		BuildTestMtdData(mtd);
+		mtd->sector_bytes = mtdcases[i].sector_bytes;
+		mtd->drive_sectors = mtdcases[i].drive_sectors;
+		mtd->flash_block_bytes = mtdcases[i].flash_block_bytes;
+		mtd->flash_page_bytes = mtdcases[i].flash_page_bytes;
+		if(mtdcases[i].expected_retval != MtdCheckParameters(mtd)) {
+			printf("i=%d\n",i);
+		}
+		EXPECT(mtdcases[i].expected_retval == MtdCheckParameters(mtd));
+	}
+
 	return TEST_OK;
 }
 
@@ -695,6 +798,9 @@ static int ValidEntryTest(void)
 	GptData *gpt = GetEmptyGptData();
 	GptHeader *h1 = (GptHeader *)gpt->primary_header;
 	GptEntry *e1 = (GptEntry *)(gpt->primary_entries);
+	MtdData *mtd = GetEmptyMtdData();
+	MtdDiskLayout *mh = &mtd->primary;
+	MtdDiskPartition *me = mh->partitions;
 
 	/* error case: entry.StartingLBA < header.FirstUsableLBA */
 	BuildTestGptData(gpt);
@@ -702,18 +808,35 @@ static int ValidEntryTest(void)
 	RefreshCrc32(gpt);
 	EXPECT(GPT_ERROR_OUT_OF_REGION == CheckEntries(e1, h1));
 
+	BuildTestMtdData(mtd);
+	if (mh->first_lba > 0) {
+		me[0].starting_lba = mh->first_lba - 1;
+		mh->crc32 = MtdHeaderCrc(mh);
+		EXPECT(GPT_ERROR_OUT_OF_REGION == MtdCheckEntries(me, mh));
+	}
+
 	/* error case: entry.EndingLBA > header.LastUsableLBA */
 	BuildTestGptData(gpt);
 	e1[2].ending_lba = h1->last_usable_lba + 1;
 	RefreshCrc32(gpt);
 	EXPECT(GPT_ERROR_OUT_OF_REGION == CheckEntries(e1, h1));
 
+	BuildTestMtdData(mtd);
+	me[0].ending_lba = mh->last_lba + 1;
+	mh->crc32 = MtdHeaderCrc(mh);
+	EXPECT(GPT_ERROR_OUT_OF_REGION == MtdCheckEntries(me, mh));
+
 	/* error case: entry.StartingLBA > entry.EndingLBA */
 	BuildTestGptData(gpt);
 	e1[3].starting_lba = e1[3].ending_lba + 1;
 	RefreshCrc32(gpt);
 	EXPECT(GPT_ERROR_OUT_OF_REGION == CheckEntries(e1, h1));
 
+	BuildTestMtdData(mtd);
+	me[0].starting_lba = me[0].ending_lba + 1;
+	mh->crc32 = MtdHeaderCrc(mh);
+	EXPECT(GPT_ERROR_OUT_OF_REGION == MtdCheckEntries(me, mh));
+
 	/* case: non active entry should be ignored. */
 	BuildTestGptData(gpt);
 	Memset(&e1[1].type, 0, sizeof(e1[1].type));
@@ -721,6 +844,12 @@ static int ValidEntryTest(void)
 	RefreshCrc32(gpt);
 	EXPECT(0 == CheckEntries(e1, h1));
 
+	BuildTestMtdData(mtd);
+	me[0].flags = 0;
+	me[0].starting_lba = me[0].ending_lba + 1;
+	mh->crc32 = MtdHeaderCrc(mh);
+	EXPECT(GPT_SUCCESS == MtdCheckEntries(me, mh));
+
 	return TEST_OK;
 }
 
@@ -729,6 +858,9 @@ static int OverlappedPartitionTest(void) {
 	GptData *gpt = GetEmptyGptData();
 	GptHeader *h = (GptHeader *)gpt->primary_header;
 	GptEntry *e = (GptEntry *)gpt->primary_entries;
+	MtdData *mtd = GetEmptyMtdData();
+	MtdDiskLayout *mh = &mtd->primary;
+	MtdDiskPartition *me = mh->partitions;
 	int i, j;
 
 	struct {
@@ -781,20 +913,29 @@ static int OverlappedPartitionTest(void) {
 
 	for (i = 0; i < ARRAY_SIZE(cases); ++i) {
 		BuildTestGptData(gpt);
+		BuildTestMtdData(mtd);
+		Memset(mh->partitions, 0, sizeof(mh->partitions));
 		ZeroEntries(gpt);
 		for(j = 0; j < ARRAY_SIZE(cases[0].entries); ++j) {
 			if (!cases[i].entries[j].starting_lba)
 				break;
 
-			if (cases[i].entries[j].active)
+			if (cases[i].entries[j].active) {
 				Memcpy(&e[j].type, &guid_kernel, sizeof(Guid));
+				me[j].flags =
+					MTD_PARTITION_TYPE_CHROMEOS_KERNEL << MTD_ATTRIBUTE_TYPE_OFFSET;
+			}
 			SetGuid(&e[j].unique, j);
 			e[j].starting_lba = cases[i].entries[j].starting_lba;
 			e[j].ending_lba = cases[i].entries[j].ending_lba;
+			me[j].starting_lba = cases[i].entries[j].starting_lba;
+			me[j].ending_lba = cases[i].entries[j].ending_lba;
+
 		}
 		RefreshCrc32(gpt);
 
 		EXPECT(cases[i].overlapped == CheckEntries(e, h));
+		EXPECT(cases[i].overlapped == MtdCheckEntries(me, mh));
 	}
 	return TEST_OK;
 }
@@ -1001,6 +1142,8 @@ static int EntryAttributeGetSetTest(void)
 {
 	GptData *gpt = GetEmptyGptData();
 	GptEntry *e = (GptEntry *)(gpt->primary_entries);
+	MtdData *mtd = GetEmptyMtdData();
+	MtdDiskPartition *m = &mtd->primary.partitions[0];
 
 	e->attrs.whole = 0x0000000000000000ULL;
 	SetEntrySuccessful(e, 1);
@@ -1011,6 +1154,15 @@ static int EntryAttributeGetSetTest(void)
 	EXPECT(0xFEFFFFFFFFFFFFFFULL == e->attrs.whole);
 	EXPECT(0 == GetEntrySuccessful(e));
 
+	m->flags = 0;
+	MtdSetEntrySuccessful(m, 1);
+	EXPECT(0x00000100 == m->flags);
+	EXPECT(1 == MtdGetEntrySuccessful(m));
+	m->flags = ~0;
+	MtdSetEntrySuccessful(m, 0);
+	EXPECT(0xFFFFFEFF == m->flags);
+	EXPECT(0 == MtdGetEntrySuccessful(m));
+
 	e->attrs.whole = 0x0000000000000000ULL;
 	SetEntryTries(e, 15);
 	EXPECT(15 == GetEntryTries(e));
@@ -1020,6 +1172,15 @@ static int EntryAttributeGetSetTest(void)
 	EXPECT(0xFF0FFFFFFFFFFFFFULL == e->attrs.whole);
 	EXPECT(0 == GetEntryTries(e));
 
+	m->flags = 0;
+	MtdSetEntryTries(m, 15);
+	EXPECT(0x000000F0 == m->flags);
+	EXPECT(15 == MtdGetEntryTries(m));
+	m->flags = ~0;
+	MtdSetEntryTries(m, 0);
+	EXPECT(0xFFFFFF0F == m->flags);
+	EXPECT(0 == MtdGetEntryTries(m));
+
 	e->attrs.whole = 0x0000000000000000ULL;
 	SetEntryPriority(e, 15);
 	EXPECT(0x000F000000000000ULL == e->attrs.whole);
@@ -1029,6 +1190,15 @@ static int EntryAttributeGetSetTest(void)
 	EXPECT(0xFFF0FFFFFFFFFFFFULL == e->attrs.whole);
 	EXPECT(0 == GetEntryPriority(e));
 
+	m->flags = 0;
+	MtdSetEntryPriority(m, 15);
+	EXPECT(0x0000000F == m->flags);
+	EXPECT(15 == MtdGetEntryPriority(m));
+	m->flags = ~0;
+	MtdSetEntryPriority(m, 0);
+	EXPECT(0xFFFFFFF0 == m->flags);
+	EXPECT(0 == MtdGetEntryPriority(m));
+
 	e->attrs.whole = 0xFFFFFFFFFFFFFFFFULL;
 	EXPECT(1 == GetEntrySuccessful(e));
 	EXPECT(15 == GetEntryPriority(e));
@@ -1068,6 +1238,11 @@ static void FreeEntry(GptEntry *e)
 	Memset(&e->type, 0, sizeof(Guid));
 }
 
+static void MtdFreeEntry(MtdDiskPartition *e)
+{
+	MtdSetEntryType(e, MTD_PARTITION_TYPE_UNUSED);
+}
+
 /* Set up an entry. */
 static void FillEntry(GptEntry *e, int is_kernel,
                       int priority, int successful, int tries)
@@ -1078,6 +1253,16 @@ static void FillEntry(GptEntry *e, int is_kernel,
 	SetEntryTries(e, tries);
 }
 
+static void MtdFillEntry(MtdDiskPartition *e, int is_kernel,
+                      int priority, int successful, int tries)
+{
+	MtdSetEntryType(e, is_kernel ? MTD_PARTITION_TYPE_CHROMEOS_KERNEL :
+									MTD_PARTITION_TYPE_CHROMEOS_FIRMWARE);
+	MtdSetEntryPriority(e, priority);
+	MtdSetEntrySuccessful(e, successful);
+	MtdSetEntryTries(e, tries);
+}
+
 /*
  * Invalidate all kernel entries and expect GptNextKernelEntry() cannot find
  * any usable kernel entry.
@@ -1097,6 +1282,20 @@ static int NoValidKernelEntryTest(void)
 	return TEST_OK;
 }
 
+static int MtdNoValidKernelEntryTest(void)
+{
+	MtdData *mtd = GetEmptyMtdData();
+	MtdDiskPartition *e1 = mtd->primary.partitions;
+
+	BuildTestMtdData(mtd);
+	MtdSetEntryPriority(e1 + KERNEL_A, 0);
+	MtdFreeEntry(e1 + KERNEL_B);
+	EXPECT(GPT_ERROR_NO_VALID_KERNEL ==
+	       MtdNextKernelEntry(mtd, NULL, NULL));
+
+	return TEST_OK;
+}
+
 static int GetNextNormalTest(void)
 {
 	GptData *gpt = GetEmptyGptData();
@@ -1184,6 +1383,169 @@ static int GetNextTriesTest(void)
 	return TEST_OK;
 }
 
+static int MtdGetNextNormalTest(void)
+{
+	MtdData *mtd = GetEmptyMtdData();
+	MtdDiskPartition *e1 = mtd->primary.partitions;
+	uint64_t start, size;
+
+	/* Normal case - both kernels successful */
+	BuildTestMtdData(mtd);
+	MtdFillEntry(e1 + KERNEL_A, 1, 2, 1, 0);
+	MtdFillEntry(e1 + KERNEL_B, 1, 2, 1, 0);
+	mtd->primary.crc32 = MtdHeaderCrc(&mtd->primary);
+	MtdInit(mtd);
+
+	EXPECT(GPT_SUCCESS == MtdNextKernelEntry(mtd, &start, &size));
+	EXPECT(KERNEL_A == mtd->current_kernel);
+	EXPECT(34 == start);
+	EXPECT(100 == size);
+
+	EXPECT(GPT_SUCCESS == MtdNextKernelEntry(mtd, &start, &size));
+	EXPECT(KERNEL_B == mtd->current_kernel);
+	EXPECT(134 == start);
+	EXPECT(99 == size);
+
+	EXPECT(GPT_ERROR_NO_VALID_KERNEL ==
+	       MtdNextKernelEntry(mtd, &start, &size));
+	EXPECT(-1 == mtd->current_kernel);
+
+	/* Call as many times as you want; you won't get another kernel... */
+	EXPECT(GPT_ERROR_NO_VALID_KERNEL ==
+	       MtdNextKernelEntry(mtd, &start, &size));
+	EXPECT(-1 == mtd->current_kernel);
+
+	return TEST_OK;
+}
+
+static int MtdGetNextPrioTest(void)
+{
+	MtdData *mtd = GetEmptyMtdData();
+	MtdDiskPartition *e1 = mtd->primary.partitions;
+	uint64_t start, size;
+
+	/* Priority 3, 4, 0, 4 - should boot order B, Y, A */
+	BuildTestMtdData(mtd);
+	MtdFillEntry(e1 + KERNEL_A, 1, 3, 1, 0);
+	MtdFillEntry(e1 + KERNEL_B, 1, 4, 1, 0);
+	MtdFillEntry(e1 + KERNEL_X, 1, 0, 1, 0);
+	MtdFillEntry(e1 + KERNEL_Y, 1, 4, 1, 0);
+	mtd->primary.crc32 = MtdHeaderCrc(&mtd->primary);
+	MtdInit(mtd);
+
+	EXPECT(GPT_SUCCESS == MtdNextKernelEntry(mtd, &start, &size));
+	EXPECT(KERNEL_B == mtd->current_kernel);
+	EXPECT(GPT_SUCCESS == MtdNextKernelEntry(mtd, &start, &size));
+	EXPECT(KERNEL_Y == mtd->current_kernel);
+	EXPECT(GPT_SUCCESS == MtdNextKernelEntry(mtd, &start, &size));
+	EXPECT(KERNEL_A == mtd->current_kernel);
+	EXPECT(GPT_ERROR_NO_VALID_KERNEL ==
+	       MtdNextKernelEntry(mtd, &start, &size));
+
+	return TEST_OK;
+}
+
+static int MtdGetNextTriesTest(void)
+{
+	MtdData *mtd = GetEmptyMtdData();
+	MtdDiskPartition *e1 = mtd->primary.partitions;
+	uint64_t start, size;
+
+	/* Tries=nonzero is attempted just like success, but tries=0 isn't */
+	BuildTestMtdData(mtd);
+	MtdFillEntry(e1 + KERNEL_A, 1, 2, 1, 0);
+	MtdFillEntry(e1 + KERNEL_B, 1, 3, 0, 0);
+	MtdFillEntry(e1 + KERNEL_X, 1, 4, 0, 1);
+	MtdFillEntry(e1 + KERNEL_Y, 1, 0, 0, 5);
+	mtd->primary.crc32 = MtdHeaderCrc(&mtd->primary);
+	MtdInit(mtd);
+
+	EXPECT(GPT_SUCCESS == MtdNextKernelEntry(mtd, &start, &size));
+	EXPECT(KERNEL_X == mtd->current_kernel);
+	EXPECT(GPT_SUCCESS == MtdNextKernelEntry(mtd, &start, &size));
+	EXPECT(KERNEL_A == mtd->current_kernel);
+	EXPECT(GPT_ERROR_NO_VALID_KERNEL ==
+	       MtdNextKernelEntry(mtd, &start, &size));
+
+	return TEST_OK;
+}
+
+static int MtdUpdateTest() {
+	MtdData *mtd = GetEmptyMtdData();
+	MtdDiskPartition *e = &mtd->primary.partitions[0];
+	uint64_t start, size;
+
+	BuildTestMtdData(mtd);
+
+	/* Tries=nonzero is attempted just like success, but tries=0 isn't */
+	MtdFillEntry(e + KERNEL_A, 1, 4, 1, 0);
+	MtdFillEntry(e + KERNEL_B, 1, 3, 0, 2);
+	MtdFillEntry(e + KERNEL_X, 1, 2, 0, 2);
+	mtd->primary.crc32 = MtdHeaderCrc(&mtd->primary);
+	mtd->modified = 0;
+	EXPECT(GPT_SUCCESS == MtdInit(mtd));
+
+	/* Successful kernel */
+	EXPECT(GPT_SUCCESS == MtdNextKernelEntry(mtd, &start, &size));
+	EXPECT(KERNEL_A == mtd->current_kernel);
+	EXPECT(1 == MtdGetEntrySuccessful(e + KERNEL_A));
+	EXPECT(4 == MtdGetEntryPriority(e + KERNEL_A));
+	EXPECT(0 == MtdGetEntryTries(e + KERNEL_A));
+	/* Trying successful kernel changes nothing */
+	EXPECT(GPT_SUCCESS == MtdUpdateKernelEntry(mtd, GPT_UPDATE_ENTRY_TRY));
+	EXPECT(1 == MtdGetEntrySuccessful(e + KERNEL_A));
+	EXPECT(4 == MtdGetEntryPriority(e + KERNEL_A));
+	EXPECT(0 == MtdGetEntryTries(e + KERNEL_A));
+	EXPECT(0 == mtd->modified);
+	/* Marking it bad also does not update it. */
+	EXPECT(GPT_SUCCESS == MtdUpdateKernelEntry(mtd, GPT_UPDATE_ENTRY_BAD));
+	EXPECT(1 == MtdGetEntrySuccessful(e + KERNEL_A));
+	EXPECT(4 == MtdGetEntryPriority(e + KERNEL_A));
+	EXPECT(0 == MtdGetEntryTries(e + KERNEL_A));
+	EXPECT(0 == mtd->modified);
+
+	/* Kernel with tries */
+	EXPECT(GPT_SUCCESS == MtdNextKernelEntry(mtd, &start, &size));
+	EXPECT(KERNEL_B == mtd->current_kernel);
+	EXPECT(0 == MtdGetEntrySuccessful(e + KERNEL_B));
+	EXPECT(3 == MtdGetEntryPriority(e + KERNEL_B));
+	EXPECT(2 == MtdGetEntryTries(e + KERNEL_B));
+	/* Marking it bad clears it */
+	EXPECT(GPT_SUCCESS == MtdUpdateKernelEntry(mtd, GPT_UPDATE_ENTRY_BAD));
+	EXPECT(0 == MtdGetEntrySuccessful(e + KERNEL_B));
+	EXPECT(0 == MtdGetEntryPriority(e + KERNEL_B));
+	EXPECT(0 == MtdGetEntryTries(e + KERNEL_B));
+	/* And that's caused the mtd to need updating */
+	EXPECT(1 == mtd->modified);
+
+	/* Another kernel with tries */
+	EXPECT(GPT_SUCCESS == MtdNextKernelEntry(mtd, &start, &size));
+	EXPECT(KERNEL_X == mtd->current_kernel);
+	EXPECT(0 == MtdGetEntrySuccessful(e + KERNEL_X));
+	EXPECT(2 == MtdGetEntryPriority(e + KERNEL_X));
+	EXPECT(2 == MtdGetEntryTries(e + KERNEL_X));
+	/* Trying it uses up a try */
+	EXPECT(GPT_SUCCESS == MtdUpdateKernelEntry(mtd, GPT_UPDATE_ENTRY_TRY));
+	EXPECT(0 == MtdGetEntrySuccessful(e + KERNEL_X));
+	EXPECT(2 == MtdGetEntryPriority(e + KERNEL_X));
+	EXPECT(1 == MtdGetEntryTries(e + KERNEL_X));
+	/* Trying it again marks it inactive */
+	EXPECT(GPT_SUCCESS == MtdUpdateKernelEntry(mtd, GPT_UPDATE_ENTRY_TRY));
+	EXPECT(0 == MtdGetEntrySuccessful(e + KERNEL_X));
+	EXPECT(0 == MtdGetEntryPriority(e + KERNEL_X));
+	EXPECT(0 == MtdGetEntryTries(e + KERNEL_X));
+
+	/* Can't update if entry isn't a kernel, or there isn't an entry */
+	MtdSetEntryType(e + KERNEL_X, MTD_PARTITION_TYPE_UNUSED);
+	EXPECT(GPT_ERROR_INVALID_UPDATE_TYPE ==
+	       MtdUpdateKernelEntry(mtd, GPT_UPDATE_ENTRY_BAD));
+	mtd->current_kernel = CGPT_KERNEL_ENTRY_NOT_FOUND;
+	EXPECT(GPT_ERROR_INVALID_UPDATE_TYPE ==
+	       MtdUpdateKernelEntry(mtd, GPT_UPDATE_ENTRY_BAD));
+
+	return TEST_OK;
+}
+
 static int GptUpdateTest(void)
 {
 	GptData *gpt = GetEmptyGptData();
@@ -1290,6 +1652,20 @@ static int UpdateInvalidKernelTypeTest(void)
 	return TEST_OK;
 }
 
+static int MtdUpdateInvalidKernelTypeTest(void)
+{
+	MtdData *mtd = GetEmptyMtdData();
+
+	BuildTestMtdData(mtd);
+	/* anything, but not CGPT_KERNEL_ENTRY_NOT_FOUND */
+	mtd->current_kernel = 0;
+	/* any invalid update_type value */
+	EXPECT(GPT_ERROR_INVALID_UPDATE_TYPE ==
+	       MtdUpdateKernelEntry(mtd, 99));
+
+	return TEST_OK;
+}
+
 /* Test duplicate UniqueGuids can be detected. */
 static int DuplicateUniqueGuidTest(void)
 {
@@ -1570,6 +1946,7 @@ int main(int argc, char *argv[])
 	} test_cases[] = {
 		{ TEST_CASE(StructSizeTest), },
 		{ TEST_CASE(TestBuildTestGptData), },
+		{ TEST_CASE(TestBuildTestMtdData), },
 		{ TEST_CASE(ParameterTests), },
 		{ TEST_CASE(HeaderCrcTest), },
 		{ TEST_CASE(HeaderSameTest), },
@@ -1587,13 +1964,19 @@ int main(int argc, char *argv[])
 		{ TEST_CASE(OverlappedPartitionTest), },
 		{ TEST_CASE(SanityCheckTest), },
 		{ TEST_CASE(NoValidKernelEntryTest), },
+		{ TEST_CASE(MtdNoValidKernelEntryTest), },
 		{ TEST_CASE(EntryAttributeGetSetTest), },
 		{ TEST_CASE(EntryTypeTest), },
 		{ TEST_CASE(GetNextNormalTest), },
 		{ TEST_CASE(GetNextPrioTest), },
 		{ TEST_CASE(GetNextTriesTest), },
+		{ TEST_CASE(MtdGetNextNormalTest), },
+		{ TEST_CASE(MtdGetNextPrioTest), },
+		{ TEST_CASE(MtdGetNextTriesTest), },
 		{ TEST_CASE(GptUpdateTest), },
+		{ TEST_CASE(MtdUpdateTest), },
 		{ TEST_CASE(UpdateInvalidKernelTypeTest), },
+		{ TEST_CASE(MtdUpdateInvalidKernelTypeTest), },
 		{ TEST_CASE(DuplicateUniqueGuidTest), },
 		{ TEST_CASE(TestCrc32TestVectors), },
 		{ TEST_CASE(GetKernelGuidTest), },