vboot/ui: add screen stack functionality

Add a stack storing previous screen states.  When the user clicks
"Back" or presses ESC, revert to the previous state.

In order to deal with the possibility of a UI cycle (repeatedly
selecting the same sequence of screens) which would eventually
use up all available memory, re-use existing target screen states
within the stack.  In other words, when switching to a specific
screen which already exists in the stack, pop until that screen
is reached, rather than creating a duplicate stack entry.

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

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

3 files changed, 139 insertions, 134 deletions

diff --git a/tests/vb2_ui_action_tests.c b/tests/vb2_ui_action_tests.c
index a664c191..b32bafe5 100644
--- a/tests/vb2_ui_action_tests.c
+++ b/tests/vb2_ui_action_tests.c
@@ -44,7 +44,6 @@ static struct vb2_gbb_header gbb;
 static int mock_calls_until_shutdown;
 
 static struct vb2_ui_context mock_ui_context;
-static struct vb2_screen_state *mock_state;
 
 static struct display_call mock_displayed[64];
 static int mock_displayed_count;
@@ -70,9 +69,9 @@ static vb2_error_t mock_action_countdown(struct vb2_ui_context *ui)
 	return VB2_REQUEST_UI_CONTINUE;
 }
 
-static vb2_error_t mock_action_change_screen(struct vb2_ui_context *ui)
+static vb2_error_t mock_action_screen_change(struct vb2_ui_context *ui)
 {
-	return vb2_ui_change_screen(ui, MOCK_SCREEN_BASE);
+	return vb2_ui_screen_change(ui, MOCK_SCREEN_BASE);
 }
 
 static vb2_error_t mock_action_base(struct vb2_ui_context *ui)
@@ -307,8 +306,11 @@ static void reset_common_data(void)
 	/* Mock ui_context based on mock screens */
 	memset(&mock_ui_context, 0, sizeof(mock_ui_context));
 	mock_ui_context.ctx = ctx;
-	mock_ui_context.state.screen = &mock_screen_temp;
-	mock_state = &mock_ui_context.state;
+
+	if (!mock_ui_context.state)
+		mock_ui_context.state = malloc(sizeof(*mock_ui_context.state));
+	memset(mock_ui_context.state, 0, sizeof(*mock_ui_context.state));
+	mock_ui_context.state->screen = &mock_screen_temp;
 
 	/* For vb2ex_display_ui */
 	memset(mock_displayed, 0, sizeof(mock_displayed));
@@ -448,52 +450,57 @@ static void menu_prev_tests(void)
 
 	/* Valid action */
 	reset_common_data();
-	mock_state->screen = &mock_screen_menu;
-	mock_state->selected_item = 2;
+	mock_ui_context.state->screen = &mock_screen_menu;
+	mock_ui_context.state->selected_item = 2;
 	mock_ui_context.key = VB_KEY_UP;
 	TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
 		"valid action");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 1, MOCK_IGNORE);
+	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 1,
+			MOCK_IGNORE);
 
 	/* Valid action with mask */
 	reset_common_data();
-	mock_state->screen = &mock_screen_menu;
-	mock_state->selected_item = 2;
-	mock_state->disabled_item_mask = 0x0a;  /* 0b01010 */
+	mock_ui_context.state->screen = &mock_screen_menu;
+	mock_ui_context.state->selected_item = 2;
+	mock_ui_context.state->disabled_item_mask = 0x0a;  /* 0b01010 */
 	mock_ui_context.key = VB_KEY_UP;
 	TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
 		"valid action with mask");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 0, MOCK_IGNORE);
+	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 0,
+			MOCK_IGNORE);
 
 	/* Invalid action (blocked) */
 	reset_common_data();
-	mock_state->screen = &mock_screen_menu;
-	mock_state->selected_item = 0;
+	mock_ui_context.state->screen = &mock_screen_menu;
+	mock_ui_context.state->selected_item = 0;
 	mock_ui_context.key = VB_KEY_UP;
 	TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
 		"invalid action (blocked)");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 0, MOCK_IGNORE);
+	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 0,
+			MOCK_IGNORE);
 
 	/* Invalid action (blocked by mask) */
 	reset_common_data();
-	mock_state->screen = &mock_screen_menu;
-	mock_state->selected_item = 2;
-	mock_state->disabled_item_mask = 0x0b;  /* 0b01011 */
+	mock_ui_context.state->screen = &mock_screen_menu;
+	mock_ui_context.state->selected_item = 2;
+	mock_ui_context.state->disabled_item_mask = 0x0b;  /* 0b01011 */
 	mock_ui_context.key = VB_KEY_UP;
 	TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
 		"invalid action (blocked by mask)");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
+	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 2,
+			MOCK_IGNORE);
 
 	/* Ignore volume-up when not DETACHABLE */
 	if (!DETACHABLE) {
 		reset_common_data();
-		mock_state->screen = &mock_screen_menu;
-		mock_state->selected_item = 2;
+		mock_ui_context.state->screen = &mock_screen_menu;
+		mock_ui_context.state->selected_item = 2;
 		mock_ui_context.key = VB_BUTTON_VOL_UP_SHORT_PRESS;
 		TEST_EQ(vb2_ui_menu_prev(&mock_ui_context),
 			VB2_REQUEST_UI_CONTINUE,
 			"ignore volume-up when not DETACHABLE");
-		screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
+		screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 2,
+				MOCK_IGNORE);
 	}
 
 	VB2_DEBUG("...done.\n");
@@ -505,52 +512,57 @@ static void menu_next_tests(void)
 
 	/* Valid action */
 	reset_common_data();
-	mock_state->screen = &mock_screen_menu;
-	mock_state->selected_item = 2;
+	mock_ui_context.state->screen = &mock_screen_menu;
+	mock_ui_context.state->selected_item = 2;
 	mock_ui_context.key = VB_KEY_DOWN;
 	TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
 		"valid action");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 3, MOCK_IGNORE);
+	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 3,
+			MOCK_IGNORE);
 
 	/* Valid action with mask */
 	reset_common_data();
-	mock_state->screen = &mock_screen_menu;
-	mock_state->selected_item = 2;
-	mock_state->disabled_item_mask = 0x0a;  /* 0b01010 */
+	mock_ui_context.state->screen = &mock_screen_menu;
+	mock_ui_context.state->selected_item = 2;
+	mock_ui_context.state->disabled_item_mask = 0x0a;  /* 0b01010 */
 	mock_ui_context.key = VB_KEY_DOWN;
 	TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
 		"valid action with mask");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 4, MOCK_IGNORE);
+	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 4,
+			MOCK_IGNORE);
 
 	/* Invalid action (blocked) */
 	reset_common_data();
-	mock_state->screen = &mock_screen_menu;
-	mock_state->selected_item = 4;
+	mock_ui_context.state->screen = &mock_screen_menu;
+	mock_ui_context.state->selected_item = 4;
 	mock_ui_context.key = VB_KEY_DOWN;
 	TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
 		"invalid action (blocked)");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 4, MOCK_IGNORE);
+	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 4,
+			MOCK_IGNORE);
 
 	/* Invalid action (blocked by mask) */
 	reset_common_data();
-	mock_state->screen = &mock_screen_menu;
-	mock_state->selected_item = 2;
-	mock_state->disabled_item_mask = 0x1a;  /* 0b11010 */
+	mock_ui_context.state->screen = &mock_screen_menu;
+	mock_ui_context.state->selected_item = 2;
+	mock_ui_context.state->disabled_item_mask = 0x1a;  /* 0b11010 */
 	mock_ui_context.key = VB_KEY_DOWN;
 	TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
 		"invalid action (blocked by mask)");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
+	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 2,
+			MOCK_IGNORE);
 
 	/* Ignore volume-down when not DETACHABLE */
 	if (!DETACHABLE) {
 		reset_common_data();
-		mock_state->screen = &mock_screen_menu;
-		mock_state->selected_item = 2;
+		mock_ui_context.state->screen = &mock_screen_menu;
+		mock_ui_context.state->selected_item = 2;
 		mock_ui_context.key = VB_BUTTON_VOL_DOWN_SHORT_PRESS;
 		TEST_EQ(vb2_ui_menu_next(&mock_ui_context),
 			VB2_REQUEST_UI_CONTINUE,
 			"ignore volume-down when not DETACHABLE");
-		screen_state_eq(mock_state, MOCK_SCREEN_MENU, 2, MOCK_IGNORE);
+		screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 2,
+				MOCK_IGNORE);
 	}
 
 	VB2_DEBUG("...done.\n");
@@ -562,26 +574,28 @@ static void menu_select_tests(void)
 
 	/* select action with no item screen */
 	reset_common_data();
-	mock_state->screen = &mock_screen_base;
+	mock_ui_context.state->screen = &mock_screen_base;
 	mock_ui_context.key = VB_KEY_ENTER;
 	TEST_EQ(vb2_ui_menu_select(&mock_ui_context),
 		VB2_REQUEST_UI_CONTINUE,
 		"vb2_ui_menu_select with no item screen");
-	screen_state_eq(mock_state, MOCK_SCREEN_BASE, 0, MOCK_IGNORE);
+	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_BASE, 0,
+			MOCK_IGNORE);
 
 	/* Try to select an item with a target (item 2) */
 	reset_common_data();
-	mock_state->screen = &mock_screen_menu;
-	mock_state->selected_item = 2;
+	mock_ui_context.state->screen = &mock_screen_menu;
+	mock_ui_context.state->selected_item = 2;
 	mock_ui_context.key = VB_KEY_ENTER;
 	TEST_EQ(vb2_ui_menu_select(&mock_ui_context),
 		VB2_REQUEST_UI_CONTINUE, "select an item with a target");
-	screen_state_eq(mock_state, MOCK_SCREEN_TARGET2, 0, MOCK_IGNORE);
+	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_TARGET2, 0,
+			MOCK_IGNORE);
 
 	/* Try to select an item with an action (item 3) */
 	reset_common_data();
-	mock_state->screen = &mock_screen_menu;
-	mock_state->selected_item = 3;
+	mock_ui_context.state->screen = &mock_screen_menu;
+	mock_ui_context.state->selected_item = 3;
 	mock_ui_context.key = VB_KEY_ENTER;
 	TEST_EQ(vb2_ui_menu_select(&mock_ui_context),
 		VB2_SUCCESS, "select an item with an action");
@@ -589,24 +603,26 @@ static void menu_select_tests(void)
 
 	/* Try to select an item with neither targets nor actions (item 4) */
 	reset_common_data();
-	mock_state->screen = &mock_screen_menu;
-	mock_state->selected_item = 4;
+	mock_ui_context.state->screen = &mock_screen_menu;
+	mock_ui_context.state->selected_item = 4;
 	mock_ui_context.key = VB_KEY_ENTER;
 	TEST_EQ(vb2_ui_menu_select(&mock_ui_context),
 		VB2_REQUEST_UI_CONTINUE,
 		"select an item with neither targets nor actions");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 4, MOCK_IGNORE);
+	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 4,
+			MOCK_IGNORE);
 
 	/* Ignore power button short press when not DETACHABLE */
 	if (!DETACHABLE) {
 		reset_common_data();
-		mock_state->screen = &mock_screen_menu;
-		mock_state->selected_item = 1;
+		mock_ui_context.state->screen = &mock_screen_menu;
+		mock_ui_context.state->selected_item = 1;
 		mock_ui_context.key = VB_BUTTON_POWER_SHORT_PRESS;
 		TEST_EQ(vb2_ui_menu_select(&mock_ui_context),
 			VB2_REQUEST_UI_CONTINUE,
 			"ignore power button short press when not DETACHABLE");
-		screen_state_eq(mock_state, MOCK_SCREEN_MENU, 1, MOCK_IGNORE);
+		screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 1,
+				MOCK_IGNORE);
 	}
 
 	VB2_DEBUG("...done.\n");
@@ -628,7 +644,7 @@ static void manual_recovery_action_tests(void)
 	set_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
 	TEST_EQ(manual_recovery_action(&mock_ui_context),
 		VB2_REQUEST_UI_CONTINUE, "NO_DISK_FOUND");
-	screen_state_eq(mock_state, VB2_SCREEN_RECOVERY_SELECT,
+	screen_state_eq(mock_ui_context.state, VB2_SCREEN_RECOVERY_SELECT,
 			MOCK_IGNORE, MOCK_IGNORE);
 
 	/* NO_DISK_FOUND -> INVALID_KERNEL -> SUCCESS */
@@ -643,7 +659,7 @@ static void manual_recovery_action_tests(void)
 	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
 	TEST_EQ(manual_recovery_action(&mock_ui_context),
 		VB2_SUCCESS, "SUCCESS");
-	screen_state_eq(mock_state, VB2_SCREEN_RECOVERY_INVALID,
+	screen_state_eq(mock_ui_context.state, VB2_SCREEN_RECOVERY_INVALID,
 			MOCK_IGNORE, MOCK_IGNORE);
 
 	/* INVALID_KERNEL */
@@ -652,7 +668,7 @@ static void manual_recovery_action_tests(void)
 		       VB_DISK_FLAG_REMOVABLE);
 	TEST_EQ(manual_recovery_action(&mock_ui_context),
 		VB2_REQUEST_UI_CONTINUE, "INVALID_KERNEL");
-	screen_state_eq(mock_state, VB2_SCREEN_RECOVERY_INVALID,
+	screen_state_eq(mock_ui_context.state, VB2_SCREEN_RECOVERY_INVALID,
 			MOCK_IGNORE, MOCK_IGNORE);
 
 	/* INVALID_KERNEL -> NO_DISK_FOUND -> SUCCESS */
@@ -667,7 +683,7 @@ static void manual_recovery_action_tests(void)
 	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
 	TEST_EQ(manual_recovery_action(&mock_ui_context),
 		VB2_SUCCESS, "SUCCESS");
-	screen_state_eq(mock_state, VB2_SCREEN_RECOVERY_SELECT,
+	screen_state_eq(mock_ui_context.state, VB2_SCREEN_RECOVERY_SELECT,
 			MOCK_IGNORE, MOCK_IGNORE);
 
 	VB2_DEBUG("...done.\n");
@@ -717,11 +733,11 @@ static void ui_loop_tests(void)
 
 	/* Global action can change screen */
 	reset_common_data();
-	TEST_EQ(ui_loop(ctx, VB2_SCREEN_BLANK, mock_action_change_screen),
+	TEST_EQ(ui_loop(ctx, VB2_SCREEN_BLANK, mock_action_screen_change),
 		VB2_REQUEST_SHUTDOWN, "global action can change screen");
 	DISPLAYED_PASS();
-	DISPLAYED_EQ("change to mock_screen_base", MOCK_IGNORE, MOCK_IGNORE,
-		     MOCK_IGNORE, MOCK_IGNORE);
+	DISPLAYED_EQ("change to mock_screen_base", MOCK_SCREEN_BASE,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	DISPLAYED_NO_EXTRA();
 
 	/*
diff --git a/tests/vb2_ui_tests.c b/tests/vb2_ui_tests.c
index a56e0a4a..47c70466 100644
--- a/tests/vb2_ui_tests.c
+++ b/tests/vb2_ui_tests.c
@@ -42,7 +42,7 @@ static struct vb2_shared_data *sd;
 static struct vb2_gbb_header gbb;
 
 static struct vb2_ui_context mock_ui_context;
-static struct vb2_screen_state *mock_state;
+static struct vb2_screen_state mock_state;
 
 static struct display_call mock_displayed[64];
 static int mock_displayed_count;
@@ -248,7 +248,7 @@ static void reset_common_data(enum reset_type t)
 	/* Mock ui_context based on real screens */
 	memset(&mock_ui_context, 0, sizeof(mock_ui_context));
 	mock_ui_context.ctx = ctx;
-	mock_state = &mock_ui_context.state;
+	mock_ui_context.state = &mock_state;
 
 	/* For vb2ex_display_ui */
 	memset(mock_displayed, 0, sizeof(mock_displayed));
@@ -653,6 +653,9 @@ static void developer_tests(void)
 		"if dev mode is disabled, goes to to_norm screen repeatedly");
 	DISPLAYED_EQ("to_norm", VB2_SCREEN_DEVELOPER_TO_NORM,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	DISPLAYED_PASS();
+	DISPLAYED_EQ("to_norm", VB2_SCREEN_DEVELOPER_TO_NORM,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	DISPLAYED_NO_EXTRA();
 
 	/* Select to_norm in dev menu and confirm */
@@ -1040,10 +1043,11 @@ static void developer_screen_tests(void)
 	add_mock_keypress(VB_KEY_ENTER);
 	/* #1: Return to secure mode */
 	add_mock_keypress(VB_KEY_ESC);
-	add_mock_keypress(VB_KEY_UP);
+	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
 	/* #2: Boot internal */
 	add_mock_keypress(VB_KEY_ESC);
+	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
 	TEST_EQ(vb2_developer_menu(ctx), VB2_SUCCESS,
 		"dev mode screen");
@@ -1061,6 +1065,7 @@ static void developer_screen_tests(void)
 	DISPLAYED_EQ("#1: return to secure mode", VB2_SCREEN_DEVELOPER_TO_NORM,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	/* #2: Boot internal */
+	DISPLAYED_PASS();
 	DISPLAYED_EQ("dev mode", VB2_SCREEN_DEVELOPER_MODE,
 		     MOCK_IGNORE, 2, MOCK_IGNORE);
 	VB2_DEBUG("#2: boot internal (no extra screen)\n");
@@ -1088,7 +1093,7 @@ static void developer_screen_tests(void)
 	/* End of menu */
 	add_mock_keypress(VB_KEY_ESC);
 	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_DOWN);  /* Blocked */
 	TEST_EQ(vb2_developer_menu(ctx), VB2_REQUEST_SHUTDOWN,
 		"dev mode screen");
 	/* #4: Advanced options */
@@ -1100,9 +1105,8 @@ static void developer_screen_tests(void)
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	/* End of menu */
 	DISPLAYED_PASS();
-	DISPLAYED_PASS();
 	DISPLAYED_EQ("end of menu", VB2_SCREEN_DEVELOPER_MODE,
-		     MOCK_IGNORE, 4, MOCK_IGNORE);
+		     MOCK_IGNORE, 5, MOCK_IGNORE);
 	DISPLAYED_NO_EXTRA();
 
 	/* Advanced options screen */
@@ -1117,12 +1121,8 @@ static void developer_screen_tests(void)
 	/* #2: Back */
 	add_mock_keypress(VB_KEY_ESC);
 	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_ENTER);
 	add_mock_keypress(VB_KEY_ENTER);
 	/* End of menu */
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
 	extend_calls_until_shutdown();
 	TEST_EQ(vb2_developer_menu(ctx), VB2_REQUEST_SHUTDOWN,
@@ -1130,8 +1130,8 @@ static void developer_screen_tests(void)
 	DISPLAYED_PASS();
 	DISPLAYED_PASS();
 	DISPLAYED_PASS();
-	/* #0: Language menu */
 	DISPLAYED_PASS();
+	/* #0: Language menu */
 	DISPLAYED_EQ("advanced options", VB2_SCREEN_ADVANCED_OPTIONS,
 		     MOCK_IGNORE, 0, 0x2);
 	DISPLAYED_EQ("#0: language menu", VB2_SCREEN_LANGUAGE_SELECT,
@@ -1139,15 +1139,11 @@ static void developer_screen_tests(void)
 	/* #1: (Disabled) */
 	/* #2: Back */
 	DISPLAYED_PASS();
-	DISPLAYED_PASS();
-	DISPLAYED_PASS();
 	DISPLAYED_EQ("advanced options", VB2_SCREEN_ADVANCED_OPTIONS,
 		     MOCK_IGNORE, 2, 0x2);
 	DISPLAYED_EQ("#2: back", VB2_SCREEN_DEVELOPER_MODE,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	/* End of menu */
-	DISPLAYED_PASS();
-	DISPLAYED_PASS();
 	DISPLAYED_EQ("end of menu", VB2_SCREEN_ADVANCED_OPTIONS,
 		     MOCK_IGNORE, 2, MOCK_IGNORE);
 	DISPLAYED_NO_EXTRA();
@@ -1164,6 +1160,7 @@ static void broken_recovery_screen_tests(void)
 	add_mock_keypress(VB_KEY_ENTER);
 	/* #1: Advanced options */
 	add_mock_keypress(VB_KEY_ESC);
+	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
 	/* End of menu */
 	add_mock_keypress(VB_KEY_ESC);
@@ -1177,6 +1174,7 @@ static void broken_recovery_screen_tests(void)
 	DISPLAYED_EQ("#0: language menu", VB2_SCREEN_LANGUAGE_SELECT,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	/* #1: Advanced options */
+	DISPLAYED_PASS();
 	DISPLAYED_EQ("broken screen", VB2_SCREEN_RECOVERY_BROKEN,
 		     MOCK_IGNORE, 1, 0x0);
 	DISPLAYED_EQ("#1: advanced options", VB2_SCREEN_ADVANCED_OPTIONS,
@@ -1195,7 +1193,7 @@ static void broken_recovery_screen_tests(void)
 	/* #1: (Disabled) */
 	/* #2: Back */
 	add_mock_keypress(VB_KEY_ESC);
-	add_mock_keypress(VB_KEY_ENTER);
+	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
 	/* End of menu */
 	add_mock_keypress(VB_KEY_ENTER);
@@ -1233,6 +1231,7 @@ static void manual_recovery_screen_tests(void)
 	add_mock_keypress(VB_KEY_ENTER);
 	/* #1: Phone recovery */
 	add_mock_keypress(VB_KEY_ESC);
+	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
 	/* #2: External disk recovery */
 	add_mock_keypress(VB_KEY_ESC);
@@ -1241,12 +1240,11 @@ static void manual_recovery_screen_tests(void)
 	/* #3: Advanced options */
 	add_mock_keypress(VB_KEY_ESC);
 	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
 	/* End of menu */
 	add_mock_keypress(VB_KEY_ESC);
 	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_DOWN);  /* Blocked */
 	extend_calls_until_shutdown();
 	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
 		"recovery select screen");
@@ -1257,6 +1255,7 @@ static void manual_recovery_screen_tests(void)
 	DISPLAYED_EQ("#0: language menu", VB2_SCREEN_LANGUAGE_SELECT,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	/* #1: Phone recovery */
+	DISPLAYED_PASS();
 	DISPLAYED_EQ("recovery select", VB2_SCREEN_RECOVERY_SELECT,
 		     MOCK_IGNORE, 1, 0x0);
 	DISPLAYED_EQ("#1: phone recovery", VB2_SCREEN_RECOVERY_PHONE_STEP1,
@@ -1269,16 +1268,14 @@ static void manual_recovery_screen_tests(void)
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	/* #3: Advanced options */
 	DISPLAYED_PASS();
-	DISPLAYED_PASS();
 	DISPLAYED_EQ("recovery select", VB2_SCREEN_RECOVERY_SELECT,
 		     MOCK_IGNORE, 3, 0x0);
 	DISPLAYED_EQ("#3: advanced options", VB2_SCREEN_ADVANCED_OPTIONS,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	/* End of menu */
 	DISPLAYED_PASS();
-	DISPLAYED_PASS();
 	DISPLAYED_EQ("end of menu", VB2_SCREEN_RECOVERY_SELECT,
-		     MOCK_IGNORE, 3, MOCK_IGNORE);
+		     MOCK_IGNORE, 4, MOCK_IGNORE);
 	DISPLAYED_NO_EXTRA();
 
 	/* Advanced options screen */
@@ -1292,19 +1289,12 @@ static void manual_recovery_screen_tests(void)
 	/* #1: Enable dev mode */
 	add_mock_keypress(VB_KEY_ESC);
 	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_ENTER);
 	add_mock_keypress(VB_KEY_ENTER);
 	/* #2: Back */
 	add_mock_keypress(VB_KEY_ESC);
 	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_ENTER);
-	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
 	/* End of menu */
-	add_mock_keypress(VB_KEY_DOWN);
-	add_mock_keypress(VB_KEY_DOWN);
 	add_mock_keypress(VB_KEY_ENTER);
 	add_mock_keypress(VB_KEY_DOWN);
 	extend_calls_until_shutdown();
@@ -1321,25 +1311,18 @@ static void manual_recovery_screen_tests(void)
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	/* #1: Enable dev mode */
 	DISPLAYED_PASS();
-	DISPLAYED_PASS();
-	DISPLAYED_PASS();
 	DISPLAYED_EQ("advanced options", VB2_SCREEN_ADVANCED_OPTIONS,
 		     MOCK_IGNORE, 1, 0x0);
 	DISPLAYED_EQ("#1: enable dev mode", VB2_SCREEN_RECOVERY_TO_DEV,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	/* #2: Back */
 	DISPLAYED_PASS();
-	DISPLAYED_PASS();
-	DISPLAYED_PASS();
-	DISPLAYED_PASS();
 	DISPLAYED_EQ("advanced options", VB2_SCREEN_ADVANCED_OPTIONS,
 		     MOCK_IGNORE, 2, 0x0);
 	DISPLAYED_EQ("#2: back", VB2_SCREEN_RECOVERY_SELECT,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	/* End of menu */
 	DISPLAYED_PASS();
-	DISPLAYED_PASS();
-	DISPLAYED_PASS();
 	DISPLAYED_EQ("end of menu", VB2_SCREEN_ADVANCED_OPTIONS,
 		     MOCK_IGNORE, 2, 0x0);
 	DISPLAYED_NO_EXTRA();
diff --git a/tests/vb2_ui_utility_tests.c b/tests/vb2_ui_utility_tests.c
index 533aa514..dbdebd3e 100644
--- a/tests/vb2_ui_utility_tests.c
+++ b/tests/vb2_ui_utility_tests.c
@@ -33,7 +33,6 @@ static uint32_t mock_locale_count;
 static int mock_shutdown_request;
 
 static struct vb2_ui_context mock_ui_context;
-static struct vb2_screen_state *mock_state;
 
 /* Mock actions */
 static uint32_t mock_action_called;
@@ -70,7 +69,7 @@ struct vb2_menu_item mock_screen_menu_items[] = {
 		.text = "option 4",
 	},
 };
-const struct vb2_screen_info mock_screen_menu = {
+struct vb2_screen_info mock_screen_menu = {
 	.id = MOCK_SCREEN_MENU,
 	.name = "mock_screen_menu: screen with 5 options",
 	.menu = {
@@ -78,7 +77,7 @@ const struct vb2_screen_info mock_screen_menu = {
 		.items = mock_screen_menu_items,
 	},
 };
-const struct vb2_screen_info mock_screen_root = {
+struct vb2_screen_info mock_screen_root = {
 	.id = MOCK_SCREEN_ROOT,
 	.name = "mock_screen_root",
 };
@@ -121,10 +120,19 @@ static void reset_common_data(void)
 	/* Mock ui_context based on mock screens */
 	memset(&mock_ui_context, 0, sizeof(mock_ui_context));
 	mock_ui_context.power_button = VB2_POWER_BUTTON_HELD_SINCE_BOOT;
-	mock_state = &mock_ui_context.state;
 
 	/* For mock actions */
 	mock_action_called = 0;
+
+	/* Reset init and action functions */
+	mock_screen_blank.init = NULL;
+	mock_screen_blank.action = NULL;
+	mock_screen_base.init = NULL;
+	mock_screen_base.action = NULL;
+	mock_screen_menu.init = NULL;
+	mock_screen_menu.action = NULL;
+	mock_screen_root.init = NULL;
+	mock_screen_root.action = NULL;
 }
 
 /* Mock functions */
@@ -260,57 +268,56 @@ static void check_shutdown_request_tests(void)
 	VB2_DEBUG("...done.\n");
 }
 
-static void vb2_ui_change_root_tests(void)
+static void screen_stack_tests(void)
 {
-	VB2_DEBUG("Testing vb2_ui_change_root...\n");
+	VB2_DEBUG("Testing screen stack functionality...\n");
 
-	/* Back to root screen */
+	/* Change to screen which does not exist */
 	reset_common_data();
-	mock_ui_context.root_screen = &mock_screen_root;
-	mock_ui_context.key = VB_KEY_ESC;
-	TEST_EQ(vb2_ui_change_root(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
-		"back to root screen");
-	screen_state_eq(mock_state, MOCK_SCREEN_ROOT, MOCK_IGNORE, MOCK_IGNORE);
-
-	VB2_DEBUG("...done.\n");
-}
-
-static void change_screen_tests(void)
-{
-	VB2_DEBUG("Testing change_screen...\n");
+	TEST_EQ(vb2_ui_screen_change(&mock_ui_context, MOCK_NO_SCREEN),
+		VB2_REQUEST_UI_CONTINUE,
+		"change to screen which does not exist");
+	TEST_PTR_EQ(mock_ui_context.state, NULL, "  stack is empty");
 
-	/* Changing screen will clear screen state */
+	/* Screen back with empty stack */
 	reset_common_data();
-	mock_state->screen = &mock_screen_menu;
-	mock_state->selected_item = 2;
-	mock_state->disabled_item_mask = 0x10;
-	TEST_EQ(vb2_ui_change_screen(&mock_ui_context, MOCK_SCREEN_BASE),
+	TEST_EQ(vb2_ui_screen_back(&mock_ui_context),
 		VB2_REQUEST_UI_CONTINUE,
-		"change_screen will clear screen state");
-	screen_state_eq(mock_state, MOCK_SCREEN_BASE, 0, 0);
+		"screen back with empty stack");
+	TEST_PTR_EQ(mock_ui_context.state, NULL, "  stack is empty");
 
-	/* Change to screen which does not exist */
+	/* Back to previous screen, restoring the state */
 	reset_common_data();
-	mock_state->screen = &mock_screen_menu;
-	TEST_EQ(vb2_ui_change_screen(&mock_ui_context, MOCK_NO_SCREEN),
-		VB2_REQUEST_UI_CONTINUE,
-		"change to screen which does not exist");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, MOCK_IGNORE);
+	mock_screen_base.init = mock_action_base;
+	vb2_ui_screen_change(&mock_ui_context, MOCK_SCREEN_ROOT);
+	vb2_ui_screen_change(&mock_ui_context, MOCK_SCREEN_BASE);
+	mock_ui_context.state->selected_item = 2;
+	mock_ui_context.state->disabled_item_mask = 0x10;
+	vb2_ui_screen_change(&mock_ui_context, MOCK_SCREEN_MENU);
+	TEST_EQ(vb2_ui_screen_back(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"back to previous screen");
+	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_BASE, 2, 0x10);
+	TEST_EQ(mock_action_called, 1, "  action called once");
 
-	/* Change to screen with init */
+	/* Change to target screen already in stack, restoring the state */
 	reset_common_data();
 	mock_screen_base.init = mock_action_base;
-	TEST_EQ(vb2_ui_change_screen(&mock_ui_context, MOCK_SCREEN_BASE),
-		VB2_SUCCESS, "change to screen with init");
+	vb2_ui_screen_change(&mock_ui_context, MOCK_SCREEN_ROOT);
+	vb2_ui_screen_change(&mock_ui_context, MOCK_SCREEN_BASE);
+	mock_ui_context.state->selected_item = 2;
+	mock_ui_context.state->disabled_item_mask = 0x10;
+	vb2_ui_screen_change(&mock_ui_context, MOCK_SCREEN_MENU);
+	TEST_EQ(vb2_ui_screen_change(&mock_ui_context, MOCK_SCREEN_BASE),
+		VB2_REQUEST_UI_CONTINUE, "change to target in stack");
+	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_BASE, 2, 0x10);
 	TEST_EQ(mock_action_called, 1, "  action called once");
 
 	/* Change to screen without init; using default init() */
 	reset_common_data();
-	mock_state->screen = &mock_screen_base;
-	TEST_EQ(vb2_ui_change_screen(&mock_ui_context, MOCK_SCREEN_MENU),
+	TEST_EQ(vb2_ui_screen_change(&mock_ui_context, MOCK_SCREEN_MENU),
 		VB2_REQUEST_UI_CONTINUE,
 		"change to screen with language selection");
-	screen_state_eq(mock_state, MOCK_SCREEN_MENU,
+	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU,
 			1,	/* Since index 0 is the language selection */
 			0);
 
@@ -350,8 +357,7 @@ static void get_language_menu_tests(void)
 int main(void)
 {
 	check_shutdown_request_tests();
-	vb2_ui_change_root_tests();
-	change_screen_tests();
+	screen_stack_tests();
 	get_language_menu_tests();
 
 	return gTestSuccess ? 0 : 255;