brillo_uefi_x86_64: Add A/B selection logic.

Boot loaders used for Brillo must support A/B selection that allows for
selecting from multiple boot slots (typically 2, perhaps up to 4). This
provides capabilities of background system updates and the presence of
redundant partitions as backups.  A/B flow lets the bootloader fall
back on valid slots when the threshold attempts to boot newly
downloaded slots is reached -- thus reverting to a known safe state.

This CL provides both the implementation of the brillo x86-64 uefi boot
loader which is aware of multiple slots and a posix test harness to
unit test the a/b flow logic of this boot loader.

Additional A/B flow logic description may be found in the Brillo Boot
Loader Requirements.

BUG=29072323
TEST=Unit tests for A/B logic plus manual testing in qemu.

Change-Id: I69358a9a845de89d04f84dd58c2bd803ff522e30

15 files changed, 1681 insertions, 271 deletions

diff --git a/brillo_uefi_x86_64/boot_loader/Android.mk b/brillo_uefi_x86_64/boot_loader/Android.mk
index 36399fe..58f4259 100644
--- a/brillo_uefi_x86_64/boot_loader/Android.mk
+++ b/brillo_uefi_x86_64/boot_loader/Android.mk
@@ -47,8 +47,8 @@ LOCAL_CFLAGS := $(bub_common_cflags) -fno-stack-protector -DBUB_ENABLE_DEBUG -DB
 LOCAL_LDFLAGS := $(bub_common_ldflags)
 LOCAL_C_INCLUDES :=
 LOCAL_SRC_FILES := \
-    bub_sysdeps_posix.c \
     bub_ab_flow.c \
+    bub_sysdeps_posix.c \
     bub_util.c
 include $(BUILD_HOST_STATIC_LIBRARY)
 
@@ -61,8 +61,9 @@ LOCAL_CFLAGS := $(bub_common_cflags) -DBUB_ENABLE_DEBUG -DBUB_COMPILATION
 LOCAL_CPPFLAGS := $(bub_common_cppflags)
 LOCAL_LDFLAGS := $(bub_common_ldflags)
 LOCAL_C_INCLUDES := \
-    $(LOCAL_PATH)/bub_sysdeps.h \
     $(LOCAL_PATH)/bub_ab_flow.h \
+    $(LOCAL_PATH)/bub_image_util.h \
+    $(LOCAL_PATH)/bub_sysdeps.h \
     $(LOCAL_PATH)/bub_util.h \
     external/gtest/include
 LOCAL_STATIC_LIBRARIES := \
@@ -72,6 +73,7 @@ LOCAL_STATIC_LIBRARIES := \
 LOCAL_SHARED_LIBRARIES := \
     libchrome
 LOCAL_SRC_FILES := \
-    bub_ab_flow_unittest.cc
+    bub_ab_flow_unittest.cc \
+    bub_image_util.cc
 LOCAL_LDLIBS_linux := -lrt
 include $(BUILD_HOST_NATIVE_TEST)
diff --git a/brillo_uefi_x86_64/boot_loader/Makefile b/brillo_uefi_x86_64/boot_loader/Makefile
index 1d134e5..8988c2d 100644
--- a/brillo_uefi_x86_64/boot_loader/Makefile
+++ b/brillo_uefi_x86_64/boot_loader/Makefile
@@ -23,7 +23,12 @@ EFI_CC          = gcc
 EFI_LD          = ld
 EFI_OBJCOPY     = objcopy
 
-EFI_SRC_FILES   = bub_main.c bub_boot_kernel.c bub_sysdeps_uefi.c bub_util.c
+EFI_SRC_FILES   = bub_ab_flow.c \
+                  bub_boot_kernel.c \
+                  bub_main.c \
+                  bub_ops_uefi.c \
+                  bub_sysdeps_uefi.c \
+                  bub_util.c
 EFI_OBJ_FILES   = $(patsubst %.c,%.o,$(EFI_SRC_FILES))
 EFI_TARGET      = brillo_boot_loader.efi
 EFI_SHARED_OBJ  = $(patsubst %.efi,%.so,$(EFI_TARGET))
diff --git a/brillo_uefi_x86_64/boot_loader/brillo_boot_loader.efi b/brillo_uefi_x86_64/boot_loader/brillo_boot_loader.efi
index 6ce6a7e..7c0f7e8 100755
--- a/brillo_uefi_x86_64/boot_loader/brillo_boot_loader.efi
+++ b/brillo_uefi_x86_64/boot_loader/brillo_boot_loader.efi
diff --git a/brillo_uefi_x86_64/boot_loader/bub_ab_flow.c b/brillo_uefi_x86_64/boot_loader/bub_ab_flow.c
index a3fbaf8..799b0c2 100644
--- a/brillo_uefi_x86_64/boot_loader/bub_ab_flow.c
+++ b/brillo_uefi_x86_64/boot_loader/bub_ab_flow.c
@@ -15,7 +15,199 @@
  */
 
 #include "bub_ab_flow.h"
+#include "bub_util.h"
 
- BubAbFlowResult bub_ab_flow(BubOps* ops, char** out_selected_suffix) {
- 	return BUB_AB_FLOW_RESULT_OK;
- }
\ No newline at end of file
+static const uint8_t magic[] = BUB_BOOT_CTRL_MAGIC;
+static const char* bub_suffixes[2] = {"_a", "_b"};
+
+static int normalize_slot(BubSlotData *slot) {
+  if (slot->priority > 0) {
+    if (slot->tries_remaining == 0 && !slot->successful_boot) {
+      bub_memset(slot, 0, sizeof(BubSlotData));
+      return 1;
+    }
+    if (slot->tries_remaining > 0 && slot->successful_boot) {
+      slot->successful_boot = 0;
+      return 1;
+    }
+  } else if (slot->tries_remaining != 0 || slot->successful_boot) {
+    bub_memset(slot, 0, sizeof(BubSlotData));
+    return 1;
+  }
+
+  return 0;
+}
+
+static int slot_is_bootable(const BubSlotData *slot) {
+  return (slot->successful_boot || slot->tries_remaining > 0);
+}
+
+static int reset_metadata(BubOps* ops) {
+  BubAbData metadata;
+
+  bub_memset(&metadata, 0, sizeof(BubAbData));
+  bub_memcpy(&metadata.magic, magic, sizeof(metadata.magic));
+  metadata.major_version = BUB_MAJOR_VERSION;
+  metadata.minor_version = BUB_MINOR_VERSION;
+  metadata.slots[0].priority = 15;
+  metadata.slots[0].tries_remaining = 7;
+  metadata.slots[0].successful_boot = 0;
+  metadata.slots[1].priority = 15;
+  metadata.slots[1].tries_remaining = 7;
+  metadata.slots[1].successful_boot = 0;
+  return bub_write_ab_data_to_misc(ops, &metadata);
+}
+
+BubAbFlowResult bub_ab_flow(BubOps* ops,
+                            char* out_selected_suffix,
+                            size_t suffix_num_bytes) {
+  bub_assert(out_selected_suffix != NULL);
+  bub_assert(suffix_num_bytes >= 3);
+
+  BubAbFlowResult ab_err;
+  BubAbData ab_ctl;
+  int target_slot_index_to_boot = -1;
+  int new_metadata = 0;
+
+  // No selection has been made yet.
+  bub_memset(out_selected_suffix, 0, 3);
+
+  ab_err = bub_read_ab_data_from_misc(ops, &ab_ctl);
+  if (ab_err != BUB_AB_FLOW_RESULT_OK) {
+    if (ab_err == BUB_AB_FLOW_ERROR_INVALID_METADATA) {
+      bub_warning("Reseting metadata.\n");
+      if (!reset_metadata(ops)) {
+        bub_warning("Unable to reset metadata.\n");
+      }
+    }
+    return ab_err;
+  }
+
+  // Ensure only proper slot states exist.
+  if (normalize_slot(&ab_ctl.slots[0])) {
+    bub_warning("State of slot A was normalized.\n");
+    new_metadata = 1;
+  }
+  if (normalize_slot(&ab_ctl.slots[1])) {
+    bub_warning("State of slot B was normalized.\n");
+    new_metadata = 1;
+  }
+
+  if (slot_is_bootable(&ab_ctl.slots[0]) && slot_is_bootable(&ab_ctl.slots[1]))
+  {
+    if (ab_ctl.slots[1].priority > ab_ctl.slots[0].priority)
+      target_slot_index_to_boot = 1;
+    else
+      target_slot_index_to_boot = 0;
+  } else if (slot_is_bootable(&ab_ctl.slots[0])) {
+    // Choose slot A.
+    target_slot_index_to_boot = 0;
+  } else if (slot_is_bootable(&ab_ctl.slots[1])) {
+    // Choose slot B.
+    target_slot_index_to_boot = 1;
+  } else {
+    if (new_metadata && !bub_write_ab_data_to_misc(ops, &ab_ctl))
+      return BUB_AB_FLOW_ERROR_WRITE_METADATA;
+    // If neither was chosen, there are no valid slots.
+    bub_warning("No valid slot found.\n");
+    return BUB_AB_FLOW_ERROR_NO_VALID_SLOTS;
+  }
+
+  // Found usable slot to attempt boot on. Decrement tries remaining and write
+  // out the new AB metadata if the slots is in "Updated" state.
+  if (ab_ctl.slots[target_slot_index_to_boot].tries_remaining > 0) {
+    ab_ctl.slots[target_slot_index_to_boot].tries_remaining--;
+    new_metadata = 1;
+  }
+
+  if (new_metadata)
+    if (!bub_write_ab_data_to_misc(ops, &ab_ctl))
+      return BUB_AB_FLOW_ERROR_WRITE_METADATA;
+
+  // Write selected suffix to caller's pointer.
+  bub_memcpy(out_selected_suffix, bub_suffixes[target_slot_index_to_boot], 3);
+
+  return BUB_AB_FLOW_RESULT_OK;
+}
+
+int bub_ab_mark_as_invalid(BubOps* ops, const char* invalid_suffix) {
+  bub_assert(invalid_suffix != NULL);
+  bub_assert(bub_strlen(invalid_suffix) >= 2);
+
+  BubAbData ab_ctl;
+  BubAbFlowResult ab_err;
+  unsigned int i;
+
+  ab_err = bub_read_ab_data_from_misc(ops, &ab_ctl);
+  if (ab_err != BUB_AB_FLOW_RESULT_OK)
+    return ab_err;
+
+  // Find suffix in small list.
+  for (i = 0; i < 2; ++i)
+  {
+    if (!bub_memcmp(bub_suffixes[i], invalid_suffix, 3)) {
+      // Found the corresponding index to invalid_suffix. Invalidate the slot
+      // and write out.
+      bub_memset(&ab_ctl.slots[i], 0, sizeof(BubSlotData));
+      return bub_write_ab_data_to_misc(ops, &ab_ctl);
+    }
+  }
+
+  bub_warning("Could not find requested slot.\n");
+  return 0;
+}
+
+BubAbFlowResult bub_read_ab_data_from_misc(BubOps* ops, BubAbData* data) {
+  BubIOResult io_result;
+  size_t num_bytes_read;
+  uint32_t crc;
+
+  io_result = ops->read_from_partition(ops, "misc", data, 0,
+                                    sizeof(BubAbData), &num_bytes_read);
+  if (io_result != BUB_IO_RESULT_OK || num_bytes_read != sizeof(BubAbData)) {
+    bub_warning("Could not read metadata from misc partition.\n");
+    return BUB_AB_FLOW_ERROR_READ_METADATA;
+  }
+
+  if(bub_memcmp(data->magic, magic, 4) != 0) {
+    bub_warning("AB metadata magic did not match.\n");
+    return BUB_AB_FLOW_ERROR_INVALID_METADATA;
+  }
+
+  crc = bub_be32toh(data->crc32);
+  data->crc32 = 0;
+  if (crc != bub_crc32((uint8_t *)data, sizeof(BubAbData))) {
+    bub_warning("AB metadata crc is invalid.\n");
+    return BUB_AB_FLOW_ERROR_INVALID_METADATA;
+  }
+
+  // Assign host byte order to necessary variables needed here.
+  data->crc32 = crc;
+
+  return BUB_AB_FLOW_RESULT_OK;
+}
+
+int bub_write_ab_data_to_misc(BubOps* ops, const BubAbData* data) {
+  BubIOResult io_result;
+  BubAbData metadata;
+
+  bub_memcpy(&metadata, data, sizeof(BubAbData));
+
+  // Assign big endian order to necessary variables here.
+
+  // Calculate crc assign back to crc field, maintaining big endianness.
+  metadata.crc32 = 0;
+  metadata.crc32 = bub_be32toh(bub_crc32((uint8_t*)&metadata,
+                                          sizeof(BubAbData)));
+  io_result = ops->write_to_partition(ops,
+                                      "misc",
+                                      &metadata,
+                                      0,
+                                      sizeof(BubAbData));
+  if (io_result != BUB_IO_RESULT_OK) {
+    bub_warning("Could not write to misc partition.\n");
+    return 0;
+  }
+
+  return 1;
+}
\ No newline at end of file
diff --git a/brillo_uefi_x86_64/boot_loader/bub_ab_flow.h b/brillo_uefi_x86_64/boot_loader/bub_ab_flow.h
index ab297d0..9cfd94c 100644
--- a/brillo_uefi_x86_64/boot_loader/bub_ab_flow.h
+++ b/brillo_uefi_x86_64/boot_loader/bub_ab_flow.h
@@ -17,13 +17,123 @@
 #ifndef BUB_AB_FLOW_H_
 #define BUB_AB_FLOW_H_
 
+#ifdef __cplusplus
+extern "C" {
+#endif
+
 #include "bub_ops.h"
 
+/* Magic for the Brillo Uefi metadata header */
+#define BUB_BOOT_CTRL_MAGIC {'B', 'U', 'E', 'F'}
+
+/* The current major and minor versions used for AB metadata structs */
+#define BUB_MAJOR_VERSION 1
+#define BUB_MINOR_VERSION 0
+
+#define BUB_BLOCK_SIZE 512
+#define BUB_AB_DATA_SIZE 64
+#define BUB_SUFFIX_SIZE 3
+
 typedef enum {
-	BUB_AB_FLOW_RESULT_OK,
-	BUB_AB_FLOW_RESULT_ERROR
+  BUB_AB_FLOW_RESULT_OK,
+  BUB_AB_FLOW_ERROR_INPUT,
+  BUB_AB_FLOW_ERROR_INVALID_METADATA,
+  BUB_AB_FLOW_ERROR_NO_VALID_SLOTS,
+  BUB_AB_FLOW_ERROR_READ_METADATA,
+  BUB_AB_FLOW_ERROR_WRITE_METADATA
 } BubAbFlowResult;
 
-BubAbFlowResult bub_ab_flow(BubOps* ops, char** out_selected_suffix);
+typedef struct {
+    // Slot priority with 15 meaning highest priority, 1 lowest
+    // priority and 0 the slot is unbootable.
+    uint8_t priority: 4;
+    // Number of times left attempting to boot this slot.
+    uint8_t tries_remaining: 3;
+    // 1 if this slot has booted successfully, 0 otherwise.
+    uint8_t successful_boot: 1;
+    // Reserved for further use.
+    uint8_t reserved[7];
+} __attribute__((__packed__)) BubSlotData;
+
+/* Bootloader Control BubAbData
+ *
+ * This struct is used to manage A/B metadata. Big-endian order is used.
+ */
+typedef struct {
+    // Bootloader Control AB magic number (see BUB_BOOT_CTRL_MAGIC).
+    uint8_t magic[4];
+    // Major version number.
+    uint8_t major_version;
+    // Minor version number.
+    uint8_t minor_version;
+    // Reserved for slot alignment.
+    uint8_t reserved1[2];
+    // Per-slot information.
+    BubSlotData slots[2];
+    // Reserved for further use.
+    uint8_t reserved2[36];
+    // CRC32 of all 60 bytes preceding this field.
+    uint32_t crc32;
+} __attribute__((__packed__)) BubAbData;
+
+
+#if defined(__GNUC__) && __GNUC__ >= 4 && __GNUC_MINOR__ >= 6
+ _Static_assert(sizeof(BubAbData) == BUB_AB_DATA_SIZE,
+                "BubAbData has wrong size!");
+#endif
+
+/* A/B flow logic for Brillo booting. Reads A/B metadata from the 'misc'
+ * partition and validates it. Chooses a bootable slot based on its state. Upon
+ * finding the bootable slot, its "tries remaining" attribute is decremented
+ * and thus A/B metadata may be modified in the "misc" partition. The suffix of
+ * a bootable slot, including a terminating NUL-byte, will be written in
+ * |out_selected_suffix| on success. Caller must specify the size of the
+ * |out_selected_suffix| in |suffix_num_bytes| which must be at least 3
+ * otherwise aborting the program. If BUB_AB_FLOW_INVALID_AB_METADATA is
+ * returned, metadata on disk will be reset to an 'updating' state where both
+ * slots have tries remaining to reattempt booting.
+ *
+ * @return: BUB_AB_FLOW_RESULT_OK on success. BUB_AB_FLOW_INVALID_AB_METADATA
+ *          if AB metadata is invalid. BUB_AB_FLOW_RESULT_ERROR if no available
+ *          memory for allocationor no valid slot found.
+ *
+ */
+BubAbFlowResult bub_ab_flow(BubOps* ops,
+                            char* out_selected_suffix,
+                            size_t suffix_num_bytes);
+
+/* Marks a boot slot with |invalid_suffix| invalid by assigning zero to its
+ * priority, tries_remaining, and successful_boot member variables. Caller
+ * must pass |invalid_suffix| as a NUL_terminated string with length 2.
+ *
+ * @return: 0 on success. 1 on failure.
+ */
+int bub_ab_mark_as_invalid(BubOps* ops, const char *invalid_suffix);
+
+/* Helper function to read and check validity of AB metadata using |ops|
+ * read_from_partition method.  Will read from "misc" partition and assign
+ * fields to |data| in host byte order. Checks magic field matches expected
+ * value and calculates crc32.
+ *
+ * @return: BUB_AB_FLOW_ERROR_READ_METADATA on i/o error.
+ *          BUB_AB_FLOW_ERROR_INVALID_METADATA if magic or crc are not
+ *          correct. BUB_AB_FLOW_RESULT_OK on success.
+ *
+ */
+BubAbFlowResult bub_read_ab_data_from_misc(BubOps* ops, BubAbData* data);
+
+/* Helper function to write AB metadata using |ops| write_to_partition method.
+ * Will write to "misc" partition and assign fields to |data| in big-endian byte
+ * order. Calculates crc32 as well.
+ *
+ * @return: 0 on i/o error, 1 on success.
+ *
+ */
+int bub_write_ab_data_to_misc(BubOps* ops, const BubAbData* data);
+
+
+#ifdef __cplusplus
+}
+#endif
 
 #endif /* BUB_AB_FLOW_H_ */
\ No newline at end of file
diff --git a/brillo_uefi_x86_64/boot_loader/bub_ab_flow_unittest.cc b/brillo_uefi_x86_64/boot_loader/bub_ab_flow_unittest.cc
index a7885a1..be57900 100644
--- a/brillo_uefi_x86_64/boot_loader/bub_ab_flow_unittest.cc
+++ b/brillo_uefi_x86_64/boot_loader/bub_ab_flow_unittest.cc
@@ -14,10 +14,33 @@
  * limitations under the License.
  */
 
-#include <gtest/gtest.h>
-#include "bub_sysdeps.h"
-#include "bub_ab_flow.h"
-#include "bub_util.h"
+#include "bub_image_util.h"
+
+#define ab_init(magic,                                                        \
+                a_priority, a_tries_remaining, a_successful_boot,             \
+                b_priority, b_tries_remaining, b_successful_boot)             \
+  do {                                                                        \
+    BubAbData init;                                                           \
+    ops_.write_ab_metadata(&init, (uint8_t[4])magic,                          \
+                           a_priority, a_tries_remaining, a_successful_boot,  \
+                           b_priority, b_tries_remaining, b_successful_boot); \
+    GenerateMiscImage(&init);                                                 \
+  } while(0)
+
+#define test_ab_flow(a_priority, a_tries_remaining, a_successful_boot,        \
+                     b_priority, b_tries_remaining, b_successful_boot,        \
+                     expected_result, expected_suffix)                        \
+  do {                                                                        \
+    char suffix[BUB_SUFFIX_SIZE] = {0};                                       \
+    BubAbData ab_result;                                                      \
+    ops_.write_ab_metadata(&ab_result, (uint8_t[4])BUB_BOOT_CTRL_MAGIC,       \
+                           a_priority, a_tries_remaining, a_successful_boot,  \
+                           b_priority, b_tries_remaining, b_successful_boot); \
+    EXPECT_EQ(expected_result,                                                \
+              bub_ab_flow((BubOps*)ops_.bub_ops(), suffix, BUB_SUFFIX_SIZE)); \
+    EXPECT_EQ(0, bub_memcmp(expected_suffix, suffix, BUB_SUFFIX_SIZE));       \
+    EXPECT_EQ(0, CompareMiscImage(ab_result));                                \
+  } while (0)
 
 static int converted_utf8_ucs2(const char* data,
                                const char* raw_bytes,
@@ -52,3 +75,404 @@ TEST(UtilTest, Utf8toUcs2) {
   // UTF-8 5 bytes encoding case.
   EXPECT_NE(0, converted_utf8_ucs2("👦", "\x66\xF4", 5));
 }
+
+TEST_F(AbTest, NoValidSlots) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 0, 0, 0, 0, 0, 0);
+
+  test_ab_flow(
+    0, 0, 0, 0, 0, 0,                   // Expected A/B state.
+    BUB_AB_FLOW_ERROR_NO_VALID_SLOTS,   // Expected A/B result.
+    "\0\0");                            // Expected A/B suffix.
+}
+
+TEST_F(AbTest, InvalidMetadataMagicInvalidSlots) {
+  char suffix[BUB_SUFFIX_SIZE] = {0};
+  BubAbData ab_init;
+  BubAbData ab_result;
+
+  ops_.write_ab_metadata(&ab_init, (uint8_t[4]){'N','O','P','E'},
+                         0, 0, 0, 0, 0, 0);
+  ops_.write_ab_metadata(&ab_result, (uint8_t[4])BUB_BOOT_CTRL_MAGIC,
+                         15, 7, 0, 15, 7, 0);
+  GenerateMiscImage(&ab_init);
+
+  // Invalid metadata should be found by ab flow at this point. It should reset
+  // each slot to an 'updating' state on disk so that we may reattempt boot.
+  EXPECT_EQ(BUB_AB_FLOW_ERROR_INVALID_METADATA,
+            bub_ab_flow((BubOps*)ops_.bub_ops(), suffix, BUB_SUFFIX_SIZE));
+  EXPECT_EQ(0, bub_memcmp((char[BUB_SUFFIX_SIZE]){0,0,0},
+                          suffix,
+                          BUB_SUFFIX_SIZE));
+  EXPECT_EQ(0, CompareMiscImage(ab_result));
+
+  // Run again to check slot A is selected with one less try remaining.
+  test_ab_flow(
+    15, 6, 0, 15, 7, 0,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");
+}
+
+TEST_F(AbTest, InvalidMetadataMagicValidSlots) {
+  char suffix[BUB_SUFFIX_SIZE] = {0};
+  BubAbData ab_init;
+  BubAbData ab_result;
+
+  ops_.write_ab_metadata(&ab_init, (uint8_t[4]){'N','O','P','E'},
+                         15, 0, 1, 14, 0, 1);
+  ops_.write_ab_metadata(&ab_result, (uint8_t[4])BUB_BOOT_CTRL_MAGIC,
+                         15, 7, 0, 15, 7, 0);
+  GenerateMiscImage(&ab_init);
+
+  // Invalid metadata should be found by ab flow at this point. It should reset
+  // each slot to an 'updating' state on disk so that we may reattempt boot.
+  EXPECT_EQ(BUB_AB_FLOW_ERROR_INVALID_METADATA,
+            bub_ab_flow((BubOps*)ops_.bub_ops(), suffix, BUB_SUFFIX_SIZE));
+  EXPECT_EQ(0, bub_memcmp((char[BUB_SUFFIX_SIZE]){0,0,0},
+                          suffix,
+                          BUB_SUFFIX_SIZE));
+  EXPECT_EQ(0, CompareMiscImage(ab_result));
+
+  // Run again to check slot A is selected with one less try remaining.
+  test_ab_flow(
+    15, 6, 0, 15, 7, 0,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+}
+
+TEST_F(AbTest, SingleSuccessfulSlot) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 14, 0, 1, 0, 0, 0);
+
+  test_ab_flow(
+    14, 0, 1, 0, 0, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+}
+
+TEST_F(AbTest, SingleTryingSlot) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 14, 3, 0, 0, 0, 0);
+
+  test_ab_flow(
+    14, 2, 0, 0, 0, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+}
+
+TEST_F(AbTest, TwoValidSlotsA) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 15, 0, 1, 14, 0, 1);
+
+  test_ab_flow(
+    15, 0, 1, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+}
+
+TEST_F(AbTest, TwoValidSlotsB) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 14, 0, 1, 15, 0, 1);
+
+  test_ab_flow(
+    14, 0, 1, 15, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_b");                 // Expected A/B suffix.
+}
+
+TEST_F(AbTest, TryingFallback) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 15, 7, 0, 14, 0, 1);
+
+  // Decremented our expected tries_remaining for slot a as we run ab_flow
+  test_ab_flow(
+    15, 6, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 5, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 4, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 3, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 2, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 1, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 0, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  // Should revert to slot b.
+  test_ab_flow(
+    0, 0, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_b");                 // Expected A/B suffix.
+}
+
+TEST_F(AbTest, TryingNoFallbackRecovery) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 15, 7, 0, 0, 0, 0);
+
+  // Decremented our expected tries_remaining for slot a as we run ab_flow
+  test_ab_flow(
+    15, 6, 0, 0, 0, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 5, 0, 0, 0, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 4, 0, 0, 0, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 3, 0, 0, 0, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 2, 0, 0, 0, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 1, 0, 0, 0, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 0, 0, 0, 0, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    0, 0, 0, 0, 0, 0,                // Expected A/B state.
+    BUB_AB_FLOW_ERROR_NO_VALID_SLOTS, // Expected A/B result.
+    "\0\0");                          // Expected A/B suffix.
+}
+
+TEST_F(AbTest, SingleTryingSuccess) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 15, 7, 0, 14, 0, 1);
+
+  // Decremented our expected tries_remaining for slot a. Boot was a success
+  // on our 6th try and we reboot 2 more times to make sure we stick to the
+  // same slot.
+
+  test_ab_flow(
+    15, 6, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 5, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 4, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 3, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 2, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 1, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 0, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  // Boot Control HAL should do this.
+  ab_init(BUB_BOOT_CTRL_MAGIC, 15, 0, 1, 14, 0, 1);
+
+  test_ab_flow(
+    15, 0, 1, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  // Should not have changed still.
+  test_ab_flow(
+    15, 0, 1, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+}
+
+TEST_F(AbTest, TwoTryingRecovery) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 15, 7, 0, 14, 7, 0);
+
+  // Decremented our expected tries_remaining for slot a.
+  test_ab_flow(
+    15, 6, 0, 14, 7, 0,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 5, 0, 14, 7, 0,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 4, 0, 14, 7, 0,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 3, 0, 14, 7, 0,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 2, 0, 14, 7, 0,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 1, 0, 14, 7, 0,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    15, 0, 0, 14, 7, 0,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  // At this point a should have run out of tries, so we expect the other
+  // updating slot to be chosen.
+
+  // Decremented our expected tries_remaining for slot b.
+  test_ab_flow(
+    0, 0, 0, 14, 6, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_b");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    0, 0, 0, 14, 5, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_b");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    0, 0, 0, 14, 4, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_b");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    0, 0, 0, 14, 3, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_b");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    0, 0, 0, 14, 2, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_b");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    0, 0, 0, 14, 1, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_b");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    0, 0, 0, 14, 0, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_b");                 // Expected A/B suffix.
+
+  test_ab_flow(
+    0, 0, 0, 0, 0, 0,                 // Expected A/B state.
+    BUB_AB_FLOW_ERROR_NO_VALID_SLOTS, // Expected A/B result.
+    "\0\0");                          // Expected A/B suffix.
+}
+
+TEST_F(AbTest, MarkedInvalidFallback) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 15, 0, 1, 14, 0, 1);
+
+  // Initially selects slot a.
+  test_ab_flow(
+    15, 0, 1, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+
+  // Invalidate slot a. We expect this slot to be all zero values with slot b
+  // unchanged.
+  EXPECT_TRUE(bub_ab_mark_as_invalid((BubOps*)ops_.bub_ops(), "_a"));
+
+  // Should select slot b now.
+  test_ab_flow(
+     0, 0, 0, 14, 0, 1,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_b");                 // Expected A/B suffix.
+}
+
+TEST_F(AbTest, ValidAndInvalidHigherPriority) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 14, 0, 1, 15, 0, 0);
+
+  // Normalizes and selects slot a.
+  test_ab_flow(
+    14, 0, 1, 0, 0, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_a");                 // Expected A/B suffix.
+}
+
+TEST_F(AbTest, ValidAndUpdatingBadSuccessfulBoot) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 14, 0, 1, 15, 7, 1);
+
+  // Normalizes and selects slot b.
+  test_ab_flow(
+    14, 0, 1, 15, 6, 0,    // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_b");                 // Expected A/B suffix.
+}
+
+TEST_F(AbTest, InvalidBadTriesRemainingAndValid) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 0, 7, 0, 14, 0, 1);
+
+  // Normalizes and selects slot b.
+  test_ab_flow(
+    0, 0, 0, 14, 0, 1,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_b");                 // Expected A/B suffix.
+}
+
+TEST_F(AbTest, InvalidBadSuccessfulBootandValid) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 0, 0, 1, 14, 0, 1);
+
+  // Normalizes and selects slot b.
+  test_ab_flow(
+    0, 0, 0, 14, 0, 1,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_b");                 // Expected A/B suffix.
+}
+
+TEST_F(AbTest, InvalidTriesBootAndUpdatingBadSuccessfulBoot) {
+  ab_init(BUB_BOOT_CTRL_MAGIC, 0, 7, 1, 15, 7, 1);
+
+  // Normalizes and selects slot b.
+  test_ab_flow(
+    0, 0, 0, 15, 6, 0,     // Expected A/B state.
+    BUB_AB_FLOW_RESULT_OK, // Expected A/B result.
+    "_b");                 // Expected A/B suffix.
+}
\ No newline at end of file
diff --git a/brillo_uefi_x86_64/boot_loader/bub_boot_kernel.c b/brillo_uefi_x86_64/boot_loader/bub_boot_kernel.c
index 0befed7..59a63d8 100644
--- a/brillo_uefi_x86_64/boot_loader/bub_boot_kernel.c
+++ b/brillo_uefi_x86_64/boot_loader/bub_boot_kernel.c
@@ -32,50 +32,13 @@
 #include "bub_boot_kernel.h"
 #include "bub_sysdeps.h"
 
-// GPT related constants
-#define GPT_REVISION 0x00010000
-#define GPT_MAGIC "EFI PART"
-#define GPT_MIN_SIZE 92
-#define GPT_ENTRIES_LBA 2
-#define BLOCK_SIZE 512
-#define ENTRIES_PER_BLOCK 4
-#define ENTRY_NAME_LEN 36
-#define MAX_GPT_ENTRIES 128
-
-typedef struct {
-  UINT8   signature[8];
-  UINT32  revision;
-  UINT32  header_size;
-  UINT32  header_crc32;
-  UINT32  reserved;
-  UINT64  header_lba;
-  UINT64  alternate_header_lba;
-  UINT64  first_usable_lba;
-  UINT64  last_usable_lba;
-  UINT8   disk_guid[16];
-  UINT64  entry_lba;
-  UINT32  entry_count;
-  UINT32  entry_size;
-  UINT32  entry_crc32;
-  UINT8   reserved2[420];
-} GPTHeader;
-
-typedef struct {
-  UINT8   type_GUID[16];
-  UINT8   unique_GUID[16];
-  UINT64  first_lba;
-  UINT64  last_lba;
-  UINT64  flags;
-  CHAR16  name[ENTRY_NAME_LEN];  // UTF-16LE encoding, NULL terminated
-} GPTEntry;
-
 /*
  * Note: The below header definitions are taken from
  *       system/core/mkbootimg/bootimg.h
  */
 typedef struct boot_img_hdr boot_img_hdr;
 
-#define BOOT_MAGIC "ANDROID!"
+#define BOOT_MAGIC {'A','N','D','R','O','I','D','!'}
 #define BOOT_MAGIC_SIZE 8
 #define BOOT_NAME_SIZE 16
 #define BOOT_ARGS_SIZE 512
@@ -143,15 +106,6 @@ struct boot_img_hdr
 **    else: jump to kernel_addr
 */
 
-#define IMG_SIZE(entry, block) \
-  (entry->last_lba - entry->first_lba) * block->Media->BlockSize
-
-#define SIZE_BLOCK_ALIGN(bytes, block_size) \
-  ((bytes + block_size - 1) / block_size) * block_size
-
-#define OFFSET_BLOCK_ALIGN(bytes, block_size) \
-  (bytes / block_size) * block_size
-
 /* uefi_call_wrapper's second arguments is the number of argumets for the
  * called function
  */
@@ -163,21 +117,24 @@ struct boot_img_hdr
 
 /* Helper method to get the parent path to the current |walker| path given the
  * initial path, |init|. Resulting path is stored in |next|.  Caller is
- * responsible for freeing |next|.
+ * responsible for freeing |next|. Stores allocated bytes for |next| in
+ * |out_bytes|.
  *
  * @return EFI_STATUS Standard UEFI error code, EFI_SUCCESS on success.
  */
 static EFI_STATUS walk_path(IN EFI_DEVICE_PATH *init,
                             IN EFI_DEVICE_PATH *walker,
-                            OUT EFI_DEVICE_PATH **next) {
-  EFI_STATUS err;
-
+                            OUT EFI_DEVICE_PATH **next,
+                            OUT UINTN* out_bytes) {
   // Number of bytes from initial path to current walker.
   UINTN walker_bytes = (UINT8 *)NextDevicePathNode(walker) - (UINT8 *)init;
+  *out_bytes = sizeof(EFI_DEVICE_PATH) + walker_bytes;
 
-  *next = (EFI_DEVICE_PATH*)bub_malloc_(sizeof(EFI_DEVICE_PATH) + walker_bytes);
-  if (*next == NULL)
+  *next = (EFI_DEVICE_PATH*)bub_malloc_(*out_bytes);
+  if (*next == NULL) {
+    *out_bytes = 0;
     return EFI_NOT_FOUND;
+  }
 
   // Copy in the previous paths.
   bub_memcpy((*next), init, walker_bytes);
@@ -203,12 +160,12 @@ static EFI_STATUS validate_gpt(const IN GPTHeader *gpth) {
     bub_warning("GPT header too small.\n");
     return EFI_NOT_FOUND;
   }
-  if (gpth->header_size > BLOCK_SIZE) {
+  if (gpth->header_size > BUB_BLOCK_SIZE) {
     bub_warning("GPT header too big.\n");
     return EFI_NOT_FOUND;
   }
 
-  GPTHeader gpth_tmp = {0};
+  GPTHeader gpth_tmp = {{0}};
   bub_memcpy(&gpth_tmp, gpth, sizeof(GPTHeader));
   UINT32 gpt_header_crc = gpth_tmp.header_crc32;
   gpth_tmp.header_crc32 = 0;
@@ -232,109 +189,6 @@ static EFI_STATUS validate_gpt(const IN GPTHeader *gpth) {
   return EFI_SUCCESS;
 }
 
-/* Looks through |block_io| to search for a GPT entry named |partition_name|, a
- * NUULL-terminated string. Allocates a GPTEntry struct for |entry_buf|. Caller
- * is responsible for freeing |entry_buf|.
- *
- * @return EFI_NOT_FOUND on error, EFI SUCCESS on success.
- */
-static EFI_STATUS PartitionEntryByName(IN EFI_BLOCK_IO *block_io,
-                                       const char* partition_name,
-                                       GPTEntry** entry_buf) {
-  EFI_STATUS err;
-  GPTHeader* gpt_header = NULL;
-  GPTEntry all_gpt_entries[MAX_GPT_ENTRIES];
-  CHAR16* partition_name_ucs2 = NULL;
-  UINTN partition_name_bytes;
-
-  gpt_header = (GPTHeader*)bub_malloc_(sizeof(GPTHeader));
-  if (gpt_header == NULL) {
-    bub_warning("Could not allocate for GPT header\n");
-    return EFI_NOT_FOUND;
-  }
-
-  *entry_buf = (GPTEntry*)bub_malloc_(sizeof(GPTEntry) * ENTRIES_PER_BLOCK);
-  if (entry_buf == NULL) {
-    bub_warning("Could not allocate for partition entry\n");
-    bub_free(gpt_header);
-    return EFI_NOT_FOUND;
-  }
-
-  err = uefi_call_wrapper(block_io->ReadBlocks, NUM_ARGS_READ_BLOCKS,
-                          block_io,
-                          block_io->Media->MediaId,
-                          1,
-                          sizeof(GPTHeader),
-                          gpt_header);
-  if (EFI_ERROR(err)) {
-    bub_warning("Could not ReadBlocks for gpt header\n");
-    bub_free(gpt_header);
-    return EFI_NOT_FOUND;
-  }
-
-  partition_name_bytes = bub_strlen(partition_name) + 1;
-  partition_name_ucs2 =
-    bub_calloc(sizeof(CHAR16) * partition_name_bytes);
-  if (partition_name_ucs2 == NULL) {
-    bub_warning ("Could not allocate for ucs2 partition name\n");
-    bub_free(gpt_header);
-    return EFI_NOT_FOUND;
-  }
-  if (utf8_to_ucs2(partition_name,
-                   partition_name_bytes,
-                   partition_name_ucs2,
-                   sizeof(CHAR16) * partition_name_bytes)) {
-    bub_warning("Could not convert partition name to UCS-2\n");
-    bub_free(gpt_header);
-    bub_free(partition_name_ucs2);
-    return EFI_NOT_FOUND;
-  }
-
-#ifdef BUB_ENABLE_DEBUG
-  Print(L"\nENTRY: %d, BLOCKSIZE: %d, GPTEntry: %d\n",
-            gpt_header->entry_count,
-            block_io->Media->BlockSize,
-            sizeof(GPTEntry));
-#endif
-
-  // Block-aligned bytes for entries.
-  UINTN entries_num_bytes = block_io->Media->BlockSize *
-                            (MAX_GPT_ENTRIES / ENTRIES_PER_BLOCK);
-
-  err = uefi_call_wrapper(block_io->ReadBlocks, NUM_ARGS_READ_BLOCKS,
-                          block_io,
-                          block_io->Media->MediaId,
-                          GPT_ENTRIES_LBA,
-                          entries_num_bytes,
-                          &all_gpt_entries);
-  if (EFI_ERROR(err)) {
-    bub_warning("Could not ReadBlocks for GPT header\n");
-    bub_free(gpt_header);
-    bub_free(partition_name_ucs2);
-    return EFI_NOT_FOUND;
-  }
-
-  // Find matching partition name.
-  UINT8 i;
-  for (i = 0; i < gpt_header->entry_count; ++i)
-    if (!bub_memcmp(all_gpt_entries[i].name,
-                    partition_name_ucs2,
-                    sizeof(CHAR16) * bub_strlen(partition_name))) {
-#ifdef BUB_ENABLE_DEBUG
-      Print(L"Boot Partition Name is: %s\n", all_gpt_entries[i].name);
-      Print(L"Boot Partition LBA is: %d\n", all_gpt_entries[i].first_lba);
-#endif
-      bub_memcpy((*entry_buf), &all_gpt_entries[i], sizeof(GPTEntry));
-      bub_free(partition_name_ucs2);
-      bub_free(gpt_header);
-      return EFI_SUCCESS;
-    }
-
-  bub_free(partition_name_ucs2);
-  bub_free(gpt_header);
-  return EFI_NOT_FOUND;
-}
-
 /* Allocates a pool of memory in the EfiLoaderData region for the LoadOptions
  * member of |loaded_image|.  The LoadOptions member is needed by next boot
  * stage. In the case of Linux kernel images, the EFI_STUB is this stage. The
@@ -391,14 +245,18 @@ static EFI_STATUS LoadParameters(EFI_HANDLE image,
  *
  * @return EFI_STATUS EFI_NOT_FOUND on fail, EFI_SUCCESS otherwise.
  */
-static EFI_STATUS getDiskBlockIO(IN EFI_HANDLE* disk_handle,
+static EFI_STATUS getDiskBlockIo(IN EFI_HANDLE* block_handle,
                                  OUT EFI_BLOCK_IO** block_io,
-                                 OUT EFI_DEVICE_PATH** block_path) {
+                                 OUT EFI_DISK_IO** disk_io,
+                                 OUT EFI_DEVICE_PATH** io_path) {
   EFI_STATUS err;
+  EFI_HANDLE disk_handle;
+  UINTN path_bytes;
+  EFI_DEVICE_PATH *disk_path;
   EFI_DEVICE_PATH *walker_path;
   EFI_DEVICE_PATH *init_path;
   GPTHeader gpt_header = {{0}};
-  init_path = DevicePathFromHandle(disk_handle);
+  init_path = DevicePathFromHandle(block_handle);
 
 #ifdef BUB_ENABLE_DEBUG
   Print(L"Initial Device Path: %s\n", DevicePathToStr(init_path));
@@ -415,41 +273,66 @@ static EFI_STATUS getDiskBlockIO(IN EFI_HANDLE* disk_handle,
     Print(L"DevicePathType: %x\n", DevicePathType(walker_path));
 #endif
 
-    err = walk_path(init_path, walker_path, &(*block_path));
+    err = walk_path(init_path, walker_path, &(*io_path), &path_bytes);
     if (EFI_ERROR(err)) {
       bub_warning("Cannot walk device path.\n");
       return EFI_NOT_FOUND;
     }
 
 #ifdef BUB_ENABLE_DEBUG
-    Print(L"Walking Device Path   : %s\n", DevicePathToStr(*block_path));
+    Print(L"Walking Device Path: %s\n", DevicePathToStr(*io_path));
 #endif
+    disk_path = (EFI_DEVICE_PATH*)bub_malloc_(path_bytes);
+    bub_memcpy(disk_path, *io_path, path_bytes);
     err = uefi_call_wrapper(BS->LocateDevicePath, NUM_ARGS_LOCATE_DEVICE_PATH,
                             &BlockIoProtocol,
-                            &(*block_path),
-                            &disk_handle);
+                            &(*io_path),
+                            &block_handle);
     if (EFI_ERROR(err)) {
       bub_warning("LocateDevicePath, BLOCK_IO_PROTOCOL.\n");
-      bub_free(*block_path);
+      bub_free(*io_path);
+      bub_free(disk_path);
+      continue;
+    }
+    err = uefi_call_wrapper(BS->LocateDevicePath, NUM_ARGS_LOCATE_DEVICE_PATH,
+                            &DiskIoProtocol,
+                            &disk_path,
+                            &disk_handle);
+    if (EFI_ERROR(err)) {
+      bub_warning("LocateDevicePath, DISK_IO_PROTOCOL.\n");
+      bub_free(*io_path);
+      bub_free(disk_path);
       continue;
     }
 
-    // Handle Block i/o
-    // Attempt to get handle on device, must be BlockIo type.
+    // Handle Block and Disk i/o.
+    // Attempt to get handle on device, must be Block/Disk Io type.
     err = uefi_call_wrapper(BS->HandleProtocol, NUM_ARGS_HANDLE_PROTOCOL,
-                            disk_handle,
+                            block_handle,
                             &BlockIoProtocol,
                             (VOID **)&(*block_io));
     if (EFI_ERROR(err)) {
-      bub_warning("HandleProtocol, BLOCK_IO_PROTOCOL.\n");
-      bub_free(*block_path);
+      bub_warning("Cannot get handle on block device.\n");
+      bub_free(*io_path);
+      bub_free(disk_path);
+      continue;
+    }
+    err = uefi_call_wrapper(BS->HandleProtocol, NUM_ARGS_HANDLE_PROTOCOL,
+                            disk_handle,
+                            &DiskIoProtocol,
+                            (VOID **)&(*disk_io));
+    if (EFI_ERROR(err)) {
+      bub_warning("Cannot get handle on disk device.\n");
+      bub_free(*io_path);
+      bub_free(disk_path);
       continue;
     }
 
     if ((*block_io)->Media->LogicalPartition ||
         !(*block_io)->Media->MediaPresent) {
       bub_warning("Logical partion or No Media Present, continue...\n");
-      bub_free(*block_path);
+      bub_free(*io_path);
+      bub_free(disk_path);
       continue;
     }
 
@@ -462,18 +345,21 @@ static EFI_STATUS getDiskBlockIO(IN EFI_HANDLE* disk_handle,
 
     if (EFI_ERROR(err)) {
       bub_warning("ReadBlocks, Block Media error.\n");
-      bub_free(*block_path);
+      bub_free(*io_path);
+      bub_free(disk_path);
       continue;
     }
 
     err = validate_gpt(&gpt_header);
     if (EFI_ERROR(err)) {
       bub_warning("Invalid GPTHeader\n");
-      bub_free(*block_path);
+      bub_free(*io_path);
+      bub_free(disk_path);
       continue;
     }
 
 #ifdef BUB_ENABLE_DEBUG
+    Print(L"Walking Device Path3   : %s\n", DevicePathToStr(disk_path));
     Print(L"Validated GPT\n");
 #endif
     return EFI_SUCCESS;
@@ -483,26 +369,129 @@ static EFI_STATUS getDiskBlockIO(IN EFI_HANDLE* disk_handle,
   return EFI_NOT_FOUND;
 }
 
-int bub_boot_kernel(EFI_HANDLE app_image, const char* boot_partition_name) {
+EFI_STATUS bub_partition_entry_by_name(IN EFI_BLOCK_IO *block_io,
+                                       const char* partition_name,
+                                       GPTEntry** entry_buf) {
+  EFI_STATUS err;
+  GPTHeader* gpt_header = NULL;
+  GPTEntry all_gpt_entries[MAX_GPT_ENTRIES];
+  CHAR16* partition_name_ucs2 = NULL;
+  UINTN partition_name_bytes;
+
+  gpt_header = (GPTHeader*)bub_malloc_(sizeof(GPTHeader));
+  if (gpt_header == NULL) {
+    bub_warning("Could not allocate for GPT header\n");
+    return EFI_NOT_FOUND;
+  }
+
+  *entry_buf = (GPTEntry*)bub_malloc_(sizeof(GPTEntry) * ENTRIES_PER_BLOCK);
+  if (entry_buf == NULL) {
+    bub_warning("Could not allocate for partition entry\n");
+    bub_free(gpt_header);
+    return EFI_NOT_FOUND;
+  }
+
+  err = uefi_call_wrapper(block_io->ReadBlocks, NUM_ARGS_READ_BLOCKS,
+                          block_io,
+                          block_io->Media->MediaId,
+                          1,
+                          sizeof(GPTHeader),
+                          gpt_header);
+  if (EFI_ERROR(err)) {
+    bub_warning("Could not ReadBlocks for gpt header\n");
+    bub_free(gpt_header);
+    bub_free(*entry_buf);
+    *entry_buf = NULL;
+    return EFI_NOT_FOUND;
+  }
+
+  partition_name_bytes = bub_strlen(partition_name) + 1;
+  partition_name_ucs2 =
+    bub_calloc(sizeof(CHAR16) * partition_name_bytes);
+  if (partition_name_ucs2 == NULL) {
+    bub_warning ("Could not allocate for ucs2 partition name\n");
+    bub_free(gpt_header);
+    bub_free(*entry_buf);
+    *entry_buf = NULL;
+    return EFI_NOT_FOUND;
+  }
+  if (utf8_to_ucs2(partition_name,
+                   partition_name_bytes,
+                   partition_name_ucs2,
+                   sizeof(CHAR16) * partition_name_bytes)) {
+    bub_warning("Could not convert partition name to UCS-2\n");
+    bub_free(gpt_header);
+    bub_free(partition_name_ucs2);
+    bub_free(*entry_buf);
+    *entry_buf = NULL;
+    return EFI_NOT_FOUND;
+  }
+
+#ifdef BUB_ENABLE_DEBUG
+  Print(L"\nENTRY: %d, BLOCKSIZE: %d, GPTEntry: %d\n",
+            gpt_header->entry_count,
+            block_io->Media->BlockSize,
+            sizeof(GPTEntry));
+#endif
+
+  // Block-aligned bytes for entries.
+  UINTN entries_num_bytes = block_io->Media->BlockSize *
+                            (MAX_GPT_ENTRIES / ENTRIES_PER_BLOCK);
+
+  err = uefi_call_wrapper(block_io->ReadBlocks, NUM_ARGS_READ_BLOCKS,
+                          block_io,
+                          block_io->Media->MediaId,
+                          GPT_ENTRIES_LBA,
+                          entries_num_bytes,
+                          &all_gpt_entries);
+  if (EFI_ERROR(err)) {
+    bub_warning("Could not ReadBlocks for GPT header\n");
+    bub_free(gpt_header);
+    bub_free(partition_name_ucs2);
+    bub_free(*entry_buf);
+    *entry_buf = NULL;
+    return EFI_NOT_FOUND;
+  }
+
+  // Find matching partition name.
+  UINT8 i;
+  for (i = 0; i < gpt_header->entry_count; ++i)
+    if ((StrLen(partition_name_ucs2) ==
+         StrLen(all_gpt_entries[i].name)) &&
+        !bub_memcmp(all_gpt_entries[i].name,
+                    partition_name_ucs2,
+                    sizeof(CHAR16) * bub_strlen(partition_name))) {
+#ifdef BUB_ENABLE_DEBUG
+      Print(L"Requested Partition: %s\n", partition_name_ucs2);
+      Print(L"Found Partition Name is: %s\n", all_gpt_entries[i].name);
+      Print(L"Found Partition LBA is: %d\n", all_gpt_entries[i].first_lba);
+#endif
+      bub_memcpy((*entry_buf), &all_gpt_entries[i], sizeof(GPTEntry));
+      bub_free(partition_name_ucs2);
+      bub_free(gpt_header);
+      return EFI_SUCCESS;
+    }
+
+  bub_free(partition_name_ucs2);
+  bub_free(gpt_header);
+  bub_free(*entry_buf);
+  *entry_buf = NULL;
+  return EFI_NOT_FOUND;
+}
+
+int bub_init(MyBubOps* bub, EFI_HANDLE app_image) {
   EFI_STATUS err;
   EFI_LOADED_IMAGE *loaded_app_image = NULL;
   EFI_GUID loaded_image_protocol = LOADED_IMAGE_PROTOCOL;
-  EFI_BLOCK_IO* block_io;
-  EFI_DEVICE_PATH* block_io_path;
-  GPTEntry* partition_entry;
-  UINT8* boot_buf = NULL;
-  UINT8* kernel_buf = NULL;
-  boot_img_hdr* head_buf = NULL;
-  EFI_HANDLE kernel_image;
-  EFI_LOADED_IMAGE *loaded_kernel_image = NULL;
 
+  bub->efi_image_handle = app_image;
   err = uefi_call_wrapper(BS->HandleProtocol, NUM_ARGS_HANDLE_PROTOCOL,
                           app_image,
                           &loaded_image_protocol,
                           (VOID **) &loaded_app_image);
   if (EFI_ERROR(err)) {
     bub_warning("HandleProtocol, LOADED_IMAGE_PROTOCOL.\n");
-    return 1;
+    return 0;
   }
 
 #ifdef BUB_ENABLE_DEBUG
@@ -515,64 +504,55 @@ int bub_boot_kernel(EFI_HANDLE app_image, const char* boot_partition_name) {
   Print(L"LoadOptionsSize   : %d\n", loaded_app_image->LoadOptionsSize);
 #endif
 
-  // Get parent device block i/o.
-  err = getDiskBlockIO(loaded_app_image->DeviceHandle,
-                       &block_io,
-                       &block_io_path);
+  // Get parent device disk and block i/o.
+  err = getDiskBlockIo(loaded_app_image->DeviceHandle,
+                       &bub->block_io,
+                       &bub->disk_io,
+                       &bub->path);
   if (EFI_ERROR(err)) {
-    bub_warning("Could not block device handle.\n");
-    return 1;
+    bub_warning("Could not acquire block or disk device handle.\n");
+    return 0;
   }
 
+  bub->parent.read_from_partition = bub_read_from_partition;
+  bub->parent.write_to_partition = bub_write_to_partition;
 
-  // Get lba of partition based on name.
-  err = PartitionEntryByName(block_io, boot_partition_name ,&partition_entry);
-  if (EFI_ERROR(err)) {
-    bub_warning("Could not find Image LBA offset.\n");
-    return 1;
-  }
-
-#ifdef BUB_ENABLE_DEBUG
-  Print(L"Block IO media block size: %d\n", block_io->Media->BlockSize);
-  Print(L"Kernel Image LBA: 0x%x\n", partition_entry->first_lba);
-  Print(L"Kernel size:  0x%x\n", IMG_SIZE(partition_entry,block_io));
-#endif
-
-  boot_buf = (UINT8*)bub_malloc_(IMG_SIZE(partition_entry, block_io));
-  if (boot_buf == NULL) {
-    bub_warning("Could not allocate for boot buffer.\n");
-    return 1;
-  }
-
-#ifdef BUB_ENABLE_DEBUG
-  Print(L"Loading Boot Image...\n");
-#endif
-
-  err = uefi_call_wrapper(block_io->ReadBlocks, NUM_ARGS_READ_BLOCKS,
-                          block_io,
-                          block_io->Media->MediaId,
-                          partition_entry->first_lba,
-                          IMG_SIZE(partition_entry, block_io),
-                          boot_buf);
+  return 1;
+}
 
-  if (EFI_ERROR(err)) {
-    bub_warning("Could not read load partition image.\n");
-    return 1;
-  }
+BubBootResult bub_boot_kernel(MyBubOps* bub, const char* boot_partition_name) {
+  EFI_STATUS err;
+  GPTEntry* partition_entry;
+  UINT8* kernel_buf = NULL;
+  boot_img_hdr* head_buf = NULL;
+  UINTN num_bytes_read;
+  EFI_HANDLE kernel_image;
+  EFI_LOADED_IMAGE *loaded_kernel_image = NULL;
 
-  head_buf = (boot_img_hdr *)boot_buf;
   err = uefi_call_wrapper(BS->AllocatePool, NUM_ARGS_ALLOCATE_POOL,
                           EfiLoaderCode,
-                          head_buf->kernel_size,
-                          &kernel_buf);
+                          sizeof(boot_img_hdr),
+                          &head_buf);
   if (EFI_ERROR(err)) {
     bub_warning("Could not allocate for kernel buffer.\n");
-    return 1;
+    return BUB_BOOT_ERROR_OOM;
+  }
+
+  if (bub->parent.read_from_partition((BubOps *)bub,
+                                      boot_partition_name,
+                                      head_buf,
+                                      0,
+                                      sizeof(boot_img_hdr), &num_bytes_read)) {
+    bub_warning("Could not read boot image header.\n");
+    return BUB_BOOT_ERROR_IO;
   }
 
 #ifdef BUB_ENABLE_DEBUG
   // Print Header info
-  Print(L"kernel size:      %x\n", head_buf->kernel_size);
+  UINT8 i = 0;
+  Print(L"magic:              %c", head_buf->magic[0]);
+  for (i = 1; i < 8; ++i)     Print(L"%c", head_buf->magic[i]);
+  Print(L"\nkernel size:        %x\n", head_buf->kernel_size);
   Print(L"kernel_addr:        %x\n", head_buf->kernel_addr);
   Print(L"ramdisk_size:       %x\n", head_buf->ramdisk_size);
   Print(L"ramdisk_addr:       %x\n", head_buf->ramdisk_addr);
@@ -581,49 +561,87 @@ int bub_boot_kernel(EFI_HANDLE app_image, const char* boot_partition_name) {
   Print(L"tags_addr:          %x\n", head_buf->tags_addr);
   Print(L"page_size:          %x\n", head_buf->page_size);
   Print(L"os_version:         %x\n", head_buf->os_version);
+  Print(L"id:                 %x\n", head_buf->id[0]);
+  for (i = 1; i < 8; ++i)     Print(L" %x ", head_buf->id[i]);
+  Print(L"\n");
+#endif
+
+  // Retrieve Gpt partition data to check address boundaries.
+  err = bub_partition_entry_by_name(bub->block_io,
+                                    boot_partition_name,
+                                    &partition_entry);
+  if (EFI_ERROR(err)) {
+    bub_warning("Could not find boot partition GPT entry.\n");
+    return BUB_BOOT_ERROR_IO;
+  }
+
+#ifdef BUB_ENABLE_DEBUG
+  Print(L"Block IO media block size: %d\n", bub->block_io->Media->BlockSize);
+  Print(L"Kernel Image LBA: 0x%x\n", partition_entry->first_lba);
+  Print(L"Kernel size:  0x%x\n", IMG_SIZE(partition_entry,bub->block_io));
 #endif
 
+  // Check boot image header magic field.
+  if (bub_memcmp((UINT8[8])BOOT_MAGIC, head_buf->magic, 8)) {
+    bub_warning("Wrong boot image header magic.\n");
+    return BUB_BOOT_ERROR_PARTITION_INVALID_FORMAT;
+  }
+
   // Checks on buffer overflow.
-  if (head_buf->kernel_size > IMG_SIZE(partition_entry, block_io)) {
+  if (head_buf->kernel_size > IMG_SIZE(partition_entry, bub->block_io)) {
     bub_warning("Kernel size beyond allowed boundary.\n");
-    return 1;
+    return BUB_BOOT_ERROR_PARTITION_INVALID_FORMAT;
   }
   if (head_buf->page_size >
-      (IMG_SIZE(partition_entry, block_io) - sizeof(boot_img_hdr))) {
+      (IMG_SIZE(partition_entry, bub->block_io) - sizeof(boot_img_hdr))) {
     bub_warning("Page size too big.\n");
-    return 1;
+    return BUB_BOOT_ERROR_PARTITION_INVALID_FORMAT;
+  }
+
+  err = uefi_call_wrapper(BS->AllocatePool, NUM_ARGS_ALLOCATE_POOL,
+                          EfiLoaderCode,
+                          head_buf->kernel_size,
+                          &kernel_buf);
+  if (EFI_ERROR(err)) {
+    bub_warning("Could not allocate for kernel buffer.\n");
+    return BUB_BOOT_ERROR_OOM;
   }
 
-  bub_memcpy(kernel_buf,
-             boot_buf + head_buf->page_size,
-             head_buf->kernel_size);
+  bub_debug("Reading kernel image.\n");
+  if (bub->parent.read_from_partition((BubOps *)bub,
+                                      boot_partition_name,
+                                      kernel_buf,
+                                      head_buf->page_size,
+                                      head_buf->kernel_size,
+                                      &num_bytes_read)) {
+    bub_warning("Could not read kernel image.\n");
+    return BUB_BOOT_ERROR_IO;
+  }
 
+  bub_debug("Loading kernel image.\n");
   err = uefi_call_wrapper(BS->LoadImage, NUM_ARGS_LOAD_IMAGE,
                           FALSE,
-                          app_image,
-                          block_io_path,
+                          bub->efi_image_handle,
+                          bub->path,
                           (void *)(kernel_buf),
                           head_buf->kernel_size,
                           &kernel_image);
   if (EFI_ERROR(err)) {
-    bub_warning("Could not LOAD kernel_image.\n");
-    return 1;
+    bub_warning("Could not load kernel image.\n");
+    return BUB_BOOT_ERROR_LOAD_KERNEL;
   }
-
-#ifdef BUB_ENABLE_DEBUG
-  Print(L"LOADED Kernel Image.\n");
-#endif
+  bub_debug("Loaded kernel image.\n");
 
   // Load parameters
   err = LoadParameters(kernel_image, head_buf, &loaded_kernel_image);
   if (EFI_ERROR(err))
-    return 1;
+    return BUB_BOOT_ERROR_PARAMETER_LOAD;
 
   err = uefi_call_wrapper(BS->StartImage, 3, kernel_image, NULL, NULL);
   if (EFI_ERROR(err)) {
-    bub_warning("Could not START kernel_image.\n");
-    return 1;
+    bub_warning("Could not start kernel image.\n");
+    return BUB_BOOT_ERROR_START_KERNEL;
   }
 
-  return 0;
+  return BUB_BOOT_RESULT_OK;
 }
\ No newline at end of file
diff --git a/brillo_uefi_x86_64/boot_loader/bub_boot_kernel.h b/brillo_uefi_x86_64/boot_loader/bub_boot_kernel.h
index 7ae85a3..2ae7535 100644
--- a/brillo_uefi_x86_64/boot_loader/bub_boot_kernel.h
+++ b/brillo_uefi_x86_64/boot_loader/bub_boot_kernel.h
@@ -26,17 +26,124 @@
 
 #include <efi.h>
 #include <efilib.h>
+#include "bub_ops.h"
 
 // For printing debug statements.
 #define BUB_ENABLE_DEBUG
 
+typedef enum {
+  BUB_BOOT_RESULT_OK,
+  BUB_BOOT_ERROR_OOM,
+  BUB_BOOT_ERROR_IO,
+  BUB_BOOT_ERROR_PARTITION_INVALID_FORMAT,
+  BUB_BOOT_ERROR_LOAD_KERNEL,
+  BUB_BOOT_ERROR_PARAMETER_LOAD,
+  BUB_BOOT_ERROR_START_KERNEL,
+} BubBootResult;
+
+// GPT related constants
+#define GPT_REVISION 0x00010000
+#define GPT_MAGIC "EFI PART"
+#define GPT_MIN_SIZE 92
+#define GPT_ENTRIES_LBA 2
+#define BUB_BLOCK_SIZE 512
+#define ENTRIES_PER_BLOCK 4
+#define ENTRY_NAME_LEN 36
+#define MAX_GPT_ENTRIES 128
+
+typedef struct {
+  UINT8   signature[8];
+  UINT32  revision;
+  UINT32  header_size;
+  UINT32  header_crc32;
+  UINT32  reserved;
+  UINT64  header_lba;
+  UINT64  alternate_header_lba;
+  UINT64  first_usable_lba;
+  UINT64  last_usable_lba;
+  UINT8   disk_guid[16];
+  UINT64  entry_lba;
+  UINT32  entry_count;
+  UINT32  entry_size;
+  UINT32  entry_crc32;
+  UINT8   reserved2[420];
+} GPTHeader;
+
+typedef struct {
+  UINT8   type_GUID[16];
+  UINT8   unique_GUID[16];
+  UINT64  first_lba;
+  UINT64  last_lba;
+  UINT64  flags;
+  CHAR16  name[ENTRY_NAME_LEN];
+} GPTEntry;
+
+
+#define IMG_SIZE(entry, block) \
+  (entry->last_lba - entry->first_lba) * block->Media->BlockSize
+
+#define SIZE_BLOCK_ALIGN(bytes, BUB_BLOCK_SIZE) \
+  ((bytes + BUB_BLOCK_SIZE - 1) / BUB_BLOCK_SIZE) * BUB_BLOCK_SIZE
+
+#define OFFSET_BLOCK_ALIGN(bytes, BUB_BLOCK_SIZE) \
+  (bytes / BUB_BLOCK_SIZE) * BUB_BLOCK_SIZE
+
+typedef struct {
+  BubOps parent;
+  EFI_HANDLE efi_image_handle;
+  EFI_DEVICE_PATH* path;
+  EFI_BLOCK_IO* block_io;
+  EFI_DISK_IO* disk_io;
+  // EFI_STATUS (*PopulateMiscPartition)(MyBubOps* self);
+} MyBubOps;
+
+BubIOResult bub_read_from_partition(BubOps* ops,
+                                    const char* partition_name,
+                                    void* buf,
+                                    int64_t offset_from_partition,
+                                    size_t num_bytes,
+                                    size_t* out_num_read);
+
+BubIOResult bub_write_to_partition(BubOps* ops,
+                                   const char* partition_name,
+                                   const void* buf,
+                                   int64_t offset_from_partition,
+                                   size_t num_bytes);
+
+/* Allocates memory for and assigns to member variables of |bub|. Also assigns
+ * the Brillo Uefi-specific read_from_partition and write_to_partition
+ * functions to its BubOps parent. |app_image| must be the EFI main-specific
+ * (the current currently running program) handle.
+ *
+ * @return int 0 on failure. non-zero on success.
+ */
+int bub_init(MyBubOps* bub, EFI_HANDLE app_image);
+
 /* Boots a UEFI kernel image given a |boot_partition_name| string belonging to a
  * bootable partition entry. The partition must be on the same block device as
  * the current UEFI application, |app_image|. |app_image| is given at the entry
  * point, efi_main(), of the UEFI application.
  *
- * @return int 1 upon failure. 0 on success.
+ * @return BUB_BOOT_ERROR_OOM on allocation,
+ *         BUB_BOOT_ERROR_IO on read/write error,
+ *         BUB_BOOT_ERROR_PARTITION_INVALID_FORMAT on bad magic or bad size
+ *           boundaries,
+ *         BUB_BOOT_ERROR_LOAD_KERNEL if unable to load kernel into memory
+ *         BUB_BOOT_ERROR_PARAMETER_LOAD if unable to load kernel parameters to
+ *          the EFI_STUB,
+ *         BUB_BOOT_ERROR_START_KERNEL if unable to execute kernel,
+ *         BUB_BOOT_RESULT_OK on success.
+ */
+BubBootResult bub_boot_kernel(MyBubOps* bub, const char* boot_partition_name);
+
+/* Looks through |block_io| to search for a GPT entry named |partition_name|, a
+ * NULL-terminated string. Allocates a GPTEntry struct for |entry_buf|. Caller
+ * is responsible for freeing |entry_buf|.
+ *
+ * @return EFI_NOT_FOUND on error, EFI SUCCESS on success.
  */
-int bub_boot_kernel(EFI_HANDLE app_image, const char* boot_partition_name);
+EFI_STATUS bub_partition_entry_by_name(IN EFI_BLOCK_IO *block_io,
+                                       const char* partition_name,
+                                       GPTEntry** entry_buf);
 
 #endif /* BUB_BOOT_KERNEL_H_ */
\ No newline at end of file
diff --git a/brillo_uefi_x86_64/boot_loader/bub_image_util.cc b/brillo_uefi_x86_64/boot_loader/bub_image_util.cc
new file mode 100644
index 0000000..e8e5a15
--- /dev/null
+++ b/brillo_uefi_x86_64/boot_loader/bub_image_util.cc
@@ -0,0 +1,258 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// NOTE: See avb_slot_verify_unittest.cc for orginal reference to similar
+//       partition testing.
+
+#include "bub_image_util.h"
+
+static BubIOResult my_ops_read_from_partition(BubOps* ops,
+                                              const char* partition, void* buf,
+                                              int64_t offset, size_t num_bytes,
+                                              size_t* out_num_read) {
+  return ((MyBubOps*)ops)
+      ->my_ops->read_from_partition(partition, buf, offset, num_bytes,
+                                    out_num_read);
+}
+
+static BubIOResult my_ops_write_to_partition(BubOps* ops, const char* partition,
+                                             const void* buf, int64_t offset,
+                                             size_t num_bytes) {
+  return ((MyBubOps*)ops)
+      ->my_ops->write_to_partition(partition, buf, offset, num_bytes);
+}
+
+
+void MyOps::set_partition_dir(const base::FilePath& partition_dir) {
+  partition_dir_ = partition_dir;
+}
+
+BubIOResult MyOps::read_from_partition(const char* partition, void* buf,
+                                int64_t offset, size_t num_bytes,
+                                size_t* out_num_read) {
+  base::FilePath path =
+      partition_dir_.Append(std::string(partition)).AddExtension("img");
+
+  if (offset < 0) {
+    int64_t file_size;
+    if (!base::GetFileSize(path, &file_size)) {
+      fprintf(stderr, "Error getting size of file '%s'\n",
+              path.value().c_str());
+      return BUB_IO_RESULT_ERROR_IO;
+    }
+    offset = file_size - (-offset);
+  }
+
+  int fd = open(path.value().c_str(), O_RDONLY);
+  if (fd < 0) {
+    fprintf(stderr, "Error opening file '%s': %s\n", path.value().c_str(),
+                strerror(errno));
+    return BUB_IO_RESULT_ERROR_IO;
+  }
+  if (lseek(fd, offset, SEEK_SET) != offset) {
+    fprintf(stderr, "Error seeking to pos %zd in file %s: %s\n", offset,
+                path.value().c_str(), strerror(errno));
+    close(fd);
+    return BUB_IO_RESULT_ERROR_IO;
+  }
+  ssize_t num_read = read(fd, buf, num_bytes);
+  if (num_read < 0) {
+    fprintf(stderr, "Error reading %zd bytes from pos %" PRId64
+                " in file %s: %s\n",
+                num_bytes, offset, path.value().c_str(), strerror(errno));
+    close(fd);
+    return BUB_IO_RESULT_ERROR_IO;
+  }
+  close(fd);
+
+  if (out_num_read != NULL) {
+    *out_num_read = num_read;
+  }
+
+  return BUB_IO_RESULT_OK;
+}
+
+BubIOResult MyOps::write_to_partition(const char* partition, const void* buf,
+                               int64_t offset, size_t num_bytes) {
+  base::FilePath path =
+      partition_dir_.Append(std::string(partition)).AddExtension("img");
+
+  if (offset < 0) {
+    int64_t file_size;
+    if (!base::GetFileSize(path, &file_size)) {
+      fprintf(stderr, "Error getting size of file '%s'\n",
+              path.value().c_str());
+      return BUB_IO_RESULT_ERROR_IO;
+    }
+    offset = file_size - (-offset);
+  }
+
+  int fd = open(path.value().c_str(), O_WRONLY);
+  if (fd < 0) {
+    fprintf(stderr, "Error opening file '%s': %s\n", path.value().c_str(),
+                strerror(errno));
+    return BUB_IO_RESULT_ERROR_IO;
+  }
+  if (lseek(fd, offset, SEEK_SET) != offset) {
+    fprintf(stderr, "Error seeking to pos %zd in file %s: %s\n", offset,
+                path.value().c_str(), strerror(errno));
+    close(fd);
+    return BUB_IO_RESULT_ERROR_IO;
+  }
+  ssize_t num_written = write(fd, buf, num_bytes);
+  if (num_written < 0) {
+    fprintf(stderr, "Error writing %zd bytes at pos %"
+                    PRId64 " in file %s: %s\n",
+            num_bytes, offset, path.value().c_str(), strerror(errno));
+    close(fd);
+    return BUB_IO_RESULT_ERROR_IO;
+  }
+  close(fd);
+
+  return BUB_IO_RESULT_OK;
+}
+
+void MyOps::write_ab_metadata(BubAbData* ab,
+                              const uint8_t* magic,
+                              uint8_t a_priority,
+                              uint8_t a_tries_remaining,
+                              uint8_t a_successful_boot,
+                              uint8_t b_priority,
+                              uint8_t b_tries_remaining,
+                              uint8_t b_successful_boot) {
+  bub_memset(ab, 0, sizeof(BubAbData));
+  bub_memcpy(ab->magic, magic, sizeof(ab->magic));
+  ab->major_version = BUB_MAJOR_VERSION;
+  ab->minor_version = BUB_MINOR_VERSION;
+  ab->slots[0].priority = a_priority;
+  ab->slots[0].tries_remaining = a_tries_remaining;
+  ab->slots[0].successful_boot = a_successful_boot;
+  ab->slots[1].priority = b_priority;
+  ab->slots[1].tries_remaining = b_tries_remaining;
+  ab->slots[1].successful_boot = b_successful_boot;
+}
+
+base::FilePath MyOps::make_metadata_image(const BubAbData* ab_metadata,
+                                          const char* name) {
+  // Generate a 1025 KiB file with known content.
+  std::vector<uint8_t> image;
+  image.resize(sizeof(BubAbData));
+  BubAbData ab_metadata_be;
+  uint8_t* image_data = (uint8_t*)&ab_metadata_be;
+
+  bub_memcpy(&ab_metadata_be, ab_metadata, sizeof(BubAbData));
+
+  // Byte swap all necessary variables here.
+
+  ab_metadata_be.crc32 = 0;
+  ab_metadata_be.crc32 =
+    bub_be32toh(bub_crc32((uint8_t*)&ab_metadata_be, sizeof(BubAbData)));
+
+  for (size_t n = 0; n < sizeof(BubAbData); n++) {
+    image[n] = image_data[n];
+  }
+  base::FilePath image_path = partition_dir_.Append(name);
+  EXPECT_EQ(sizeof(BubAbData),
+            static_cast<const size_t>(base::WriteFile(
+                image_path, reinterpret_cast<const char*>(image.data()),
+                image.size())));
+  return image_path;
+}
+
+void AbTest::SetUp() {
+  base::FilePath ret;
+  char* buf = strdup("/tmp/bub-tests.XXXXXX");
+  ASSERT_TRUE(mkdtemp(buf) != nullptr);
+  testdir_ = base::FilePath(buf);
+  ops_.set_partition_dir(testdir_);
+  free(buf);
+}
+
+MyOps::MyOps() {
+  bub_ops_ = new MyBubOps;
+  bub_ops_->parent.read_from_partition = my_ops_read_from_partition;
+  bub_ops_->parent.write_to_partition = my_ops_write_to_partition;
+  bub_ops_->my_ops = this;
+}
+
+MyOps::~MyOps() { delete bub_ops_; }
+
+void AbTest::GenerateMiscImage(const BubAbData* ab_metadata) {
+  ops_.make_metadata_image(ab_metadata, "misc.img");
+}
+
+int AbTest::CompareMiscImage(BubAbData ab_expected) {
+  const uint8_t A = 0, B = 1;
+  size_t num_bytes_read;
+  BubAbData ab_expected_be;
+  BubAbData* ab_actual = (BubAbData*)bub_calloc(sizeof(BubAbData));
+
+  bub_memcpy(&ab_expected_be, &ab_expected, sizeof(BubAbData));
+
+  // Byte swap all necessary variables here.
+
+  ab_expected_be.crc32 = 0;
+  ab_expected_be.crc32 =
+    bub_be32toh(bub_crc32((uint8_t*)&ab_expected_be, sizeof(BubAbData)));
+
+  if ((ops_.bub_ops_)->parent.read_from_partition((BubOps *)ops_.bub_ops_,
+                                                  "misc", ab_actual, 0,
+                                                  sizeof(BubAbData),
+                                                  &num_bytes_read)) {
+    fprintf(stderr, "Could not read from misc partition.\n");
+    bub_free(ab_actual);
+    return 1;
+  }
+  if (num_bytes_read != sizeof(BubAbData)) {
+    fprintf(stderr, "Bad misc partition read.\n");
+    bub_free(ab_actual);
+    return 1;
+  }
+
+  // Check magic and version numbers.
+  EXPECT_EQ(0, bub_memcmp(&ab_expected_be, ab_actual, 8));
+
+  // Check slots values.
+  EXPECT_EQ(ab_expected_be.slots[A].priority,
+            ab_actual->slots[A].priority);
+  EXPECT_EQ(ab_expected_be.slots[A].tries_remaining,
+            ab_actual->slots[A].tries_remaining);
+  EXPECT_EQ(ab_expected_be.slots[A].successful_boot,
+            ab_actual->slots[A].successful_boot);
+  EXPECT_EQ(0, bub_memcmp(ab_expected_be.slots[A].reserved,
+                          ab_actual->slots[A].reserved,
+                          sizeof(ab_actual->slots[A].reserved)));
+  EXPECT_EQ(ab_expected_be.slots[B].priority,
+            ab_actual->slots[B].priority);
+  EXPECT_EQ(ab_expected_be.slots[B].tries_remaining,
+            ab_actual->slots[B].tries_remaining);
+  EXPECT_EQ(ab_expected_be.slots[B].successful_boot,
+            ab_actual->slots[B].successful_boot);
+  EXPECT_EQ(0, bub_memcmp(&ab_expected_be.slots[B].reserved,
+                          ab_actual->slots[B].reserved,
+                          sizeof(ab_actual->slots[A].reserved)));
+
+  // Check reserved and crc bytes.
+  // TODO: Compute and compare crc value here.
+  EXPECT_EQ(0, bub_memcmp(ab_actual->reserved2,
+                          ab_expected_be.reserved2,
+                          sizeof(ab_expected_be.reserved2)));
+
+  EXPECT_EQ(ab_expected_be.crc32, ab_actual->crc32);
+
+  bub_free(ab_actual);
+  return 0;
+}
diff --git a/brillo_uefi_x86_64/boot_loader/bub_image_util.h b/brillo_uefi_x86_64/boot_loader/bub_image_util.h
new file mode 100644
index 0000000..4adba75
--- /dev/null
+++ b/brillo_uefi_x86_64/boot_loader/bub_image_util.h
@@ -0,0 +1,96 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef BUB_IMAGE_UTIL_H_
+#define BUB_IMAGE_UTIL_H_
+
+#include <fcntl.h>
+#include <gtest/gtest.h>
+#include <base/files/file_util.h>
+
+#include "bub_sysdeps.h"
+#include "bub_ab_flow.h"
+#include "bub_util.h"
+
+struct MyBubOps;
+typedef struct MyBubOps MyBubOps;
+
+class MyOps {
+ public:
+  MyOps();
+  ~MyOps();
+
+  BubOps* bub_ops() { return (BubOps*)bub_ops_; }
+  void set_partition_dir(const base::FilePath& partition_dir);
+  BubIOResult read_from_partition(const char* partition, void* buf,
+                                  int64_t offset, size_t num_bytes,
+                                  size_t* out_num_read);
+  BubIOResult write_to_partition(const char* partition, const void* buf,
+                                 int64_t offset, size_t num_bytes);
+
+  /* Assigns to |ab| metadata using |magic| and [a,b]_*| parameters. This
+   * function does not swap byte order nor does it calculate the crc.
+   */
+  void write_ab_metadata(BubAbData* ab,
+                         const uint8_t* magic,
+                         uint8_t a_priority,
+                         uint8_t a_tries_remaining,
+                         uint8_t a_successful_boot,
+                         uint8_t b_priority,
+                         uint8_t b_tries_remaining,
+                         uint8_t b_successful_boot);
+
+  /* Writes out a misc.img file in a temp directory using |ab_metadata|.
+   * Byte swapping is done prior to writing to ensure the big endianness
+   * expected in the Misc partition.
+   */
+  base::FilePath make_metadata_image(const BubAbData* ab_metadata,
+                                   	 const char* name);
+
+  MyBubOps* bub_ops_;
+  base::FilePath partition_dir_;
+};
+
+struct MyBubOps {
+  BubOps parent;
+  MyOps* my_ops;
+};
+
+class AbTest : public ::testing::Test {
+  public:
+    AbTest() {}
+
+    // Create temporary directory to stash images in.
+    void SetUp() override;
+
+    /* Wrapper function to generate test misc image by calling MyOps
+     * make_metadata_image.
+     */
+    void GenerateMiscImage(const BubAbData* ab_metadata);
+
+    /* Tests expected vs actual contents of ab metadata found in the Misc
+     * partition. Byte swapping to big endianness and crc for |ab_expected|
+     * is done prior to test comparisons.
+     */
+    int CompareMiscImage(BubAbData ab_expected);
+
+    // Temporary directory created in SetUp().
+    base::FilePath testdir_;
+
+    MyOps ops_;
+};
+
+#endif /* BUB_IMAGE_UTIL_H_ */
\ No newline at end of file
diff --git a/brillo_uefi_x86_64/boot_loader/bub_main.c b/brillo_uefi_x86_64/boot_loader/bub_main.c
index 5889b18..54e6f62 100644
--- a/brillo_uefi_x86_64/boot_loader/bub_main.c
+++ b/brillo_uefi_x86_64/boot_loader/bub_main.c
@@ -24,23 +24,46 @@
 
 #include <efi.h>
 #include <efilib.h>
+#include "bub_ab_flow.h"
 #include "bub_boot_kernel.h"
 #include "bub_sysdeps.h"
 
 
 EFI_STATUS EFIAPI efi_main (EFI_HANDLE ImageHandle,
                             EFI_SYSTEM_TABLE* SystemTable) {
-  int err;
+  MyBubOps ops;
+  BubAbFlowResult ab_result;
+  BubBootResult boot_result;
+  char slot_suffix[BUB_SUFFIX_SIZE] = {0};
+  char boot_name[7] = "boot\0\0\0";
 
   InitializeLib(ImageHandle, SystemTable);
   bub_print("Brillo UEFI A/B BOOT LOADER\n");
 
-  err = bub_boot_kernel(ImageHandle, "boot_a");
-  if (err) {
-    bub_error("Error loading kernel.\n");
-    uefi_call_wrapper(BS->Stall, 1, 15 * 1000 * 1000);
-    return EFI_LOAD_ERROR;
-  }
+  if (!bub_init(&ops, ImageHandle))
+    bub_error("Could not initialize Brillo Uefi object.");
+
+  // Attempt AB flow and boot.  Invalidate metadata for slots having bad
+  // partition format.
+  do {
+    ab_result = bub_ab_flow((BubOps *)&ops, slot_suffix, BUB_SUFFIX_SIZE);
+    if (ab_result != BUB_AB_FLOW_RESULT_OK)
+      bub_error("Could not choose A/B slot.\n");
+
+    bub_memcpy(boot_name + 4, slot_suffix, BUB_SUFFIX_SIZE);
+
+    boot_result = bub_boot_kernel(&ops, boot_name);
+    if (boot_result == BUB_BOOT_ERROR_PARTITION_INVALID_FORMAT) {
+      bub_warning("Marking slot as invalid.\n");
+
+      if (bub_ab_mark_as_invalid((MyBubOps *)&ops, slot_suffix))
+        bub_error("Could not mark slot invalid.");
+
+    }
+    else if (boot_result != BUB_BOOT_RESULT_OK)
+      bub_error("Error loading kernel.\n");
+
+  } while (boot_result == BUB_BOOT_ERROR_PARTITION_INVALID_FORMAT);
 
   return EFI_SUCCESS;
 }
\ No newline at end of file
diff --git a/brillo_uefi_x86_64/boot_loader/bub_ops_uefi.c b/brillo_uefi_x86_64/boot_loader/bub_ops_uefi.c
new file mode 100644
index 0000000..232f125
--- /dev/null
+++ b/brillo_uefi_x86_64/boot_loader/bub_ops_uefi.c
@@ -0,0 +1,135 @@
+/*
+ * Copyright (C) 2016 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <efi.h>
+#include <efilib.h>
+#include "bub_sysdeps.h"
+#include "bub_boot_kernel.h"
+
+BubIOResult bub_read_from_partition(BubOps* ops,
+                                    const char* partition_name,
+                                    void* buf,
+                                    int64_t offset_from_partition,
+                                    size_t num_bytes,
+                                    size_t* out_num_read) {
+  bub_assert(partition_name != NULL);
+  bub_assert(buf != NULL);
+  bub_assert(out_num_read != NULL);
+
+  EFI_STATUS err;
+  GPTEntry *partition_entry;
+  UINT64 partition_size;
+  MyBubOps* bub = (MyBubOps*)ops;
+
+  err = bub_partition_entry_by_name(bub->block_io,
+                                    partition_name,
+                                    &partition_entry);
+  if (EFI_ERROR(err))
+    return BUB_IO_RESULT_ERROR_NO_SUCH_PARTITION;
+
+  partition_size = IMG_SIZE(partition_entry, bub->block_io);
+
+  if (offset_from_partition < 0) {
+    if ((-offset_from_partition) > partition_size) {
+      bub_warning("Offset outside range.\n");
+      bub_free(partition_entry);
+      return BUB_IO_RESULT_ERROR_RANGE_OUTSIDE_PARTITION;
+    }
+    offset_from_partition = partition_size - (-offset_from_partition);
+  }
+
+  // Check if num_bytes goes beyond partition end. If so, don't read beyond
+  // this boundary -- do a partial I/O instead.
+  if (num_bytes > partition_size - offset_from_partition)
+    *out_num_read = partition_size - offset_from_partition;
+  else
+    *out_num_read = num_bytes;
+
+  err = uefi_call_wrapper(bub->disk_io->ReadDisk, 5,
+                          bub->disk_io,
+                          bub->block_io->Media->MediaId,
+                          (partition_entry->first_lba *
+                            bub->block_io->Media->BlockSize) +
+                            offset_from_partition,
+                          *out_num_read,
+                          buf);
+  if (EFI_ERROR(err)) {
+    bub_warning("Could not read from Disk.\n");
+    *out_num_read = 0;
+    bub_free(partition_entry);
+    return BUB_IO_RESULT_ERROR_IO;
+  }
+
+  bub_free(partition_entry);
+  return BUB_IO_RESULT_OK;
+}
+
+BubIOResult bub_write_to_partition(BubOps* ops,
+                                   const char* partition_name,
+                                   const void* buf,
+                                   int64_t offset_from_partition,
+                                   size_t num_bytes) {
+  bub_assert(partition_name != NULL);
+  bub_assert(buf != NULL);
+
+  EFI_STATUS err;
+  GPTEntry *partition_entry;
+  UINT64 partition_size;
+  MyBubOps* bub = (MyBubOps*)ops;
+
+  err = bub_partition_entry_by_name(bub->block_io,
+                                    partition_name,
+                                    &partition_entry);
+  if (EFI_ERROR(err))
+    return BUB_IO_RESULT_ERROR_NO_SUCH_PARTITION;
+
+  partition_size = IMG_SIZE(partition_entry, bub->block_io);
+
+  if (offset_from_partition < 0) {
+    if ((-offset_from_partition) > partition_size) {
+      bub_warning("Offset outside range.\n");
+      bub_free(partition_entry);
+      return BUB_IO_RESULT_ERROR_RANGE_OUTSIDE_PARTITION;
+    }
+    offset_from_partition = partition_size - (-offset_from_partition);
+  }
+
+  // Check if num_bytes goes beyond partition end. If so, error out -- no
+  // partial I/O.
+  if (num_bytes > partition_size - offset_from_partition) {
+    bub_warning("Cannot write beyond partition boundary.\n");
+    bub_free(partition_entry);
+    return BUB_IO_RESULT_ERROR_RANGE_OUTSIDE_PARTITION;
+  }
+
+  err = uefi_call_wrapper(bub->disk_io->WriteDisk, 5,
+                          bub->disk_io,
+                          bub->block_io->Media->MediaId,
+                          (partition_entry->first_lba *
+                            bub->block_io->Media->BlockSize) +
+                            offset_from_partition,
+                          num_bytes,
+                          buf);
+
+  if (EFI_ERROR(err)) {
+    bub_warning("Could not write to Disk.\n");
+    bub_free(partition_entry);
+    return BUB_IO_RESULT_ERROR_IO;
+  }
+
+  bub_free(partition_entry);
+  return BUB_IO_RESULT_OK;
+}
\ No newline at end of file
diff --git a/brillo_uefi_x86_64/boot_loader/bub_sysdeps_uefi.c b/brillo_uefi_x86_64/boot_loader/bub_sysdeps_uefi.c
index 051a748..37df25f 100644
--- a/brillo_uefi_x86_64/boot_loader/bub_sysdeps_uefi.c
+++ b/brillo_uefi_x86_64/boot_loader/bub_sysdeps_uefi.c
@@ -23,7 +23,7 @@ int bub_memcmp(const void* src1, const void* src2, size_t n) {
 }
 
 int bub_strcmp(const char* s1, const char* s2) {
-  return (int)StrCmp(s1, s2);
+  return (int)strcmpa(s1, s2);
 }
 
 void* bub_memcpy(void* dest, const void* src, size_t n) {
diff --git a/brillo_uefi_x86_64/boot_loader/bub_util.c b/brillo_uefi_x86_64/boot_loader/bub_util.c
index fa7294b..ed1e1f4 100644
--- a/brillo_uefi_x86_64/boot_loader/bub_util.c
+++ b/brillo_uefi_x86_64/boot_loader/bub_util.c
@@ -57,4 +57,35 @@ void* bub_calloc(size_t size) {
     bub_memset(x, 0, size);
 
   return x;
+}
+
+uint32_t bub_be32toh(uint32_t in) {
+  uint8_t* d = (uint8_t*)&in;
+  uint32_t ret;
+  ret = ((uint32_t)d[0]) << 24;
+  ret |= ((uint32_t)d[1]) << 16;
+  ret |= ((uint32_t)d[2]) << 8;
+  ret |= ((uint32_t)d[3]);
+  return ret;
+}
+
+uint64_t bub_be64toh(uint64_t in) {
+  uint8_t* d = (uint8_t*)&in;
+  uint64_t ret;
+  ret = ((uint64_t)d[0]) << 56;
+  ret |= ((uint64_t)d[1]) << 48;
+  ret |= ((uint64_t)d[2]) << 40;
+  ret |= ((uint64_t)d[3]) << 32;
+  ret |= ((uint64_t)d[4]) << 24;
+  ret |= ((uint64_t)d[5]) << 16;
+  ret |= ((uint64_t)d[6]) << 8;
+  ret |= ((uint64_t)d[7]);
+  return ret;
+}
+
+uint32_t bub_crc32(const uint8_t* data, size_t data_size) {
+
+  /* TODO: Determine crc polynomial and implement here. */
+
+  return 0;
 }
\ No newline at end of file
diff --git a/brillo_uefi_x86_64/boot_loader/bub_util.h b/brillo_uefi_x86_64/boot_loader/bub_util.h
index 78978d8..74b32f3 100644
--- a/brillo_uefi_x86_64/boot_loader/bub_util.h
+++ b/brillo_uefi_x86_64/boot_loader/bub_util.h
@@ -43,6 +43,15 @@ int utf8_to_ucs2(const uint8_t* utf8_data,
  */
 void* bub_calloc(size_t size);
 
+/* Converts a 32-bit unsigned integer from big-endian to host byte order. */
+uint32_t bub_be32toh(uint32_t in);
+
+/* Converts a 64-bit unsigned integer from big-endian to host byte order. */
+uint64_t bub_be64toh(uint64_t in);
+
+/* Calculates and returns crc32 value of |data| given byte size, |data_size|. */
+uint32_t bub_crc32(const uint8_t* data, size_t data_size);
+
 #ifdef __cplusplus
 }
 #endif