vboot_api_kernel.c: Remove kparams_ptr

Pass VbSelectAndLoadKernelParams kparams as a function argument instead
of using global variable kparams_ptr. Remove VbSelectAndLoadKernel and
replace its tests with the unit tests for vb2_set_boot_mode,
vb2api_kernel_phase2, vb2api_kernel_finalize, and vb2api_normal_boot.

BUG=b:172339016
BRANCH=none
TEST=make runtests

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

21 files changed, 490 insertions, 594 deletions

diff --git a/Makefile b/Makefile
index 2c6f70c2..b6ab96b4 100644
--- a/Makefile
+++ b/Makefile
@@ -714,7 +714,6 @@ TEST_NAMES = \
 	tests/gpt_misc_tests \
 	tests/sha_benchmark \
 	tests/subprocess_tests \
-	tests/vboot_api_kernel4_tests \
 	tests/vboot_api_kernel_tests \
 	tests/vboot_kernel_tests \
 	tests/vboot_kernel2_tests \
@@ -748,7 +747,9 @@ TEST2X_NAMES = \
 	tests/vb2_host_key_tests \
 	tests/vb2_host_nvdata_flashrom_tests \
 	tests/vb2_kernel_tests \
+	tests/vb2_kernel2_tests \
 	tests/vb2_misc_tests \
+	tests/vb2_misc2_tests \
 	tests/vb2_nvstorage_tests \
 	tests/vb2_rsa_utility_tests \
 	tests/vb2_recovery_reasons_tests \
@@ -1285,7 +1286,6 @@ ifeq ($(filter-out 0,${MOCK_TPM})$(filter-out 0,${TPM2_MODE}),)
 # tlcl_tests only works when MOCK_TPM is disabled
 	${RUNTEST} ${BUILD_RUN}/tests/tlcl_tests
 endif
-	${RUNTEST} ${BUILD_RUN}/tests/vboot_api_kernel4_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vboot_api_kernel_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vboot_kernel_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vboot_kernel2_tests
@@ -1303,7 +1303,9 @@ run2tests: install_for_test
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_gbb_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_host_key_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_kernel_tests
+	${RUNTEST} ${BUILD_RUN}/tests/vb2_kernel2_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_misc_tests
+	${RUNTEST} ${BUILD_RUN}/tests/vb2_misc2_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_nvstorage_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_rsa_utility_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_secdata_firmware_tests
diff --git a/firmware/2lib/2auxfw_sync.c b/firmware/2lib/2auxfw_sync.c
index eaea1d46..a60165c1 100644
--- a/firmware/2lib/2auxfw_sync.c
+++ b/firmware/2lib/2auxfw_sync.c
@@ -48,6 +48,7 @@ static vb2_error_t auxfw_sync_check_update(struct vb2_context *ctx,
 	return vb2ex_auxfw_check(severity);
 }
 
+test_mockable
 vb2_error_t vb2api_auxfw_sync(struct vb2_context *ctx)
 {
 	enum vb2_auxfw_update_severity fw_update = VB2_AUXFW_NO_UPDATE;
diff --git a/firmware/2lib/2ec_sync.c b/firmware/2lib/2ec_sync.c
index e75313b4..b1bc7ebd 100644
--- a/firmware/2lib/2ec_sync.c
+++ b/firmware/2lib/2ec_sync.c
@@ -338,6 +338,7 @@ static vb2_error_t ec_sync_phase2(struct vb2_context *ctx)
 	return sync_ec(ctx);
 }
 
+test_mockable
 vb2_error_t vb2api_ec_sync(struct vb2_context *ctx)
 {
 	struct vb2_shared_data *sd = vb2_get_sd(ctx);
diff --git a/firmware/2lib/2kernel.c b/firmware/2lib/2kernel.c
index bec9c507..30d6bbe6 100644
--- a/firmware/2lib/2kernel.c
+++ b/firmware/2lib/2kernel.c
@@ -37,7 +37,8 @@ static int vb2_reset_nv_requests(struct vb2_context *ctx)
 	return need_reboot;
 }
 
-vb2_error_t vb2api_normal_boot(struct vb2_context *ctx)
+vb2_error_t vb2api_normal_boot(struct vb2_context *ctx,
+			       VbSelectAndLoadKernelParams *kparams)
 {
 	struct vb2_shared_data *sd = vb2_get_sd(ctx);
 	uint32_t max_rollforward = vb2_nv_get(ctx,
@@ -51,7 +52,7 @@ vb2_error_t vb2api_normal_boot(struct vb2_context *ctx)
 		return VB2_REQUEST_REBOOT;
 	}
 
-	vb2_error_t rv = VbTryLoadKernel(ctx, VB_DISK_FLAG_FIXED);
+	vb2_error_t rv = VbTryLoadKernel(ctx, VB_DISK_FLAG_FIXED, kparams);
 
 	VB2_DEBUG("Checking if TPM kernel version needs advancing\n");
 
diff --git a/firmware/2lib/2rsa.c b/firmware/2lib/2rsa.c
index dcd8bad0..eb07b04c 100644
--- a/firmware/2lib/2rsa.c
+++ b/firmware/2lib/2rsa.c
@@ -14,7 +14,6 @@
 #include "2rsa_private.h"
 #include "2sha.h"
 #include "2sysincludes.h"
-#include "vboot_test.h"
 
 /**
  * a[] -= mod
diff --git a/firmware/2lib/2secdata_kernel.c b/firmware/2lib/2secdata_kernel.c
index 0d4208fa..9973cd46 100644
--- a/firmware/2lib/2secdata_kernel.c
+++ b/firmware/2lib/2secdata_kernel.c
@@ -11,7 +11,6 @@
 #include "2secdata.h"
 #include "2secdata_struct.h"
 #include "2sysincludes.h"
-#include "vboot_test.h"
 
 #define MAJOR_VER(x) (((x) & 0xf0) >> 4)
 #define MINOR_VER(x) ((x) & 0x0f)
@@ -213,6 +212,7 @@ uint32_t vb2_secdata_kernel_get(struct vb2_context *ctx,
 	return 0;
 }
 
+test_mockable
 void vb2_secdata_kernel_set(struct vb2_context *ctx,
 			    enum vb2_secdata_kernel_param param,
 			    uint32_t value)
@@ -273,6 +273,7 @@ void vb2_secdata_kernel_set(struct vb2_context *ctx,
 	VB2_REC_OR_DIE(ctx, "%s\n", msg);
 }
 
+test_mockable
 const uint8_t *vb2_secdata_kernel_get_ec_hash(struct vb2_context *ctx)
 {
 	struct vb2_shared_data *sd = vb2_get_sd(ctx);
diff --git a/firmware/2lib/2stub.c b/firmware/2lib/2stub.c
index ef0bbd8a..51d2eb20 100644
--- a/firmware/2lib/2stub.c
+++ b/firmware/2lib/2stub.c
@@ -160,33 +160,6 @@ vb2_error_t vb2ex_auxfw_finalize(struct vb2_context *ctx)
 }
 
 /*****************************************************************************/
-/* UI-related stubs */
-
-__attribute__((weak))
-vb2_error_t vb2ex_broken_screen_ui(struct vb2_context *ctx)
-{
-	return VB2_SUCCESS;
-}
-
-__attribute__((weak))
-vb2_error_t vb2ex_manual_recovery_ui(struct vb2_context *ctx)
-{
-	return VB2_SUCCESS;
-}
-
-__attribute__((weak))
-vb2_error_t vb2ex_developer_ui(struct vb2_context *ctx)
-{
-	return VB2_SUCCESS;
-}
-
-__attribute__((weak))
-vb2_error_t vb2ex_diagnostic_ui(struct vb2_context *ctx)
-{
-	return VB2_SUCCESS;
-}
-
-/*****************************************************************************/
 /* Timer-related stubs */
 
 __attribute__((weak))
diff --git a/firmware/2lib/include/2api.h b/firmware/2lib/include/2api.h
index ff89a29d..a833da7c 100644
--- a/firmware/2lib/include/2api.h
+++ b/firmware/2lib/include/2api.h
@@ -28,6 +28,7 @@
 #include "2return_codes.h"
 #include "2rsa.h"
 #include "2secdata_struct.h"
+#include "vboot_api.h"
 
 #define _VB2_TRY_IMPL(expr, ctx, recovery_reason, ...) do { \
 	vb2_error_t _vb2_try_rv = (expr); \
@@ -165,14 +166,14 @@ enum vb2_context_flags {
 
 	/*
 	 * System supports EC software sync.  Caller may set this flag at any
-	 * time before calling VbSelectAndLoadKernel().
+	 * time before calling vb2api_kernel_phase2().
 	 */
 	VB2_CONTEXT_EC_SYNC_SUPPORTED = (1 << 15),
 
 	/*
 	 * EC software sync is slow to update; warning screen should be
 	 * displayed.  Caller may set this flag at any time before calling
-	 * VbSelectAndLoadKernel().  Deprecated as part of chromium:1038259.
+	 * vb2api_kernel_phase2().  Deprecated as part of chromium:1038259.
 	 */
 	VB2_CONTEXT_DEPRECATED_EC_SYNC_SLOW = (1 << 16),
 
@@ -847,9 +848,11 @@ vb2_error_t vb2api_kernel_phase2(struct vb2_context *ctx);
  * Handle a normal boot.
  *
  * @param ctx		Vboot context.
+ * @param kparams	Params specific to loading the kernel
  * @return VB2_SUCCESS, or error code on error.
  */
-vb2_error_t vb2api_normal_boot(struct vb2_context *ctx);
+vb2_error_t vb2api_normal_boot(struct vb2_context *ctx,
+			       VbSelectAndLoadKernelParams *kparams);
 
 /**
  * Finalize for kernel verification stage.
@@ -1484,54 +1487,6 @@ vb2_error_t vb2ex_ec_battery_cutoff(void);
 
 /*****************************************************************************/
 /* Functions for firmware UI. */
-/* TODO(b/172339016): Remove vb2ex_*_ui(). */
-
-/**
- * UI for a non-manual recovery ("BROKEN").
- *
- * Enter the broken screen UI, which shows that an unrecoverable error was
- * encountered last boot. Wait for the user to physically reset or shut down.
- *
- * @param ctx		Vboot context
- * @return VB2_SUCCESS, or non-zero error code.
- */
-vb2_error_t vb2ex_broken_screen_ui(struct vb2_context *ctx);
-
-/**
- * UI for a manual recovery-mode boot.
- *
- * Enter the recovery menu, which prompts the user to insert recovery media,
- * navigate the step-by-step recovery, or enter developer mode if allowed.
- *
- * @param ctx		Vboot context
- * @return VB2_SUCCESS, or non-zero error code.
- */
-vb2_error_t vb2ex_manual_recovery_ui(struct vb2_context *ctx);
-
-/**
- * UI for a developer-mode boot.
- *
- * Enter the developer menu, which provides options to switch out of developer
- * mode, boot from external media, use legacy bootloader, or boot Chrome OS from
- * disk.
- *
- * If a timeout occurs, take the default boot action.
- *
- * @param ctx		Vboot context
- * @return VB2_SUCCESS, or non-zero error code.
- */
-vb2_error_t vb2ex_developer_ui(struct vb2_context *ctx);
-
-/**
- * UI for a diagnostic tools boot.
- *
- * Enter the diagnostic tools menu, which provides debug information and
- * diagnostic tests of various hardware components.
- *
- * @param ctx		Vboot context
- * @return VB2_SUCCESS, or non-zero error code.
- */
-vb2_error_t vb2ex_diagnostic_ui(struct vb2_context *ctx);
 
 /**
  * Get the vboot debug info.
diff --git a/firmware/include/vboot_api.h b/firmware/include/vboot_api.h
index ce3ac2d6..4dba0a06 100644
--- a/firmware/include/vboot_api.h
+++ b/firmware/include/vboot_api.h
@@ -45,15 +45,14 @@ typedef struct VbSharedDataHeader VbSharedDataHeader;
 typedef void *VbExDiskHandle_t;
 
 typedef struct VbSelectAndLoadKernelParams {
-	/* Inputs to VbSelectAndLoadKernel() */
+	/* Inputs to VbTryLoadKernel() */
 	/* Destination buffer for kernel (normally at 0x100000 on x86) */
 	void *kernel_buffer;
 	/* Size of kernel buffer in bytes */
 	uint32_t kernel_buffer_size;
 
 	/*
-	 * Outputs from VbSelectAndLoadKernel(); valid only if it returns
-	 * success.
+	 * Outputs from VbTryLoadKernel(); valid only if it returns success.
 	 */
 	/* Handle of disk containing loaded kernel */
 	VbExDiskHandle_t disk_handle;
@@ -70,14 +69,6 @@ typedef struct VbSelectAndLoadKernelParams {
 } VbSelectAndLoadKernelParams;
 
 /**
- * Select and loads the kernel.
- *
- * Returns VB2_SUCCESS if success, non-zero if error; on error, caller
- * should reboot. */
-vb2_error_t VbSelectAndLoadKernel(struct vb2_context *ctx,
-				  VbSelectAndLoadKernelParams *kparams);
-
-/**
  * Attempt loading a kernel from the specified type(s) of disks.
  *
  * If successful, sets kparams.disk_handle to the disk for the kernel and
@@ -85,9 +76,11 @@ vb2_error_t VbSelectAndLoadKernel(struct vb2_context *ctx,
  *
  * @param ctx			Vboot context
  * @param disk_flags		Flags to pass to VbExDiskGetInfo()
+ * @param kparams               Params specific to loading the kernel
  * @return VB2_SUCCESS or the most specific VB2_ERROR_LK error.
  */
-vb2_error_t VbTryLoadKernel(struct vb2_context *ctx, uint32_t disk_flags);
+vb2_error_t VbTryLoadKernel(struct vb2_context *ctx, uint32_t disk_flags,
+			    VbSelectAndLoadKernelParams *kparams);
 
 /* miniOS flags */
 
@@ -106,10 +99,12 @@ vb2_error_t VbTryLoadKernel(struct vb2_context *ctx, uint32_t disk_flags);
  *
  * @param ctx			Vboot context
  * @param minios_flags		Flags for miniOS
+ * @param kparams               Params specific to loading the kernel
  * @return VB2_SUCCESS or the most specific VB2_ERROR_LK error.
  */
 vb2_error_t VbTryLoadMiniOsKernel(struct vb2_context *ctx,
-				  uint32_t minios_flags);
+				  uint32_t minios_flags,
+				  VbSelectAndLoadKernelParams *kparams);
 
 /*****************************************************************************/
 /* Disk access (previously in boot_device.h) */
diff --git a/firmware/lib/include/vboot_struct.h b/firmware/lib/include/vboot_struct.h
index 374dfd40..edcb2b45 100644
--- a/firmware/lib/include/vboot_struct.h
+++ b/firmware/lib/include/vboot_struct.h
@@ -137,7 +137,7 @@ typedef struct VbSharedDataHeader {
 	uint8_t reserved4[7];
 	/* Flags from firmware keyblock */
 	uint64_t fw_keyblock_flags;
-	/* Kernel TPM version at start of VbSelectAndLoadKernel() */
+	/* Kernel TPM version at start of vb2api_kernel_phase1 */
 	uint32_t kernel_version_tpm_start;
 	/* Kernel lowest version found */
 	uint32_t kernel_version_lowest;
diff --git a/firmware/lib/include/vboot_test.h b/firmware/lib/include/vboot_test.h
deleted file mode 100644
index fb1f52ad..00000000
--- a/firmware/lib/include/vboot_test.h
+++ /dev/null
@@ -1,17 +0,0 @@
-/* Copyright 2019 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.
- *
- * This header is for APIs that are only used by test code.
- */
-
-#ifndef VBOOT_REFERENCE_TEST_API_H_
-#define VBOOT_REFERENCE_TEST_API_H_
-
-/****************************************************************************
- * vboot_api_kernel.c */
-
-struct VbSelectAndLoadKernelParams;
-struct VbSelectAndLoadKernelParams **VbApiKernelGetParamsPtr(void);
-
-#endif  /* VBOOT_REFERENCE_TEST_API_H_ */
diff --git a/firmware/lib/vboot_api_kernel.c b/firmware/lib/vboot_api_kernel.c
index f4e5146a..0fdc2d33 100644
--- a/firmware/lib/vboot_api_kernel.c
+++ b/firmware/lib/vboot_api_kernel.c
@@ -15,18 +15,6 @@
 #include "load_kernel_fw.h"
 #include "vboot_api.h"
 #include "vboot_struct.h"
-#include "vboot_test.h"
-
-/* Global variables */
-static VbSelectAndLoadKernelParams *kparams_ptr;
-
-#ifdef CHROMEOS_ENVIRONMENT
-/* Global variable accessor for unit tests */
-struct VbSelectAndLoadKernelParams **VbApiKernelGetParamsPtr(void)
-{
-	return &kparams_ptr;
-}
-#endif
 
 static int is_valid_disk(VbDiskInfo *info, uint32_t disk_flags)
 {
@@ -40,7 +28,8 @@ static int is_valid_disk(VbDiskInfo *info, uint32_t disk_flags)
 
 static vb2_error_t VbTryLoadKernelImpl(struct vb2_context *ctx,
 				       uint32_t disk_flags, int minios,
-				       uint32_t minios_flags)
+				       uint32_t minios_flags,
+				       VbSelectAndLoadKernelParams *kparams)
 {
 	vb2_error_t rv = VB2_ERROR_LK_NO_DISK_FOUND;
 	VbDiskInfo* disk_info = NULL;
@@ -48,11 +37,8 @@ static vb2_error_t VbTryLoadKernelImpl(struct vb2_context *ctx,
 	uint32_t i;
 	vb2_error_t new_rv;
 
-	/* TODO: Should have been set by VbSelectAndLoadKernel. Remove when
-	   this global is no longer needed. */
-	VB2_ASSERT(kparams_ptr);
-
-	kparams_ptr->disk_handle = NULL;
+	VB2_ASSERT(kparams);
+	kparams->disk_handle = NULL;
 
 	/* Find disks */
 	if (VB2_SUCCESS != VbExDiskGetInfo(&disk_info, &disk_count, disk_flags))
@@ -70,14 +56,14 @@ static vb2_error_t VbTryLoadKernelImpl(struct vb2_context *ctx,
 				  disk_info[i].flags);
 			continue;
 		}
-		kparams_ptr->disk_handle = disk_info[i].handle;
+		kparams->disk_handle = disk_info[i].handle;
 
 		if (minios) {
-			new_rv = LoadMiniOsKernel(ctx, kparams_ptr,
+			new_rv = LoadMiniOsKernel(ctx, kparams,
 						  &disk_info[i], minios_flags);
 			VB2_DEBUG("LoadMiniOsKernel() = %#x\n", new_rv);
 		} else {
-			new_rv = LoadKernel(ctx, kparams_ptr, &disk_info[i]);
+			new_rv = LoadKernel(ctx, kparams, &disk_info[i]);
 			VB2_DEBUG("LoadKernel() = %#x\n", new_rv);
 		}
 
@@ -118,69 +104,20 @@ static vb2_error_t VbTryLoadKernelImpl(struct vb2_context *ctx,
 }
 
 test_mockable
-vb2_error_t VbTryLoadKernel(struct vb2_context *ctx, uint32_t disk_flags)
+vb2_error_t VbTryLoadKernel(struct vb2_context *ctx, uint32_t disk_flags,
+			    VbSelectAndLoadKernelParams *kparams)
 {
 	ctx->flags &= ~VB2_CONTEXT_DISABLE_TPM;
-	return VbTryLoadKernelImpl(ctx, disk_flags, 0, 0);
+	return VbTryLoadKernelImpl(ctx, disk_flags, 0, 0, kparams);
 }
 
 test_mockable
 vb2_error_t VbTryLoadMiniOsKernel(struct vb2_context *ctx,
-				  uint32_t minios_flags)
+				  uint32_t minios_flags,
+				  VbSelectAndLoadKernelParams *kparams)
 {
-	VB2_TRY(VbTryLoadKernelImpl(ctx, VB_DISK_FLAG_FIXED, 1, minios_flags));
+	VB2_TRY(VbTryLoadKernelImpl(ctx, VB_DISK_FLAG_FIXED, 1, minios_flags,
+				    kparams));
 	ctx->flags |= VB2_CONTEXT_DISABLE_TPM;
 	return VB2_SUCCESS;
 }
-
-vb2_error_t VbSelectAndLoadKernel(struct vb2_context *ctx,
-				  VbSelectAndLoadKernelParams *kparams)
-{
-	/* TODO: Send this argument through subsequent function calls, rather
-	   than relying on a global to pass it to VbTryLoadKernel. */
-	kparams_ptr = kparams;
-
-	VB2_TRY(vb2api_kernel_phase1(ctx));
-	VB2_TRY(vb2api_kernel_phase2(ctx));
-
-	switch (ctx->boot_mode) {
-	case VB2_BOOT_MODE_MANUAL_RECOVERY:
-		/* Manual recovery boot.  This has UI. */
-		VB2_TRY(vb2ex_manual_recovery_ui(ctx));
-		break;
-	case VB2_BOOT_MODE_BROKEN_SCREEN:
-		/*
-		 * In EFS2, recovery mode can be entered even when battery is
-		 * drained or damaged. EC-RO sets NO_BOOT flag in such case and
-		 * uses PD power to boot AP.
-		 *
-		 * TODO: Inform user why recovery failed to start.
-		 */
-		if (ctx->flags & VB2_CONTEXT_NO_BOOT)
-			VB2_DEBUG("NO_BOOT in RECOVERY mode\n");
-
-		/* Broken screen.  This has UI. */
-		VB2_TRY(vb2ex_broken_screen_ui(ctx));
-		break;
-	case VB2_BOOT_MODE_DIAGNOSTICS:
-		/* Diagnostic boot.  This has UI. */
-		VB2_TRY(vb2ex_diagnostic_ui(ctx));
-		/*
-		 * The diagnostic menu should either boot a rom, or
-		 * return either of reboot or shutdown.
-		 */
-		return VB2_REQUEST_REBOOT;
-	case VB2_BOOT_MODE_DEVELOPER:
-		/* Developer boot.  This has UI. */
-		VB2_TRY(vb2ex_developer_ui(ctx));
-		break;
-	case VB2_BOOT_MODE_NORMAL:
-		/* Normal boot */
-		VB2_TRY(vb2api_normal_boot(ctx));
-		break;
-	default:
-		return VB2_ERROR_ESCAPE_NO_BOOT;
-	}
-
-	return vb2api_kernel_finalize(ctx);
-}
diff --git a/tests/vb2_kernel2_tests.c b/tests/vb2_kernel2_tests.c
new file mode 100644
index 00000000..f42c0f78
--- /dev/null
+++ b/tests/vb2_kernel2_tests.c
@@ -0,0 +1,163 @@
+/* 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 "host_common.h"
+#include "load_kernel_fw.h"
+#include "test_common.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 enum vb2_boot_mode *boot_mode;
+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");
+
+	boot_mode = (enum vb2_boot_mode *)&ctx->boot_mode;
+	*boot_mode = 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();
+	mock_vbtlk_expect_fixed = 0;
+	vb2_nv_set(ctx, VB2_NV_DISPLAY_REQUEST, 1);
+	TEST_EQ(vb2api_normal_boot(ctx, &kparams), VB2_REQUEST_REBOOT,
+		"vb2api_normal_boot() reboot to reset NVRAM display request");
+	TEST_EQ(vb2_nv_get(ctx, VB2_NV_DISPLAY_REQUEST), 0,
+		"  display request reset");
+
+	reset_common_data();
+	mock_vbtlk_expect_fixed = 0;
+	vb2_nv_set(ctx, VB2_NV_DIAG_REQUEST, 1);
+	TEST_EQ(vb2api_normal_boot(ctx, &kparams), VB2_REQUEST_REBOOT,
+		"vb2api_normal_boot() reboot to reset NVRAM diag request");
+	TEST_EQ(vb2_nv_get(ctx, VB2_NV_DIAG_REQUEST), 0,
+		"  diag request reset");
+
+	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 c467fa9d..d5e6b55a 100644
--- a/tests/vb2_kernel_tests.c
+++ b/tests/vb2_kernel_tests.c
@@ -24,6 +24,7 @@ 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 enum vb2_boot_mode *boot_mode;
+static VbSelectAndLoadKernelParams kparams;
 
 /* Mocked function data */
 
@@ -35,14 +36,16 @@ static struct {
 
 static int mock_read_res_fail_on_call;
 static int mock_secdata_fwmp_check_retval;
-static int mock_vbtlk_expect_fixed;
-static int mock_vbtlk_expect_removable;
-static vb2_error_t mock_vbtlk_retval;
+static int mock_commit_data_called;
+static int mock_ec_sync_called;
+static int mock_ec_sync_retval;
+static int mock_battery_cutoff_called;
 
 /* Type of test to reset for */
 enum reset_type {
 	FOR_PHASE1,
-	FOR_NORMAL_BOOT,
+	FOR_PHASE2,
+	FOR_FINALIZE,
 };
 
 static void reset_common_data(enum reset_type t)
@@ -51,6 +54,8 @@ static void reset_common_data(enum reset_type t)
 
 	memset(workbuf, 0xaa, sizeof(workbuf));
 
+	memset(&kparams, 0, sizeof(kparams));
+
 	TEST_SUCC(vb2api_init(workbuf, sizeof(workbuf), &ctx),
 		  "vb2api_init failed");
 
@@ -63,10 +68,10 @@ static void reset_common_data(enum reset_type t)
 
 	mock_read_res_fail_on_call = 0;
 	mock_secdata_fwmp_check_retval = VB2_SUCCESS;
-	mock_vbtlk_expect_fixed = 0;
-	mock_vbtlk_expect_removable = 0;
-	mock_vbtlk_retval = VB2_SUCCESS;
-
+	mock_commit_data_called = 0;
+	mock_ec_sync_called = 0;
+	mock_ec_sync_retval = VB2_SUCCESS;
+	mock_battery_cutoff_called = 0;
 
 	/* Recovery key in mock GBB */
 	memset(&mock_gbb, 0, sizeof(mock_gbb));
@@ -86,8 +91,6 @@ static void reset_common_data(enum reset_type t)
 	boot_mode = (enum vb2_boot_mode *)&ctx->boot_mode;
 	if (t == FOR_PHASE1)
 		*boot_mode = VB2_BOOT_MODE_BROKEN_SCREEN;
-	else if (t == FOR_NORMAL_BOOT)
-		*boot_mode = VB2_BOOT_MODE_NORMAL;
 	else
 		*boot_mode = VB2_BOOT_MODE_UNDEFINED;
 
@@ -117,6 +120,34 @@ vb2_error_t vb2api_secdata_fwmp_check(struct vb2_context *c, uint8_t *size)
 	return mock_secdata_fwmp_check_retval;
 }
 
+vb2_error_t vb2api_ec_sync(struct vb2_context *c)
+{
+	mock_ec_sync_called = 1;
+	return mock_ec_sync_retval;
+}
+
+vb2_error_t vb2api_auxfw_sync(struct vb2_context *c)
+{
+	return VB2_SUCCESS;
+}
+
+vb2_error_t vb2ex_ec_battery_cutoff(void)
+{
+	TEST_EQ(mock_ec_sync_called, 1,
+		"  battery cutoff must happen after EC sync");
+	mock_battery_cutoff_called = 1;
+	return VB2_SUCCESS;
+}
+
+const uint8_t *vb2_secdata_kernel_get_ec_hash(struct vb2_context *c)
+{
+	/*
+	 * Return NULL to prevent EC reboot due to
+	 * VB2_SD_FLAG_ECSYNC_HMIR_UPDATED.
+	 */
+	return NULL;
+}
+
 struct vb2_gbb_header *vb2_get_gbb(struct vb2_context *c)
 {
 	return &mock_gbb.h;
@@ -148,24 +179,10 @@ vb2_error_t vb2ex_read_resource(struct vb2_context *c,
 	return VB2_SUCCESS;
 }
 
-vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t disk_flags)
+vb2_error_t vb2ex_commit_data(struct vb2_context *c)
 {
-	/*
-	 * 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),
-		"  VbTryLoadKernel unexpected fixed disk call");
-	TEST_EQ(!!mock_vbtlk_expect_removable,
-		!!(disk_flags & VB_DISK_FLAG_REMOVABLE),
-		"  VbTryLoadKernel unexpected removable disk call");
-
-	return mock_vbtlk_retval;
+	mock_commit_data_called = 1;
+	return VB2_SUCCESS;
 }
 
 /* Tests */
@@ -304,37 +321,125 @@ static void phase1_tests(void)
 		"phase1 fw preamble");
 }
 
-static void normal_boot_tests(void)
+static void phase2_tests(void)
 {
-	reset_common_data(FOR_NORMAL_BOOT);
-	mock_vbtlk_expect_fixed = 1;
-	TEST_EQ(vb2api_normal_boot(ctx), VB2_SUCCESS,
-		"vb2api_normal_boot() returns VB2_SUCCESS");
-
-	reset_common_data(FOR_NORMAL_BOOT);
-	mock_vbtlk_expect_fixed = 1;
-	mock_vbtlk_retval = VB2_ERROR_MOCK;
-	TEST_EQ(vb2api_normal_boot(ctx), VB2_ERROR_MOCK,
-		"vb2api_normal_boot() returns VB2_ERROR_MOCK");
-
-	reset_common_data(FOR_NORMAL_BOOT);
-	vb2_nv_set(ctx, VB2_NV_DISPLAY_REQUEST, 1);
-	TEST_EQ(vb2api_normal_boot(ctx), VB2_REQUEST_REBOOT,
-		"vb2api_normal_boot() reboot to reset NVRAM display request");
-	TEST_EQ(vb2_nv_get(ctx, VB2_NV_DISPLAY_REQUEST), 0,
-		"  display request reset");
-
-	reset_common_data(FOR_NORMAL_BOOT);
+	reset_common_data(FOR_PHASE2);
+	*boot_mode = VB2_BOOT_MODE_NORMAL;
+	TEST_SUCC(vb2api_kernel_phase2(ctx), "Normal mode");
+	TEST_EQ(mock_ec_sync_called, 1, "  EC sync");
+
+	reset_common_data(FOR_PHASE2);
+	*boot_mode = VB2_BOOT_MODE_DEVELOPER;
+	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	TEST_SUCC(vb2api_kernel_phase2(ctx), "Developer mode");
+	TEST_EQ(mock_ec_sync_called, 1, "  EC sync");
+
+	reset_common_data(FOR_PHASE2);
+	*boot_mode = VB2_BOOT_MODE_DIAGNOSTICS;
+	TEST_SUCC(vb2api_kernel_phase2(ctx), "Diagnostics mode");
+	TEST_EQ(mock_ec_sync_called, 1, "  EC sync");
+
+	/* Commit data for recovery mode */
+	reset_common_data(FOR_PHASE2);
+	*boot_mode = VB2_BOOT_MODE_MANUAL_RECOVERY;
+	ctx->flags |= VB2_CONTEXT_RECOVERY_MODE;
+	sd->recovery_reason = VB2_RECOVERY_RO_MANUAL;
+	TEST_SUCC(vb2api_kernel_phase2(ctx), "Manual recovery mode");
+	TEST_EQ(mock_commit_data_called, 1, "  commit data");
+	TEST_EQ(mock_ec_sync_called, 0, "  EC sync");
+
+	reset_common_data(FOR_PHASE2);
+	*boot_mode = VB2_BOOT_MODE_BROKEN_SCREEN;
+	ctx->flags |= VB2_CONTEXT_RECOVERY_MODE;
+	sd->recovery_reason = 123;
+	TEST_SUCC(vb2api_kernel_phase2(ctx), "Broken screen mode");
+	TEST_EQ(mock_commit_data_called, 1, "  commit data");
+	TEST_EQ(mock_ec_sync_called, 0, "  EC sync");
+
+	/* Boot recovery - memory retraining */
+	reset_common_data(FOR_PHASE2);
+	*boot_mode = VB2_BOOT_MODE_MANUAL_RECOVERY;
+	ctx->flags |= VB2_CONTEXT_RECOVERY_MODE;
+	sd->recovery_reason = VB2_RECOVERY_TRAIN_AND_REBOOT;
+	TEST_EQ(vb2api_kernel_phase2(ctx), VB2_REQUEST_REBOOT,
+		"Recovery train and reboot");
+
+	/* Clear VB2_NV_DIAG_REQUEST */
+	reset_common_data(FOR_PHASE2);
+	*boot_mode = VB2_BOOT_MODE_NORMAL;
 	vb2_nv_set(ctx, VB2_NV_DIAG_REQUEST, 1);
-	TEST_EQ(vb2api_normal_boot(ctx), VB2_REQUEST_REBOOT,
-		"vb2api_normal_boot() reboot to reset NVRAM diag request");
+	TEST_SUCC(vb2api_kernel_phase2(ctx), "Normal mode with DIAG_REQUEST");
 	TEST_EQ(vb2_nv_get(ctx, VB2_NV_DIAG_REQUEST), 0,
-		"  diag request reset");}
+		"  clear VB2_NV_DIAG_REQUEST");
+	TEST_EQ(mock_commit_data_called, 1, "  commit data");
+
+	reset_common_data(FOR_PHASE2);
+	*boot_mode = VB2_BOOT_MODE_DIAGNOSTICS;
+	vb2_nv_set(ctx, VB2_NV_DIAG_REQUEST, 1);
+	TEST_SUCC(vb2api_kernel_phase2(ctx), "Diagnostics mode");
+	TEST_EQ(vb2_nv_get(ctx, VB2_NV_DIAG_REQUEST), 0,
+		"  clear VB2_NV_DIAG_REQUEST");
+	TEST_EQ(mock_commit_data_called, 1, "  commit data");
+
+	/* Battery cutoff called after EC sync */
+	reset_common_data(FOR_PHASE2);
+	*boot_mode = VB2_BOOT_MODE_NORMAL;
+	vb2_nv_set(ctx, VB2_NV_BATTERY_CUTOFF_REQUEST, 1);
+	TEST_EQ(vb2api_kernel_phase2(ctx), VB2_REQUEST_SHUTDOWN,
+		"Set VB2_NV_BATTERY_CUTOFF_REQUEST");
+	TEST_EQ(mock_battery_cutoff_called, 1,
+		"  battery_cutoff called after EC sync");
+
+	/* Return EC sync error */
+	reset_common_data(FOR_PHASE2);
+	*boot_mode = VB2_BOOT_MODE_NORMAL;
+	mock_ec_sync_retval = VB2_ERROR_MOCK;
+	TEST_EQ(vb2api_kernel_phase2(ctx), VB2_ERROR_MOCK,
+		"Return EC sync error");
+
+	/* Undefined boot mode */
+	reset_common_data(FOR_PHASE2);
+	*boot_mode = VB2_BOOT_MODE_UNDEFINED;
+	TEST_EQ(vb2api_kernel_phase2(ctx), VB2_ERROR_ESCAPE_NO_BOOT,
+		"Undefined boot mode");
+}
+
+static void finalize_tests(void)
+{
+	/* NO_BOOT with EC sync support */
+	reset_common_data(FOR_FINALIZE);
+	ctx->flags |= VB2_CONTEXT_NO_BOOT;
+	ctx->flags |= VB2_CONTEXT_EC_SYNC_SUPPORTED;
+	TEST_EQ(vb2api_kernel_finalize(ctx), VB2_ERROR_ESCAPE_NO_BOOT,
+		"Recovery for NO_BOOT escape");
+	TEST_EQ(vb2_nv_get(ctx, VB2_NV_RECOVERY_REQUEST),
+		VB2_RECOVERY_ESCAPE_NO_BOOT, "  recovery_reason");
+
+	/* NO_BOOT with EC sync disabled */
+	reset_common_data(FOR_FINALIZE);
+	ctx->flags |= VB2_CONTEXT_NO_BOOT;
+	ctx->flags |= VB2_CONTEXT_EC_SYNC_SUPPORTED;
+	mock_gbb.h.flags |= VB2_GBB_FLAG_DISABLE_EC_SOFTWARE_SYNC;
+	TEST_SUCC(vb2api_kernel_finalize(ctx),
+		  "NO_BOOT ignored with gbb DISABLE_EC_SOFTWARE_SYNC");
+
+	/* Normal case with EC sync support */
+	reset_common_data(FOR_FINALIZE);
+	ctx->flags |= VB2_CONTEXT_EC_SYNC_SUPPORTED;
+	TEST_SUCC(vb2api_kernel_finalize(ctx), "Disable VB2_CONTEXT_NO_BOOT");
+
+	/* NO_BOOT without EC sync support */
+	reset_common_data(FOR_FINALIZE);
+	ctx->flags |= VB2_CONTEXT_NO_BOOT;
+	TEST_SUCC(vb2api_kernel_finalize(ctx),
+		  "Disable VB2_CONTEXT_EC_SYNC_SUPPORTED");
+}
 
 int main(int argc, char* argv[])
 {
 	phase1_tests();
-	normal_boot_tests();
+	phase2_tests();
+	finalize_tests();
 
 	return gTestSuccess ? 0 : 255;
 }
diff --git a/tests/vb2_keyblock_fuzzer.c b/tests/vb2_keyblock_fuzzer.c
index 2aa83e13..cd211c47 100644
--- a/tests/vb2_keyblock_fuzzer.c
+++ b/tests/vb2_keyblock_fuzzer.c
@@ -9,7 +9,6 @@
 #include "2rsa.h"
 #include "2rsa_private.h"
 #include "2secdata.h"
-#include "vboot_test.h"
 
 static struct vb2_context *ctx;
 static uint8_t workbuf[VB2_FIRMWARE_WORKBUF_RECOMMENDED_SIZE]
diff --git a/tests/vb2_misc2_tests.c b/tests/vb2_misc2_tests.c
new file mode 100644
index 00000000..a4ed5789
--- /dev/null
+++ b/tests/vb2_misc2_tests.c
@@ -0,0 +1,128 @@
+/* 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 vb2_set_boot_mode.
+ */
+
+#include "2api.h"
+#include "2common.h"
+#include "2misc.h"
+#include "2nvstorage.h"
+#include "test_common.h"
+
+/* Mock data */
+
+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 struct vb2_gbb_header gbb;
+
+static int mock_diagnostic_ui_enabled;
+
+/* Mock functions */
+
+struct vb2_gbb_header *vb2_get_gbb(struct vb2_context *c)
+{
+	return &gbb;
+}
+
+int vb2api_diagnostic_ui_enabled(struct vb2_context *c)
+{
+	return mock_diagnostic_ui_enabled;
+}
+
+static void reset_common_data(void)
+{
+	memset(workbuf, 0xaa, sizeof(workbuf));
+
+	memset(&gbb, 0, sizeof(gbb));
+
+	TEST_SUCC(vb2api_init(workbuf, sizeof(workbuf), &ctx),
+		  "vb2api_init failed");
+	sd = vb2_get_sd(ctx);
+
+	vb2_nv_init(ctx);
+
+	mock_diagnostic_ui_enabled = 0;
+}
+
+static void set_boot_mode_tests(void)
+{
+	/* Normal boot */
+	reset_common_data();
+	vb2_set_boot_mode(ctx);
+	TEST_EQ(ctx->boot_mode, VB2_BOOT_MODE_NORMAL, "Normal boot");
+
+	/* Check that NV_DIAG_REQUEST triggers diagnostic mode */
+	reset_common_data();
+	mock_diagnostic_ui_enabled = 1;
+	vb2_nv_set(ctx, VB2_NV_DIAG_REQUEST, 1);
+	vb2_set_boot_mode(ctx);
+	TEST_EQ(ctx->boot_mode, VB2_BOOT_MODE_DIAGNOSTICS,
+		"Normal boot with diag UI enabled");
+
+	reset_common_data();
+	vb2_nv_set(ctx, VB2_NV_DIAG_REQUEST, 1);
+	vb2_set_boot_mode(ctx);
+	TEST_EQ(ctx->boot_mode, VB2_BOOT_MODE_NORMAL,
+		"Normal boot with diag UI disabled");
+
+	/* Developer boot */
+	reset_common_data();
+	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	sd->flags |= VB2_SD_FLAG_DEV_MODE_ENABLED;
+	vb2_set_boot_mode(ctx);
+	TEST_EQ(ctx->boot_mode, VB2_BOOT_MODE_DEVELOPER, "Dev boot");
+
+	/* Recovery boot */
+	reset_common_data();
+	sd->recovery_reason = 123;
+	ctx->flags |= VB2_CONTEXT_RECOVERY_MODE;
+	vb2_set_boot_mode(ctx);
+	TEST_EQ(ctx->boot_mode, VB2_BOOT_MODE_BROKEN_SCREEN, "Broken screen");
+
+	reset_common_data();
+	sd->recovery_reason = VB2_RECOVERY_RO_MANUAL;
+	ctx->flags |= VB2_CONTEXT_RECOVERY_MODE;
+	gbb.flags |= VB2_GBB_FLAG_FORCE_MANUAL_RECOVERY;
+	vb2_set_boot_mode(ctx);
+	TEST_EQ(ctx->boot_mode, VB2_BOOT_MODE_MANUAL_RECOVERY,
+		"Manual recovery: forced by GBB flags");
+
+	reset_common_data();
+	sd->recovery_reason = VB2_RECOVERY_RO_MANUAL;
+	ctx->flags |= VB2_CONTEXT_FORCE_RECOVERY_MODE;
+	ctx->flags |= VB2_CONTEXT_EC_TRUSTED;
+	vb2_set_boot_mode(ctx);
+	TEST_EQ(ctx->boot_mode, VB2_BOOT_MODE_MANUAL_RECOVERY,
+		"Manual recovery: physical rec switch");
+
+	reset_common_data();
+	ctx->flags |= VB2_CONTEXT_EC_TRUSTED;
+	vb2_set_boot_mode(ctx);
+	TEST_NEQ(ctx->boot_mode, VB2_BOOT_MODE_MANUAL_RECOVERY,
+		 "VB2_CONTEXT_FORCE_RECOVERY_MODE is not set");
+
+	reset_common_data();
+	ctx->flags |= VB2_CONTEXT_FORCE_RECOVERY_MODE;
+	ctx->flags |= VB2_CONTEXT_NO_BOOT;
+	ctx->flags |= VB2_CONTEXT_EC_TRUSTED;
+	vb2_set_boot_mode(ctx);
+	TEST_NEQ(ctx->boot_mode, VB2_BOOT_MODE_MANUAL_RECOVERY,
+		 "Block manual recovery if NO_BOOT");
+
+	reset_common_data();
+	ctx->flags |= VB2_CONTEXT_FORCE_RECOVERY_MODE;
+	vb2_set_boot_mode(ctx);
+	TEST_NEQ(ctx->boot_mode, VB2_BOOT_MODE_MANUAL_RECOVERY,
+		 "Block manual recovery for untrusted EC");
+}
+
+int main(void)
+{
+	set_boot_mode_tests();
+
+	return gTestSuccess ? 0 : 255;
+}
diff --git a/tests/vb2_preamble_fuzzer.c b/tests/vb2_preamble_fuzzer.c
index 186ee9fa..92c240e7 100644
--- a/tests/vb2_preamble_fuzzer.c
+++ b/tests/vb2_preamble_fuzzer.c
@@ -9,7 +9,6 @@
 #include "2rsa.h"
 #include "2rsa_private.h"
 #include "2secdata.h"
-#include "vboot_test.h"
 
 static struct vb2_context *ctx;
 static uint8_t workbuf[VB2_FIRMWARE_WORKBUF_RECOMMENDED_SIZE]
diff --git a/tests/vb2_rsa_utility_tests.c b/tests/vb2_rsa_utility_tests.c
index 4ad5a95b..18347348 100644
--- a/tests/vb2_rsa_utility_tests.c
+++ b/tests/vb2_rsa_utility_tests.c
@@ -14,7 +14,6 @@
 #include "rsa_padding_test.h"
 #include "test_common.h"
 #include "vboot_api.h"
-#include "vboot_test.h"
 
 /**
  * Test RSA utility funcs
diff --git a/tests/vb2_secdata_kernel_tests.c b/tests/vb2_secdata_kernel_tests.c
index dc41f1a1..5ac63ee5 100644
--- a/tests/vb2_secdata_kernel_tests.c
+++ b/tests/vb2_secdata_kernel_tests.c
@@ -13,7 +13,6 @@
 #include "2secdata_struct.h"
 #include "2sysincludes.h"
 #include "test_common.h"
-#include "vboot_test.h"
 
 static uint8_t workbuf[VB2_FIRMWARE_WORKBUF_RECOMMENDED_SIZE]
 	__attribute__((aligned(VB2_WORKBUF_ALIGN)));
diff --git a/tests/vboot_api_kernel4_tests.c b/tests/vboot_api_kernel4_tests.c
deleted file mode 100644
index 6f210d38..00000000
--- a/tests/vboot_api_kernel4_tests.c
+++ /dev/null
@@ -1,342 +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 vboot_api_kernel, part 4 - select and load kernel
- */
-
-#include "2api.h"
-#include "2common.h"
-#include "2misc.h"
-#include "2nvstorage.h"
-#include "2secdata.h"
-#include "2sysincludes.h"
-#include "host_common.h"
-#include "load_kernel_fw.h"
-#include "test_common.h"
-#include "tlcl.h"
-#include "tss_constants.h"
-#include "vboot_struct.h"
-#include "vboot_test.h"
-
-/* Mock data */
-
-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 VbSelectAndLoadKernelParams kparams;
-static struct vb2_gbb_header gbb;
-
-static uint32_t kernel_version;
-static uint32_t new_version;
-static vb2_error_t vbboot_retval;
-static int commit_data_called;
-static vb2_error_t secdata_kernel_init_retval;
-static vb2_error_t secdata_fwmp_init_retval;
-static vb2_error_t kernel_phase1_retval;
-static uint32_t current_recovery_reason;
-static uint32_t expected_recovery_reason;
-
-static int mock_diagnostic_ui_enabled;
-
-static void reset_common_data(void)
-{
-	memset(&kparams, 0, sizeof(kparams));
-
-	memset(&gbb, 0, sizeof(gbb));
-	gbb.major_version = VB2_GBB_MAJOR_VER;
-	gbb.minor_version = VB2_GBB_MINOR_VER;
-	gbb.flags = 0;
-
-	TEST_SUCC(vb2api_init(workbuf, sizeof(workbuf), &ctx),
-		  "vb2api_init failed");
-	sd = vb2_get_sd(ctx);
-	sd->flags |= VB2_SD_FLAG_DISPLAY_AVAILABLE;
-	ctx->flags |= VB2_CONTEXT_NO_SECDATA_FWMP;
-	sd->preamble_size = 1;
-
-	vb2_nv_init(ctx);
-	vb2_nv_set(ctx, VB2_NV_KERNEL_MAX_ROLLFORWARD, 0xffffffff);
-	commit_data_called = 0;
-
-	kernel_version = new_version = 0x10002;
-	vbboot_retval = VB2_SUCCESS;
-	secdata_kernel_init_retval = VB2_SUCCESS;
-	secdata_fwmp_init_retval = VB2_SUCCESS;
-	kernel_phase1_retval = VB2_SUCCESS;
-	current_recovery_reason = 0;
-	expected_recovery_reason = 0;
-
-	mock_diagnostic_ui_enabled = 0;
-
-	sd->status |= VB2_SD_STATUS_SECDATA_KERNEL_INIT;
-	sd->status |= VB2_SD_STATUS_SECDATA_FWMP_INIT;
-}
-
-static void test_slk(vb2_error_t retval, int recovery_reason, const char *desc)
-{
-	if (sd->recovery_reason)
-		ctx->flags |= VB2_CONTEXT_RECOVERY_MODE;
-
-	expected_recovery_reason = recovery_reason;
-	/* The VbSelectAndLoadKernel directly leverages the value at
-	   ctx->boot_mode, so we have to call vb2_set_boot_mode first. */
-	vb2_set_boot_mode(ctx);
-	TEST_EQ(VbSelectAndLoadKernel(ctx, &kparams), retval, desc);
-	TEST_EQ(vb2_nv_get(ctx, VB2_NV_RECOVERY_REQUEST),
-		recovery_reason, "  recovery reason");
-}
-
-/* Mock functions */
-
-struct vb2_gbb_header *vb2_get_gbb(struct vb2_context *c)
-{
-	return &gbb;
-}
-
-vb2_error_t vb2api_kernel_phase1(struct vb2_context *c)
-{
-	sd->kernel_version_secdata = kernel_version;
-	sd->kernel_version = kernel_version;
-	return kernel_phase1_retval;
-}
-
-vb2_error_t vb2ex_commit_data(struct vb2_context *c)
-{
-	current_recovery_reason = vb2_nv_get(ctx, VB2_NV_RECOVERY_REQUEST);
-	commit_data_called = 1;
-	return VB2_SUCCESS;
-}
-
-vb2_error_t vb2_secdata_fwmp_init(struct vb2_context *c)
-{
-	return secdata_fwmp_init_retval;
-}
-
-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)
-{
-	sd->kernel_version = new_version;
-
-	if (vbboot_retval == -1)
-		return VB2_ERROR_MOCK;
-
-	return vbboot_retval;
-}
-
-static vb2_error_t boot_dev(struct vb2_context *c)
-{
-	if (vbboot_retval == -2)
-		return VB2_ERROR_MOCK;
-
-	return vbboot_retval;
-}
-
-vb2_error_t vb2ex_developer_ui(struct vb2_context *c)
-{
-	return boot_dev(c);
-}
-
-static void rec_check(struct vb2_context *c)
-{
-	TEST_EQ(current_recovery_reason, expected_recovery_reason,
-		"  recovery reason");
-	TEST_TRUE(commit_data_called, "  commit data");
-}
-
-vb2_error_t vb2ex_manual_recovery_ui(struct vb2_context *c)
-{
-	rec_check(c);
-	if (vbboot_retval == -3)
-		return VB2_ERROR_MOCK;
-	return vbboot_retval;
-}
-
-vb2_error_t vb2ex_broken_screen_ui(struct vb2_context *c)
-{
-	rec_check(c);
-	if (vbboot_retval == -4)
-		return VB2_ERROR_MOCK;
-	return vbboot_retval;
-}
-
-vb2_error_t vb2ex_diagnostic_ui(struct vb2_context *c)
-{
-	if (vbboot_retval == -5)
-		return VB2_ERROR_MOCK;
-
-	return vbboot_retval;
-}
-
-int vb2api_diagnostic_ui_enabled(struct vb2_context *c)
-{
-	return mock_diagnostic_ui_enabled;
-}
-
-vb2_error_t vb2ex_tpm_set_mode(enum vb2_tpm_mode mode_val)
-{
-	return VB2_SUCCESS;
-}
-
-/* Tests */
-
-static void select_and_load_kernel_tests(void)
-{
-	/* Normal boot */
-	reset_common_data();
-	test_slk(0, 0, "Normal");
-	TEST_EQ(kernel_version, 0x10002, "  version");
-	TEST_NEQ(sd->status & VB2_SD_STATUS_EC_SYNC_COMPLETE, 0,
-		 "  EC sync complete");
-	TEST_FALSE(commit_data_called, "  no commit data");
-
-	/* Check EC sync toggling */
-	reset_common_data();
-	ctx->flags |= VB2_CONTEXT_EC_SYNC_SUPPORTED;
-	gbb.flags |= VB2_GBB_FLAG_DISABLE_EC_SOFTWARE_SYNC;
-	test_slk(0, 0, "EC sync disabled by GBB");
-	TEST_NEQ(sd->status & VB2_SD_STATUS_EC_SYNC_COMPLETE, 0,
-		 "  EC sync complete");
-
-	reset_common_data();
-	ctx->flags |= VB2_CONTEXT_EC_SYNC_SUPPORTED;
-	test_slk(0, 0, "Normal with EC sync");
-	TEST_NEQ(sd->status & VB2_SD_STATUS_EC_SYNC_COMPLETE, 0,
-		 "  EC sync complete");
-
-	reset_common_data();
-	new_version = 0x20003;
-	test_slk(0, 0, "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_slk(0, 0, "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_slk(0, 0, "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_slk(0, 0, "Max roll forward can't rollback");
-	TEST_EQ(kernel_version, 0x10002, "  version");
-
-	/* Boot normal */
-	reset_common_data();
-	vbboot_retval = -1;
-	test_slk(VB2_ERROR_MOCK, 0, "Normal boot bad");
-
-	/* Check that NV_DIAG_REQUEST triggers diagnostic UI */
-	reset_common_data();
-	mock_diagnostic_ui_enabled = 1;
-	vb2_nv_set(ctx, VB2_NV_DIAG_REQUEST, 1);
-	vbboot_retval = -5;
-	test_slk(VB2_ERROR_MOCK, 0,
-		 "Normal boot with diag enabled");
-	TEST_EQ(vb2_nv_get(ctx, VB2_NV_DIAG_REQUEST), 0,
-		"  diag not requested");
-
-	reset_common_data();
-	vb2_nv_set(ctx, VB2_NV_DIAG_REQUEST, 1);
-	test_slk(0, 0, "Normal boot with diag disabled (reboot to unset)");
-
-	/* Boot normal - phase1 failure */
-	reset_common_data();
-	kernel_phase1_retval = VB2_ERROR_MOCK;
-	test_slk(VB2_ERROR_MOCK, 0, "Normal phase1 failure");
-
-	/* Recovery - VB2_ERROR_ESCAPE_NO_BOOT */
-	reset_common_data();
-	ctx->flags |= VB2_CONTEXT_NO_BOOT;
-	ctx->flags |= VB2_CONTEXT_EC_SYNC_SUPPORTED;
-	test_slk(VB2_ERROR_ESCAPE_NO_BOOT,
-		 VB2_RECOVERY_ESCAPE_NO_BOOT, "Recovery for NO_BOOT escape");
-
-	/* Boot normal - VB2_ERROR_ESCAPE_NO_BOOT */
-	reset_common_data();
-	ctx->flags |= VB2_CONTEXT_NO_BOOT;
-	ctx->flags |= VB2_CONTEXT_EC_SYNC_SUPPORTED;
-	gbb.flags |= VB2_GBB_FLAG_DISABLE_EC_SOFTWARE_SYNC;
-	test_slk(VB2_SUCCESS, 0, "DISABLE_EC_SOFTWARE_SYNC ignores NO_BOOT");
-
-	/* Boot dev */
-	reset_common_data();
-	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
-	vbboot_retval = -2;
-	test_slk(VB2_ERROR_MOCK, 0, "Dev boot bad");
-	TEST_FALSE(commit_data_called, "  no commit data");
-
-	reset_common_data();
-	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
-	new_version = 0x20003;
-	test_slk(0, 0, "Dev doesn't roll forward");
-	TEST_EQ(kernel_version, 0x10002, "  version");
-
-	/* Boot dev - phase1 failure */
-	reset_common_data();
-	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
-	kernel_phase1_retval = VB2_ERROR_MOCK;
-	test_slk(VB2_ERROR_MOCK, 0, "Dev phase1 failure");
-
-	/* Boot recovery */
-	reset_common_data();
-	sd->recovery_reason = 123;
-	vbboot_retval = -4;
-	test_slk(VB2_ERROR_MOCK, 0, "Recovery boot bad");
-	TEST_TRUE(commit_data_called, "  commit data");
-
-	reset_common_data();
-	sd->recovery_reason = 123;
-	new_version = 0x20003;
-	test_slk(0, 0, "Recovery doesn't roll forward");
-	TEST_EQ(kernel_version, 0x10002, "  version");
-	TEST_TRUE(commit_data_called, "  commit data");
-
-	/* Boot recovery - phase1 failure */
-	reset_common_data();
-	sd->recovery_reason = 123;
-	kernel_phase1_retval = VB2_ERROR_MOCK;
-	test_slk(VB2_ERROR_MOCK, 0, "Recovery phase1 failure");
-	TEST_FALSE(commit_data_called, "  no commit data");
-
-	/* Boot recovery - memory retraining */
-	reset_common_data();
-	sd->recovery_reason = VB2_RECOVERY_TRAIN_AND_REBOOT;
-	test_slk(VB2_REQUEST_REBOOT, 0, "Recovery train and reboot");
-	TEST_FALSE(commit_data_called, "  no commit data");
-
-	/* Boot BROKEN recovery */
-	reset_common_data();
-	sd->recovery_reason = 123;
-	vb2_nv_set(ctx, VB2_NV_RECOVERY_SUBCODE, 13);
-	test_slk(0, 0, "BROKEN recovery");
-	TEST_TRUE(commit_data_called, "  commit data");
-
-	/* Boot manual recovery */
-	reset_common_data();
-	sd->recovery_reason = VB2_RECOVERY_RO_MANUAL;
-	vb2_nv_set(ctx, VB2_NV_RECOVERY_SUBCODE, 13);
-	ctx->flags &= VB2_CONTEXT_FORCE_RECOVERY_MODE;
-	test_slk(0, 0, "Manual recovery");
-	TEST_TRUE(commit_data_called, "  commit data");
-}
-
-int main(void)
-{
-	select_and_load_kernel_tests();
-
-	return gTestSuccess ? 0 : 255;
-}
diff --git a/tests/vboot_api_kernel_tests.c b/tests/vboot_api_kernel_tests.c
index b2d4236e..f885cd86 100644
--- a/tests/vboot_api_kernel_tests.c
+++ b/tests/vboot_api_kernel_tests.c
@@ -12,7 +12,6 @@
 #include "load_kernel_fw.h"
 #include "test_common.h"
 #include "vboot_api.h"
-#include "vboot_test.h"
 
 #define MAX_TEST_DISKS 10
 #define DEFAULT_COUNT -1
@@ -452,8 +451,7 @@ static void ResetMocks(test_case_t *test_case)
 	TEST_SUCC(vb2api_init(workbuf, sizeof(workbuf), &ctx),
 		  "vb2api_init failed");
 
-	memset(&kparams, 0, sizeof(VbSelectAndLoadKernelParams));
-	*VbApiKernelGetParamsPtr() = &kparams;
+	memset(&kparams, 0, sizeof(kparams));
 
 
 	memset(&mock_disks, 0, sizeof(mock_disks));
@@ -591,7 +589,7 @@ static void VbTryLoadKernelTest(void)
 		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),
+		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");
@@ -618,7 +616,7 @@ static void VbTryLoadMiniOsKernelTest(void)
 		printf("Test case: %s ...\n", minios_tests[i].name);
 		ResetMocks(&minios_tests[i]);
 		ctx->flags = t->ctx_flags;
-		TEST_EQ(VbTryLoadMiniOsKernel(ctx, 0),
+		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");