vboot/ui: move invalid_disk_last into vb2_context_ui

BUG=b:146399181
TEST=make clean && make runtests
BRANCH=none

Change-Id: I055499580325a43dcba3e1d037bac60deff30af8
Signed-off-by: Joel Kitching <kitching@google.com>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2210015
Commit-Queue: Joel Kitching <kitching@chromium.org>
Tested-by: Joel Kitching <kitching@chromium.org>
Reviewed-by: Yu-Ping Wu <yupingso@chromium.org>

5 files changed, 45 insertions, 71 deletions

diff --git a/firmware/2lib/2ui.c b/firmware/2lib/2ui.c
index e839f017..adf07867 100644
--- a/firmware/2lib/2ui.c
+++ b/firmware/2lib/2ui.c
@@ -19,11 +19,6 @@
 #define KEY_DELAY_MS 20  /* Delay between key scans in UI loops */
 
 /*****************************************************************************/
-/* Global variables */
-
-int invalid_disk_last = -1;
-
-/*****************************************************************************/
 /* Utility functions */
 
 /**
@@ -217,8 +212,14 @@ vb2_error_t (*input_action_lookup(int key))(struct vb2_ui_context *ui)
 
 vb2_error_t vb2_ui_change_screen(struct vb2_ui_context *ui, enum vb2_screen id)
 {
-	const struct vb2_screen_info *new_screen_info = vb2_get_screen_info(id);
+	const struct vb2_screen_info *new_screen_info;
 
+	if (ui->state.screen && ui->state.screen->id == id) {
+		VB2_DEBUG("WARNING: Already on screen %#x; ignoring\n", id);
+		return VB2_REQUEST_UI_CONTINUE;
+	}
+
+	new_screen_info = vb2_get_screen_info(id);
 	if (new_screen_info == NULL) {
 		VB2_DEBUG("ERROR: Screen entry %#x not found; ignoring\n", id);
 		return VB2_REQUEST_UI_CONTINUE;
@@ -340,19 +341,18 @@ vb2_error_t vb2_manual_recovery_menu(struct vb2_context *ctx)
 
 vb2_error_t try_recovery_action(struct vb2_ui_context *ui)
 {
-	int invalid_disk;
+	/* See if we have a recovery kernel available yet. */
 	vb2_error_t rv = VbTryLoadKernel(ui->ctx, VB_DISK_FLAG_REMOVABLE);
-
 	if (rv == VB2_SUCCESS)
 		return rv;
 
 	/* If disk validity state changed, switch to appropriate screen. */
-	invalid_disk = rv != VB2_ERROR_LK_NO_DISK_FOUND;
-	if (invalid_disk_last != invalid_disk) {
-		invalid_disk_last = invalid_disk;
-		return vb2_ui_change_screen(ui, invalid_disk ?
-					    VB2_SCREEN_RECOVERY_INVALID :
-					    VB2_SCREEN_RECOVERY_SELECT);
+	if (ui->recovery_rv != rv) {
+		ui->recovery_rv = rv;
+		return vb2_ui_change_screen(ui,
+					    rv == VB2_ERROR_LK_NO_DISK_FOUND ?
+					    VB2_SCREEN_RECOVERY_SELECT :
+					    VB2_SCREEN_RECOVERY_INVALID);
 	}
 
 	return VB2_REQUEST_UI_CONTINUE;
diff --git a/firmware/2lib/include/2ui.h b/firmware/2lib/include/2ui.h
index 847ec827..0d4df42f 100644
--- a/firmware/2lib/include/2ui.h
+++ b/firmware/2lib/include/2ui.h
@@ -60,13 +60,16 @@ struct vb2_ui_context {
 	uint32_t key;
 	int key_trusted;
 
-	/* For developer mode screen. */
+	/* For check_shutdown_request. */
+	enum vb2_power_button power_button;
+
+	/* For developer mode. */
 	int disable_timer;
 	uint64_t start_time;
 	int beep_count;
 
-	/* For check_shutdown_request. */
-	enum vb2_power_button power_button;
+	/* For manual recovery. */
+	vb2_error_t recovery_rv;
 };
 
 vb2_error_t vb2_ui_change_screen(struct vb2_ui_context *ui, enum vb2_screen id);
diff --git a/firmware/2lib/include/2ui_private.h b/firmware/2lib/include/2ui_private.h
index f7af05e1..891c0414 100644
--- a/firmware/2lib/include/2ui_private.h
+++ b/firmware/2lib/include/2ui_private.h
@@ -12,8 +12,6 @@
 
 vb2_error_t check_shutdown_request(struct vb2_ui_context *ui);
 
-extern int invalid_disk_last;
-
 struct input_action {
 	int key;
 	vb2_error_t (*action)(struct vb2_ui_context *ui);
diff --git a/tests/vb2_ui_action_tests.c b/tests/vb2_ui_action_tests.c
index 0641d671..f438f01f 100644
--- a/tests/vb2_ui_action_tests.c
+++ b/tests/vb2_ui_action_tests.c
@@ -249,9 +249,6 @@ static void reset_common_data(void)
 	/* For check_shutdown_request */
 	mock_calls_until_shutdown = 10;
 
-	/* For try_recovery_action */
-	invalid_disk_last = -1;
-
 	/* Mock ui_context based on mock screens */
 	mock_ui_context = (struct vb2_ui_context){
 		.ctx = ctx,
@@ -650,79 +647,58 @@ static void try_recovery_action_tests(void)
 {
 	VB2_DEBUG("Testing try recovery action...\n");
 
-	/* Success on the first try */
+	/* SUCCESS */
 	reset_common_data();
 	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
 	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS,
-		"success on the first try");
+		"SUCCESS");
 	TEST_EQ(mock_get_screen_info_last, -1, "  no change_screen");
 
-	/* No disk found on the first try */
+	/* NO_DISK_FOUND */
 	reset_common_data();
 	set_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
 	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"no disk found on the first try");
+		"NO_DISK_FOUND");
 	TEST_EQ(mock_get_screen_info_last, VB2_SCREEN_RECOVERY_SELECT,
 		"  recovery select screen");
 
-	/* Invalid disk on the first try */
+	/* NO_DISK_FOUND -> INVALID_KERNEL -> SUCCESS */
 	reset_common_data();
-	set_mock_vbtlk(VB2_ERROR_MOCK, VB_DISK_FLAG_REMOVABLE);
+	set_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
 	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"invalid on the first try");
+		"NO_DISK_FOUND");
+	set_mock_vbtlk(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+		       VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"INVALID_KERNEL");
+	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS, "SUCCESS");
 	TEST_EQ(mock_get_screen_info_last, VB2_SCREEN_RECOVERY_INVALID,
 		"  recovery invalid screen");
 
-	/* Success, last == 0 */
+	/* INVALID_KERNEL */
 	reset_common_data();
-	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
-	invalid_disk_last = 0;
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS,
-		"success, last == 0");
-	TEST_EQ(mock_get_screen_info_last, -1, "  no change_screen");
-
-	/* No disk found, last == 0 */
-	reset_common_data();
-	set_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
-	invalid_disk_last = 0;
+	set_mock_vbtlk(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+		       VB_DISK_FLAG_REMOVABLE);
 	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"no disk found, last == 0");
-	TEST_EQ(mock_get_screen_info_last, -1, "  no change_screen");
-
-	/* Invalid disk, last == 0 */
-	reset_common_data();
-	set_mock_vbtlk(VB2_ERROR_MOCK, VB_DISK_FLAG_REMOVABLE);
-	invalid_disk_last = 0;
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"invalid, last == 0");
+		"INVALID_KERNEL");
 	TEST_EQ(mock_get_screen_info_last, VB2_SCREEN_RECOVERY_INVALID,
 		"  recovery invalid screen");
 
-	/* Success, last == 1 */
-	reset_common_data();
-	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
-	invalid_disk_last = 1;
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS,
-		"success, last == 1");
-	TEST_EQ(mock_get_screen_info_last, -1, "  no change_screen");
-
-	/* No disk found, last == 1 */
+	/* INVALID_KERNEL -> NO_DISK_FOUND -> SUCCESS */
 	reset_common_data();
+	set_mock_vbtlk(VB2_ERROR_LK_INVALID_KERNEL_FOUND,
+		       VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"INVALID_KERNEL");
 	set_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
-	invalid_disk_last = 1;
 	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"no disk found, last == 1");
+		"NO_DISK_FOUND");
+	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS, "SUCCESS");
 	TEST_EQ(mock_get_screen_info_last, VB2_SCREEN_RECOVERY_SELECT,
 		"  recovery select screen");
 
-	/* Invalid disk, last == 1 */
-	reset_common_data();
-	set_mock_vbtlk(VB2_ERROR_MOCK, VB_DISK_FLAG_REMOVABLE);
-	invalid_disk_last = 1;
-	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"invalid, last == 1");
-	TEST_EQ(mock_get_screen_info_last, -1, "  no change_screen");
-
 	VB2_DEBUG("...done.\n");
 }
 
diff --git a/tests/vb2_ui_tests.c b/tests/vb2_ui_tests.c
index 4f36051c..abb05f9c 100644
--- a/tests/vb2_ui_tests.c
+++ b/tests/vb2_ui_tests.c
@@ -164,9 +164,6 @@ static void reset_common_data(enum reset_type t)
 	sd = vb2_get_sd(ctx);
 	sd->status |= VB2_SD_STATUS_SECDATA_KERNEL_INIT;
 
-	/* For try_recovery_action */
-	invalid_disk_last = -1;
-
 	/* Mock ui_context based on real screens */
 	mock_ui_context = (struct vb2_ui_context){
 		.ctx = ctx,