2load_kernel.c: Expose load kernel as vb2_api

Move VbTryLoadKernel() and VbTryLoadMiniOsKernel() to depthcharge.
In addition, expose LoadKernel() and LoadMiniOsKernel() as
vb2api_load_kernel() and vb2api_load_minos_kernel(), respectively.
Since we remove VbTryLoadKernel() in vb2api_normal_boot(), rename
vb2api_normal_boot() to check_kernel_version() and call it in
vb2api_kernel_finalize().

Also, rename VbSelectAndLoadKernelParams, VbDiskInfo, and
VbExDiskHandle_t as vb2_kernel_params, vb2_disk_info, and
vb2ex_disk_handle_t, respectively.

BUG=b:172339016
BRANCH=none
TEST=make runtests
TEST=FW_NAME=tomato emerge-cherry depthcharge

Cq-Depend: chromium:3758201
Signed-off-by: Hsin-Te Yuan <yuanhsinte@google.com>
Change-Id: Iaf182ca81797fee24c6104dac9d4770ff75c7fb2
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/3755923
Reviewed-by: Yu-Ping Wu <yupingso@chromium.org>
Reviewed-by: Julius Werner <jwerner@chromium.org>

10 files changed, 244 insertions, 956 deletions

diff --git a/tests/cgpt_fuzzer.c b/tests/cgpt_fuzzer.c
index d440ee82..6762bc38 100644
--- a/tests/cgpt_fuzzer.c
+++ b/tests/cgpt_fuzzer.c
@@ -25,7 +25,7 @@ struct GptDataParams {
 
 static struct MockDisk mock_disk;
 
-vb2_error_t VbExDiskRead(VbExDiskHandle_t h, uint64_t lba_start,
+vb2_error_t VbExDiskRead(vb2ex_disk_handle_t h, uint64_t lba_start,
 			 uint64_t lba_count, void *buffer)
 {
 	size_t lba_size = mock_disk.size >> mock_disk.sector_shift;
diff --git a/tests/gpt_misc_tests.c b/tests/gpt_misc_tests.c
index 30632991..e7ffccf7 100644
--- a/tests/gpt_misc_tests.c
+++ b/tests/gpt_misc_tests.c
@@ -29,7 +29,7 @@ static char call_log[4096];
 static int disk_read_to_fail;
 static int disk_write_to_fail;
 
-static VbExDiskHandle_t handle;
+static vb2ex_disk_handle_t handle;
 static uint8_t mock_disk[MOCK_SECTOR_SIZE * MOCK_SECTOR_COUNT];
 static GptHeader *mock_gpt_primary =
 	(GptHeader*)&mock_disk[MOCK_SECTOR_SIZE * 1];
@@ -91,7 +91,7 @@ static void ResetMocks(void)
 
 /* Mocks */
 
-vb2_error_t VbExDiskRead(VbExDiskHandle_t h, uint64_t lba_start,
+vb2_error_t VbExDiskRead(vb2ex_disk_handle_t h, uint64_t lba_start,
 			 uint64_t lba_count, void *buffer)
 {
 	LOGCALL("VbExDiskRead(h, %d, %d)\n", (int)lba_start, (int)lba_count);
@@ -105,7 +105,7 @@ vb2_error_t VbExDiskRead(VbExDiskHandle_t h, uint64_t lba_start,
 	return VB2_SUCCESS;
 }
 
-vb2_error_t VbExDiskWrite(VbExDiskHandle_t h, uint64_t lba_start,
+vb2_error_t VbExDiskWrite(vb2ex_disk_handle_t h, uint64_t lba_start,
 			  uint64_t lba_count, const void *buffer)
 {
 	LOGCALL("VbExDiskWrite(h, %d, %d)\n", (int)lba_start, (int)lba_count);
@@ -207,7 +207,7 @@ static void ReadWriteGptTest(void)
 	 * Invalidate primary GPT header and check that it is
 	 * repaired by GptRepair().
 	 *
-	 * This would normally be called by LoadKernel()->GptInit()
+	 * This would normally be called by vb2api_load_kernel()->GptInit()
 	 * but this callback is mocked in these tests.
 	 */
 	ResetMocks();
@@ -232,7 +232,7 @@ static void ReadWriteGptTest(void)
 	 * Invalidate secondary GPT header and check that it can be
 	 * repaired by GptRepair().
 	 *
-	 * This would normally be called by LoadKernel()->GptInit()
+	 * This would normally be called by vb2api_load_kernel()->GptInit()
 	 * but this callback is mocked in these tests.
 	 */
 	ResetMocks();
diff --git a/tests/vb2_auxfw_sync_tests.c b/tests/vb2_auxfw_sync_tests.c
index bb236be8..b9c869ab 100644
--- a/tests/vb2_auxfw_sync_tests.c
+++ b/tests/vb2_auxfw_sync_tests.c
@@ -15,7 +15,6 @@
 #include "2sysincludes.h"
 #include "common/tests.h"
 #include "host_common.h"
-#include "load_kernel_fw.h"
 #include "vboot_struct.h"
 
 /* Mock data */
diff --git a/tests/vb2_ec_sync_tests.c b/tests/vb2_ec_sync_tests.c
index 4afc5ef0..f317561a 100644
--- a/tests/vb2_ec_sync_tests.c
+++ b/tests/vb2_ec_sync_tests.c
@@ -12,7 +12,6 @@
 #include "2sysincludes.h"
 #include "common/tests.h"
 #include "host_common.h"
-#include "load_kernel_fw.h"
 #include "vboot_struct.h"
 
 /* Mock data */
diff --git a/tests/vb2_kernel2_tests.c b/tests/vb2_kernel2_tests.c
deleted file mode 100644
index 04defe50..00000000
--- a/tests/vb2_kernel2_tests.c
+++ /dev/null
@@ -1,146 +0,0 @@
-/* Copyright 2022 The ChromiumOS Authors.
- * Use of this source code is governed by a BSD-style license that can be
- * found in the LICENSE file.
- *
- * Tests for vb2api_normal_boot.
- */
-
-#include "2api.h"
-#include "2common.h"
-#include "2misc.h"
-#include "2nvstorage.h"
-#include "2secdata.h"
-#include "2sysincludes.h"
-#include "common/boot_mode.h"
-#include "common/tests.h"
-#include "host_common.h"
-#include "load_kernel_fw.h"
-#include "tlcl.h"
-#include "tss_constants.h"
-#include "vboot_struct.h"
-
-/* Common context for tests */
-static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
-	__attribute__((aligned(VB2_WORKBUF_ALIGN)));
-static struct vb2_context *ctx;
-static struct vb2_shared_data *sd;
-static uint32_t kernel_version;
-static uint32_t new_version;
-static VbSelectAndLoadKernelParams kparams;
-
-/* Mocked function data */
-static struct vb2_gbb_header mock_gbb;
-static int mock_vbtlk_expect_fixed;
-static int mock_vbtlk_expect_removable;
-static vb2_error_t mock_vbtlk_retval;
-
-static void reset_common_data(void)
-{
-	memset(&kparams, 0, sizeof(kparams));
-
-	memset(&mock_gbb, 0, sizeof(mock_gbb));
-	mock_gbb.major_version = VB2_GBB_MAJOR_VER;
-	mock_gbb.minor_version = VB2_GBB_MINOR_VER;
-	mock_gbb.flags = 0;
-
-	mock_vbtlk_expect_fixed = 1;
-	mock_vbtlk_expect_removable = 0;
-	mock_vbtlk_retval = VB2_SUCCESS;
-
-	TEST_SUCC(vb2api_init(workbuf, sizeof(workbuf), &ctx),
-		  "vb2api_init failed");
-
-	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_NORMAL);
-
-	sd = vb2_get_sd(ctx);
-
-	vb2_nv_init(ctx);
-	vb2_nv_set(ctx, VB2_NV_KERNEL_MAX_ROLLFORWARD, 0xffffffff);
-
-	kernel_version = new_version = 0x10002;
-
-	sd->kernel_version_secdata = kernel_version;
-	sd->kernel_version = kernel_version;
-}
-
-/* Mock functions */
-
-struct vb2_gbb_header *vb2_get_gbb(struct vb2_context *c)
-{
-	return &mock_gbb;
-}
-
-void vb2_secdata_kernel_set(struct vb2_context *c,
-			    enum vb2_secdata_kernel_param param,
-			    uint32_t value)
-{
-	kernel_version = value;
-}
-
-vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t disk_flags,
-			    VbSelectAndLoadKernelParams *kpa)
-{
-	/*
-	 * TODO: Currently we don't have a good way of testing for an ordered
-	 * sequence of VB_DISK_FLAG_FIXED and then VB_DISK_FLAG_REMOVABLE.  If
-	 * both are set, then just assume success.
-	 */
-	if (mock_vbtlk_expect_fixed && mock_vbtlk_expect_removable)
-		return mock_vbtlk_retval;
-
-	TEST_EQ(!!mock_vbtlk_expect_fixed,
-		!!(disk_flags & VB_DISK_FLAG_FIXED),
-		"  unexpected fixed disk call");
-	TEST_EQ(!!mock_vbtlk_expect_removable,
-		!!(disk_flags & VB_DISK_FLAG_REMOVABLE),
-		"  unexpected removable disk call");
-
-	sd->kernel_version = new_version;
-
-	return mock_vbtlk_retval;
-}
-
-/* Tests */
-
-static void normal_boot_tests(void)
-{
-	reset_common_data();
-	TEST_SUCC(vb2api_normal_boot(ctx, &kparams),
-		  "vb2api_normal_boot() returns VB2_SUCCESS");
-
-	reset_common_data();
-	mock_vbtlk_retval = VB2_ERROR_MOCK;
-	TEST_EQ(vb2api_normal_boot(ctx, &kparams), VB2_ERROR_MOCK,
-		"vb2api_normal_boot() returns VB2_ERROR_MOCK");
-
-	reset_common_data();
-	new_version = 0x20003;
-	TEST_SUCC(vb2api_normal_boot(ctx, &kparams), "Roll forward");
-	TEST_EQ(kernel_version, 0x20003, "  version");
-
-	reset_common_data();
-	vb2_nv_set(ctx, VB2_NV_FW_RESULT, VB2_FW_RESULT_TRYING);
-	new_version = 0x20003;
-	TEST_SUCC(vb2api_normal_boot(ctx, &kparams),
-		  "Don't roll forward kernel when trying new FW");
-	TEST_EQ(kernel_version, 0x10002, "  version");
-
-	reset_common_data();
-	vb2_nv_set(ctx, VB2_NV_KERNEL_MAX_ROLLFORWARD, 0x30005);
-	new_version = 0x40006;
-	TEST_SUCC(vb2api_normal_boot(ctx, &kparams), "Limit max roll forward");
-	TEST_EQ(kernel_version, 0x30005, "  version");
-
-	reset_common_data();
-	vb2_nv_set(ctx, VB2_NV_KERNEL_MAX_ROLLFORWARD, 0x10001);
-	new_version = 0x40006;
-	TEST_SUCC(vb2api_normal_boot(ctx, &kparams),
-		  "Max roll forward can't rollback");
-	TEST_EQ(kernel_version, 0x10002, "  version");
-}
-
-int main(void)
-{
-	normal_boot_tests();
-	return gTestSuccess ? 0 : 255;
-}
diff --git a/tests/vb2_kernel_tests.c b/tests/vb2_kernel_tests.c
index ca945e95..6d32eb2d 100644
--- a/tests/vb2_kernel_tests.c
+++ b/tests/vb2_kernel_tests.c
@@ -14,8 +14,6 @@
 #include "2sysincludes.h"
 #include "common/boot_mode.h"
 #include "common/tests.h"
-#include "vboot_struct.h"
-#include "vboot_api.h"
 
 /* Common context for tests */
 static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
@@ -24,7 +22,7 @@ static struct vb2_context *ctx;
 static struct vb2_shared_data *sd;
 static struct vb2_fw_preamble *fwpre;
 static const char fw_kernel_key_data[36] = "Test kernel key data";
-static VbSelectAndLoadKernelParams kparams;
+static struct vb2_kernel_params kparams;
 
 /* Mocked function data */
 
@@ -40,6 +38,9 @@ static int mock_commit_data_called;
 static int mock_ec_sync_called;
 static int mock_ec_sync_retval;
 static int mock_battery_cutoff_called;
+static int mock_kernel_flag;
+static int mock_kernel_flag_set;
+static int mock_kernel_version;
 
 /* Type of test to reset for */
 enum reset_type {
@@ -62,16 +63,15 @@ static void reset_common_data(enum reset_type t)
 	sd = vb2_get_sd(ctx);
 	vb2_nv_init(ctx);
 
-	vb2api_secdata_kernel_create(ctx);
-	vb2_secdata_kernel_init(ctx);
-	vb2_secdata_kernel_set(ctx, VB2_SECDATA_KERNEL_VERSIONS, 0x20002);
-
 	mock_read_res_fail_on_call = 0;
 	mock_secdata_fwmp_check_retval = VB2_SUCCESS;
 	mock_commit_data_called = 0;
 	mock_ec_sync_called = 0;
 	mock_ec_sync_retval = VB2_SUCCESS;
 	mock_battery_cutoff_called = 0;
+	mock_kernel_flag = 0;
+	mock_kernel_flag_set = 0;
+	mock_kernel_version = 0x10002;
 
 	/* Recovery key in mock GBB */
 	memset(&mock_gbb, 0, sizeof(mock_gbb));
@@ -86,6 +86,8 @@ static void reset_common_data(enum reset_type t)
 	mock_gbb.h.recovery_key_size =
 		mock_gbb.recovery_key.key_offset +
 		mock_gbb.recovery_key.key_size;
+	mock_gbb.h.major_version = VB2_GBB_MAJOR_VER;
+	mock_gbb.h.minor_version = VB2_GBB_MINOR_VER;
 
 	if (t == FOR_PHASE1) {
 		uint8_t *kdata;
@@ -103,6 +105,11 @@ static void reset_common_data(enum reset_type t)
 		sd->preamble_size = sizeof(*fwpre) + k->key_size;
 		vb2_set_workbuf_used(ctx,
 				     sd->preamble_offset + sd->preamble_size);
+	} else if (t == FOR_FINALIZE) {
+		SET_BOOT_MODE(ctx, VB2_BOOT_MODE_NORMAL);
+		vb2_nv_set(ctx, VB2_NV_KERNEL_MAX_ROLLFORWARD, 0xffffffff);
+		sd->kernel_version_secdata = mock_kernel_version;
+		sd->kernel_version = mock_kernel_version;
 	}
 };
 
@@ -178,6 +185,37 @@ vb2_error_t vb2ex_commit_data(struct vb2_context *c)
 	return VB2_SUCCESS;
 }
 
+void vb2_secdata_kernel_set(struct vb2_context *c,
+			    enum vb2_secdata_kernel_param param,
+			    uint32_t value)
+{
+	switch (param) {
+	case VB2_SECDATA_KERNEL_FLAGS:
+		mock_kernel_flag = value;
+		mock_kernel_flag_set = 1;
+		break;
+	case VB2_SECDATA_KERNEL_VERSIONS:
+		mock_kernel_version = value;
+		break;
+	default:
+		vb2ex_abort();
+	}
+}
+
+uint32_t vb2_secdata_kernel_get(struct vb2_context *c,
+				enum vb2_secdata_kernel_param param)
+{
+	switch (param) {
+	case VB2_SECDATA_KERNEL_FLAGS:
+		return mock_kernel_flag;
+	case VB2_SECDATA_KERNEL_VERSIONS:
+		return mock_kernel_version;
+	default:
+		vb2ex_abort();
+	}
+	return 0;
+}
+
 /* Tests */
 
 static void phase1_tests(void)
@@ -204,7 +242,7 @@ static void phase1_tests(void)
 	TEST_EQ(k->key_size, sizeof(fw_kernel_key_data), "  key_size");
 	TEST_EQ(memcmp((uint8_t *)k + k->key_offset, fw_kernel_key_data,
 		       k->key_size), 0, "  key data");
-	TEST_EQ(sd->kernel_version_secdata, 0x20002,
+	TEST_EQ(sd->kernel_version_secdata, 0x10002,
 		"  secdata_kernel version");
 
 	/* Test successful call in recovery mode */
@@ -230,7 +268,7 @@ static void phase1_tests(void)
 	TEST_EQ(memcmp((uint8_t *)k + k->key_offset,
 		       mock_gbb.recovery_key_data, k->key_size), 0,
 		"  key data");
-	TEST_EQ(sd->kernel_version_secdata, 0x20002,
+	TEST_EQ(sd->kernel_version_secdata, 0x10002,
 		"  secdata_kernel version");
 
 	/* Test flags for experimental features in non-recovery path */
@@ -249,22 +287,10 @@ static void phase1_tests(void)
 	/*
 	 * Test flags are unchanged for experimental features in recovery path
 	 */
-
-	/* Set 8 bits to 0 */
 	reset_common_data(FOR_PHASE1);
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_BROKEN_SCREEN, 123);
-	vb2_secdata_kernel_set(ctx, VB2_SECDATA_KERNEL_FLAGS, 0);
 	TEST_SUCC(vb2api_kernel_phase1(ctx), "phase1 rec good");
-	TEST_EQ(vb2_secdata_kernel_get(ctx, VB2_SECDATA_KERNEL_FLAGS), 0,
-		"VB2_SECDATA_KERNEL_FLAGS remains unchanged in recovery path");
-
-	/* Set 8 bits to 1 */
-	reset_common_data(FOR_PHASE1);
-	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_BROKEN_SCREEN, 123);
-	vb2_secdata_kernel_set(ctx, VB2_SECDATA_KERNEL_FLAGS, UINT8_MAX);
-	TEST_SUCC(vb2api_kernel_phase1(ctx), "phase1 rec good");
-	TEST_EQ(vb2_secdata_kernel_get(ctx, VB2_SECDATA_KERNEL_FLAGS),
-	        UINT8_MAX,
+	TEST_EQ(mock_kernel_flag_set, 0,
 		"VB2_SECDATA_KERNEL_FLAGS remains unchanged in recovery path");
 
 	/* Bad secdata_fwmp causes failure in normal mode only */
@@ -404,6 +430,34 @@ static void phase2_tests(void)
 
 static void finalize_tests(void)
 {
+	/* Kernel version roll forward */
+	reset_common_data(FOR_FINALIZE);
+	sd->kernel_version = 0x20003;
+	TEST_EQ(vb2api_kernel_finalize(ctx), VB2_SUCCESS,
+		"Kernel version roll forward");
+	TEST_EQ(mock_kernel_version, 0x20003, "  kernel version");
+
+	reset_common_data(FOR_FINALIZE);
+	vb2_nv_set(ctx, VB2_NV_FW_RESULT, VB2_FW_RESULT_TRYING);
+	sd->kernel_version = 0x20003;
+	TEST_EQ(vb2api_kernel_finalize(ctx), VB2_SUCCESS,
+		"Don't roll forward kernel when trying new FW");
+	TEST_EQ(mock_kernel_version, 0x10002, "  kernel version");
+
+	reset_common_data(FOR_FINALIZE);
+	vb2_nv_set(ctx, VB2_NV_KERNEL_MAX_ROLLFORWARD, 0x30005);
+	sd->kernel_version = 0x40006;
+	TEST_EQ(vb2api_kernel_finalize(ctx), VB2_SUCCESS,
+		"Limit max roll forward");
+	TEST_EQ(mock_kernel_version, 0x30005, "  kernel version");
+
+	reset_common_data(FOR_FINALIZE);
+	vb2_nv_set(ctx, VB2_NV_KERNEL_MAX_ROLLFORWARD, 0x10001);
+	sd->kernel_version = 0x40006;
+	TEST_EQ(vb2api_kernel_finalize(ctx), VB2_SUCCESS,
+		"Max roll forward can't rollback");
+	TEST_EQ(mock_kernel_version, 0x10002, "  kernel version");
+
 	/* NO_BOOT with EC sync support */
 	reset_common_data(FOR_FINALIZE);
 	ctx->flags |= VB2_CONTEXT_NO_BOOT;
diff --git a/tests/vboot_kernel2_tests.c b/tests/vb2_load_kernel2_tests.c
index 1f2f7a85..dfb06fbf 100644
--- a/tests/vboot_kernel2_tests.c
+++ b/tests/vb2_load_kernel2_tests.c
@@ -12,7 +12,6 @@
 #include "2secdata.h"
 #include "common/boot_mode.h"
 #include "common/tests.h"
-#include "load_kernel_fw.h"
 #include "vboot_api.h"
 
 #define MAX_MOCK_KERNELS 10
@@ -21,7 +20,7 @@
 /* Internal struct to simulate a stream for sector-based disks */
 struct disk_stream {
 	/* Disk handle */
-	VbExDiskHandle_t handle;
+	vb2ex_disk_handle_t handle;
 
 	/* Next sector to read */
 	uint64_t sector;
@@ -49,8 +48,8 @@ static struct vb2_workbuf wb;
 static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
 	__attribute__((aligned(VB2_WORKBUF_ALIGN)));
 
-static VbSelectAndLoadKernelParams lkp;
-static VbDiskInfo disk_info;
+static struct vb2_kernel_params lkp;
+static struct vb2_disk_info disk_info;
 static struct vb2_keyblock kbh;
 static struct vb2_kernel_preamble kph;
 static uint8_t kernel_buffer[80000];
@@ -90,12 +89,11 @@ static void reset_common_data(void)
 	memset(&lkp, 0, sizeof(lkp));
 	lkp.kernel_buffer = kernel_buffer;
 	lkp.kernel_buffer_size = sizeof(kernel_buffer);
-	lkp.disk_handle = (VbExDiskHandle_t)1;
 
 	memset(&disk_info, 0, sizeof(disk_info));
 	disk_info.bytes_per_lba = 512;
 	disk_info.lba_count = 1024;
-	disk_info.handle = lkp.disk_handle;
+	disk_info.handle = (vb2ex_disk_handle_t)1;
 
 	memset(&kbh, 0, sizeof(kbh));
 	kbh.data_key.key_version = 2;
@@ -120,7 +118,7 @@ static void reset_common_data(void)
 
 /* Mocks */
 
-vb2_error_t VbExStreamOpen(VbExDiskHandle_t handle, uint64_t lba_start,
+vb2_error_t VbExStreamOpen(vb2ex_disk_handle_t handle, uint64_t lba_start,
 			   uint64_t lba_count, VbExStream_t *stream)
 {
 	struct disk_stream *s;
@@ -258,38 +256,44 @@ static void load_minios_kernel_tests(void)
 	disk_info.bytes_per_lba = KBUF_SIZE;
 	disk_info.lba_count = 1;
 	add_mock_kernel(0, VB2_SUCCESS);
-	TEST_SUCC(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_SUCC(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		  "{valid kernel}");
+	TEST_PTR_EQ(lkp.disk_handle, disk_info.handle,
+		    "  fill disk_handle when success");
 
 	reset_common_data();
 	disk_info.bytes_per_lba = KBUF_SIZE;
 	disk_info.lba_count = 1;
-	TEST_EQ(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_EQ(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		VB2_ERROR_LK_NO_KERNEL_FOUND, "{no kernel}");
 
 	reset_common_data();
 	disk_info.bytes_per_lba = KBUF_SIZE;
 	disk_info.lba_count = 2;
 	add_mock_kernel(1, VB2_SUCCESS);
-	TEST_SUCC(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_SUCC(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		  "{no kernel, valid kernel}");
 	TEST_EQ(cur_kernel->sector, 1, "  select kernel");
+	TEST_PTR_EQ(lkp.disk_handle, disk_info.handle,
+		    "  fill disk_handle when success");
 
 	reset_common_data();
 	disk_info.bytes_per_lba = KBUF_SIZE;
 	disk_info.lba_count = 2;
 	add_mock_kernel(0, VB2_ERROR_MOCK);
 	add_mock_kernel(1, VB2_SUCCESS);
-	TEST_SUCC(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_SUCC(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		  "{invalid kernel, valid kernel}");
 	TEST_EQ(cur_kernel->sector, 1, "  select second kernel");
+	TEST_PTR_EQ(lkp.disk_handle, disk_info.handle,
+		    "  fill disk_handle when success");
 
 	reset_common_data();
 	disk_info.bytes_per_lba = KBUF_SIZE;
 	disk_info.lba_count = 2;
 	add_mock_kernel(0, VB2_ERROR_MOCK);
 	add_mock_kernel(1, VB2_ERROR_MOCK);
-	TEST_EQ(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_EQ(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		VB2_ERROR_LK_NO_KERNEL_FOUND,
 		"{invalid kernel, invalid kernel}");
 
@@ -298,9 +302,11 @@ static void load_minios_kernel_tests(void)
 	disk_info.lba_count = 2;
 	add_mock_kernel(0, VB2_SUCCESS);
 	add_mock_kernel(1, VB2_SUCCESS);
-	TEST_SUCC(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_SUCC(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		  "{valid kernel, valid kernel} minios_priority=0");
 	TEST_EQ(cur_kernel->sector, 0, "  select first kernel");
+	TEST_PTR_EQ(lkp.disk_handle, disk_info.handle,
+		    "  fill disk_handle when success");
 
 	reset_common_data();
 	disk_info.bytes_per_lba = KBUF_SIZE;
@@ -308,19 +314,23 @@ static void load_minios_kernel_tests(void)
 	add_mock_kernel(0, VB2_SUCCESS);
 	add_mock_kernel(1, VB2_SUCCESS);
 	vb2_nv_set(ctx, VB2_NV_MINIOS_PRIORITY, 1);
-	TEST_SUCC(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_SUCC(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		  "{valid kernel, valid kernel} minios_priority=1");
 	TEST_EQ(cur_kernel->sector, 1, "  select second kernel");
+	TEST_PTR_EQ(lkp.disk_handle, disk_info.handle,
+		    "  fill disk_handle when success");
 
 	reset_common_data();
 	disk_info.bytes_per_lba = KBUF_SIZE;
 	disk_info.lba_count = 2;
 	add_mock_kernel(0, VB2_SUCCESS);
 	add_mock_kernel(1, VB2_SUCCESS);
-	TEST_SUCC(LoadMiniOsKernel(ctx, &lkp, &disk_info,
-				   VB_MINIOS_FLAG_NON_ACTIVE),
+	TEST_SUCC(vb2api_load_minios_kernel(ctx, &lkp, &disk_info,
+					    VB2_MINIOS_FLAG_NON_ACTIVE),
 		  "{valid kernel, valid kernel} minios_priority=0 non-active");
 	TEST_EQ(cur_kernel->sector, 1, "  select second kernel");
+	TEST_PTR_EQ(lkp.disk_handle, disk_info.handle,
+		    "  fill disk_handle when success");
 
 	reset_common_data();
 	disk_info.bytes_per_lba = KBUF_SIZE;
@@ -328,8 +338,8 @@ static void load_minios_kernel_tests(void)
 	add_mock_kernel(0, VB2_ERROR_MOCK);
 	add_mock_kernel(1, VB2_SUCCESS);
 	vb2_nv_set(ctx, VB2_NV_MINIOS_PRIORITY, 1);
-	TEST_EQ(LoadMiniOsKernel(ctx, &lkp, &disk_info,
-				 VB_MINIOS_FLAG_NON_ACTIVE),
+	TEST_EQ(vb2api_load_minios_kernel(ctx, &lkp, &disk_info,
+					  VB2_MINIOS_FLAG_NON_ACTIVE),
 		VB2_ERROR_LK_NO_KERNEL_FOUND,
 		"{invalid kernel, valid kernel} minios_priority=1 non-active");
 
@@ -337,7 +347,7 @@ static void load_minios_kernel_tests(void)
 	disk_info.bytes_per_lba = VB2_KEYBLOCK_MAGIC_SIZE;
 	disk_info.lba_count = 4;
 	add_mock_kernel(1, VB2_SUCCESS);
-	TEST_EQ(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_EQ(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		VB2_ERROR_LK_NO_KERNEL_FOUND,
 		"valid kernel header near start of disk (disk too small)");
 
@@ -345,7 +355,7 @@ static void load_minios_kernel_tests(void)
 	disk_info.bytes_per_lba = VB2_KEYBLOCK_MAGIC_SIZE;
 	disk_info.lba_count = 1000;
 	add_mock_kernel(999, VB2_SUCCESS);
-	TEST_EQ(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_EQ(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		VB2_ERROR_LK_NO_KERNEL_FOUND,
 		"valid kernel header near end of disk");
 
@@ -353,36 +363,46 @@ static void load_minios_kernel_tests(void)
 	disk_info.bytes_per_lba = 1024;
 	disk_info.lba_count = 128;
 	add_mock_kernel(63, VB2_SUCCESS);
-	TEST_SUCC(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_SUCC(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		  "start/end overlap assuming >128 MB search range (start)");
+	TEST_PTR_EQ(lkp.disk_handle, disk_info.handle,
+		    "  fill disk_handle when success");
 
 	reset_common_data();
 	disk_info.bytes_per_lba = 1024;
 	disk_info.lba_count = 128;
 	add_mock_kernel(64, VB2_SUCCESS);
-	TEST_SUCC(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_SUCC(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		  "start/end overlap assuming >128 MB search range (end)");
+	TEST_PTR_EQ(lkp.disk_handle, disk_info.handle,
+		    "  fill disk_handle when success");
 
 	reset_common_data();
 	disk_info.bytes_per_lba = 128;
 	disk_info.lba_count = 1024;
 	add_mock_kernel(3, VB2_SUCCESS);
-	TEST_SUCC(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_SUCC(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		  "kernel at last sector in batch assuming 512 KB batches");
+	TEST_PTR_EQ(lkp.disk_handle, disk_info.handle,
+		    "  fill disk_handle when success");
 
 	reset_common_data();
 	disk_info.bytes_per_lba = 256;
 	disk_info.lba_count = 1024;
 	add_mock_kernel(3, VB2_SUCCESS);
-	TEST_SUCC(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_SUCC(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		  "kernel at last sector in batch assuming 1 MB batches");
+	TEST_PTR_EQ(lkp.disk_handle, disk_info.handle,
+		    "  fill disk_handle when success");
 
 	reset_common_data();
 	disk_info.bytes_per_lba = 512;
 	disk_info.lba_count = 1024;
 	add_mock_kernel(3, VB2_SUCCESS);
-	TEST_SUCC(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_SUCC(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		  "kernel at last sector in batch assuming 2 MB batches");
+	TEST_PTR_EQ(lkp.disk_handle, disk_info.handle,
+		    "  fill disk_handle when success");
 
 	reset_common_data();
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_DEVELOPER_0
@@ -391,8 +411,10 @@ static void load_minios_kernel_tests(void)
 	disk_info.bytes_per_lba = KBUF_SIZE;
 	disk_info.lba_count = 2;
 	add_mock_kernel(0, VB2_SUCCESS);
-	TEST_SUCC(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_SUCC(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		  "kernel with minios keyblock flag");
+	TEST_PTR_EQ(lkp.disk_handle, disk_info.handle,
+		    "  fill disk_handle when success");
 
 	reset_common_data();
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_DEVELOPER_0
@@ -401,7 +423,7 @@ static void load_minios_kernel_tests(void)
 	disk_info.bytes_per_lba = KBUF_SIZE;
 	disk_info.lba_count = 2;
 	add_mock_kernel(0, VB2_SUCCESS);
-	TEST_EQ(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_EQ(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		VB2_ERROR_LK_NO_KERNEL_FOUND,
 		"kernel with !minios keyblock flag");
 
@@ -411,7 +433,7 @@ static void load_minios_kernel_tests(void)
 	add_mock_kernel(0, VB2_SUCCESS);
 	sd->kernel_version_secdata = 5 << 24;
 	kph.kernel_version = 4;
-	TEST_EQ(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_EQ(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		VB2_ERROR_LK_NO_KERNEL_FOUND,
 		"kernel version too old");
 
@@ -421,7 +443,7 @@ static void load_minios_kernel_tests(void)
 	add_mock_kernel(0, VB2_SUCCESS);
 	sd->kernel_version_secdata = 5 << 24;
 	kph.kernel_version = 0x100;
-	TEST_EQ(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
+	TEST_EQ(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
 		VB2_ERROR_LK_NO_KERNEL_FOUND,
 		"kernel version greater than 0xff");
 
@@ -431,8 +453,10 @@ static void load_minios_kernel_tests(void)
 	add_mock_kernel(0, VB2_SUCCESS);
 	sd->kernel_version_secdata = 5 << 24;
 	kph.kernel_version = 6;
-	TEST_SUCC(LoadMiniOsKernel(ctx, &lkp, &disk_info, 0),
-		 "newer kernel version");
+	TEST_SUCC(vb2api_load_minios_kernel(ctx, &lkp, &disk_info, 0),
+		  "newer kernel version");
+	TEST_PTR_EQ(lkp.disk_handle, disk_info.handle,
+		    "  fill disk_handle when success");
 }
 
 int main(void)
diff --git a/tests/vboot_kernel_tests.c b/tests/vb2_load_kernel_tests.c
index e4162fbf..bb6f3487 100644
--- a/tests/vboot_kernel_tests.c
+++ b/tests/vb2_load_kernel_tests.c
@@ -16,7 +16,6 @@
 #include "common/boot_mode.h"
 #include "common/tests.h"
 #include "gpt.h"
-#include "load_kernel_fw.h"
 #include "vboot_api.h"
 
 /* Mock kernel partition */
@@ -41,8 +40,8 @@ static int unpack_key_fail;
 static int gpt_flag_external;
 
 static struct vb2_gbb_header gbb;
-static VbSelectAndLoadKernelParams lkp;
-static VbDiskInfo disk_info;
+static struct vb2_kernel_params lkp;
+static struct vb2_disk_info disk_info;
 static struct vb2_keyblock kbh;
 static struct vb2_kernel_preamble kph;
 static struct vb2_secdata_fwmp *fwmp;
@@ -76,13 +75,12 @@ static void ResetMocks(void)
 	memset(&lkp, 0, sizeof(lkp));
 	lkp.kernel_buffer = kernel_buffer;
 	lkp.kernel_buffer_size = sizeof(kernel_buffer);
-	lkp.disk_handle = (VbExDiskHandle_t)1;
 
 	memset(&disk_info, 0, sizeof(disk_info));
 	disk_info.bytes_per_lba = 512;
 	disk_info.streaming_lba_count = 1024;
 	disk_info.lba_count = 1024;
-	disk_info.handle = lkp.disk_handle;
+	disk_info.handle = (vb2ex_disk_handle_t)1;
 
 	memset(&kbh, 0, sizeof(kbh));
 	kbh.data_key.key_version = 2;
@@ -155,7 +153,7 @@ vb2_error_t vb2_gbb_read_recovery_key(struct vb2_context *c,
 	return VB2_SUCCESS;
 }
 
-vb2_error_t VbExDiskRead(VbExDiskHandle_t h, uint64_t lba_start,
+vb2_error_t VbExDiskRead(vb2ex_disk_handle_t h, uint64_t lba_start,
 			 uint64_t lba_count, void *buffer)
 {
 	if ((int)lba_start == disk_read_to_fail)
@@ -164,7 +162,7 @@ vb2_error_t VbExDiskRead(VbExDiskHandle_t h, uint64_t lba_start,
 	return VB2_SUCCESS;
 }
 
-int AllocAndReadGptData(VbExDiskHandle_t disk_handle, GptData *gptdata)
+int AllocAndReadGptData(vb2ex_disk_handle_t disk_handle, GptData *gptdata)
 {
 	return GPT_SUCCESS;
 }
@@ -196,7 +194,7 @@ int GptUpdateKernelEntry(GptData *gpt, uint32_t update_type)
 	return GPT_SUCCESS;
 }
 
-int WriteAndFreeGptData(VbExDiskHandle_t disk_handle, GptData *gptdata)
+int WriteAndFreeGptData(vb2ex_disk_handle_t disk_handle, GptData *gptdata)
 {
 	return GPT_SUCCESS;
 }
@@ -278,31 +276,34 @@ void vb2api_fail(struct vb2_context *c, uint8_t reason, uint8_t subcode)
 	TEST_TRUE(0, "  called vb2api_fail()");
 }
 
-static void TestLoadKernel(int expect_retval, const char *test_name)
+static void test_load_kernel(vb2_error_t expect_retval, const char *test_name)
 {
-	TEST_EQ(LoadKernel(ctx, &lkp, &disk_info), expect_retval, test_name);
+	TEST_EQ(vb2api_load_kernel(ctx, &lkp, &disk_info), expect_retval,
+		test_name);
+	if (expect_retval == VB2_SUCCESS)
+		TEST_PTR_EQ(lkp.disk_handle, disk_info.handle,
+			    "  fill disk_handle when success");
 }
 
 /**
- * Trivial invalid calls to LoadKernel()
+ * Trivial invalid calls to vb2api_load_kernel()
  */
-static void InvalidParamsTest(void)
+static void invalid_params_tests(void)
 {
 	ResetMocks();
 	gpt_init_fail = 1;
-	TestLoadKernel(VB2_ERROR_LK_NO_KERNEL_FOUND, "Bad GPT");
+	test_load_kernel(VB2_ERROR_LK_NO_KERNEL_FOUND, "Bad GPT");
 
 	/* This causes the stream open call to fail */
 	ResetMocks();
-	lkp.disk_handle = NULL;
 	disk_info.handle = NULL;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND, "Bad disk handle");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND, "Bad disk handle");
 }
 
-static void LoadKernelTest(void)
+static void load_kernel_tests(void)
 {
 	ResetMocks();
-	TestLoadKernel(0, "First kernel good");
+	test_load_kernel(VB2_SUCCESS, "First kernel good");
 	TEST_EQ(lkp.partition_number, 1, "  part num");
 	TEST_EQ(lkp.bootloader_address, 0xbeadd008, "  bootloader addr");
 	TEST_EQ(lkp.bootloader_size, 0x1234, "  bootloader size");
@@ -313,42 +314,42 @@ static void LoadKernelTest(void)
 	ResetMocks();
 	mock_parts[1].start = 300;
 	mock_parts[1].size = 150;
-	TestLoadKernel(0, "Two good kernels");
+	test_load_kernel(VB2_SUCCESS, "Two good kernels");
 	TEST_EQ(lkp.partition_number, 1, "  part num");
 	TEST_EQ(mock_part_next, 1, "  didn't read second one");
 
 	/* Fail if no kernels found */
 	ResetMocks();
 	mock_parts[0].size = 0;
-	TestLoadKernel(VB2_ERROR_LK_NO_KERNEL_FOUND, "No kernels");
+	test_load_kernel(VB2_ERROR_LK_NO_KERNEL_FOUND, "No kernels");
 
 	/* Skip kernels which are too small */
 	ResetMocks();
 	mock_parts[0].size = 10;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND, "Too small");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND, "Too small");
 
 	ResetMocks();
 	disk_read_to_fail = 100;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Fail reading kernel start");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Fail reading kernel start");
 
 	ResetMocks();
 	keyblock_verify_fail = 1;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Fail key block sig");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Fail key block sig");
 
 	/* In dev mode, fail if hash is bad too */
 	ResetMocks();
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
 	keyblock_verify_fail = 2;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Fail key block dev hash");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Fail key block dev hash");
 
 	/* But just bad sig is ok */
 	ResetMocks();
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
 	keyblock_verify_fail = 1;
-	TestLoadKernel(0, "Succeed keyblock dev sig");
+	test_load_kernel(VB2_SUCCESS, "Succeed keyblock dev sig");
 	TEST_EQ(sd->flags & VB2_SD_FLAG_KERNEL_SIGNED, 0, "  use hash");
 
 	/* In dev mode and requiring signed kernel, fail if sig is bad */
@@ -356,37 +357,37 @@ static void LoadKernelTest(void)
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
 	vb2_nv_set(ctx, VB2_NV_DEV_BOOT_SIGNED_ONLY, 1);
 	keyblock_verify_fail = 1;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Fail key block dev sig");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Fail key block dev sig");
 
 	ResetMocks();
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
 	fwmp->flags |= VB2_SECDATA_FWMP_DEV_ENABLE_OFFICIAL_ONLY;
 	keyblock_verify_fail = 1;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Fail key block dev sig fwmp");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Fail key block dev sig fwmp");
 
 	/* Check keyblock flags */
 	ResetMocks();
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_0
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_1
 		| VB2_KEYBLOCK_FLAG_MINIOS_0;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Keyblock dev flag mismatch");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Keyblock dev flag mismatch");
 
 	ResetMocks();
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_1
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_0
 		| VB2_KEYBLOCK_FLAG_MINIOS_0;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Keyblock rec flag mismatch");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Keyblock rec flag mismatch");
 
 	ResetMocks();
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_0
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_0
 		| VB2_KEYBLOCK_FLAG_MINIOS_1;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Keyblock minios flag mismatch");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Keyblock minios flag mismatch");
 
 	ResetMocks();
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_MANUAL_RECOVERY,
@@ -394,8 +395,8 @@ static void LoadKernelTest(void)
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_1
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_1
 		| VB2_KEYBLOCK_FLAG_MINIOS_0;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Keyblock recdev flag mismatch");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Keyblock recdev flag mismatch");
 
 	ResetMocks();
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_MANUAL_RECOVERY,
@@ -403,7 +404,7 @@ static void LoadKernelTest(void)
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_1
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_0
 		| VB2_KEYBLOCK_FLAG_MINIOS_0;
-	TestLoadKernel(0, "Keyblock rec flag okay");
+	test_load_kernel(VB2_SUCCESS, "Keyblock rec flag okay");
 
 	ResetMocks();
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_MANUAL_RECOVERY,
@@ -412,8 +413,8 @@ static void LoadKernelTest(void)
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_1
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_0
 		| VB2_KEYBLOCK_FLAG_MINIOS_0;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Keyblock rec!dev flag mismatch");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Keyblock rec!dev flag mismatch");
 
 	ResetMocks();
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_MANUAL_RECOVERY,
@@ -422,7 +423,7 @@ static void LoadKernelTest(void)
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_1
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_1
 		| VB2_KEYBLOCK_FLAG_MINIOS_0;
-	TestLoadKernel(0, "Keyblock recdev flag okay");
+	test_load_kernel(VB2_SUCCESS, "Keyblock recdev flag okay");
 
 	/* Check keyblock flags (dev mode + signed kernel required) */
 	ResetMocks();
@@ -431,8 +432,8 @@ static void LoadKernelTest(void)
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_1
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_0
 		| VB2_KEYBLOCK_FLAG_MINIOS_0;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Keyblock dev flag mismatch (signed kernel required)");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Keyblock dev flag mismatch (signed kernel required)");
 
 	ResetMocks();
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
@@ -440,8 +441,8 @@ static void LoadKernelTest(void)
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_1
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_0
 		| VB2_KEYBLOCK_FLAG_MINIOS_0;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Keyblock dev flag mismatch (signed kernel required)");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Keyblock dev flag mismatch (signed kernel required)");
 
 	ResetMocks();
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
@@ -449,8 +450,8 @@ static void LoadKernelTest(void)
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_0
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_0
 		| VB2_KEYBLOCK_FLAG_MINIOS_1;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Keyblock dev flag mismatch (signed kernel required)");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Keyblock dev flag mismatch (signed kernel required)");
 
 	ResetMocks();
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
@@ -458,91 +459,92 @@ static void LoadKernelTest(void)
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_0
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_1
 		| VB2_KEYBLOCK_FLAG_MINIOS_0;
-	TestLoadKernel(0, "Keyblock dev flag okay (signed kernel required)");
+	test_load_kernel(VB2_SUCCESS,
+			 "Keyblock dev flag okay (signed kernel required)");
 
 	/* Check kernel key version */
 	ResetMocks();
 	kbh.data_key.key_version = 1;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Keyblock kernel key rollback");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Keyblock kernel key rollback");
 
 	ResetMocks();
 	kbh.data_key.key_version = 0x10000;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Keyblock kernel key version too big");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Keyblock kernel key version too big");
 
 	ResetMocks();
 	kbh.data_key.key_version = 3;
-	TestLoadKernel(0, "Keyblock version roll forward");
+	test_load_kernel(VB2_SUCCESS, "Keyblock version roll forward");
 	TEST_EQ(sd->kernel_version, 0x30001, "  SD version");
 
 	ResetMocks();
 	kbh.data_key.key_version = 3;
 	mock_parts[1].start = 300;
 	mock_parts[1].size = 150;
-	TestLoadKernel(0, "Two kernels roll forward");
+	test_load_kernel(VB2_SUCCESS, "Two kernels roll forward");
 	TEST_EQ(mock_part_next, 2, "  read both");
 	TEST_EQ(sd->kernel_version, 0x30001, "  SD version");
 
 	ResetMocks();
 	kbh.data_key.key_version = 1;
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
-	TestLoadKernel(0, "Key version ignored in dev mode");
+	test_load_kernel(VB2_SUCCESS, "Key version ignored in dev mode");
 
 	ResetMocks();
 	kbh.data_key.key_version = 1;
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_MANUAL_RECOVERY,
 		      VB2_RECOVERY_RO_MANUAL);
-	TestLoadKernel(0, "Key version ignored in rec mode");
+	test_load_kernel(VB2_SUCCESS, "Key version ignored in rec mode");
 
 	ResetMocks();
 	unpack_key_fail = 2;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND, "Bad data key");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND, "Bad data key");
 
 	ResetMocks();
 	preamble_verify_fail = 1;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND, "Bad preamble");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND, "Bad preamble");
 
 	ResetMocks();
 	kph.kernel_version = 0;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Kernel version rollback");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Kernel version rollback");
 
 	ResetMocks();
 	kph.kernel_version = 0;
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
-	TestLoadKernel(0, "Kernel version ignored in dev mode");
+	test_load_kernel(VB2_SUCCESS, "Kernel version ignored in dev mode");
 
 	ResetMocks();
 	kph.kernel_version = 0;
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_MANUAL_RECOVERY,
 		      VB2_RECOVERY_RO_MANUAL);
-	TestLoadKernel(0, "Kernel version ignored in rec mode");
+	test_load_kernel(VB2_SUCCESS, "Kernel version ignored in rec mode");
 
 	/* Check kernel version (dev mode + signed kernel required) */
 	ResetMocks();
 	kbh.data_key.key_version = 0;
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
 	vb2_nv_set(ctx, VB2_NV_DEV_BOOT_SIGNED_ONLY, 1);
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Keyblock key version checked in dev mode "
-		       "(signed kernel required)");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Keyblock key version checked in dev mode "
+			 "(signed kernel required)");
 
 	ResetMocks();
 	kbh.data_key.key_version = 0;
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
 	fwmp->flags |= VB2_SECDATA_FWMP_DEV_ENABLE_OFFICIAL_ONLY;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Keyblock key version checked in dev mode "
-		       "(signed kernel required)");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Keyblock key version checked in dev mode "
+			 "(signed kernel required)");
 
 	/* Check developer key hash - bad */
 	ResetMocks();
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
 	fwmp->flags |= VB2_SECDATA_FWMP_DEV_USE_KEY_HASH;
 	fwmp->dev_key_hash[0]++;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Fail keyblock dev fwmp hash");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Fail keyblock dev fwmp hash");
 
 	/* Check developer key hash - bad (recovery mode) */
 	ResetMocks();
@@ -550,72 +552,73 @@ static void LoadKernelTest(void)
 		      VB2_RECOVERY_RO_MANUAL);
 	fwmp->flags |= VB2_SECDATA_FWMP_DEV_USE_KEY_HASH;
 	fwmp->dev_key_hash[0]++;
-	TestLoadKernel(0, "Bad keyblock dev fwmp hash ignored in rec mode");
+	test_load_kernel(VB2_SUCCESS,
+			 "Bad keyblock dev fwmp hash ignored in rec mode");
 
 	/* Check developer key hash - good */
 	ResetMocks();
 	SET_BOOT_MODE(ctx, VB2_BOOT_MODE_DEVELOPER);
 	fwmp->flags |= VB2_SECDATA_FWMP_DEV_USE_KEY_HASH;
-	TestLoadKernel(0, "Good keyblock dev fwmp hash");
+	test_load_kernel(VB2_SUCCESS, "Good keyblock dev fwmp hash");
 
 	ResetMocks();
 	kph.preamble_size |= 0x07;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Kernel body offset");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Kernel body offset");
 
 	ResetMocks();
 	kph.preamble_size += 65536;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Kernel body offset huge");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Kernel body offset huge");
 
 	/* Check getting kernel load address from header */
 	ResetMocks();
 	kph.body_load_address = (size_t)kernel_buffer;
 	lkp.kernel_buffer = NULL;
-	TestLoadKernel(0, "Get load address from preamble");
+	test_load_kernel(VB2_SUCCESS, "Get load address from preamble");
 	TEST_PTR_EQ(lkp.kernel_buffer, kernel_buffer, "  address");
 	/* Size is rounded up to nearest sector */
 	TEST_EQ(lkp.kernel_buffer_size, 70144, "  size");
 
 	ResetMocks();
 	lkp.kernel_buffer_size = 8192;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Kernel too big for buffer");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Kernel too big for buffer");
 
 	ResetMocks();
 	mock_parts[0].size = 130;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Kernel too big for partition");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Kernel too big for partition");
 
 	ResetMocks();
 	kph.body_signature.data_size = 8192;
-	TestLoadKernel(0, "Kernel tiny");
+	test_load_kernel(VB2_SUCCESS, "Kernel tiny");
 
 	ResetMocks();
 	disk_read_to_fail = 228;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-		       "Fail reading kernel data");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+			 "Fail reading kernel data");
 
 	ResetMocks();
 	verify_data_fail = 1;
-	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND, "Bad data");
+	test_load_kernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND, "Bad data");
 
 	/* Check that EXTERNAL_GPT flag makes it down */
 	ResetMocks();
-	disk_info.flags |= VB_DISK_FLAG_EXTERNAL_GPT;
-	TestLoadKernel(0, "Succeed external GPT");
+	disk_info.flags |= VB2_DISK_FLAG_EXTERNAL_GPT;
+	test_load_kernel(VB2_SUCCESS, "Succeed external GPT");
 	TEST_EQ(gpt_flag_external, 1, "GPT was external");
 
 	/* Check recovery from unreadble primary GPT */
 	ResetMocks();
 	disk_read_to_fail = 1;
-	TestLoadKernel(0, "Can't read disk");
+	test_load_kernel(VB2_SUCCESS, "Can't read disk");
 }
 
 int main(void)
 {
-	InvalidParamsTest();
-	LoadKernelTest();
+	invalid_params_tests();
+	load_kernel_tests();
 
 	return gTestSuccess ? 0 : 255;
 }
diff --git a/tests/vboot_api_kernel_tests.c b/tests/vboot_api_kernel_tests.c
deleted file mode 100644
index 5c1d0182..00000000
--- a/tests/vboot_api_kernel_tests.c
+++ /dev/null
@@ -1,644 +0,0 @@
-/* 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.
- *
- * Tests for VbTryLoadKernel()
- */
-
-#include "2common.h"
-#include "2misc.h"
-#include "2nvstorage.h"
-#include "2sysincludes.h"
-#include "common/tests.h"
-#include "load_kernel_fw.h"
-#include "vboot_api.h"
-
-#define MAX_TEST_DISKS 10
-#define DEFAULT_COUNT -1
-
-typedef struct {
-	uint64_t bytes_per_lba;
-	uint64_t lba_count;
-	uint32_t flags;
-	const char *diskname;
-} disk_desc_t;
-
-typedef struct {
-	const char *name;
-
-	/* inputs for test case */
-	uint32_t ctx_flags;
-	uint32_t want_flags;
-	vb2_error_t diskgetinfo_return_val;
-	disk_desc_t disks_to_provide[MAX_TEST_DISKS];
-	int disk_count_to_return;
-	vb2_error_t loadkernel_return_val[MAX_TEST_DISKS];
-	uint8_t external_expected[MAX_TEST_DISKS];
-
-	/* outputs from test */
-	uint32_t expected_recovery_request_val;
-	const char *expected_to_find_disk;
-	const char *expected_to_load_disk;
-	uint32_t expected_return_val;
-
-} test_case_t;
-
-/****************************************************************************/
-/* Test cases */
-
-static const char pickme[] = "correct choice";
-#define DONT_CARE ((const char *)42)
-
-test_case_t normal_tests[] = {
-	{
-		.name = "first drive (removable)",
-		.ctx_flags = 0,
-		.want_flags = VB_DISK_FLAG_REMOVABLE | VB_DISK_FLAG_FIXED,
-		.disks_to_provide = {
-			{4096, 100, VB_DISK_FLAG_REMOVABLE, pickme},
-			{4096, 100, VB_DISK_FLAG_FIXED, "holygrail"},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {0},
-		.external_expected = {0},
-
-		.expected_recovery_request_val = VB2_RECOVERY_NOT_REQUESTED,
-		.expected_to_find_disk = pickme,
-		.expected_to_load_disk = pickme,
-		.expected_return_val = VB2_SUCCESS
-	},
-	{
-		.name = "first drive (fixed)",
-		.ctx_flags = 0,
-		.want_flags = VB_DISK_FLAG_REMOVABLE | VB_DISK_FLAG_FIXED,
-		.disks_to_provide = {
-			{4096, 100, VB_DISK_FLAG_FIXED, pickme},
-			{4096, 100, VB_DISK_FLAG_REMOVABLE, "holygrail"},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {0},
-		.external_expected = {0},
-
-		.expected_recovery_request_val = VB2_RECOVERY_NOT_REQUESTED,
-		.expected_to_find_disk = pickme,
-		.expected_to_load_disk = pickme,
-		.expected_return_val = VB2_SUCCESS
-	},
-	{
-		.name = "first removable drive",
-		.ctx_flags = 0,
-		.want_flags = VB_DISK_FLAG_REMOVABLE,
-		.disks_to_provide = {
-			/* too small */
-			{512,  10,   VB_DISK_FLAG_REMOVABLE, 0},
-			/* wrong LBA */
-			{511, 100,  VB_DISK_FLAG_REMOVABLE, 0},
-			/* not a power of 2 */
-			{2047, 100,  VB_DISK_FLAG_REMOVABLE, 0},
-			/* wrong type */
-			{512,  100,  VB_DISK_FLAG_FIXED, 0},
-			/* wrong flags */
-			{512,  100,  0, 0},
-			/* still wrong flags */
-			{512,  100,  -1, 0},
-			{4096, 100,  VB_DISK_FLAG_REMOVABLE, pickme},
-			/* already got one */
-			{512,  100,  VB_DISK_FLAG_REMOVABLE, "holygrail"},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {0},
-		.external_expected = {0},
-
-		.expected_recovery_request_val = VB2_RECOVERY_NOT_REQUESTED,
-		.expected_to_find_disk = pickme,
-		.expected_to_load_disk = pickme,
-		.expected_return_val = VB2_SUCCESS
-	},
-	{
-		.name = "first removable drive (skip external GPT)",
-		.ctx_flags = 0,
-		.want_flags = VB_DISK_FLAG_REMOVABLE,
-		.disks_to_provide = {
-			/* too small */
-			{512,  10,   VB_DISK_FLAG_REMOVABLE, 0},
-			/* wrong LBA */
-			{511, 100,  VB_DISK_FLAG_REMOVABLE, 0},
-			/* not a power of 2 */
-			{2047, 100,  VB_DISK_FLAG_REMOVABLE, 0},
-			/* wrong type */
-			{512,  100,  VB_DISK_FLAG_FIXED, 0},
-			/* wrong flags */
-			{512,  100,  0, 0},
-			/* still wrong flags */
-			{512,  100,  -1, 0},
-			{512,  100,
-			 VB_DISK_FLAG_REMOVABLE | VB_DISK_FLAG_EXTERNAL_GPT,
-			 pickme},
-			/* already got one */
-			{512,  100,  VB_DISK_FLAG_REMOVABLE, "holygrail"},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {0, 0},
-		.external_expected = {1, 0},
-
-		.expected_recovery_request_val = VB2_RECOVERY_NOT_REQUESTED,
-		.expected_to_find_disk = pickme,
-		.expected_to_load_disk = pickme,
-		.expected_return_val = VB2_SUCCESS
-	},
-	{
-		.name = "second removable drive",
-		.ctx_flags = 0,
-		.want_flags = VB_DISK_FLAG_REMOVABLE,
-		.disks_to_provide = {
-			/* wrong flags */
-			{512,  100,  0, 0},
-			{512,  100,  VB_DISK_FLAG_REMOVABLE, "not yet"},
-			{512,  100,  VB_DISK_FLAG_REMOVABLE, pickme},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {VB2_ERROR_LK_INVALID_KERNEL_FOUND, 0},
-
-		.expected_recovery_request_val = VB2_RECOVERY_NOT_REQUESTED,
-		.expected_to_find_disk = pickme,
-		.expected_to_load_disk = pickme,
-		.expected_return_val = VB2_SUCCESS
-	},
-	{
-		.name = "first fixed drive",
-		.ctx_flags = 0,
-		.want_flags = VB_DISK_FLAG_FIXED,
-		.disks_to_provide = {
-			/* too small */
-			{512,   10,  VB_DISK_FLAG_FIXED, 0},
-			/* wrong LBA */
-			{511, 100,  VB_DISK_FLAG_FIXED, 0},
-			/* not a power of 2 */
-			{2047, 100,  VB_DISK_FLAG_REMOVABLE, 0},
-			/* wrong type */
-			{512,  100,  VB_DISK_FLAG_REMOVABLE, 0},
-			/* wrong flags */
-			{512,  100,  0, 0},
-			/* still wrong flags */
-			{512,  100,  -1, 0},
-			/* flags */
-			{512,  100,  VB_DISK_FLAG_REMOVABLE|VB_DISK_FLAG_FIXED,
-			 0},
-			{512,  100,  VB_DISK_FLAG_FIXED, pickme},
-			/* already got one */
-			{512,  100,  VB_DISK_FLAG_FIXED, "holygrail"},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {0},
-
-		.expected_recovery_request_val = VB2_RECOVERY_NOT_REQUESTED,
-		.expected_to_find_disk = pickme,
-		.expected_to_load_disk = pickme,
-		.expected_return_val = VB2_SUCCESS
-	},
-	{
-		.name = "no drives at all",
-		.ctx_flags = 0,
-		.want_flags = VB_DISK_FLAG_FIXED,
-		.disks_to_provide = {},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-
-		.expected_recovery_request_val = VB2_RECOVERY_RW_NO_DISK,
-		.expected_to_find_disk = 0,
-		.expected_to_load_disk = 0,
-		.expected_return_val = VB2_ERROR_LK_NO_DISK_FOUND
-	},
-	{
-		.name = "VbExDiskGetInfo() error",
-		.ctx_flags = 0,
-		.want_flags = VB_DISK_FLAG_FIXED,
-		.disks_to_provide = {
-			{512,  10, VB_DISK_FLAG_REMOVABLE, 0},
-			{512, 100, VB_DISK_FLAG_FIXED, 0},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_ERROR_UNKNOWN,
-
-		.expected_recovery_request_val = VB2_RECOVERY_RW_NO_DISK,
-		.expected_to_find_disk = 0,
-		.expected_to_load_disk = 0,
-		.expected_return_val = VB2_ERROR_LK_NO_DISK_FOUND,
-	},
-	{
-		.name = "invalid kernel",
-		.ctx_flags = 0,
-		.want_flags = VB_DISK_FLAG_FIXED,
-		.disks_to_provide = {
-			/* too small */
-			{512,   10,  VB_DISK_FLAG_FIXED, 0},
-			/* wrong LBA */
-			{511, 100,  VB_DISK_FLAG_FIXED, 0},
-			/* not a power of 2 */
-			{2047, 100,  VB_DISK_FLAG_FIXED, 0},
-			/* wrong type */
-			{512,  100,  VB_DISK_FLAG_REMOVABLE, 0},
-			/* wrong flags */
-			{512,  100,  0, 0},
-			/* still wrong flags */
-			{512,  100,  -1, 0},
-			/* doesn't load */
-			{512,  100,  VB_DISK_FLAG_FIXED, "corrupted kernel"},
-			/* doesn't load */
-			{512,  100,  VB_DISK_FLAG_FIXED, "stateful partition"},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-					  VB2_ERROR_LK_NO_KERNEL_FOUND},
-
-		.expected_recovery_request_val = VB2_RECOVERY_RW_INVALID_OS,
-		.expected_to_find_disk = DONT_CARE,
-		.expected_to_load_disk = 0,
-		.expected_return_val = VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-	},
-	{
-		.name = "invalid kernel, order flipped",
-		.ctx_flags = 0,
-		.want_flags = VB_DISK_FLAG_FIXED,
-		.disks_to_provide = {
-			{512, 1000, VB_DISK_FLAG_FIXED, "stateful partition"},
-			{512, 1000, VB_DISK_FLAG_FIXED, "corrupted kernel"},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {VB2_ERROR_LK_NO_KERNEL_FOUND,
-					  VB2_ERROR_LK_INVALID_KERNEL_FOUND},
-
-		.expected_recovery_request_val = VB2_RECOVERY_RW_INVALID_OS,
-		.expected_to_find_disk = DONT_CARE,
-		.expected_to_load_disk = 0,
-		.expected_return_val = VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-	},
-	{
-		.name = "no Chrome OS partitions",
-		.ctx_flags = 0,
-		.want_flags = VB_DISK_FLAG_FIXED,
-		.disks_to_provide = {
-			{512, 100, VB_DISK_FLAG_FIXED, "stateful partition"},
-			{512, 1000, VB_DISK_FLAG_FIXED, "Chrubuntu"},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {VB2_ERROR_LK_NO_KERNEL_FOUND,
-					  VB2_ERROR_LK_NO_KERNEL_FOUND},
-
-		.expected_recovery_request_val = VB2_RECOVERY_RW_NO_KERNEL,
-		.expected_to_find_disk = DONT_CARE,
-		.expected_to_load_disk = 0,
-		.expected_return_val = VB2_ERROR_LK_NO_KERNEL_FOUND,
-	},
-	{
-		.name = "invalid kernel (removable)",
-		.ctx_flags = 0,
-		.want_flags = VB_DISK_FLAG_REMOVABLE,
-		.disks_to_provide = {
-			{512,  100,  VB_DISK_FLAG_REMOVABLE, "corrupted"},
-			{512,  100,  VB_DISK_FLAG_REMOVABLE, "data"},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-					  VB2_ERROR_LK_NO_KERNEL_FOUND},
-
-		.expected_recovery_request_val = VB2_RECOVERY_RW_INVALID_OS,
-		.expected_to_find_disk = DONT_CARE,
-		.expected_to_load_disk = 0,
-		.expected_return_val = VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-	},
-	{
-		.name = "invalid kernel (removable, rec mode)",
-		.ctx_flags = VB2_CONTEXT_RECOVERY_MODE,
-		.want_flags = VB_DISK_FLAG_REMOVABLE,
-		.disks_to_provide = {
-			{512,  100,  VB_DISK_FLAG_REMOVABLE, "corrupted"},
-			{512,  100,  VB_DISK_FLAG_REMOVABLE, "data"},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-					  VB2_ERROR_LK_NO_KERNEL_FOUND},
-
-		.expected_recovery_request_val = VB2_RECOVERY_NOT_REQUESTED,
-		.expected_to_find_disk = DONT_CARE,
-		.expected_to_load_disk = 0,
-		.expected_return_val = VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-	},
-	{
-		.name = "invalid kernel (removable, dev mode)",
-		.ctx_flags = VB2_CONTEXT_DEVELOPER_MODE,
-		.want_flags = VB_DISK_FLAG_REMOVABLE,
-		.disks_to_provide = {
-			{512,  100,  VB_DISK_FLAG_REMOVABLE, "corrupted"},
-			{512,  100,  VB_DISK_FLAG_REMOVABLE, "data"},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-					  VB2_ERROR_LK_NO_KERNEL_FOUND},
-
-		.expected_recovery_request_val = VB2_RECOVERY_NOT_REQUESTED,
-		.expected_to_find_disk = DONT_CARE,
-		.expected_to_load_disk = 0,
-		.expected_return_val = VB2_ERROR_LK_INVALID_KERNEL_FOUND,
-	},
-	{
-		.name = "no kernel (removable)",
-		.ctx_flags = 0,
-		.want_flags = VB_DISK_FLAG_REMOVABLE,
-		.disks_to_provide = {
-			{512,  100,  VB_DISK_FLAG_REMOVABLE, "data"},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {VB2_ERROR_LK_NO_KERNEL_FOUND},
-
-		.expected_recovery_request_val = VB2_RECOVERY_RW_NO_KERNEL,
-		.expected_to_find_disk = DONT_CARE,
-		.expected_to_load_disk = 0,
-		.expected_return_val = VB2_ERROR_LK_NO_KERNEL_FOUND,
-	},
-	{
-		.name = "no kernel (removable, rec mode)",
-		.ctx_flags = VB2_CONTEXT_RECOVERY_MODE,
-		.want_flags = VB_DISK_FLAG_REMOVABLE,
-		.disks_to_provide = {
-			{512,  100,  VB_DISK_FLAG_REMOVABLE, "data"},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {VB2_ERROR_LK_NO_KERNEL_FOUND},
-
-		.expected_recovery_request_val = VB2_RECOVERY_NOT_REQUESTED,
-		.expected_to_find_disk = DONT_CARE,
-		.expected_to_load_disk = 0,
-		.expected_return_val = VB2_ERROR_LK_NO_KERNEL_FOUND,
-	},
-};
-
-test_case_t minios_tests[] = {
-	{
-		.name = "pick first fixed drive",
-		.ctx_flags = 0,
-		.disks_to_provide = {
-			{4096, 100, VB_DISK_FLAG_REMOVABLE, 0},
-			{4096, 100, VB_DISK_FLAG_FIXED, pickme},
-			{4096, 100, VB_DISK_FLAG_FIXED, "holygrail"},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {0},
-		.external_expected = {0},
-
-		.expected_recovery_request_val = VB2_RECOVERY_NOT_REQUESTED,
-		.expected_to_find_disk = pickme,
-		.expected_to_load_disk = pickme,
-		.expected_return_val = VB2_SUCCESS
-	},
-	{
-		.name = "skip failed fixed drive",
-		.ctx_flags = 0,
-		.disks_to_provide = {
-			{4096, 100, VB_DISK_FLAG_FIXED, "holygrail"},
-			{4096, 100, VB_DISK_FLAG_FIXED, pickme},
-		},
-		.disk_count_to_return = DEFAULT_COUNT,
-		.diskgetinfo_return_val = VB2_SUCCESS,
-		.loadkernel_return_val = {VB2_ERROR_LK_INVALID_KERNEL_FOUND, 0},
-		.external_expected = {0},
-
-		.expected_recovery_request_val = VB2_RECOVERY_NOT_REQUESTED,
-		.expected_to_find_disk = pickme,
-		.expected_to_load_disk = pickme,
-		.expected_return_val = VB2_SUCCESS
-	},
-};
-
-/****************************************************************************/
-
-/* Mock data */
-static VbDiskInfo mock_disks[MAX_TEST_DISKS];
-static test_case_t *t;
-static int load_kernel_calls;
-static int lk_normal_calls;
-static int lk_minios_calls;
-static uint32_t got_recovery_request_val;
-static const char *got_find_disk;
-static const char *got_load_disk;
-static uint32_t got_return_val;
-static uint32_t got_external_mismatch;
-static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
-	__attribute__((aligned(VB2_WORKBUF_ALIGN)));
-static struct vb2_context *ctx;
-static struct VbSelectAndLoadKernelParams kparams;
-
-/**
- * Reset mock data (for use before each test)
- */
-static void ResetMocks(test_case_t *test_case)
-{
-	TEST_SUCC(vb2api_init(workbuf, sizeof(workbuf), &ctx),
-		  "vb2api_init failed");
-
-	memset(&kparams, 0, sizeof(kparams));
-
-
-	memset(&mock_disks, 0, sizeof(mock_disks));
-	load_kernel_calls = 0;
-	lk_normal_calls = 0;
-	lk_minios_calls = 0;
-
-	got_recovery_request_val = VB2_RECOVERY_NOT_REQUESTED;
-	got_find_disk = 0;
-	got_load_disk = 0;
-	got_return_val = 0xdeadbeef;
-
-	t = test_case;
-}
-
-static int is_nonzero(const void *vptr, size_t count)
-{
-	const char *p = (const char *)vptr;
-	while (count--)
-		if (*p++)
-			return 1;
-
-	return 0;
-}
-
-/****************************************************************************/
-/* Mocked verification functions */
-
-vb2_error_t VbExDiskGetInfo(VbDiskInfo **infos_ptr, uint32_t *count,
-			    uint32_t disk_flags)
-{
-	int i;
-	int num_disks = 0;
-
-	VB2_DEBUG("My %s\n", __FUNCTION__);
-
-	*infos_ptr = mock_disks;
-
-	for(i = 0; i < MAX_TEST_DISKS; i++) {
-		if (is_nonzero(&t->disks_to_provide[i],
-			       sizeof(t->disks_to_provide[i]))) {
-			mock_disks[num_disks].bytes_per_lba =
-				t->disks_to_provide[i].bytes_per_lba;
-			mock_disks[num_disks].lba_count =
-				mock_disks[num_disks].streaming_lba_count =
-				t->disks_to_provide[i].lba_count;
-			mock_disks[num_disks].flags =
-				t->disks_to_provide[i].flags;
-			mock_disks[num_disks].handle = (VbExDiskHandle_t)
-				t->disks_to_provide[i].diskname;
-			VB2_DEBUG("  mock_disk[%d] %" PRIu64 " %" PRIu64
-				  " %#x %s\n", i,
-				  mock_disks[num_disks].bytes_per_lba,
-				  mock_disks[num_disks].lba_count,
-				  mock_disks[num_disks].flags,
-				  (mock_disks[num_disks].handle
-				   ? (char *)mock_disks[num_disks].handle
-				   : "0"));
-			num_disks++;
-		} else {
-			mock_disks[num_disks].handle =
-				(VbExDiskHandle_t)"INVALID";
-		}
-	}
-
-	if (t->disk_count_to_return >= 0)
-		*count = t->disk_count_to_return;
-	else
-		*count = num_disks;
-
-	VB2_DEBUG("  *count=%" PRIu32 "\n", *count);
-	VB2_DEBUG("  return %#x\n", t->diskgetinfo_return_val);
-
-	return t->diskgetinfo_return_val;
-}
-
-vb2_error_t VbExDiskFreeInfo(VbDiskInfo *infos,
-			   VbExDiskHandle_t preserve_handle)
-{
-	got_load_disk = (const char *)preserve_handle;
-	VB2_DEBUG("%s(): got_load_disk = %s\n", __FUNCTION__,
-		  got_load_disk ? got_load_disk : "0");
-	return VB2_SUCCESS;
-}
-
-static vb2_error_t LoadKernelImpl(struct vb2_context *c,
-				  VbSelectAndLoadKernelParams *params,
-				  VbDiskInfo *disk_info)
-{
-	got_find_disk = (const char *)params->disk_handle;
-	VB2_DEBUG("%s(%d): got_find_disk = %s\n", __FUNCTION__,
-		  load_kernel_calls,
-		  got_find_disk ? got_find_disk : "0");
-	if (t->external_expected[load_kernel_calls] !=
-			!!(disk_info->flags & VB_DISK_FLAG_EXTERNAL_GPT))
-		got_external_mismatch++;
-	return t->loadkernel_return_val[load_kernel_calls++];
-}
-
-vb2_error_t LoadKernel(struct vb2_context *c,
-		       VbSelectAndLoadKernelParams *params,
-		       VbDiskInfo *disk_info)
-{
-	lk_normal_calls++;
-	return LoadKernelImpl(c, params, disk_info);
-}
-
-vb2_error_t LoadMiniOsKernel(struct vb2_context *c,
-			     VbSelectAndLoadKernelParams *params,
-			     VbDiskInfo *disk_info, uint32_t minios_flags)
-{
-	lk_minios_calls++;
-	return LoadKernelImpl(c, params, disk_info);
-}
-
-void vb2_nv_set(struct vb2_context *c,
-		enum vb2_nv_param param,
-		uint32_t value)
-{
-	if (param != VB2_NV_RECOVERY_REQUEST)
-		return;
-	VB2_DEBUG("%s(): got_recovery_request_val = %d (%#x)\n", __FUNCTION__,
-		  value, value);
-	got_recovery_request_val = value;
-}
-
-/****************************************************************************/
-
-static void VbTryLoadKernelTest(void)
-{
-	int i;
-	int num_tests = ARRAY_SIZE(normal_tests);
-
-	for (i = 0; i < num_tests; i++) {
-		printf("Test case: %s ...\n", normal_tests[i].name);
-		ResetMocks(&normal_tests[i]);
-		ctx->flags = t->ctx_flags;
-		TEST_EQ(VbTryLoadKernel(ctx, t->want_flags, &kparams),
-			t->expected_return_val, "  return value");
-		TEST_EQ(got_recovery_request_val,
-			t->expected_recovery_request_val, "  recovery_request");
-		if (t->expected_to_find_disk != DONT_CARE) {
-			TEST_PTR_EQ(got_find_disk, t->expected_to_find_disk,
-				    "  find disk");
-			TEST_PTR_EQ(got_load_disk, t->expected_to_load_disk,
-				    "  load disk");
-		}
-		TEST_EQ(got_external_mismatch, 0, "  external GPT errors");
-		TEST_EQ(!!(ctx->flags & VB2_CONTEXT_DISABLE_TPM), 0,
-			"  ctx flag disable_tpm");
-	}
-	TEST_EQ(lk_normal_calls, load_kernel_calls, "  LoadKernel called");
-	TEST_EQ(lk_minios_calls, 0, "  LoadMiniOsKernel not called");
-}
-
-static void VbTryLoadMiniOsKernelTest(void)
-{
-	int i;
-	int num_tests =  ARRAY_SIZE(minios_tests);
-
-	for (i = 0; i < num_tests; i++) {
-		printf("Test case: %s ...\n", minios_tests[i].name);
-		ResetMocks(&minios_tests[i]);
-		ctx->flags = t->ctx_flags;
-		TEST_EQ(VbTryLoadMiniOsKernel(ctx, 0, &kparams),
-			t->expected_return_val, "  return value");
-		TEST_EQ(got_recovery_request_val,
-			t->expected_recovery_request_val, "  recovery_request");
-		if (t->expected_to_find_disk != DONT_CARE) {
-			TEST_PTR_EQ(got_find_disk, t->expected_to_find_disk,
-				    "  find disk");
-			TEST_PTR_EQ(got_load_disk, t->expected_to_load_disk,
-				    "  load disk");
-		}
-		TEST_EQ(got_external_mismatch, 0, "  external GPT errors");
-		TEST_EQ(!!(ctx->flags & VB2_CONTEXT_DISABLE_TPM), 1,
-			"  ctx flag disable_tpm");
-	}
-	TEST_EQ(lk_normal_calls, 0, "  LoadKernel not called");
-	TEST_EQ(lk_minios_calls, load_kernel_calls,
-		"  LoadMiniOsKernel called");
-}
-
-int main(void)
-{
-	VbTryLoadKernelTest();
-	VbTryLoadMiniOsKernelTest();
-
-	return gTestSuccess ? 0 : 255;
-}
diff --git a/tests/verify_kernel.c b/tests/verify_kernel.c
index a79f33a0..8af47918 100644
--- a/tests/verify_kernel.c
+++ b/tests/verify_kernel.c
@@ -14,7 +14,6 @@
 #include "host_common.h"
 #include "util_misc.h"
 #include "vboot_api.h"
-#include "load_kernel_fw.h"
 
 static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
 	__attribute__((aligned(VB2_WORKBUF_ALIGN)));
@@ -23,13 +22,13 @@ static struct vb2_shared_data *sd;
 
 static uint8_t *diskbuf;
 
-static VbSelectAndLoadKernelParams params;
-static VbDiskInfo disk_info;
+static struct vb2_kernel_params params;
+static struct vb2_disk_info disk_info;
 
-vb2_error_t VbExDiskRead(VbExDiskHandle_t handle, uint64_t lba_start,
+vb2_error_t VbExDiskRead(vb2ex_disk_handle_t handle, uint64_t lba_start,
 			 uint64_t lba_count, void *buffer)
 {
-	if (handle != (VbExDiskHandle_t)1)
+	if (handle != (vb2ex_disk_handle_t)1)
 		return VB2_ERROR_UNKNOWN;
 	if (lba_start >= disk_info.streaming_lba_count)
 		return VB2_ERROR_UNKNOWN;
@@ -40,10 +39,10 @@ vb2_error_t VbExDiskRead(VbExDiskHandle_t handle, uint64_t lba_start,
 	return VB2_SUCCESS;
 }
 
-vb2_error_t VbExDiskWrite(VbExDiskHandle_t handle, uint64_t lba_start,
+vb2_error_t VbExDiskWrite(vb2ex_disk_handle_t handle, uint64_t lba_start,
 			  uint64_t lba_count, const void *buffer)
 {
-	if (handle != (VbExDiskHandle_t)1)
+	if (handle != (vb2ex_disk_handle_t)1)
 		return VB2_ERROR_UNKNOWN;
 	if (lba_start >= disk_info.streaming_lba_count)
 		return VB2_ERROR_UNKNOWN;
@@ -87,8 +86,7 @@ int main(int argc, char *argv[])
 	}
 
 	/* Set up params */
-	params.disk_handle = (VbExDiskHandle_t)1;
-	disk_info.handle = (VbExDiskHandle_t)1;
+	disk_info.handle = (vb2ex_disk_handle_t)1;
 	disk_info.bytes_per_lba = 512;
 	disk_info.streaming_lba_count = disk_bytes / 512;
 	disk_info.lba_count = disk_info.streaming_lba_count;
@@ -123,8 +121,8 @@ int main(int argc, char *argv[])
 	}
 
 	/*
-	 * LoadKernel() cares only about VBNV_DEV_BOOT_SIGNED_ONLY, and only in
-	 * dev mode.  So just use defaults for nv storage.
+	 * vb2api_load_kernel() cares only about VBNV_DEV_BOOT_SIGNED_ONLY, and
+	 * only in dev mode.  So just use defaults for nv storage.
 	 */
 	vb2_nv_init(ctx);
 	/* We need to init kernel secdata for
@@ -134,9 +132,10 @@ int main(int argc, char *argv[])
 	vb2_secdata_kernel_init(ctx);
 
 	/* Try loading kernel */
-	rv = LoadKernel(ctx, &params, &disk_info);
+	rv = vb2api_load_kernel(ctx, &params, &disk_info);
 	if (rv != VB2_SUCCESS) {
-		fprintf(stderr, "LoadKernel() failed with code %d\n", rv);
+		fprintf(stderr, "vb2api_load_kernel() failed with code %d\n",
+			rv);
 		return 1;
 	}