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>

6 files changed, 261 insertions, 214 deletions

diff --git a/firmware/2lib/2ui.c b/firmware/2lib/2ui.c
index 4d363f68..81243ac2 100644
--- a/firmware/2lib/2ui.c
+++ b/firmware/2lib/2ui.c
@@ -95,11 +95,11 @@ const struct vb2_menu *get_menu(struct vb2_ui_context *ui)
 		.num_items = 0,
 		.items = NULL,
 	};
-	if (ui->state.screen->get_menu) {
-		menu = ui->state.screen->get_menu(ui);
+	if (ui->state->screen->get_menu) {
+		menu = ui->state->screen->get_menu(ui);
 		return menu ? menu : &empty_menu;
 	} else {
-		return &ui->state.screen->menu;
+		return &ui->state->screen->menu;
 	}
 }
 
@@ -125,7 +125,7 @@ vb2_error_t menu_navigation_action(struct vb2_ui_context *ui)
 	case VB_KEY_ENTER:
 		return vb2_ui_menu_select(ui);
 	case VB_KEY_ESC:
-		return vb2_ui_change_root(ui);
+		return vb2_ui_screen_back(ui);
 	default:
 		if (key != 0)
 			VB2_DEBUG("Pressed key %#x, trusted? %d\n",
@@ -142,13 +142,13 @@ vb2_error_t vb2_ui_menu_prev(struct vb2_ui_context *ui)
 	if (!DETACHABLE && ui->key == VB_BUTTON_VOL_UP_SHORT_PRESS)
 		return VB2_REQUEST_UI_CONTINUE;
 
-	item = ui->state.selected_item - 1;
+	item = ui->state->selected_item - 1;
 	while (item >= 0 &&
-	       ((1 << item) & ui->state.disabled_item_mask))
+	       ((1 << item) & ui->state->disabled_item_mask))
 		item--;
 	/* Only update if item is valid */
 	if (item >= 0)
-		ui->state.selected_item = item;
+		ui->state->selected_item = item;
 
 	return VB2_REQUEST_UI_CONTINUE;
 }
@@ -162,13 +162,13 @@ vb2_error_t vb2_ui_menu_next(struct vb2_ui_context *ui)
 		return VB2_REQUEST_UI_CONTINUE;
 
 	menu = get_menu(ui);
-	item = ui->state.selected_item + 1;
+	item = ui->state->selected_item + 1;
 	while (item < menu->num_items &&
-	       ((1 << item) & ui->state.disabled_item_mask))
+	       ((1 << item) & ui->state->disabled_item_mask))
 		item++;
 	/* Only update if item is valid */
 	if (item < menu->num_items)
-		ui->state.selected_item = item;
+		ui->state->selected_item = item;
 
 	return VB2_REQUEST_UI_CONTINUE;
 }
@@ -185,7 +185,7 @@ vb2_error_t vb2_ui_menu_select(struct vb2_ui_context *ui)
 	if (menu->num_items == 0)
 		return VB2_REQUEST_UI_CONTINUE;
 
-	menu_item = &menu->items[ui->state.selected_item];
+	menu_item = &menu->items[ui->state->selected_item];
 
 	if (menu_item->action) {
 		VB2_DEBUG("Menu item <%s> run action\n", menu_item->text);
@@ -193,7 +193,7 @@ vb2_error_t vb2_ui_menu_select(struct vb2_ui_context *ui)
 	} else if (menu_item->target) {
 		VB2_DEBUG("Menu item <%s> to target screen %#x\n",
 			  menu_item->text, menu_item->target);
-		return vb2_ui_change_screen(ui, menu_item->target);
+		return vb2_ui_screen_change(ui, menu_item->target);
 	}
 
 	VB2_DEBUG("Menu item <%s> no action or target screen\n",
@@ -204,28 +204,35 @@ vb2_error_t vb2_ui_menu_select(struct vb2_ui_context *ui)
 /*****************************************************************************/
 /* Screen navigation functions */
 
-vb2_error_t vb2_ui_change_root(struct vb2_ui_context *ui)
+vb2_error_t vb2_ui_screen_back(struct vb2_ui_context *ui)
 {
-	return vb2_ui_change_screen(ui, ui->root_screen->id);
+	struct vb2_screen_state *tmp;
+
+	if (ui->state && ui->state->prev) {
+		tmp = ui->state->prev;
+		free(ui->state);
+		ui->state = tmp;
+	} else {
+		VB2_DEBUG("ERROR: No previous screen; ignoring\n");
+	}
+
+	return VB2_REQUEST_UI_CONTINUE;
 }
 
 static vb2_error_t default_screen_init(struct vb2_ui_context *ui)
 {
 	const struct vb2_menu *menu = get_menu(ui);
-	ui->state.selected_item = 0;
+	ui->state->selected_item = 0;
 	if (menu->num_items > 1 && menu->items[0].is_language_select)
-		ui->state.selected_item = 1;
+		ui->state->selected_item = 1;
 	return VB2_REQUEST_UI_CONTINUE;
 }
 
-vb2_error_t vb2_ui_change_screen(struct vb2_ui_context *ui, enum vb2_screen id)
+vb2_error_t vb2_ui_screen_change(struct vb2_ui_context *ui, enum vb2_screen 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;
-	}
+	struct vb2_screen_state *cur_state;
+	int state_exists = 0;
 
 	new_screen_info = vb2_get_screen_info(id);
 	if (new_screen_info == NULL) {
@@ -233,13 +240,41 @@ vb2_error_t vb2_ui_change_screen(struct vb2_ui_context *ui, enum vb2_screen id)
 		return VB2_REQUEST_UI_CONTINUE;
 	}
 
-	memset(&ui->state, 0, sizeof(ui->state));
-	ui->state.screen = new_screen_info;
+	/* Check to see if the screen state already exists in our stack. */
+	cur_state = ui->state;
+	while (cur_state != NULL) {
+		if (cur_state->screen->id == id) {
+			state_exists = 1;
+			break;
+		}
+		cur_state = cur_state->prev;
+	}
+
+	if (state_exists) {
+		/* Pop until the requested screen is at the top of stack. */
+		while (ui->state->screen->id != id) {
+			cur_state = ui->state;
+			ui->state = cur_state->prev;
+			free(cur_state);
+		}
+	} else {
+		/* Allocate the requested screen on top of the stack. */
+		cur_state = malloc(sizeof(*ui->state));
+		memset(cur_state, 0, sizeof(*ui->state));
+		if (cur_state == NULL) {
+			VB2_DEBUG("WARNING: malloc failed; ignoring\n");
+			return VB2_REQUEST_UI_CONTINUE;
+		}
+		cur_state->prev = ui->state;
+		cur_state->screen = new_screen_info;
+		ui->state = cur_state;
+		if (ui->state->screen->init)
+			return ui->state->screen->init(ui);
+		else
+			return default_screen_init(ui);
+	}
 
-	if (ui->state.screen->init)
-		return ui->state.screen->init(ui);
-	else
-		return default_screen_init(ui);
+	return VB2_REQUEST_UI_CONTINUE;
 }
 
 /*****************************************************************************/
@@ -252,16 +287,17 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 	struct vb2_screen_state prev_state;
 	enum vb2_ui_error prev_error_code;
 	const struct vb2_menu *menu;
+	const struct vb2_screen_info *root_info;
 	uint32_t key_flags;
 	vb2_error_t rv;
 
 	memset(&ui, 0, sizeof(ui));
 	ui.ctx = ctx;
-	ui.root_screen = vb2_get_screen_info(root_screen_id);
-	if (ui.root_screen == NULL)
+	root_info = vb2_get_screen_info(root_screen_id);
+	if (root_info == NULL)
 		VB2_DIE("Root screen not found.\n");
 	ui.locale_id = vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX);
-	rv = vb2_ui_change_screen(&ui, ui.root_screen->id);
+	rv = vb2_ui_screen_change(&ui, root_screen_id);
 	if (rv != VB2_REQUEST_UI_CONTINUE)
 		return rv;
 	memset(&prev_state, 0, sizeof(prev_state));
@@ -269,19 +305,19 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 
 	while (1) {
 		/* Draw if there are state changes. */
-		if (memcmp(&prev_state, &ui.state, sizeof(ui.state)) ||
+		if (memcmp(&prev_state, ui.state, sizeof(*ui.state)) ||
 		    /* we want to redraw/beep on a transition */
 		    prev_error_code != ui.error_code) {
 
 			menu = get_menu(&ui);
 			VB2_DEBUG("<%s> menu item <%s>\n",
-				  ui.state.screen->name,
+				  ui.state->screen->name,
 				  menu->num_items ?
-				  menu->items[ui.state.selected_item].text :
+				  menu->items[ui.state->selected_item].text :
 				  "null");
-			vb2ex_display_ui(ui.state.screen->id, ui.locale_id,
-					 ui.state.selected_item,
-					 ui.state.disabled_item_mask,
+			vb2ex_display_ui(ui.state->screen->id, ui.locale_id,
+					 ui.state->selected_item,
+					 ui.state->disabled_item_mask,
 					 ui.error_code);
 			/*
 			 * Only beep if we're transitioning from no
@@ -292,7 +328,7 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 				vb2ex_beep(250, 400);
 
 			/* Update prev variables. */
-			memcpy(&prev_state, &ui.state, sizeof(ui.state));
+			memcpy(&prev_state, ui.state, sizeof(*ui.state));
 			prev_error_code = ui.error_code;
 		}
 
@@ -313,8 +349,8 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 			return rv;
 
 		/* Run screen action. */
-		if (ui.state.screen->action) {
-			rv = ui.state.screen->action(&ui);
+		if (ui.state->screen->action) {
+			rv = ui.state->screen->action(&ui);
 			if (rv != VB2_REQUEST_UI_CONTINUE)
 				return rv;
 		}
@@ -350,7 +386,7 @@ vb2_error_t developer_action(struct vb2_ui_context *ui)
 {
 	/* Developer mode keyboard shortcuts */
 	if (ui->key == VB_KEY_CTRL('S'))
-		return vb2_ui_change_screen(ui, VB2_SCREEN_DEVELOPER_TO_NORM);
+		return vb2_ui_screen_change(ui, VB2_SCREEN_DEVELOPER_TO_NORM);
 	if (ui->key == VB_KEY_CTRL('U') ||
 	    (DETACHABLE && ui->key == VB_BUTTON_VOL_UP_LONG_PRESS))
 		return vb2_ui_developer_mode_boot_external_action(ui);
@@ -402,7 +438,7 @@ vb2_error_t manual_recovery_action(struct vb2_ui_context *ui)
 		VB2_DEBUG("Recovery VbTryLoadKernel %#x --> %#x\n",
 			  ui->recovery_rv, rv);
 		ui->recovery_rv = rv;
-		return vb2_ui_change_screen(ui,
+		return vb2_ui_screen_change(ui,
 					    rv == VB2_ERROR_LK_NO_DISK_FOUND ?
 					    VB2_SCREEN_RECOVERY_SELECT :
 					    VB2_SCREEN_RECOVERY_INVALID);
@@ -411,7 +447,7 @@ vb2_error_t manual_recovery_action(struct vb2_ui_context *ui)
 	/* Manual recovery keyboard shortcuts */
 	if (ui->key == VB_KEY_CTRL('D') ||
 	    (DETACHABLE && ui->key == VB_BUTTON_VOL_UP_DOWN_COMBO_PRESS))
-		return vb2_ui_change_screen(ui, VB2_SCREEN_RECOVERY_TO_DEV);
+		return vb2_ui_screen_change(ui, VB2_SCREEN_RECOVERY_TO_DEV);
 
 	/* TODO(b/144969088): Re-implement as debug info screen. */
 	if (ui->key == '\t')
diff --git a/firmware/2lib/2ui_screens.c b/firmware/2lib/2ui_screens.c
index c43abb8f..97c9b5e5 100644
--- a/firmware/2lib/2ui_screens.c
+++ b/firmware/2lib/2ui_screens.c
@@ -32,7 +32,7 @@
 
 #define BACK_ITEM { \
 	.text = "Back", \
-	.action = vb2_ui_change_root, \
+	.action = vb2_ui_screen_back, \
 }
 
 #define ADVANCED_OPTIONS_ITEM { \
@@ -65,7 +65,7 @@ static const struct vb2_screen_info blank_screen = {
 static vb2_error_t language_select_action(struct vb2_ui_context *ui)
 {
 	vb2_error_t rv;
-	ui->locale_id = ui->state.selected_item;
+	ui->locale_id = ui->state->selected_item;
 	VB2_DEBUG("Locale changed to %u\n", ui->locale_id);
 
 	/* Write locale id back to nvdata. */
@@ -77,7 +77,7 @@ static vb2_error_t language_select_action(struct vb2_ui_context *ui)
 	if (rv && !(ui->ctx->flags & VB2_CONTEXT_RECOVERY_MODE))
 		return rv;
 
-	return vb2_ui_change_root(ui);
+	return vb2_ui_screen_back(ui);
 }
 
 const struct vb2_menu *get_language_menu(struct vb2_ui_context *ui)
@@ -117,14 +117,14 @@ static vb2_error_t language_select_init(struct vb2_ui_context *ui)
 	if (menu->num_items == 0) {
 		VB2_DEBUG("ERROR: No menu items found; "
 			  "rejecting entering language selection screen\n");
-		return vb2_ui_change_root(ui);
+		return vb2_ui_screen_back(ui);
 	}
 	if (ui->locale_id < menu->num_items) {
-		ui->state.selected_item = ui->locale_id;
+		ui->state->selected_item = ui->locale_id;
 	} else {
 		VB2_DEBUG("WARNING: Current locale not found in menu items; "
 			  "initializing selected_item to 0\n");
-		ui->state.selected_item = 0;
+		ui->state->selected_item = 0;
 	}
 	return VB2_REQUEST_UI_CONTINUE;
 }
@@ -159,12 +159,12 @@ static const struct vb2_screen_info recovery_broken_screen = {
 
 vb2_error_t advanced_options_init(struct vb2_ui_context *ui)
 {
-	ui->state.selected_item = ADVANCED_OPTIONS_ITEM_DEVELOPER_MODE;
+	ui->state->selected_item = ADVANCED_OPTIONS_ITEM_DEVELOPER_MODE;
 	if (vb2_get_sd(ui->ctx)->flags & VB2_SD_FLAG_DEV_MODE_ENABLED ||
 	    !vb2_allow_recovery(ui->ctx)) {
-		ui->state.disabled_item_mask |=
+		ui->state->disabled_item_mask |=
 			1 << ADVANCED_OPTIONS_ITEM_DEVELOPER_MODE;
-		ui->state.selected_item = ADVANCED_OPTIONS_ITEM_BACK;
+		ui->state->selected_item = ADVANCED_OPTIONS_ITEM_BACK;
 	}
 
 	return VB2_REQUEST_UI_CONTINUE;
@@ -195,12 +195,12 @@ static const struct vb2_screen_info advanced_options_screen = {
 
 vb2_error_t recovery_select_init(struct vb2_ui_context *ui)
 {
-	ui->state.selected_item = RECOVERY_SELECT_ITEM_PHONE;
+	ui->state->selected_item = RECOVERY_SELECT_ITEM_PHONE;
 	if (!vb2api_phone_recovery_ui_enabled(ui->ctx)) {
 		VB2_DEBUG("WARNING: Phone recovery not available\n");
-		ui->state.disabled_item_mask |=
+		ui->state->disabled_item_mask |=
 			1 << RECOVERY_SELECT_ITEM_PHONE;
-		ui->state.selected_item = RECOVERY_SELECT_ITEM_EXTERNAL_DISK;
+		ui->state->selected_item = RECOVERY_SELECT_ITEM_EXTERNAL_DISK;
 	}
 	return VB2_REQUEST_UI_CONTINUE;
 }
@@ -251,21 +251,21 @@ vb2_error_t recovery_to_dev_init(struct vb2_ui_context *ui)
 	if (vb2_get_sd(ui->ctx)->flags & VB2_SD_FLAG_DEV_MODE_ENABLED) {
 		/* We're in dev mode, so let user know they can't transition */
 		ui->error_code = VB2_UI_ERROR_DEV_MODE_ALREADY_ENABLED;
-		return vb2_ui_change_root(ui);
+		return vb2_ui_screen_back(ui);
 	}
 
 	if (!PHYSICAL_PRESENCE_KEYBOARD && vb2ex_physical_presence_pressed()) {
 		VB2_DEBUG("Presence button stuck?\n");
-		return vb2_ui_change_root(ui);
+		return vb2_ui_screen_back(ui);
 	}
 
-	ui->state.selected_item = RECOVERY_TO_DEV_ITEM_CONFIRM;
+	ui->state->selected_item = RECOVERY_TO_DEV_ITEM_CONFIRM;
 
 	/* Disable "Confirm" button for other physical presence types. */
 	if (!PHYSICAL_PRESENCE_KEYBOARD) {
-		ui->state.disabled_item_mask |=
+		ui->state->disabled_item_mask |=
 			1 << RECOVERY_TO_DEV_ITEM_CONFIRM;
-		ui->state.selected_item = RECOVERY_TO_DEV_ITEM_CANCEL;
+		ui->state->selected_item = RECOVERY_TO_DEV_ITEM_CANCEL;
 	}
 
 	ui->physical_presence_button_pressed = 0;
@@ -278,7 +278,7 @@ static vb2_error_t recovery_to_dev_finalize(struct vb2_ui_context *ui)
 	VB2_DEBUG("Physical presence confirmed!\n");
 
 	/* Sanity check, should never happen. */
-	if (ui->state.screen->id != VB2_SCREEN_RECOVERY_TO_DEV ||
+	if (ui->state->screen->id != VB2_SCREEN_RECOVERY_TO_DEV ||
 	    (vb2_get_sd(ui->ctx)->flags & VB2_SD_FLAG_DEV_MODE_ENABLED) ||
 	    !vb2_allow_recovery(ui->ctx)) {
 		VB2_DEBUG("ERROR: Dev transition sanity check failed\n");
@@ -306,7 +306,7 @@ vb2_error_t recovery_to_dev_action(struct vb2_ui_context *ui)
 
 	if (ui->key == ' ') {
 		VB2_DEBUG("SPACE means cancel dev mode transition\n");
-		return vb2_ui_change_root(ui);
+		return vb2_ui_screen_back(ui);
 	}
 
 	/* Keyboard physical presence case covered by "Confirm" action. */
@@ -335,7 +335,7 @@ static const struct vb2_menu_item recovery_to_dev_items[] = {
 	},
 	[RECOVERY_TO_DEV_ITEM_CANCEL] = {
 		.text = "Cancel",
-		.action = vb2_ui_change_root,
+		.action = vb2_ui_screen_back,
 	},
 	POWER_OFF_ITEM,
 };
@@ -438,27 +438,27 @@ vb2_error_t developer_mode_init(struct vb2_ui_context *ui)
 	enum vb2_dev_default_boot_target default_boot =
 		vb2api_get_dev_default_boot_target(ui->ctx);
 
-	/* Get me outta here! */
+	/* TODO(b/159579189): Split this case into a separate root screen */
 	if (!vb2_dev_boot_allowed(ui->ctx))
-		vb2_ui_change_screen(ui, VB2_SCREEN_DEVELOPER_TO_NORM);
+		vb2_ui_screen_change(ui, VB2_SCREEN_DEVELOPER_TO_NORM);
 
 	/* Don't show "Return to secure mode" button if GBB forces dev mode. */
 	if (vb2_get_gbb(ui->ctx)->flags & VB2_GBB_FLAG_FORCE_DEV_SWITCH_ON)
-		ui->state.disabled_item_mask |=
+		ui->state->disabled_item_mask |=
 			1 << DEVELOPER_MODE_ITEM_RETURN_TO_SECURE;
 
 	/* Don't show "Boot from external disk" button if not allowed. */
 	if (!vb2_dev_boot_external_allowed(ui->ctx))
-		ui->state.disabled_item_mask |=
+		ui->state->disabled_item_mask |=
 			1 << DEVELOPER_MODE_ITEM_BOOT_EXTERNAL;
 
 	/* Choose the default selection. */
 	switch (default_boot) {
 	case VB2_DEV_DEFAULT_BOOT_TARGET_EXTERNAL:
-		ui->state.selected_item = DEVELOPER_MODE_ITEM_BOOT_EXTERNAL;
+		ui->state->selected_item = DEVELOPER_MODE_ITEM_BOOT_EXTERNAL;
 		break;
 	default:
-		ui->state.selected_item = DEVELOPER_MODE_ITEM_BOOT_INTERNAL;
+		ui->state->selected_item = DEVELOPER_MODE_ITEM_BOOT_INTERNAL;
 		break;
 	}
 
@@ -498,18 +498,20 @@ vb2_error_t vb2_ui_developer_mode_boot_external_action(
 	if (rv == VB2_SUCCESS) {
 		return VB2_SUCCESS;
 	} else if (rv == VB2_ERROR_LK_NO_DISK_FOUND) {
-		if (ui->state.screen->id != VB2_SCREEN_DEVELOPER_BOOT_EXTERNAL) {
+		if (ui->state->screen->id !=
+		    VB2_SCREEN_DEVELOPER_BOOT_EXTERNAL) {
 			VB2_DEBUG("No external disk found\n");
 			ui->error_code = VB2_UI_ERROR_DEV_EXTERNAL_BOOT_FAILED;
 		}
-		return vb2_ui_change_screen(
+		return vb2_ui_screen_change(
 			ui, VB2_SCREEN_DEVELOPER_BOOT_EXTERNAL);
 	} else {
-		if (ui->state.screen->id != VB2_SCREEN_DEVELOPER_INVALID_DISK) {
+		if (ui->state->screen->id !=
+		    VB2_SCREEN_DEVELOPER_INVALID_DISK) {
 			VB2_DEBUG("Invalid external disk: %#x\n", rv);
 			ui->error_code = VB2_UI_ERROR_DEV_EXTERNAL_BOOT_FAILED;
 		}
-		return vb2_ui_change_screen(
+		return vb2_ui_screen_change(
 			ui, VB2_SCREEN_DEVELOPER_INVALID_DISK);
 	}
 }
@@ -519,6 +521,10 @@ vb2_error_t developer_mode_action(struct vb2_ui_context *ui)
 	const int use_short = vb2api_use_short_dev_screen_delay(ui->ctx);
 	uint64_t elapsed;
 
+	/* TODO(b/159579189): Split this case into a separate root screen */
+	if (!vb2_dev_boot_allowed(ui->ctx))
+		vb2_ui_screen_change(ui, VB2_SCREEN_DEVELOPER_TO_NORM);
+
 	/* Once any user interaction occurs, stop the timer. */
 	if (ui->key)
 		ui->disable_timer = 1;
@@ -600,7 +606,7 @@ static const struct vb2_menu_item developer_to_norm_items[] = {
 	},
 	{
 		.text = "Cancel",
-		.action = vb2_ui_change_root,
+		.action = vb2_ui_screen_back,
 	},
 	POWER_OFF_ITEM,
 };
diff --git a/firmware/2lib/include/2ui.h b/firmware/2lib/include/2ui.h
index 652362dc..35dddb89 100644
--- a/firmware/2lib/include/2ui.h
+++ b/firmware/2lib/include/2ui.h
@@ -56,6 +56,7 @@ struct vb2_screen_state {
 	const struct vb2_screen_info *screen;
 	uint32_t selected_item;
 	uint32_t disabled_item_mask;
+	struct vb2_screen_state *prev;
 };
 
 enum vb2_power_button {
@@ -66,8 +67,7 @@ enum vb2_power_button {
 
 struct vb2_ui_context {
 	struct vb2_context *ctx;
-	const struct vb2_screen_info *root_screen;
-	struct vb2_screen_state state;
+	struct vb2_screen_state *state;
 	uint32_t locale_id;
 	uint32_t key;
 	int key_trusted;
@@ -150,14 +150,14 @@ vb2_error_t vb2_ui_menu_select(struct vb2_ui_context *ui);
 /* Screen navigation functions */
 
 /**
- * Return back to the root screen.
+ * Return back to the previous screen.
  *
- * Return to the root screen originally provided to the ui_loop() function.
+ * If the current screen is already the root scren, the request is ignored.
  *
  * @param ui		UI context pointer
  * @return VB2_REQUEST_UI_CONTINUE, or error code on error.
  */
-vb2_error_t vb2_ui_change_root(struct vb2_ui_context *ui);
+vb2_error_t vb2_ui_screen_back(struct vb2_ui_context *ui);
 
 /**
  * Change to the given screen.
@@ -167,7 +167,7 @@ vb2_error_t vb2_ui_change_root(struct vb2_ui_context *ui);
  * @param ui		UI context pointer
  * @return VB2_REQUEST_UI_CONTINUE, or error code on error.
  */
-vb2_error_t vb2_ui_change_screen(struct vb2_ui_context *ui, enum vb2_screen id);
+vb2_error_t vb2_ui_screen_change(struct vb2_ui_context *ui, enum vb2_screen id);
 
 /*****************************************************************************/
 /* UI loops */
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;