ui: Check external disk status on selecting "Back"

(Manually cherry picked from CL:3195521)

Consider the series of user actions in developer UI:

0. Screen DEVLOPER_MODE is shown
1. Insert an invalid external disk
2. Press Ctrl+U to trigger external boot
   => INVALID_DISK screen is shown
3. Unplug the invalid external disk
   => BOOT_EXTERNAL screen is shown
4. Select the "Back" button

Then, the screen will temporarily change to INVALID_DISK because it is
the previous screen in the history stack. However, in the next UI loop,
the action function of INVALID_DISK will change the screen to
BOOT_EXTERNAL because there is no external disk plugged in. This bug
causes the user to be stuck in INVALID_DISK and INVALID_DISK screens.

To solve the problem, add init() and reinit() hooks for these two
polling screens, so that when we enter either of the screens, the status
of the external disk will be checked. If it doesn't match the meaning of
the target screen, ui_screen_back() will be called to go back to the
previous screen.

This solution elegantly handles more complex situations like this:

0. Screen DEVLOPER_MODE is shown
   Stack: [DEVELOPER_MODE]
1. Insert an invalid external disk
2. Press Ctrl+U to trigger external boot
   => INVALID_DISK screen is shown
   Stack: [DEVELOPER_MODE, INVALID_DISK]
3. Press TAB
   => DEBUG_INFO screen is shown
   Stack: [DEVELOPER_MODE, INVALID_DISK, DEBUG_INFO]
4. Unplug the invalid external disk and press Ctrl+U
   => BOOT_EXTERNAL screen is shown
   Stack: [DEVELOPER_MODE, INVALID_DISK, DEBUG_INFO, BOOT_EXTERNAL]
5. Select "Back"
   => DEBUG_INFO screen is shown
   Stack: [DEVELOPER_MODE, INVALID_DISK, DEBUG_INFO]
6. Select "Back"
   => DEVELOPER_MODE screen is shown (because INVALID_DISK is skipped)
   Stack: [DEVELOPER_MODE]

BUG=b:201510767, b:203580383
TEST=make run2tests
BRANCH=volteer

Change-Id: Ic41321ba6aca60fd13f6d33859ec171663d68145
Signed-off-by: Yu-Ping Wu <yupingso@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/3233766
Reviewed-by: Nick Vaccaro <nvaccaro@google.com>

1 files changed, 43 insertions, 0 deletions

diff --git a/firmware/2lib/2ui_screens.c b/firmware/2lib/2ui_screens.c
index e3d7067e..662f5b4c 100644
--- a/firmware/2lib/2ui_screens.c
+++ b/firmware/2lib/2ui_screens.c
@@ -885,6 +885,32 @@ static const struct vb2_screen_info developer_to_norm_screen = {
 /******************************************************************************/
 /* VB2_SCREEN_DEVELOPER_BOOT_EXTERNAL */
 
+static vb2_error_t developer_boot_external_check(struct vb2_ui_context *ui)
+{
+	if (!(ui->ctx->flags & VB2_CONTEXT_DEVELOPER_MODE) ||
+	    !vb2_dev_boot_allowed(ui->ctx) ||
+	    !vb2_dev_boot_external_allowed(ui->ctx)) {
+		VB2_DEBUG("ERROR: Dev mode external boot not allowed\n");
+		ui->error_beep = 1;
+		ui->error_code = VB2_UI_ERROR_EXTERNAL_BOOT_NOT_ENABLED;
+		return vb2_ui_screen_back(ui);
+	}
+	return VB2_SUCCESS;
+}
+
+static vb2_error_t developer_boot_external_init(struct vb2_ui_context *ui)
+{
+	vb2_error_t rv;
+
+	VB2_TRY(developer_boot_external_check(ui));
+	rv = VbTryLoadKernel(ui->ctx, VB_DISK_FLAG_REMOVABLE);
+	/* If the status of the external disk doesn't match, skip the screen. */
+	if (rv != VB2_ERROR_LK_NO_DISK_FOUND)
+		return vb2_ui_screen_back(ui);
+
+	return VB2_SUCCESS;
+}
+
 static const struct vb2_menu_item developer_boot_external_items[] = {
 	LANGUAGE_SELECT_ITEM,
 	BACK_ITEM,
@@ -896,6 +922,8 @@ static const struct vb2_screen_info developer_boot_external_screen = {
 	.name = "Developer boot from external disk",
 	.action = vb2_ui_developer_mode_boot_external_action,
 	.menu = MENU_ITEMS(developer_boot_external_items),
+	.init = developer_boot_external_init,
+	.reinit = developer_boot_external_init,
 };
 
 /******************************************************************************/
@@ -907,11 +935,26 @@ static const struct vb2_menu_item developer_invalid_disk_items[] = {
 	POWER_OFF_ITEM,
 };
 
+static vb2_error_t developer_invalid_disk_init(struct vb2_ui_context *ui)
+{
+	vb2_error_t rv;
+
+	VB2_TRY(developer_boot_external_check(ui));
+	rv = VbTryLoadKernel(ui->ctx, VB_DISK_FLAG_REMOVABLE);
+	/* If the status of the external disk doesn't match, skip the screen. */
+	if (rv == VB2_SUCCESS || rv == VB2_ERROR_LK_NO_DISK_FOUND)
+		return vb2_ui_screen_back(ui);
+
+	return VB2_SUCCESS;
+}
+
 static const struct vb2_screen_info developer_invalid_disk_screen = {
 	.id = VB2_SCREEN_DEVELOPER_INVALID_DISK,
 	.name = "Invalid external disk in dev mode",
 	.action = vb2_ui_developer_mode_boot_external_action,
 	.menu = MENU_ITEMS(developer_invalid_disk_items),
+	.init = developer_invalid_disk_init,
+	.reinit = developer_invalid_disk_init,
 };
 
 /******************************************************************************/