Reland "2api: Add vb2_boot_mode in vb2_context denoted the current boot mode"stabilize-14396.Bstabilize-14395.B

This reverts commit f0ba3fe277fdc623565eab0725aed743e24a1a44.

Reason for revert: We can try another reland after the coreboot
CL https://review.coreboot.org/c/coreboot/+/58253 is
downstreamed and the guybrush psp_verstage.bin is signed again.

Original change's description:
> Revert "2api: Add vb2_boot_mode in vb2_context denoted the current boot mode"
>
> This reverts commit 2cf76574dbf57c70da3a835f2a45d5619dc058f5.
>
> Reason for revert: b:202258389
>
> Original change's description:
> > 2api: Add vb2_boot_mode in vb2_context denoted the current boot mode
> >
> > Add enum vb2_boot_mode which denotes the most relevant boot mode between
> > normal, manual recovery, broken recovery, diagnostic, and developer.
> >
> > The boot mode constant stores in vb2_context, is exposed externally, and
> > decided in vb2api_fw_phase1.
> >
> > Split out the logic of manual recovery and broken screen. (The broken
> > recovery is a recovery boot with !vb2api_allow_recovery)
> >
> > Add the fifth boot mode, diagnostic boot mode
> >
> > A boot could match more the one boot mode, this api will pick the most
> > relevant boot mode based on the following order:
> > 1. Manual recovery boot
> > 2. Broken recovery boot
> > 3. Diagnostics boot
> > 4. Developer boot
> > 5. Normal boot
> >
> > This constant is used in:
> > * lib/vboot_api_kernel.c: VbSelectAndLoadKernel()
> > * lib/vboot_kernel.c: replace the original vb2_boot_mode
> > * coreboot/bootmode for adding elogs
> >
> > Also bump the vb2sd minor version from 0 to 1.
> >
> > BUG=b:185551931, b:177196147, b:181931817
> > BRANCH=none
> > TEST=CC=x86_64-pc-linux-gnu-clang;
> >      make clean && make runtests
> > TEST=emerge coreboot vboot_reference depthcharge
> >
> > Signed-off-by: Hsuan Ting Chen <roccochen@chromium.org>
> > Change-Id: I421e4d51c261ba2bdec996a5fb2ebccb33513fa4
> > Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2944250
> > Reviewed-by: Yu-Ping Wu <yupingso@chromium.org>
>
> Bug: b:185551931, b:177196147, b:181931817
> Change-Id: Ic184a24f5b580bcaf40d75045476ae77c674dd3d
> Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/3209725
> Bot-Commit: Rubber Stamper <rubber-stamper@appspot.gserviceaccount.com>
> Commit-Queue: Bhanu Prakash Maiya <bhanumaiya@google.com>

Bug: b:185551931, b:177196147, b:181931817
Change-Id: Ie480bf76eb0164f6e498b72e0533cd5c5762d6f6
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/3274699
Reviewed-by: Hsuan Ting Chen <roccochen@chromium.org>
Tested-by: Hsuan Ting Chen <roccochen@chromium.org>
Reviewed-by: Yu-Ping Wu <yupingso@chromium.org>
Commit-Queue: Hsuan Ting Chen <roccochen@chromium.org>

10 files changed, 153 insertions, 52 deletions

diff --git a/firmware/2lib/2api.c b/firmware/2lib/2api.c
index 593a12cf..53e8e98e 100644
--- a/firmware/2lib/2api.c
+++ b/firmware/2lib/2api.c
@@ -89,6 +89,9 @@ vb2_error_t vb2api_fw_phase1(struct vb2_context *ctx)
 	if (ctx->flags & VB2_CONTEXT_DISPLAY_INIT)
 		sd->flags |= VB2_SD_FLAG_DISPLAY_AVAILABLE;
 
+	/* Decide the boot mode */
+	vb2_set_boot_mode(ctx);
+
 	/* Return error if recovery is needed */
 	if (ctx->flags & VB2_CONTEXT_RECOVERY_MODE) {
 		/* Always clear RAM when entering recovery mode */
diff --git a/firmware/2lib/2misc.c b/firmware/2lib/2misc.c
index b1696365..c1d860f1 100644
--- a/firmware/2lib/2misc.c
+++ b/firmware/2lib/2misc.c
@@ -733,3 +733,24 @@ char *vb2api_get_debug_info(struct vb2_context *ctx)
 	buf[DEBUG_INFO_MAX_LENGTH] = '\0';
 	return buf;
 }
+
+void vb2_set_boot_mode(struct vb2_context *ctx)
+{
+	struct vb2_shared_data *sd = vb2_get_sd(ctx);
+
+	/* Cast boot mode to non-constant and assign */
+	enum vb2_boot_mode *boot_mode = (enum vb2_boot_mode *)&ctx->boot_mode;
+	*boot_mode = VB2_BOOT_MODE_NORMAL;
+
+	if (sd->recovery_reason) {
+		if (vb2api_allow_recovery(ctx))
+			*boot_mode = VB2_BOOT_MODE_MANUAL_RECOVERY;
+		else
+			*boot_mode = VB2_BOOT_MODE_BROKEN_SCREEN;
+	} else if (vb2api_diagnostic_ui_enabled(ctx) &&
+		   vb2_nv_get(ctx, VB2_NV_DIAG_REQUEST)) {
+		*boot_mode = VB2_BOOT_MODE_DIAGNOSTICS;
+	} else if (ctx->flags & VB2_CONTEXT_DEVELOPER_MODE) {
+		*boot_mode = VB2_BOOT_MODE_DEVELOPER;
+	}
+}
diff --git a/firmware/2lib/include/2api.h b/firmware/2lib/include/2api.h
index 8966999d..cd558b05 100644
--- a/firmware/2lib/include/2api.h
+++ b/firmware/2lib/include/2api.h
@@ -259,6 +259,58 @@ enum vb2_context_flags {
 	VB2_CONTEXT_DEV_BOOT_ALTFW_ALLOWED = (1 << 27),
 };
 
+/* Boot mode decided in vb2api_fw_phase1.
+ *
+ * Boot mode is a constant set by verified boot and may be read (but should not
+ * be set or cleared) by the caller.
+ * The boot modes are mutually exclusive. If a boot fulfill more than one
+ * constraints of the listing boot modes, it will be set to the most important
+ * one. The priority is the same as the listing order.
+ */
+enum vb2_boot_mode {
+	/* Undefined, The boot mode is not set. */
+	VB2_BOOT_MODE_UNDEFINED = 0,
+
+	/*
+	 * Manual recovery boot, regardless of dev mode state.
+	 *
+	 * VB2_CONTEXT_RECOVERY_MODE is set and the recovery is physically
+	 * requested (a.k.a. Manual recovery).  All other recovery requests
+	 * including manual recovery requested by a (compromised) host will end
+	 * up with a broken screen.
+	 */
+	VB2_BOOT_MODE_MANUAL_RECOVERY = 1,
+
+	/*
+	 * Broken screen.
+	 *
+	 * If a recovery boot is not a manual recovery (a.k.a. not requested
+	 * physically), the recovery is not allowed and will end up with
+	 * broken screen.
+	 */
+	VB2_BOOT_MODE_BROKEN_SCREEN = 2,
+
+	/*
+	 * Diagnostic boot.
+	 *
+	 * If diagnostic boot is enabled (a.k.a. vb2api_diagnostic_ui_enabled)
+	 * and the nvdata contains VB2_NV_DIAG_REQUEST from previous boot, it
+	 * will boot to diagnostic mode.
+	 */
+	VB2_BOOT_MODE_DIAGNOSTICS = 3,
+
+	/*
+	 * Developer boot: self-signed kernel okay.
+	 *
+	 * The developer mode switch is set (a.k.a. VB2_CONTEXT_DEVELOPER_MODE)
+	 * and we are in the developer boot mode.
+	 */
+	VB2_BOOT_MODE_DEVELOPER = 4,
+
+	/* Normal boot: kernel must be verified. */
+	VB2_BOOT_MODE_NORMAL = 5,
+};
+
 /* Helper for aligning fields in vb2_context. */
 #define VB2_PAD_STRUCT3(size, align, count) \
 	uint8_t _pad##count[align - (((size - 1) % align) + 1)]
@@ -333,6 +385,16 @@ struct vb2_context {
 	 */
 	uint8_t secdata_fwmp[VB2_SECDATA_FWMP_MAX_SIZE];
 	VB2_PAD_STRUCT(VB2_SECDATA_FWMP_MAX_SIZE, 8);
+
+	/**********************************************************************
+	 * Fields below added in struct version 3.1.
+	 */
+
+	/*
+	 * Mutually exclusive boot mode.
+	 * This constant is initialized after calling vb2api_fw_phase1().
+	 */
+	const enum vb2_boot_mode boot_mode;
 };
 
 /* Resource index for vb2ex_read_resource() */
diff --git a/firmware/2lib/include/2misc.h b/firmware/2lib/include/2misc.h
index b36e1270..3d292872 100644
--- a/firmware/2lib/include/2misc.h
+++ b/firmware/2lib/include/2misc.h
@@ -191,4 +191,19 @@ void vb2_clear_recovery(struct vb2_context *ctx);
  */
 void vb2_fill_dev_boot_flags(struct vb2_context *ctx);
 
+/**
+ * Determine and set a mutually exclusive boot mode in the vboot context.
+ *
+ * Determine the most relevant boot mode for current boot, store into
+ * ctx->boot_mode, which is a ctx field introduced in struct version 3.1.
+ *
+ * This function should be only called by vb2api_fw_phase1.
+ * The vb2api_fw_phase1 should call this function at its end phase once and all
+ * the following steps should directly access ctx->boot_mode to retrieve the
+ * most relevant boot mode.
+ *
+ * @param ctx		Vboot context.
+ */
+void vb2_set_boot_mode(struct vb2_context *ctx);
+
 #endif  /* VBOOT_REFERENCE_2MISC_H_ */
diff --git a/firmware/2lib/include/2struct.h b/firmware/2lib/include/2struct.h
index 08a1efe3..310f4bc5 100644
--- a/firmware/2lib/include/2struct.h
+++ b/firmware/2lib/include/2struct.h
@@ -94,7 +94,7 @@ enum vb2_shared_data_status {
 
 /* Current version of vb2_shared_data struct */
 #define VB2_SHARED_DATA_VERSION_MAJOR 3
-#define VB2_SHARED_DATA_VERSION_MINOR 0
+#define VB2_SHARED_DATA_VERSION_MINOR 1
 
 /* MAX_SIZE should not be changed without bumping up DATA_VERSION_MAJOR. */
 #define VB2_CONTEXT_MAX_SIZE 384
diff --git a/firmware/lib/vboot_api_kernel.c b/firmware/lib/vboot_api_kernel.c
index 54447c26..85a8f4b4 100644
--- a/firmware/lib/vboot_api_kernel.c
+++ b/firmware/lib/vboot_api_kernel.c
@@ -182,13 +182,15 @@ vb2_error_t VbSelectAndLoadKernel(struct vb2_context *ctx,
 	}
 
 	/*
-	 * If in non-manual recovery mode, save the recovery reason as subcode.
+	 * If in the broken screen, save the recovery reason as subcode.
 	 * Otherwise, clear any leftover recovery requests or subcodes.
 	 */
 	vb2_clear_recovery(ctx);
 
 	/* Select boot path */
-	if (ctx->flags & VB2_CONTEXT_RECOVERY_MODE) {
+	switch (ctx->boot_mode) {
+	case VB2_BOOT_MODE_MANUAL_RECOVERY:
+	case VB2_BOOT_MODE_BROKEN_SCREEN:
 		/* If we're in recovery mode just to do memory retraining, all
 		   we need to do is reboot. */
 		if (sd->recovery_reason == VB2_RECOVERY_TRAIN_AND_REBOOT) {
@@ -213,12 +215,12 @@ vb2_error_t VbSelectAndLoadKernel(struct vb2_context *ctx,
 			VB2_DEBUG("NO_BOOT in RECOVERY mode\n");
 
 		/* Recovery boot.  This has UI. */
-		if (vb2api_allow_recovery(ctx))
+		if (ctx->boot_mode == VB2_BOOT_MODE_MANUAL_RECOVERY)
 			VB2_TRY(vb2ex_manual_recovery_ui(ctx));
 		else
 			VB2_TRY(vb2ex_broken_screen_ui(ctx));
-	} else if (vb2api_diagnostic_ui_enabled(ctx) &&
-		   vb2_nv_get(ctx, VB2_NV_DIAG_REQUEST)) {
+		break;
+	case VB2_BOOT_MODE_DIAGNOSTICS:
 		/*
 		 * Need to clear the request flag and commit nvdata changes
 		 * immediately to avoid booting back into diagnostic tool when a
@@ -234,12 +236,16 @@ vb2_error_t VbSelectAndLoadKernel(struct vb2_context *ctx,
 		 * return either of reboot or shutdown.
 		 */
 		return VB2_REQUEST_REBOOT;
-	} else if (ctx->flags & VB2_CONTEXT_DEVELOPER_MODE) {
+	case VB2_BOOT_MODE_DEVELOPER:
 		/* Developer boot.  This has UI. */
 		VB2_TRY(vb2ex_developer_ui(ctx));
-	} else {
+		break;
+	case VB2_BOOT_MODE_NORMAL:
 		/* Normal boot */
 		VB2_TRY(vb2_normal_boot(ctx));
+		break;
+	default:
+		return VB2_ERROR_ESCAPE_NO_BOOT;
 	}
 
 	/*
diff --git a/firmware/lib/vboot_kernel.c b/firmware/lib/vboot_kernel.c
index cb83d40d..1edf4a55 100644
--- a/firmware/lib/vboot_kernel.c
+++ b/firmware/lib/vboot_kernel.c
@@ -6,6 +6,7 @@
  * (Firmware portion)
  */
 
+#include "2api.h"
 #include "2common.h"
 #include "2misc.h"
 #include "2nvstorage.h"
@@ -31,34 +32,6 @@ enum vb2_load_partition_flags {
 
 #define LOWEST_TPM_VERSION 0xffffffff
 
-enum vb2_boot_mode {
-	/* Normal boot: kernel must be verified. */
-	VB2_BOOT_MODE_NORMAL = 0,
-
-	/* Recovery boot, regardless of dev mode state. */
-	VB2_BOOT_MODE_RECOVERY = 1,
-
-	/* Developer boot: self-signed kernel okay. */
-	VB2_BOOT_MODE_DEVELOPER = 2,
-};
-
-/**
- * Return the current boot mode (normal, recovery, or dev).
- *
- * @param ctx          Vboot context
- * @return Current boot mode (see vb2_boot_mode enum).
- */
-static enum vb2_boot_mode get_boot_mode(struct vb2_context *ctx)
-{
-	if (ctx->flags & VB2_CONTEXT_RECOVERY_MODE)
-		return VB2_BOOT_MODE_RECOVERY;
-
-	if (ctx->flags & VB2_CONTEXT_DEVELOPER_MODE)
-		return VB2_BOOT_MODE_DEVELOPER;
-
-	return VB2_BOOT_MODE_NORMAL;
-}
-
 /**
  * Check if a valid keyblock is required.
  *
@@ -69,7 +42,7 @@ static enum vb2_boot_mode get_boot_mode(struct vb2_context *ctx)
 static int need_valid_keyblock(struct vb2_context *ctx)
 {
 	/* Normal and recovery modes always require official OS */
-	if (get_boot_mode(ctx) != VB2_BOOT_MODE_DEVELOPER)
+	if (ctx->boot_mode != VB2_BOOT_MODE_DEVELOPER)
 		return 1;
 
 	/* FWMP can require developer mode to use signed kernels */
@@ -264,9 +237,8 @@ static vb2_error_t vb2_verify_kernel_vblock(
 	}
 
 	/* Check for rollback of key version except in recovery mode. */
-	enum vb2_boot_mode boot_mode = get_boot_mode(ctx);
 	uint32_t key_version = keyblock->data_key.key_version;
-	if (boot_mode != VB2_BOOT_MODE_RECOVERY) {
+	if (ctx->boot_mode != VB2_BOOT_MODE_MANUAL_RECOVERY) {
 		if (key_version < (sd->kernel_version_secdata >> 16)) {
 			keyblock_valid = 0;
 			if (need_keyblock_valid) {
@@ -288,7 +260,7 @@ static vb2_error_t vb2_verify_kernel_vblock(
 	}
 
 	/* If in developer mode and using key hash, check it. */
-	if (boot_mode == VB2_BOOT_MODE_DEVELOPER &&
+	if (ctx->boot_mode == VB2_BOOT_MODE_DEVELOPER &&
 	    vb2_secdata_fwmp_get_flag(ctx, VB2_SECDATA_FWMP_DEV_USE_KEY_HASH)) {
 		VB2_TRY(vb2_verify_kernel_dev_key_hash(ctx, keyblock));
 	}
@@ -356,7 +328,7 @@ static vb2_error_t vb2_verify_kernel_vblock(
 
 	/* If not in recovery mode, check for rollback of the kernel version. */
 	if (need_keyblock_valid &&
-	    boot_mode != VB2_BOOT_MODE_RECOVERY &&
+	    ctx->boot_mode != VB2_BOOT_MODE_MANUAL_RECOVERY &&
 	    sd->kernel_version < sd->kernel_version_secdata) {
 		VB2_DEBUG("Kernel version too low.\n");
 		return VB2_ERROR_KERNEL_PREAMBLE_VERSION_ROLLBACK;
@@ -744,7 +716,7 @@ vb2_error_t LoadKernel(struct vb2_context *ctx,
 		 * non-officially-signed kernel, there's no rollback
 		 * protection, so we can stop at the first valid kernel.
 		 */
-		if (get_boot_mode(ctx) == VB2_BOOT_MODE_RECOVERY ||
+		if (ctx->boot_mode == VB2_BOOT_MODE_MANUAL_RECOVERY ||
 		    !keyblock_valid) {
 			VB2_DEBUG("In recovery mode or dev-signed kernel\n");
 			break;
diff --git a/tests/vboot_api_kernel4_tests.c b/tests/vboot_api_kernel4_tests.c
index f122eb38..c1e71afa 100644
--- a/tests/vboot_api_kernel4_tests.c
+++ b/tests/vboot_api_kernel4_tests.c
@@ -80,6 +80,9 @@ static void test_slk(vb2_error_t retval, int recovery_reason, const char *desc)
 		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");
diff --git a/tests/vboot_kernel2_tests.c b/tests/vboot_kernel2_tests.c
index bc4ea0fe..5fc66f14 100644
--- a/tests/vboot_kernel2_tests.c
+++ b/tests/vboot_kernel2_tests.c
@@ -47,6 +47,7 @@ static struct vb2_shared_data *sd;
 static struct vb2_workbuf wb;
 static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
 	__attribute__((aligned(VB2_WORKBUF_ALIGN)));
+static enum vb2_boot_mode *boot_mode;
 
 static VbSelectAndLoadKernelParams lkp;
 static VbDiskInfo disk_info;
@@ -82,6 +83,9 @@ static void reset_common_data(void)
 	vb2_secdata_kernel_init(ctx);
 	ctx->flags = VB2_CONTEXT_RECOVERY_MODE;
 
+	boot_mode = (enum vb2_boot_mode *)&ctx->boot_mode;
+	*boot_mode = VB2_BOOT_MODE_MANUAL_RECOVERY;
+
 	sd = vb2_get_sd(ctx);
 	sd->kernel_version_secdata = 0xabcdef | (1 << 24);
 
diff --git a/tests/vboot_kernel_tests.c b/tests/vboot_kernel_tests.c
index 14e7e1f2..1f77cc5c 100644
--- a/tests/vboot_kernel_tests.c
+++ b/tests/vboot_kernel_tests.c
@@ -51,6 +51,7 @@ static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
 static struct vb2_context *ctx;
 static struct vb2_shared_data *sd;
 static struct vb2_packed_key mock_key;
+static enum vb2_boot_mode *boot_mode;
 
 /**
  * Reset mock data (for use before each test)
@@ -114,6 +115,9 @@ static void ResetMocks(void)
 	fwmp = (struct vb2_secdata_fwmp *)ctx->secdata_fwmp;
 	memcpy(&fwmp->dev_key_hash, mock_digest, sizeof(fwmp->dev_key_hash));
 
+	boot_mode = (enum vb2_boot_mode *)&ctx->boot_mode;
+	*boot_mode = VB2_BOOT_MODE_NORMAL;
+
 	// TODO: more workbuf fields - flags, secdata_firmware
 
 	vb2api_secdata_kernel_create(ctx);
@@ -337,6 +341,7 @@ static void LoadKernelTest(void)
 	/* In dev mode, fail if hash is bad too */
 	ResetMocks();
 	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = VB2_BOOT_MODE_DEVELOPER;
 	keyblock_verify_fail = 2;
 	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
 		       "Fail key block dev hash");
@@ -344,6 +349,7 @@ static void LoadKernelTest(void)
 	/* But just bad sig is ok */
 	ResetMocks();
 	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = VB2_BOOT_MODE_DEVELOPER;
 	keyblock_verify_fail = 1;
 	TestLoadKernel(0, "Succeed keyblock dev sig");
 	TEST_EQ(sd->flags & VB2_SD_FLAG_KERNEL_SIGNED, 0, "  use hash");
@@ -351,6 +357,7 @@ static void LoadKernelTest(void)
 	/* In dev mode and requiring signed kernel, fail if sig is bad */
 	ResetMocks();
 	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = 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,
@@ -358,6 +365,7 @@ static void LoadKernelTest(void)
 
 	ResetMocks();
 	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = VB2_BOOT_MODE_DEVELOPER;
 	fwmp->flags |= VB2_SECDATA_FWMP_DEV_ENABLE_OFFICIAL_ONLY;
 	keyblock_verify_fail = 1;
 	TestLoadKernel(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
@@ -387,6 +395,7 @@ static void LoadKernelTest(void)
 
 	ResetMocks();
 	ctx->flags |= VB2_CONTEXT_RECOVERY_MODE;
+	*boot_mode = VB2_BOOT_MODE_MANUAL_RECOVERY;
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_1
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_1
 		| VB2_KEYBLOCK_FLAG_MINIOS_0;
@@ -395,6 +404,7 @@ static void LoadKernelTest(void)
 
 	ResetMocks();
 	ctx->flags |= VB2_CONTEXT_RECOVERY_MODE;
+	*boot_mode = VB2_BOOT_MODE_MANUAL_RECOVERY;
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_1
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_0
 		| VB2_KEYBLOCK_FLAG_MINIOS_0;
@@ -402,6 +412,7 @@ static void LoadKernelTest(void)
 
 	ResetMocks();
 	ctx->flags |= VB2_CONTEXT_RECOVERY_MODE | VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = VB2_BOOT_MODE_MANUAL_RECOVERY;
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_1
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_0
 		| VB2_KEYBLOCK_FLAG_MINIOS_0;
@@ -410,6 +421,7 @@ static void LoadKernelTest(void)
 
 	ResetMocks();
 	ctx->flags |= VB2_CONTEXT_RECOVERY_MODE | VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = VB2_BOOT_MODE_MANUAL_RECOVERY;
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_1
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_1
 		| VB2_KEYBLOCK_FLAG_MINIOS_0;
@@ -418,6 +430,7 @@ static void LoadKernelTest(void)
 	/* Check keyblock flags (dev mode + signed kernel required) */
 	ResetMocks();
 	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = VB2_BOOT_MODE_DEVELOPER;
 	vb2_nv_set(ctx, VB2_NV_DEV_BOOT_SIGNED_ONLY, 1);
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_1
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_0
@@ -427,6 +440,7 @@ static void LoadKernelTest(void)
 
 	ResetMocks();
 	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = VB2_BOOT_MODE_DEVELOPER;
 	fwmp->flags |= VB2_SECDATA_FWMP_DEV_ENABLE_OFFICIAL_ONLY;
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_1
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_0
@@ -436,6 +450,7 @@ static void LoadKernelTest(void)
 
 	ResetMocks();
 	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = VB2_BOOT_MODE_DEVELOPER;
 	fwmp->flags |= VB2_SECDATA_FWMP_DEV_ENABLE_OFFICIAL_ONLY;
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_0
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_0
@@ -445,6 +460,7 @@ static void LoadKernelTest(void)
 
 	ResetMocks();
 	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = VB2_BOOT_MODE_DEVELOPER;
 	vb2_nv_set(ctx, VB2_NV_DEV_BOOT_SIGNED_ONLY, 1);
 	kbh.keyblock_flags = VB2_KEYBLOCK_FLAG_RECOVERY_0
 		| VB2_KEYBLOCK_FLAG_DEVELOPER_1
@@ -477,12 +493,12 @@ static void LoadKernelTest(void)
 
 	ResetMocks();
 	kbh.data_key.key_version = 1;
-	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = VB2_BOOT_MODE_DEVELOPER;
 	TestLoadKernel(0, "Key version ignored in dev mode");
 
 	ResetMocks();
 	kbh.data_key.key_version = 1;
-	ctx->flags |= VB2_CONTEXT_RECOVERY_MODE;
+	*boot_mode = VB2_BOOT_MODE_MANUAL_RECOVERY;
 	TestLoadKernel(0, "Key version ignored in rec mode");
 
 	ResetMocks();
@@ -500,18 +516,18 @@ static void LoadKernelTest(void)
 
 	ResetMocks();
 	kph.kernel_version = 0;
-	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = VB2_BOOT_MODE_DEVELOPER;
 	TestLoadKernel(0, "Kernel version ignored in dev mode");
 
 	ResetMocks();
 	kph.kernel_version = 0;
-	ctx->flags |= VB2_CONTEXT_RECOVERY_MODE;
+	*boot_mode = VB2_BOOT_MODE_MANUAL_RECOVERY;
 	TestLoadKernel(0, "Kernel version ignored in rec mode");
 
 	/* Check kernel version (dev mode + signed kernel required) */
 	ResetMocks();
 	kbh.data_key.key_version = 0;
-	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = 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 "
@@ -519,7 +535,7 @@ static void LoadKernelTest(void)
 
 	ResetMocks();
 	kbh.data_key.key_version = 0;
-	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = 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 "
@@ -527,7 +543,7 @@ static void LoadKernelTest(void)
 
 	/* Check developer key hash - bad */
 	ResetMocks();
-	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = 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,
@@ -535,15 +551,14 @@ static void LoadKernelTest(void)
 
 	/* Check developer key hash - bad (recovery mode) */
 	ResetMocks();
-	ctx->flags |= VB2_CONTEXT_RECOVERY_MODE;
-	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = VB2_BOOT_MODE_MANUAL_RECOVERY;
 	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");
 
 	/* Check developer key hash - good */
 	ResetMocks();
-	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+	*boot_mode = VB2_BOOT_MODE_DEVELOPER;
 	fwmp->flags |= VB2_SECDATA_FWMP_DEV_USE_KEY_HASH;
 	TestLoadKernel(0, "Good keyblock dev fwmp hash");