vboot: Implement common UI loop

Add config DETACHABLE to control the navigation in menu UI.

Implement 4 screens:
- VB2_SCREEN_RECOVERY_SELECT
- VB2_SCREEN_RECOVERY_INVALID
- VB2_SCREEN_RECOVERY_PHONE_STEP1
- VB2_SCREEN_RECOVERY_DISK_STEP1

Handling user inputs.
- Shutdown request through VbExIsShutdownRequested.
- Navigate with up, down, and enter key.
- Navigate with volume up, volume down, and power button
  in DETACHABLE.

Implement common UI loop, currently used for manual and non-manual
recovery (developer forthcoming).

BRANCH=none
BUG=b:146399181
TEST=USE="menu_ui" emerge-nami depthcharge
TEST=USE="menu_ui detachable" emerge-nami depthcharge
TEST=make clean && make runtests
TEST=DETACHABLE=1; make clean && make runtests

Cq-Depend: chromium:2152212
Signed-off-by: Hsuan Ting Chen <roccochen@chromium.org>
Change-Id: I4e0f2cdf053f75935529826df215b06c8a9af4cc
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2117810
Reviewed-by: Yu-Ping Wu <yupingso@chromium.org>

8 files changed, 1718 insertions, 70 deletions

diff --git a/Makefile b/Makefile
index cee35bdb..b5b68b20 100644
--- a/Makefile
+++ b/Makefile
@@ -185,6 +185,13 @@ ifneq (${MOCK_TPM},)
 CFLAGS += -DMOCK_TPM
 endif
 
+# DETACHABLE indicates whether the device is a detachable or not.
+ifneq ($(filter-out 0,${DETACHABLE}),)
+CFLAGS += -DDETACHABLE=1
+else
+CFLAGS += -DDETACHABLE=0
+endif
+
 # Enable the menu-based user interface.
 ifneq ($(filter-out 0,${MENU_UI}),)
 CFLAGS += -DMENU_UI=1
@@ -372,6 +379,7 @@ FWLIB_SRCS = \
 	firmware/2lib/2sha_utility.c \
 	firmware/2lib/2tpm_bootmode.c \
 	firmware/2lib/2ui.c \
+	firmware/2lib/2ui_screens.c \
 	firmware/lib/cgptlib/cgptlib.c \
 	firmware/lib/cgptlib/cgptlib_internal.c \
 	firmware/lib/cgptlib/crc32.c \
@@ -716,6 +724,7 @@ TEST2X_NAMES = \
 	tests/vb2_sha_api_tests \
 	tests/vb2_sha_tests \
 	tests/vb2_ui_tests \
+	tests/vb2_ui_utility_tests \
 	tests/hmac_test
 
 TEST20_NAMES = \
@@ -1233,6 +1242,7 @@ run2tests: install_for_test
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_sha_api_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_sha_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb2_ui_tests
+	${RUNTEST} ${BUILD_RUN}/tests/vb2_ui_utility_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb20_api_kernel_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb20_kernel_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vb20_misc_tests
diff --git a/firmware/2lib/2ui.c b/firmware/2lib/2ui.c
index 5b38e08a..eab7c58d 100644
--- a/firmware/2lib/2ui.c
+++ b/firmware/2lib/2ui.c
@@ -12,18 +12,284 @@
 #include "2return_codes.h"
 #include "2secdata.h"
 #include "2ui.h"
+#include "2ui_private.h"
+#include "vboot_api.h"  /* For VB_SHUTDOWN_REQUEST_POWER_BUTTON */
 #include "vboot_kernel.h"
 
+#define KEY_DELAY_MS 20  /* Delay between key scans in UI loops */
+
+/*****************************************************************************/
+/* Global variables */
+
+enum power_button_state power_button;
+int invalid_disk_last = -1;
+
+/*****************************************************************************/
+/* Utility functions */
+
+/**
+ * Checks GBB flags against VbExIsShutdownRequested() shutdown request to
+ * determine if a shutdown is required.
+ *
+ * @param ctx		Context pointer
+ * @param key		Pressed key (VB_BUTTON_POWER_SHORT_PRESS)
+ * @return true if a shutdown is required, or false otherwise.
+ */
+int shutdown_required(struct vb2_context *ctx, uint32_t key)
+{
+	struct vb2_gbb_header *gbb = vb2_get_gbb(ctx);
+	uint32_t shutdown_request = VbExIsShutdownRequested();
+
+	/*
+	 * Ignore power button push until after we have seen it released.
+	 * This avoids shutting down immediately if the power button is still
+	 * being held on startup. After we've recognized a valid power button
+	 * push then don't report the event until after the button is released.
+	 */
+	if (shutdown_request & VB_SHUTDOWN_REQUEST_POWER_BUTTON) {
+		shutdown_request &= ~VB_SHUTDOWN_REQUEST_POWER_BUTTON;
+		if (power_button == POWER_BUTTON_RELEASED)
+			power_button = POWER_BUTTON_PRESSED;
+	} else {
+		if (power_button == POWER_BUTTON_PRESSED)
+			shutdown_request |= VB_SHUTDOWN_REQUEST_POWER_BUTTON;
+		power_button = POWER_BUTTON_RELEASED;
+	}
+
+	if (key == VB_BUTTON_POWER_SHORT_PRESS)
+		shutdown_request |= VB_SHUTDOWN_REQUEST_POWER_BUTTON;
+
+	/* If desired, ignore shutdown request due to lid closure. */
+	if (gbb->flags & VB2_GBB_FLAG_DISABLE_LID_SHUTDOWN)
+		shutdown_request &= ~VB_SHUTDOWN_REQUEST_LID_CLOSED;
+
+	/*
+	 * In detachables, disable shutdown due to power button.
+	 * It is used for menu selection instead.
+	 */
+	if (DETACHABLE)
+		shutdown_request &= ~VB_SHUTDOWN_REQUEST_POWER_BUTTON;
+
+	return !!shutdown_request;
+}
+
 /*****************************************************************************/
-/* Entry points */
+/* Menu navigation actions */
+
+/**
+ * Update selected_item, taking into account disabled indices (from
+ * disabled_item_mask).  The selection does not wrap, meaning that we block
+ * on the 0 or max index when we hit the top or bottom of the menu.
+ */
+vb2_error_t menu_up_action(struct vb2_ui_context *ui)
+{
+	int item;
+
+	if (!DETACHABLE && ui->key == VB_BUTTON_VOL_UP_SHORT_PRESS)
+		return VB2_REQUEST_UI_CONTINUE;
+
+	item = ui->state.selected_item - 1;
+	while (item >= 0 &&
+	       ((1 << item) & ui->state.disabled_item_mask))
+		item--;
+	/* Only update if item is valid */
+	if (item >= 0)
+		ui->state.selected_item = item;
+
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
+vb2_error_t menu_down_action(struct vb2_ui_context *ui)
+{
+	int item;
+
+	if (!DETACHABLE && ui->key == VB_BUTTON_VOL_DOWN_SHORT_PRESS)
+		return VB2_REQUEST_UI_CONTINUE;
+
+	item = ui->state.selected_item + 1;
+	while (item < ui->state.screen->num_items &&
+	       ((1 << item) & ui->state.disabled_item_mask))
+		item++;
+	/* Only update if item is valid */
+	if (item < ui->state.screen->num_items)
+		ui->state.selected_item = item;
+
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
+/**
+ * Navigate to the target screen of the current menu item selection.
+ */
+vb2_error_t menu_select_action(struct vb2_ui_context *ui)
+{
+	const struct vb2_menu_item *menu_item;
+
+	if (!DETACHABLE && ui->key == VB_BUTTON_POWER_SHORT_PRESS)
+		return VB2_REQUEST_UI_CONTINUE;
+
+	if (ui->state.screen->num_items == 0)
+		return VB2_REQUEST_UI_CONTINUE;
+
+	menu_item = &ui->state.screen->items[ui->state.selected_item];
+
+	VB2_DEBUG("Select <%s> menu item <%s>\n",
+		  ui->state.screen->name, menu_item->text);
+
+	if (menu_item->target) {
+		VB2_DEBUG("Changing to target screen %#x for menu item <%s>\n",
+			  menu_item->target, menu_item->text);
+		change_screen(ui, menu_item->target);
+	} else {
+		VB2_DEBUG("No target set for menu item <%s>\n",
+			  menu_item->text);
+	}
+
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
+/**
+ * Return back to the previous screen.
+ */
+vb2_error_t menu_back_action(struct vb2_ui_context *ui)
+{
+	change_screen(ui, ui->root_screen->id);
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
+/*****************************************************************************/
+/* Action lookup tables */
+
+static struct input_action action_table[] = {
+	{ VB_KEY_UP,				menu_up_action },
+	{ VB_KEY_DOWN,				menu_down_action },
+	{ VB_KEY_ENTER,  			menu_select_action },
+	{ VB_BUTTON_VOL_UP_SHORT_PRESS, 	menu_up_action },
+	{ VB_BUTTON_VOL_DOWN_SHORT_PRESS, 	menu_down_action },
+	{ VB_BUTTON_POWER_SHORT_PRESS, 		menu_select_action },
+	{ VB_KEY_ESC, 			 	menu_back_action },
+};
+
+vb2_error_t (*input_action_lookup(int key))(struct vb2_ui_context *ui)
+{
+	int i;
+	for (i = 0; i < ARRAY_SIZE(action_table); i++)
+		if (action_table[i].key == key)
+			return action_table[i].action;
+	return NULL;
+}
+
+/*****************************************************************************/
+/* Core UI functions */
+
+void change_screen(struct vb2_ui_context *ui, enum vb2_screen id)
+{
+	const struct vb2_screen_info *new_screen_info = vb2_get_screen_info(id);
+	int locale_id;
+	if (new_screen_info == NULL) {
+		VB2_DEBUG("ERROR: Screen entry %#x not found; ignoring\n", id);
+	} else {
+		locale_id = ui->state.locale_id;
+		memset(&ui->state, 0, sizeof(ui->state));
+		ui->state.screen = new_screen_info;
+		ui->state.locale_id = locale_id;
+	}
+}
+
+void validate_selection(struct vb2_screen_state *state)
+{
+	if ((state->selected_item == 0 && state->screen->num_items == 0) ||
+	    (state->selected_item < state->screen->num_items &&
+	     !((1 << state->selected_item) & state->disabled_item_mask)))
+		return;
+
+	/* Selection invalid; select the first available non-disabled item. */
+	state->selected_item = 0;
+	while (((1 << state->selected_item) & state->disabled_item_mask) &&
+	       state->selected_item < state->screen->num_items)
+		state->selected_item++;
+
+	/* No non-disabled items available; just choose 0. */
+	if (state->selected_item >= state->screen->num_items)
+		state->selected_item = 0;
+}
+
+vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
+		    vb2_error_t (*global_action)(struct vb2_ui_context *ui))
+{
+	struct vb2_ui_context ui;
+	struct vb2_screen_state prev_state;
+	uint32_t key;
+	uint32_t key_flags;
+	vb2_error_t (*action)(struct vb2_ui_context *ui);
+	vb2_error_t rv;
+
+	ui.ctx = ctx;
+	ui.root_screen = vb2_get_screen_info(root_screen_id);
+	if (ui.root_screen == NULL)
+		VB2_DIE("Root screen not found.\n");
+	change_screen(&ui, ui.root_screen->id);
+	memset(&prev_state, 0, sizeof(prev_state));
+
+	while (1) {
+		/* Draw if there are state changes. */
+		if (memcmp(&prev_state, &ui.state, sizeof(ui.state))) {
+			memcpy(&prev_state, &ui.state, sizeof(ui.state));
+
+			VB2_DEBUG("<%s> menu item <%s>\n",
+				  ui.state.screen->name,
+				  ui.state.screen->num_items ?
+				  ui.state.screen->items[
+				  ui.state.selected_item].text : "null");
+
+			/* TODO: Stop hard-coding the locale. */
+			vb2ex_display_ui(ui.state.screen->id, 0,
+					 ui.state.selected_item,
+					 ui.state.disabled_item_mask);
+		}
+
+		/* Check for shutdown request. */
+		key = VbExKeyboardReadWithFlags(&key_flags);
+		if (shutdown_required(ctx, key)) {
+			VB2_DEBUG("Shutdown required!\n");
+			return VB2_REQUEST_SHUTDOWN;
+		}
+
+		/* Run input action function if found. */
+		action = input_action_lookup(key);
+		if (action) {
+			ui.key = key;
+			rv = action(&ui);
+			ui.key = 0;
+			if (rv != VB2_REQUEST_UI_CONTINUE)
+				return rv;
+			validate_selection(&ui.state);
+		} else if (key) {
+			VB2_DEBUG("Pressed key %#x, trusted? %d\n", key,
+				  !!(key_flags & VB_KEY_FLAG_TRUSTED_KEYBOARD));
+		}
+
+		/* Run global action function if available. */
+		if (global_action) {
+			rv = global_action(&ui);
+			validate_selection(&ui.state);
+			if (rv != VB2_REQUEST_UI_CONTINUE)
+				return rv;
+		}
+
+		/* Delay. */
+		VbExSleepMs(KEY_DELAY_MS);
+	}
+
+	return VB2_SUCCESS;
+}
+
+/*****************************************************************************/
+/* Developer mode */
 
 vb2_error_t vb2_developer_menu(struct vb2_context *ctx)
 {
 	enum vb2_dev_default_boot default_boot;
 
-	/* TODO(roccochen): Init, wait for user, and boot. */
-	vb2ex_display_ui(VB2_SCREEN_BLANK, 0, 0, 0);
-
 	/* If dev mode was disabled, loop forever. */
 	if (!vb2_dev_boot_allowed(ctx))
 		while (1);
@@ -45,22 +311,39 @@ vb2_error_t vb2_developer_menu(struct vb2_context *ctx)
 	return VbTryLoadKernel(ctx, VB_DISK_FLAG_FIXED);
 }
 
+/*****************************************************************************/
+/* Broken recovery */
+
 vb2_error_t vb2_broken_recovery_menu(struct vb2_context *ctx)
 {
-	/* TODO(roccochen): Init and wait for user to reset or shutdown. */
-	vb2ex_display_ui(VB2_SCREEN_BLANK, 0, 0, 0);
+	return ui_loop(ctx, VB2_SCREEN_RECOVERY_BROKEN, NULL);
+}
 
-	while (1);
+/*****************************************************************************/
+/* Manual recovery */
 
-	return VB2_SUCCESS;
+vb2_error_t vb2_manual_recovery_menu(struct vb2_context *ctx)
+{
+	return ui_loop(ctx, VB2_SCREEN_RECOVERY_SELECT, try_recovery_action);
 }
 
-vb2_error_t vb2_manual_recovery_menu(struct vb2_context *ctx)
+vb2_error_t try_recovery_action(struct vb2_ui_context *ui)
 {
-	/* TODO(roccochen): Init and wait for user. */
-	vb2ex_display_ui(VB2_SCREEN_BLANK, 0 ,0, 0);
+	int invalid_disk;
+	vb2_error_t rv = VbTryLoadKernel(ui->ctx, VB_DISK_FLAG_REMOVABLE);
 
-	while (1);
+	if (rv == VB2_SUCCESS)
+		return rv;
 
-	return VB2_SUCCESS;
+	/* 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;
+		if (invalid_disk)
+			change_screen(ui, VB2_SCREEN_RECOVERY_INVALID);
+		else
+			change_screen(ui, VB2_SCREEN_RECOVERY_SELECT);
+	}
+
+	return VB2_REQUEST_UI_CONTINUE;
 }
diff --git a/firmware/2lib/2ui_screens.c b/firmware/2lib/2ui_screens.c
new file mode 100644
index 00000000..5b235c24
--- /dev/null
+++ b/firmware/2lib/2ui_screens.c
@@ -0,0 +1,107 @@
+/* Copyright 2020 The Chromium OS Authors. All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ *
+ * Firmware screen definitions.
+ */
+
+#include "2common.h"
+#include "2ui.h"
+
+#define MENU_ITEMS(a) \
+	.num_items = ARRAY_SIZE(a), \
+	.items = a
+
+static const struct vb2_menu_item empty_menu[] = { };
+
+/******************************************************************************/
+/* VB2_SCREEN_BLANK */
+
+static const struct vb2_screen_info blank_screen = {
+	.id = VB2_SCREEN_BLANK,
+	.name = "Blank",
+	MENU_ITEMS(empty_menu),
+};
+
+/******************************************************************************/
+/* VB2_SCREEN_RECOVERY_BROKEN */
+
+static const struct vb2_screen_info recovery_broken_screen = {
+	.id = VB2_SCREEN_RECOVERY_BROKEN,
+	.name = "Recover broken device",
+	MENU_ITEMS(empty_menu),
+};
+
+/******************************************************************************/
+/* VB2_SCREEN_RECOVERY_SELECT */
+
+static const struct vb2_menu_item recovery_select_items[] = {
+	{
+		.text = "Recovery using phone",
+		.target = VB2_SCREEN_RECOVERY_PHONE_STEP1,
+	},
+	{
+		.text = "Recovery using external disk",
+		.target = VB2_SCREEN_RECOVERY_DISK_STEP1,
+	},
+};
+
+static const struct vb2_screen_info recovery_select_screen = {
+	.id = VB2_SCREEN_RECOVERY_SELECT,
+	.name = "Recovery method selection",
+	MENU_ITEMS(recovery_select_items),
+};
+
+/******************************************************************************/
+/* VB2_SCREEN_RECOVERY_INVALID */
+
+static const struct vb2_screen_info recovery_invalid_screen = {
+	.id = VB2_SCREEN_RECOVERY_INVALID,
+	.name = "Invalid recovery inserted",
+	MENU_ITEMS(empty_menu),
+};
+
+/******************************************************************************/
+/* VB2_SCREEN_RECOVERY_PHONE_STEP1 */
+
+static const struct vb2_screen_info recovery_phone_step1_screen = {
+	.id = VB2_SCREEN_RECOVERY_PHONE_STEP1,
+	.name = "Phone recovery step 1",
+	MENU_ITEMS(empty_menu),
+};
+
+/******************************************************************************/
+/* VB2_SCREEN_RECOVERY_DISK_STEP1 */
+
+static const struct vb2_screen_info recovery_disk_step1_screen = {
+	.id = VB2_SCREEN_RECOVERY_DISK_STEP1,
+	.name = "Disk recovery step 1",
+	MENU_ITEMS(empty_menu),
+};
+
+/******************************************************************************/
+/*
+ * TODO(chromium:1035800): Refactor UI code across vboot and depthcharge.
+ * Currently vboot and depthcharge maintain their own copies of menus/screens.
+ * vboot detects keyboard input and controls the navigation among different menu
+ * items and screens, while depthcharge performs the actual rendering of each
+ * screen, based on the menu information passed from vboot.
+ */
+static const struct vb2_screen_info *screens[] = {
+	&blank_screen,
+	&recovery_broken_screen,
+	&recovery_select_screen,
+	&recovery_invalid_screen,
+	&recovery_phone_step1_screen,
+	&recovery_disk_step1_screen,
+};
+
+const struct vb2_screen_info *vb2_get_screen_info(enum vb2_screen id)
+{
+	int i;
+	for (i = 0; i < ARRAY_SIZE(screens); i++) {
+		if (screens[i]->id == id)
+			return screens[i];
+	}
+	return NULL;
+}
diff --git a/firmware/2lib/include/2api.h b/firmware/2lib/include/2api.h
index 13d24ca9..dc6f19ab 100644
--- a/firmware/2lib/include/2api.h
+++ b/firmware/2lib/include/2api.h
@@ -1177,10 +1177,14 @@ enum vb2_screen {
 	VB2_SCREEN_RECOVERY_BROKEN		= 0x110,
 	/* First recovery screen to select recovering from disk or phone */
 	VB2_SCREEN_RECOVERY_SELECT		= 0x200,
+	/* Invalid recovery media inserted */
+	VB2_SCREEN_RECOVERY_INVALID		= 0x201,
 	/* Recovery using disk */
 	VB2_SCREEN_RECOVERY_DISK_STEP1		= 0x210,
 	VB2_SCREEN_RECOVERY_DISK_STEP2		= 0x211,
 	VB2_SCREEN_RECOVERY_DISK_STEP3		= 0x212,
+	/* Recovery using phone */
+	VB2_SCREEN_RECOVERY_PHONE_STEP1		= 0x220,
 };
 
 /**
diff --git a/firmware/2lib/include/2ui.h b/firmware/2lib/include/2ui.h
index c5fdc1c6..f179024f 100644
--- a/firmware/2lib/include/2ui.h
+++ b/firmware/2lib/include/2ui.h
@@ -8,6 +8,56 @@
 #ifndef VBOOT_REFERENCE_2UI_H_
 #define VBOOT_REFERENCE_2UI_H_
 
+#include <2api.h>
+#include <2sysincludes.h>
+
+/*****************************************************************************/
+/* Data structures */
+
+struct vb2_screen_info {
+	/* Screen id */
+	enum vb2_screen id;
+	/* Screen name for printing to console only */
+	const char *name;
+	/* Number of menu items */
+	uint16_t num_items;
+	/* List of menu items */
+	const struct vb2_menu_item *items;
+};
+
+struct vb2_menu_item {
+	/* Text description */
+	const char *text;
+	/* Target screen */
+	enum vb2_screen target;
+};
+
+struct vb2_screen_state {
+	const struct vb2_screen_info *screen;
+	uint32_t locale_id;
+	uint32_t selected_item;
+	uint32_t disabled_item_mask;
+};
+
+struct vb2_ui_context {
+	struct vb2_context *ctx;
+	const struct vb2_screen_info *root_screen;
+	struct vb2_screen_state state;
+	uint32_t key;
+};
+
+/**
+ * Get info struct of a screen.
+ *
+ * @param screen	Screen from enum vb2_screen
+ *
+ * @return screen info struct on success, NULL on error.
+ */
+const struct vb2_screen_info *vb2_get_screen_info(enum vb2_screen id);
+
+/*****************************************************************************/
+/* UI loops */
+
 /**
  * UI for a developer-mode boot.
  *
diff --git a/firmware/2lib/include/2ui_private.h b/firmware/2lib/include/2ui_private.h
new file mode 100644
index 00000000..40ee6b51
--- /dev/null
+++ b/firmware/2lib/include/2ui_private.h
@@ -0,0 +1,41 @@
+/* Copyright 2020 The Chromium OS Authors. All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ *
+ * Private declarations for 2ui.c. Defined here for testing purposes.
+ */
+
+#include "2api.h"
+
+#ifndef VBOOT_REFERENCE_2UI_PRIVATE_H_
+#define VBOOT_REFERENCE_2UI_PRIVATE_H_
+
+enum power_button_state {
+	POWER_BUTTON_HELD_SINCE_BOOT = 0,
+	POWER_BUTTON_RELEASED,
+	POWER_BUTTON_PRESSED,  /* Must have been previously released */
+};
+extern enum power_button_state power_button;
+int shutdown_required(struct vb2_context *ctx, uint32_t key);
+
+extern int invalid_disk_last;
+
+struct input_action {
+	int key;
+	vb2_error_t (*action)(struct vb2_ui_context *ui);
+};
+
+vb2_error_t menu_up_action(struct vb2_ui_context *ui);
+vb2_error_t menu_down_action(struct vb2_ui_context *ui);
+vb2_error_t menu_select_action(struct vb2_ui_context *ui);
+vb2_error_t menu_back_action(struct vb2_ui_context *ui);
+vb2_error_t (*input_action_lookup(int key))(struct vb2_ui_context *ui);
+
+void change_screen(struct vb2_ui_context *ui, enum vb2_screen id);
+void validate_selection(struct vb2_screen_state *state);
+vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
+		    vb2_error_t (*global_action)(struct vb2_ui_context *ui));
+
+vb2_error_t try_recovery_action(struct vb2_ui_context *ui);
+
+#endif  /* VBOOT_REFERENCE_2UI_PRIVATE_H_ */
diff --git a/tests/vb2_ui_tests.c b/tests/vb2_ui_tests.c
index bf9124bc..a9984491 100644
--- a/tests/vb2_ui_tests.c
+++ b/tests/vb2_ui_tests.c
@@ -10,18 +10,33 @@
 #include "2misc.h"
 #include "2nvstorage.h"
 #include "2ui.h"
+#include "2ui_private.h"
 #include "test_common.h"
-#include "vboot_api.h"
 #include "vboot_kernel.h"
 
+/* Fixed value for ignoring some checks. */
+#define MOCK_IGNORE 0xffffu
+
 /* Mock data */
 static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
 	__attribute__((aligned(VB2_WORKBUF_ALIGN)));
 static struct vb2_context *ctx;
+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_displayed[64];
+static int mock_displayed_count;
+static int mock_displayed_i;
+
+static int mock_calls_until_shutdown;
 
-static enum vb2_screen mock_screens_displayed[64];
-static uint32_t mock_locales_displayed[64];
-static uint32_t mock_screens_count = 0;
+static uint32_t mock_key[64];
+static int mock_key_trusted[64];
+static int mock_key_count;
+static int mock_key_total;
 
 static enum vb2_dev_default_boot mock_default_boot;
 static int mock_dev_boot_allowed;
@@ -31,16 +46,34 @@ static int mock_dev_boot_usb_allowed;
 static int mock_vbexlegacy_called;
 static enum VbAltFwIndex_t mock_altfw_num;
 
-static vb2_error_t mock_vbtlk_retval[5];
-static uint32_t mock_vbtlk_expected_flag[5];
+static vb2_error_t mock_vbtlk_retval[32];
+static uint32_t mock_vbtlk_expected_flag[32];
 static int mock_vbtlk_count;
 static int mock_vbtlk_total;
 
+static void add_mock_key(uint32_t press, int trusted)
+{
+	if (mock_key_total >= ARRAY_SIZE(mock_key) ||
+	    mock_key_total >= ARRAY_SIZE(mock_key_trusted)) {
+		TEST_TRUE(0, "  mock_key ran out of entries!");
+		return;
+	}
+
+	mock_key[mock_key_total] = press;
+	mock_key_trusted[mock_key_total] = trusted;
+	mock_key_total++;
+}
+
+static void add_mock_keypress(uint32_t press)
+{
+	add_mock_key(press, 0);
+}
+
 static void add_mock_vbtlk(vb2_error_t retval, uint32_t get_info_flags)
 {
 	if (mock_vbtlk_total >= ARRAY_SIZE(mock_vbtlk_retval) ||
 	    mock_vbtlk_total >= ARRAY_SIZE(mock_vbtlk_expected_flag)) {
-		TEST_TRUE(0, "Test failed as mock_vbtlk ran out of entries!");
+		TEST_TRUE(0, "  mock_vbtlk ran out of entries!");
 		return;
 	}
 
@@ -49,24 +82,102 @@ static void add_mock_vbtlk(vb2_error_t retval, uint32_t get_info_flags)
 	mock_vbtlk_total++;
 }
 
+static void displayed_eq(const char *text,
+			 enum vb2_screen screen,
+			 uint32_t locale_id,
+			 uint32_t selected_item,
+			 uint32_t disabled_item_mask)
+{
+	char text_buf[256];
+
+	if (mock_displayed_i >= mock_displayed_count) {
+		sprintf(text_buf, "  missing screen %s", text);
+		TEST_TRUE(0, text_buf);
+		return;
+	}
+
+	if (screen != MOCK_IGNORE) {
+		sprintf(text_buf, "  screen of %s", text);
+		TEST_EQ(mock_displayed[mock_displayed_i].screen->id, screen,
+			text_buf);
+	}
+	if (locale_id != MOCK_IGNORE) {
+		sprintf(text_buf, "  locale_id of %s", text);
+		TEST_EQ(mock_displayed[mock_displayed_i].locale_id, locale_id,
+			text_buf);
+	}
+	if (selected_item != MOCK_IGNORE) {
+		sprintf(text_buf, "  selected_item of %s", text);
+		TEST_EQ(mock_displayed[mock_displayed_i].selected_item,
+			selected_item, text_buf);
+	}
+	if (disabled_item_mask != MOCK_IGNORE) {
+		sprintf(text_buf, "  disabled_item_mask of %s", text);
+		TEST_EQ(mock_displayed[mock_displayed_i].disabled_item_mask,
+			disabled_item_mask, text_buf);
+	}
+	mock_displayed_i++;
+}
+
+static void displayed_no_extra(void)
+{
+	if (mock_displayed_i == 0)
+		TEST_EQ(mock_displayed_count, 0, "  no screen");
+	else
+		TEST_EQ(mock_displayed_count, mock_displayed_i,
+			"  no extra screens");
+}
+
 /* Reset mock data (for use before each test) */
-static void reset_common_data()
+static void reset_common_data(void)
 {
 	TEST_SUCC(vb2api_init(workbuf, sizeof(workbuf), &ctx),
 		  "vb2api_init failed");
-	vb2_nv_init(ctx);
 
-	memset(mock_screens_displayed, 0, sizeof(mock_screens_displayed));
-	mock_screens_count = 0;
+	memset(&gbb, 0, sizeof(gbb));
+
+	vb2_nv_init(ctx);
 
+	sd = vb2_get_sd(ctx);
+
+	/* For global actions */
+	invalid_disk_last = -1;
+	mock_ui_context.ctx = ctx;
+	mock_ui_context.root_screen = vb2_get_screen_info(VB2_SCREEN_BLANK);
+	mock_ui_context.state.screen = vb2_get_screen_info(VB2_SCREEN_BLANK);
+	mock_ui_context.state.locale_id = 0,
+	mock_ui_context.state.selected_item = 0,
+	mock_ui_context.state.disabled_item_mask = 0,
+	mock_ui_context.key = 0;
+	mock_state = &mock_ui_context.state;
+
+	/* For vb2ex_display_ui */
+	memset(mock_displayed, 0, sizeof(mock_displayed));
+	mock_displayed_count = 0;
+	mock_displayed_i = 0;
+
+	/* For shutdown_required */
+	mock_calls_until_shutdown = 10;
+
+	/* For VbExKeyboardRead */
+	memset(mock_key, 0, sizeof(mock_key));
+	memset(mock_key_trusted, 0, sizeof(mock_key_trusted));
+	mock_key_count = 0;
+	mock_key_total = 0;
+	/* Avoid iteration #0 which has a screen change by global action */
+	add_mock_keypress(0);
+
+	/* For dev_boot* in 2misc.h */
 	mock_default_boot = VB2_DEV_DEFAULT_BOOT_DISK;
 	mock_dev_boot_allowed = 1;
 	mock_dev_boot_legacy_allowed = 0;
 	mock_dev_boot_usb_allowed = 0;
 
+	/* For VbExLegacy */
 	mock_vbexlegacy_called = 0;
 	mock_altfw_num = -100;
 
+	/* For VbTryLoadKernel */
 	memset(mock_vbtlk_retval, 0, sizeof(mock_vbtlk_retval));
 	memset(mock_vbtlk_expected_flag, 0, sizeof(mock_vbtlk_expected_flag));
 	mock_vbtlk_count = 0;
@@ -74,6 +185,67 @@ static void reset_common_data()
 }
 
 /* Mock functions */
+struct vb2_gbb_header *vb2_get_gbb(struct vb2_context *c)
+{
+	return &gbb;
+}
+
+vb2_error_t vb2ex_display_ui(enum vb2_screen screen,
+			     uint32_t locale_id,
+			     uint32_t selected_item,
+			     uint32_t disabled_item_mask)
+{
+	VB2_DEBUG("displayed %d: screen = %#x, locale_id = %u, "
+		  "selected_item = %u, disabled_item_mask = %#x\n",
+		  mock_displayed_count, screen, locale_id, selected_item,
+		  disabled_item_mask);
+
+	if (mock_displayed_count >= ARRAY_SIZE(mock_displayed)) {
+		TEST_TRUE(0, "  mock vb2ex_display_ui ran out of entries!");
+		return VB2_ERROR_MOCK;
+	}
+
+	mock_displayed[mock_displayed_count] = (struct vb2_screen_state){
+		.screen = vb2_get_screen_info(screen),
+		.locale_id = locale_id,
+		.selected_item = selected_item,
+		.disabled_item_mask = disabled_item_mask,
+	};
+	mock_displayed_count++;
+
+	return VB2_SUCCESS;
+}
+
+uint32_t VbExIsShutdownRequested(void)
+{
+	if (mock_calls_until_shutdown < 0)  /* Never request shutdown */
+		return 0;
+	if (mock_calls_until_shutdown == 0)
+		return 1;
+	mock_calls_until_shutdown--;
+
+	return 0;
+}
+
+uint32_t VbExKeyboardRead(void)
+{
+	return VbExKeyboardReadWithFlags(NULL);
+}
+
+uint32_t VbExKeyboardReadWithFlags(uint32_t *key_flags)
+{
+	if (mock_key_count < mock_key_total) {
+		if (key_flags != NULL) {
+			if (mock_key_trusted[mock_key_count])
+				*key_flags = VB_KEY_FLAG_TRUSTED_KEYBOARD;
+			else
+				*key_flags = 0;
+		}
+		return mock_key[mock_key_count++];
+	}
+
+	return 0;
+}
 
 enum vb2_dev_default_boot vb2_get_dev_boot_target(struct vb2_context *c)
 {
@@ -105,51 +277,114 @@ vb2_error_t VbExLegacy(enum VbAltFwIndex_t altfw_num)
 
 vb2_error_t VbTryLoadKernel(struct vb2_context *c, uint32_t get_info_flags)
 {
-	if (mock_vbtlk_count >= mock_vbtlk_total) {
-		TEST_TRUE(0, "  VbTryLoadKernel called too many times.");
+	if (mock_vbtlk_total == 0) {
+		TEST_TRUE(0, "  VbTryLoadKernel is not allowed!");
 		return VB2_ERROR_MOCK;
 	}
 
+	/* Return last entry if called too many times */
+	if (mock_vbtlk_count >= mock_vbtlk_total)
+		mock_vbtlk_count = mock_vbtlk_total - 1;
+
 	TEST_EQ(mock_vbtlk_expected_flag[mock_vbtlk_count], get_info_flags,
 		"  unexpected get_info_flags");
 
 	return mock_vbtlk_retval[mock_vbtlk_count++];
 }
 
-vb2_error_t vb2ex_display_ui(enum vb2_screen screen,
-			     uint32_t locale_id,
-			     uint32_t selected_item,
-			     uint32_t disabled_item_mask)
+/* Tests */
+static void try_recovery_action_tests(void)
 {
-	VB2_DEBUG("screens %d: screen = %#x, locale_id = %u\n",
-		  mock_screens_count, screen, locale_id);
+	VB2_DEBUG("Testing try recovery action...\n");
 
-	if (mock_screens_count >= ARRAY_SIZE(mock_screens_displayed) ||
-	    mock_screens_count >= ARRAY_SIZE(mock_locales_displayed)) {
-		TEST_TRUE(0, "Test failed as mock vb2ex_display_ui ran out of"
-			  " entries!");
-		return VB2_ERROR_MOCK;
-	}
+	/* Success on the first try */
+	reset_common_data();
+	add_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_SUCCESS,
+		"success on the first try");
+	TEST_EQ(mock_state->screen->id, VB2_SCREEN_BLANK,
+		"  screen remains the same");
 
-	mock_screens_displayed[mock_screens_count] = screen;
-	mock_locales_displayed[mock_screens_count] = locale_id;
-	/* TODO(roccochen): handle the rest of two arguments */
-	mock_screens_count++;
+	/* No disk found on the first try */
+	reset_common_data();
+	add_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");
+	TEST_EQ(mock_state->screen->id, VB2_SCREEN_RECOVERY_SELECT,
+		"  recovery select screen");
 
-	return VB2_SUCCESS;
-}
+	/* Invalid disk on the first try */
+	reset_common_data();
+	add_mock_vbtlk(VB2_ERROR_MOCK, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"invalid on the first try");
+	TEST_EQ(mock_state->screen->id, VB2_SCREEN_RECOVERY_INVALID,
+		"  recovery invalid screen");
 
-/* Tests */
+	/* Success, last == 0 */
+	reset_common_data();
+	add_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_state->screen->id, VB2_SCREEN_BLANK,
+		"  screen remains the same");
+
+	/* No disk found, last == 0 */
+	reset_common_data();
+	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	invalid_disk_last = 0;
+	TEST_EQ(try_recovery_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"no disk found, last == 0");
+	TEST_EQ(mock_state->screen->id, VB2_SCREEN_BLANK, "  screen no change");
+
+	/* Invalid disk, last == 0 */
+	reset_common_data();
+	add_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");
+	TEST_EQ(mock_state->screen->id, VB2_SCREEN_RECOVERY_INVALID,
+		"  recovery invalid screen");
+
+	/* Success, last == 1 */
+	reset_common_data();
+	add_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_state->screen->id, VB2_SCREEN_BLANK,
+		"  screen remains the same");
+
+	/* No disk found, last == 1 */
+	reset_common_data();
+	add_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");
+	TEST_EQ(mock_state->screen->id, VB2_SCREEN_RECOVERY_SELECT,
+		"  recovery select screen");
+
+	/* Invalid disk, last == 1 */
+	reset_common_data();
+	add_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_state->screen->id, VB2_SCREEN_BLANK, "  screen no change");
+
+	VB2_DEBUG("...done.\n");
+}
 
 static void developer_tests(void)
 {
+	VB2_DEBUG("Testing developer mode...\n");
+
 	/* Proceed */
 	reset_common_data();
 	add_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_FIXED);
 	TEST_EQ(vb2_developer_menu(ctx), VB2_SUCCESS, "proceed");
-	TEST_EQ(mock_screens_displayed[0], VB2_SCREEN_BLANK,
-		"  final blank screen");
-	TEST_EQ(mock_screens_count, 1, "  no extra screens");
+	displayed_no_extra();
 	TEST_EQ(vb2_nv_get(ctx, VB2_NV_RECOVERY_REQUEST), 0,
 		"  recovery reason");
 	TEST_EQ(mock_vbtlk_count, mock_vbtlk_total, "  used up mock_vbtlk");
@@ -161,9 +396,7 @@ static void developer_tests(void)
 	TEST_EQ(vb2_developer_menu(ctx), VB2_SUCCESS, "proceed to legacy");
 	TEST_EQ(mock_vbexlegacy_called, 1, "  try legacy");
 	TEST_EQ(mock_altfw_num, 0, "  check altfw_num");
-	TEST_EQ(mock_screens_displayed[0], VB2_SCREEN_BLANK,
-		"  final blank screen");
-	TEST_EQ(mock_screens_count, 1, "  no extra screens");
+	displayed_no_extra();
 	TEST_EQ(mock_vbtlk_count, mock_vbtlk_total, "  used up mock_vbtlk");
 
 	/* Proceed to legacy only if enabled */
@@ -173,50 +406,164 @@ static void developer_tests(void)
 	TEST_EQ(vb2_developer_menu(ctx), VB2_SUCCESS,
 		"default legacy not enabled");
 	TEST_EQ(mock_vbexlegacy_called, 0, "  not legacy");
-	TEST_EQ(mock_screens_displayed[0], VB2_SCREEN_BLANK,
-		"  final blank screen");
-	TEST_EQ(mock_screens_count, 1, "  no extra screens");
-	TEST_EQ(vb2_nv_get(ctx, VB2_NV_RECOVERY_REQUEST), 0,
-		"  no recovery");
+	displayed_no_extra();
 	TEST_EQ(mock_vbtlk_count, mock_vbtlk_total, "  used up mock_vbtlk");
 
-	/* Proceed to usb */
+	/* Proceed to USB */
 	reset_common_data();
 	add_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
 	mock_default_boot = VB2_DEV_DEFAULT_BOOT_USB;
 	mock_dev_boot_usb_allowed = 1;
-	TEST_EQ(vb2_developer_menu(ctx), VB2_SUCCESS, "proceed to usb");
-	TEST_EQ(mock_screens_displayed[0], VB2_SCREEN_BLANK,
-		"  final blank screen");
-	TEST_EQ(mock_screens_count, 1, "  no extra screens");
+	TEST_EQ(vb2_developer_menu(ctx), VB2_SUCCESS, "proceed to USB");
+	displayed_no_extra();
 	TEST_EQ(mock_vbtlk_count, mock_vbtlk_total, "  used up mock_vbtlk");
 
-	/* Proceed to usb only if enabled */
+	/* Proceed to USB only if enabled */
 	reset_common_data();
 	add_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_FIXED);
 	mock_default_boot = VB2_DEV_DEFAULT_BOOT_USB;
 	TEST_EQ(vb2_developer_menu(ctx), VB2_SUCCESS,
-		"default usb not enabled");
-	TEST_EQ(mock_screens_displayed[0], VB2_SCREEN_BLANK,
-		"  final blank screen");
-	TEST_EQ(mock_screens_count, 1, "  no extra screens");
-	TEST_EQ(vb2_nv_get(ctx, VB2_NV_RECOVERY_REQUEST), 0,
-		"  no recovery");
+		"default USB not enabled");
+	displayed_no_extra();
 	TEST_EQ(mock_vbtlk_count, mock_vbtlk_total, "  used up mock_vbtlk");
+
+	VB2_DEBUG("...done.\n");
 }
 
 static void broken_recovery_tests(void)
 {
-	/* TODO(roccochen) */
+	VB2_DEBUG("Testing broken recovery mode...\n");
+
+	/* BROKEN screen shutdown request */
+	if (!DETACHABLE) {
+		reset_common_data();
+		add_mock_keypress(VB_BUTTON_POWER_SHORT_PRESS);
+		mock_calls_until_shutdown = -1;
+		TEST_EQ(vb2_broken_recovery_menu(ctx),
+			VB2_REQUEST_SHUTDOWN,
+			"power button short pressed = shutdown");
+		displayed_eq("broken screen", VB2_SCREEN_RECOVERY_BROKEN,
+			     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+		displayed_no_extra();
+	}
+
+	/* Shortcuts that are always ignored in BROKEN */
+	reset_common_data();
+	add_mock_key(VB_KEY_CTRL('D'), 1);
+	add_mock_key(VB_KEY_CTRL('U'), 1);
+	add_mock_key(VB_KEY_CTRL('L'), 1);
+	add_mock_key(VB_BUTTON_VOL_UP_DOWN_COMBO_PRESS, 1);
+	add_mock_key(VB_BUTTON_VOL_UP_LONG_PRESS, 1);
+	add_mock_key(VB_BUTTON_VOL_DOWN_LONG_PRESS, 1);
+	TEST_EQ(vb2_broken_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
+		"Shortcuts ignored in BROKEN");
+	TEST_EQ(mock_calls_until_shutdown, 0, "  ignore all");
+	displayed_eq("broken screen", VB2_SCREEN_RECOVERY_BROKEN,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_no_extra();
+
+	VB2_DEBUG("...done.\n");
 }
 
 static void manual_recovery_tests(void)
 {
-	/* TODO(roccochen) */
+	VB2_DEBUG("Testing manual recovery mode...\n");
+
+	/* Timeout, shutdown */
+	reset_common_data();
+	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
+		"timeout, shutdown");
+	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_no_extra();
+
+	/* Power button short pressed = shutdown request */
+	if (!DETACHABLE) {
+		reset_common_data();
+		add_mock_keypress(VB_BUTTON_POWER_SHORT_PRESS);
+		add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND,
+			       VB_DISK_FLAG_REMOVABLE);
+		TEST_EQ(vb2_manual_recovery_menu(ctx),
+			VB2_REQUEST_SHUTDOWN,
+			"power button short pressed = shutdown");
+		displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+			     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+		displayed_no_extra();
+	}
+
+	/* Item 1 = phone recovery */
+	reset_common_data();
+	add_mock_keypress(VB_KEY_ENTER);
+	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
+		"phone recovery");
+	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+		     MOCK_IGNORE, 0, MOCK_IGNORE);
+	displayed_eq("phone recovery", VB2_SCREEN_RECOVERY_PHONE_STEP1,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_no_extra();
+
+	/* Item 2 = external disk recovery */
+	reset_common_data();
+	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_ENTER);
+	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
+		"external disk recovery");
+	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+		     MOCK_IGNORE, 0, MOCK_IGNORE);
+	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+		     MOCK_IGNORE, 1, MOCK_IGNORE);
+	displayed_eq("disk recovery", VB2_SCREEN_RECOVERY_DISK_STEP1,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_no_extra();
+
+	/* Boots if we have a valid image on first try */
+	reset_common_data();
+	add_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
+	add_mock_vbtlk(VB2_ERROR_MOCK, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_SUCCESS,
+		"boots if valid on first try");
+	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_no_extra();
+
+	/* Boots eventually if we get a valid image later */
+	reset_common_data();
+	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	add_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
+	add_mock_vbtlk(VB2_ERROR_MOCK, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_SUCCESS,
+		"boots after valid image appears");
+	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_no_extra();
+
+	/* Invalid image, then remove, then valid image */
+	reset_common_data();
+	add_mock_vbtlk(VB2_ERROR_MOCK, VB_DISK_FLAG_REMOVABLE);
+	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	add_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
+	add_mock_vbtlk(VB2_ERROR_MOCK, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_SUCCESS,
+		"boots after valid image appears");
+	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_INVALID,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_no_extra();
+
+	VB2_DEBUG("...done.\n");
 }
 
 int main(void)
 {
+	try_recovery_action_tests();
 	developer_tests();
 	broken_recovery_tests();
 	manual_recovery_tests();
diff --git a/tests/vb2_ui_utility_tests.c b/tests/vb2_ui_utility_tests.c
new file mode 100644
index 00000000..c6f6ff38
--- /dev/null
+++ b/tests/vb2_ui_utility_tests.c
@@ -0,0 +1,806 @@
+/* Copyright 2020 The Chromium OS Authors. All rights reserved.
+ * Use of this source code is governed by a BSD-style license that can be
+ * found in the LICENSE file.
+ *
+ * Tests for UI functions without real UI flow.
+ */
+
+#include "2api.h"
+#include "2common.h"
+#include "2misc.h"
+#include "2nvstorage.h"
+#include "2ui.h"
+#include "2ui_private.h"
+#include "test_common.h"
+#include "vboot_api.h"
+#include "vboot_kernel.h"
+
+/* Fixed value for ignoring some checks. */
+#define MOCK_IGNORE 0xffffu
+
+/* Mock screen index for testing screen utility functions. */
+#define MOCK_NO_SCREEN 0xef0
+#define MOCK_SCREEN_BASE 0xeff
+#define MOCK_SCREEN_MENU 0xfff
+#define MOCK_SCREEN_TARGET0 0xff0
+#define MOCK_SCREEN_TARGET1 0xff1
+#define MOCK_SCREEN_TARGET2 0xff2
+#define MOCK_SCREEN_TARGET3 0xff3
+#define MOCK_SCREEN_TARGET4 0xff4
+
+/* Mock data */
+static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
+	__attribute__((aligned(VB2_WORKBUF_ALIGN)));
+static struct vb2_context *ctx;
+static struct vb2_gbb_header gbb;
+
+static int mock_shutdown_request;
+
+static struct vb2_ui_context mock_ui_context;
+static struct vb2_screen_state *mock_state;
+
+static struct vb2_screen_state mock_displayed[64];
+static int mock_displayed_count;
+static int mock_displayed_i;
+
+static uint32_t mock_key[64];
+static int mock_key_trusted[64];
+static int mock_key_count;
+static int mock_key_total;
+
+/* Mocks for testing screen utility functions. */
+const struct vb2_menu_item mock_empty_menu[] = {};
+struct vb2_screen_info mock_screen_blank = {
+	.id = VB2_SCREEN_BLANK,
+	.name = "mock_screen_blank",
+	.num_items = ARRAY_SIZE(mock_empty_menu),
+	.items = mock_empty_menu,
+};
+struct vb2_screen_info mock_screen_base =
+{
+	.id = MOCK_SCREEN_BASE,
+	.name = "mock_screen_base: menuless screen",
+	.num_items = ARRAY_SIZE(mock_empty_menu),
+	.items = mock_empty_menu,
+};
+struct vb2_menu_item mock_screen_menu_items[] =
+{
+	{
+		.text = "option 0",
+		.target = MOCK_SCREEN_TARGET0,
+	},
+	{
+		.text = "option 1",
+		.target = MOCK_SCREEN_TARGET1,
+	},
+	{
+		.text = "option 2",
+		.target = MOCK_SCREEN_TARGET2,
+	},
+	{
+		.text = "option 3",
+		.target = MOCK_SCREEN_TARGET3,
+	},
+	{
+		.text = "option 4 (no target)",
+	},
+};
+const struct vb2_screen_info mock_screen_menu =
+{
+	.id = MOCK_SCREEN_MENU,
+	.name = "mock_screen_menu: screen with 5 options",
+	.num_items = ARRAY_SIZE(mock_screen_menu_items),
+	.items = mock_screen_menu_items,
+};
+const struct vb2_screen_info mock_screen_target0 =
+{
+	.id = MOCK_SCREEN_TARGET0,
+	.name = "mock_screen_target0",
+	.num_items = ARRAY_SIZE(mock_empty_menu),
+	.items = mock_empty_menu,
+};
+const struct vb2_screen_info mock_screen_target1 =
+{
+	.id = MOCK_SCREEN_TARGET1,
+	.name = "mock_screen_target1",
+	.num_items = ARRAY_SIZE(mock_empty_menu),
+	.items = mock_empty_menu,
+};
+const struct vb2_screen_info mock_screen_target2 =
+{
+	.id = MOCK_SCREEN_TARGET2,
+	.name = "mock_screen_target2",
+	.num_items = ARRAY_SIZE(mock_empty_menu),
+	.items = mock_empty_menu,
+};
+const struct vb2_screen_info mock_screen_target3 =
+{
+	.id = MOCK_SCREEN_TARGET3,
+	.name = "mock_screen_target3",
+	.num_items = ARRAY_SIZE(mock_empty_menu),
+	.items = mock_empty_menu,
+};
+const struct vb2_screen_info mock_screen_target4 =
+{
+	.id = MOCK_SCREEN_TARGET4,
+	.name = "mock_screen_target4",
+	.num_items = ARRAY_SIZE(mock_empty_menu),
+	.items = mock_empty_menu,
+};
+
+/* Actions for tests */
+static uint32_t global_action_called;
+static vb2_error_t global_action_countdown(struct vb2_ui_context *ui)
+{
+	if (++global_action_called >= 10)
+		return VB2_SUCCESS;
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
+static vb2_error_t global_action_change_screen(struct vb2_ui_context *ui)
+{
+	change_screen(ui, MOCK_SCREEN_BASE);
+	if (++global_action_called >= 10)
+		return VB2_SUCCESS;
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
+static void screen_state_eq(const struct vb2_screen_state *state,
+			    enum vb2_screen screen,
+			    uint32_t locale_id,
+			    uint32_t selected_item,
+			    uint32_t disabled_item_mask)
+{
+	if (screen != MOCK_IGNORE) {
+		if (state->screen == NULL)
+			TEST_TRUE(0, "  state.screen does not exist");
+		else
+			TEST_EQ(state->screen->id, screen, "  state.screen");
+	}
+	if (locale_id != MOCK_IGNORE)
+		TEST_EQ(state->locale_id, locale_id, "  state.locale_id");
+	if (selected_item != MOCK_IGNORE)
+		TEST_EQ(state->selected_item,
+			selected_item, "  state.selected_item");
+	if (disabled_item_mask != MOCK_IGNORE)
+		TEST_EQ(state->disabled_item_mask,
+			disabled_item_mask, "  state.disabled_item_mask");
+}
+
+static void add_mock_key(uint32_t press, int trusted)
+{
+	if (mock_key_total >= ARRAY_SIZE(mock_key) ||
+	    mock_key_total >= ARRAY_SIZE(mock_key_trusted)) {
+		TEST_TRUE(0, "  mock_key ran out of entries!");
+		return;
+	}
+
+	mock_key[mock_key_total] = press;
+	mock_key_trusted[mock_key_total] = trusted;
+	mock_key_total++;
+}
+
+static void add_mock_keypress(uint32_t press)
+{
+	add_mock_key(press, 0);
+}
+
+static void displayed_eq(const char *text,
+			 enum vb2_screen screen,
+			 uint32_t locale_id,
+			 uint32_t selected_item,
+			 uint32_t disabled_item_mask)
+{
+	char text_buf[256];
+
+	if (mock_displayed_i >= mock_displayed_count) {
+		sprintf(text_buf, "  missing screen %s", text);
+		TEST_TRUE(0, text_buf);
+		return;
+	}
+
+	if (screen != MOCK_IGNORE) {
+		sprintf(text_buf, "  screen of %s", text);
+		TEST_EQ(mock_displayed[mock_displayed_i].screen->id, screen,
+			text_buf);
+	}
+	if (locale_id != MOCK_IGNORE) {
+		sprintf(text_buf, "  locale_id of %s", text);
+		TEST_EQ(mock_displayed[mock_displayed_i].locale_id, locale_id,
+			text_buf);
+	}
+	if (selected_item != MOCK_IGNORE) {
+		sprintf(text_buf, "  selected_item of %s", text);
+		TEST_EQ(mock_displayed[mock_displayed_i].selected_item,
+			selected_item, text_buf);
+	}
+	if (disabled_item_mask != MOCK_IGNORE) {
+		sprintf(text_buf, "  disabled_item_mask of %s", text);
+		TEST_EQ(mock_displayed[mock_displayed_i].disabled_item_mask,
+			disabled_item_mask, text_buf);
+	}
+	mock_displayed_i++;
+}
+
+static void displayed_no_extra(void)
+{
+	if (mock_displayed_i == 0)
+		TEST_EQ(mock_displayed_count, 0, "  no screen");
+	else
+		TEST_EQ(mock_displayed_count, mock_displayed_i,
+			"  no extra screens");
+}
+
+/* Reset mock data (for use before each test) */
+static void reset_common_data(void)
+{
+	TEST_SUCC(vb2api_init(workbuf, sizeof(workbuf), &ctx),
+		  "vb2api_init failed");
+
+	memset(&gbb, 0, sizeof(gbb));
+
+	vb2_nv_init(ctx);
+
+	/* For shutdown_required */
+	power_button = POWER_BUTTON_HELD_SINCE_BOOT;
+	mock_shutdown_request = MOCK_IGNORE;
+
+	/* For menu actions */
+	mock_ui_context = (struct vb2_ui_context){
+		.ctx = ctx,
+		.root_screen = &mock_screen_blank,
+		.state = (struct vb2_screen_state){
+			.screen = &mock_screen_blank,
+			.locale_id = 0,
+			.selected_item = 0,
+			.disabled_item_mask = 0,
+		},
+		.key = 0,
+
+	};
+	mock_state = &mock_ui_context.state;
+
+	/* For vb2ex_display_ui */
+	memset(mock_displayed, 0, sizeof(mock_displayed));
+	mock_displayed_count = 0;
+	mock_displayed_i = 0;
+
+	/* For VbExKeyboardRead */
+	memset(mock_key, 0, sizeof(mock_key));
+	memset(mock_key_trusted, 0, sizeof(mock_key_trusted));
+	mock_key_count = 0;
+	mock_key_total = 0;
+
+	/* For global actions */
+	global_action_called = 0;
+}
+
+/* Mock functions */
+struct vb2_gbb_header *vb2_get_gbb(struct vb2_context *c)
+{
+	return &gbb;
+}
+
+uint32_t VbExIsShutdownRequested(void)
+{
+	if (mock_shutdown_request != MOCK_IGNORE)
+		return mock_shutdown_request;
+
+	return 0;
+}
+
+const struct vb2_screen_info *vb2_get_screen_info(enum vb2_screen screen)
+{
+	switch ((int)screen) {
+	case VB2_SCREEN_BLANK:
+		return &mock_screen_blank;
+	case MOCK_SCREEN_BASE:
+		return &mock_screen_base;
+	case MOCK_SCREEN_MENU:
+		return &mock_screen_menu;
+	case MOCK_SCREEN_TARGET0:
+		return &mock_screen_target0;
+	case MOCK_SCREEN_TARGET1:
+		return &mock_screen_target1;
+	case MOCK_SCREEN_TARGET2:
+		return &mock_screen_target2;
+	case MOCK_SCREEN_TARGET3:
+		return &mock_screen_target3;
+	case MOCK_SCREEN_TARGET4:
+		return &mock_screen_target4;
+	default:
+		return NULL;
+	}
+}
+
+vb2_error_t vb2ex_display_ui(enum vb2_screen screen,
+			     uint32_t locale_id,
+			     uint32_t selected_item,
+			     uint32_t disabled_item_mask)
+{
+	VB2_DEBUG("displayed %d: screen = %#x, locale_id = %u, "
+		  "selected_item = %u, disabled_item_mask = %#x\n",
+		  mock_displayed_count, screen, locale_id, selected_item,
+		  disabled_item_mask);
+
+	if (mock_displayed_count >= ARRAY_SIZE(mock_displayed)) {
+		TEST_TRUE(0, "  mock vb2ex_display_ui ran out of entries!");
+		return VB2_ERROR_MOCK;
+	}
+
+	mock_displayed[mock_displayed_count] = (struct vb2_screen_state){
+		.screen = vb2_get_screen_info(screen),
+		.locale_id = locale_id,
+		.selected_item = selected_item,
+		.disabled_item_mask = disabled_item_mask,
+	};
+	mock_displayed_count++;
+
+	return VB2_SUCCESS;
+}
+
+uint32_t VbExKeyboardRead(void)
+{
+	return VbExKeyboardReadWithFlags(NULL);
+}
+
+uint32_t VbExKeyboardReadWithFlags(uint32_t *key_flags)
+{
+	if (mock_key_count < mock_key_total) {
+		if (key_flags != NULL) {
+			if (mock_key_trusted[mock_key_count])
+				*key_flags = VB_KEY_FLAG_TRUSTED_KEYBOARD;
+			else
+				*key_flags = 0;
+		}
+		return mock_key[mock_key_count++];
+	}
+
+	return 0;
+}
+
+/* Tests */
+static void shutdown_required_tests(void)
+{
+	VB2_DEBUG("Testing shutdown_required...\n");
+
+	/* Release, press, hold, and release */
+	if (!DETACHABLE) {
+		reset_common_data();
+		mock_shutdown_request = 0;
+		TEST_EQ(shutdown_required(ctx, 0), 0,
+			"release, press, hold, and release");
+		mock_shutdown_request = VB_SHUTDOWN_REQUEST_POWER_BUTTON;
+		TEST_EQ(shutdown_required(ctx, 0), 0, "  press");
+		TEST_EQ(shutdown_required(ctx, 0), 0, "  hold");
+		mock_shutdown_request = 0;
+		TEST_EQ(shutdown_required(ctx, 0), 1, "  release");
+	}
+
+	/* Press is ignored because we may held since boot */
+	if (!DETACHABLE) {
+		reset_common_data();
+		mock_shutdown_request = VB_SHUTDOWN_REQUEST_POWER_BUTTON;
+		TEST_EQ(shutdown_required(ctx, 0), 0, "press is ignored");
+	}
+
+	/* Power button short press from key */
+	if (!DETACHABLE) {
+		reset_common_data();
+		mock_shutdown_request = 0;
+		TEST_EQ(shutdown_required(ctx, VB_BUTTON_POWER_SHORT_PRESS), 1,
+			"power button short press");
+	}
+
+	/* Lid closure = shutdown request anyway */
+	reset_common_data();
+	mock_shutdown_request = VB_SHUTDOWN_REQUEST_LID_CLOSED;
+	TEST_EQ(shutdown_required(ctx, 0), 1, "lid closure");
+	TEST_EQ(shutdown_required(ctx, 'A'), 1, "  lidsw + random key");
+
+	/* Lid ignored by GBB flags */
+	reset_common_data();
+	gbb.flags |= VB2_GBB_FLAG_DISABLE_LID_SHUTDOWN;
+	mock_shutdown_request = VB_SHUTDOWN_REQUEST_LID_CLOSED;
+	TEST_EQ(shutdown_required(ctx, 0), 0, "lid ignored");
+	if (!DETACHABLE) {  /* Power button works for non DETACHABLE */
+		mock_shutdown_request = VB_SHUTDOWN_REQUEST_LID_CLOSED |
+					VB_SHUTDOWN_REQUEST_POWER_BUTTON;
+		TEST_EQ(shutdown_required(ctx, 0), 0, "  lidsw + pwdsw");
+		mock_shutdown_request = 0;
+		TEST_EQ(shutdown_required(ctx, 0), 1, "  pwdsw release");
+	}
+
+	/* Lid ignored; power button short pressed */
+	if (!DETACHABLE) {
+		reset_common_data();
+		gbb.flags |= VB2_GBB_FLAG_DISABLE_LID_SHUTDOWN;
+		mock_shutdown_request = VB_SHUTDOWN_REQUEST_LID_CLOSED;
+		TEST_EQ(shutdown_required(ctx, VB_BUTTON_POWER_SHORT_PRESS), 1,
+			"lid ignored; power button short pressed");
+	}
+
+	/* DETACHABLE ignore power button */
+	if (DETACHABLE) {
+		/* Flag pwdsw */
+		reset_common_data();
+		mock_shutdown_request = VB_SHUTDOWN_REQUEST_POWER_BUTTON;
+		TEST_EQ(shutdown_required(ctx, 0), 0,
+			"DETACHABLE: ignore pwdsw");
+		mock_shutdown_request = 0;
+		TEST_EQ(shutdown_required(ctx, 0), 0,
+			"  ignore on release");
+
+		/* Power button short press */
+		reset_common_data();
+		mock_shutdown_request = 0;
+		TEST_EQ(shutdown_required(
+		    ctx, VB_BUTTON_POWER_SHORT_PRESS), 0,
+		    "DETACHABLE: ignore power button short press");
+	}
+
+	VB2_DEBUG("...done.\n");
+}
+
+static void menu_action_tests(void)
+{
+	int i, target_id;
+	char test_name[256];
+
+	VB2_DEBUG("Testing menu actions...\n");
+
+	/* Valid menu_up_action */
+	reset_common_data();
+	mock_state->screen = &mock_screen_menu;
+	mock_state->selected_item = 2;
+	mock_ui_context.key = VB_KEY_UP;
+	TEST_EQ(menu_up_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"valid menu_up_action");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 1,
+			MOCK_IGNORE);
+
+	/* Valid menu_up_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.key = VB_KEY_UP;
+	TEST_EQ(menu_up_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"valid menu_up_action with mask");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 0,
+			MOCK_IGNORE);
+
+	/* Invalid menu_up_action (blocked) */
+	reset_common_data();
+	mock_state->screen = &mock_screen_menu;
+	mock_state->selected_item = 0;
+	mock_ui_context.key = VB_KEY_UP;
+	TEST_EQ(menu_up_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"invalid menu_up_action (blocked)");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 0,
+			MOCK_IGNORE);
+
+	/* Invalid menu_up_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.key = VB_KEY_UP;
+	TEST_EQ(menu_up_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"invalid menu_up_action (blocked by mask)");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 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.key = VB_BUTTON_VOL_UP_SHORT_PRESS;
+		TEST_EQ(menu_up_action(&mock_ui_context),
+			VB2_REQUEST_UI_CONTINUE,
+			"ignore volume-up when not DETACHABLE");
+		screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 2,
+				MOCK_IGNORE);
+	}
+
+	/* Valid menu_down_action */
+	reset_common_data();
+	mock_state->screen = &mock_screen_menu;
+	mock_state->selected_item = 2;
+	mock_ui_context.key = VB_KEY_DOWN;
+	TEST_EQ(menu_down_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"valid menu_down_action");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 3,
+			MOCK_IGNORE);
+
+	/* Valid menu_down_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.key = VB_KEY_DOWN;
+	TEST_EQ(menu_down_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"valid menu_down_action with mask");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 4,
+			MOCK_IGNORE);
+
+	/* Invalid menu_down_action (blocked) */
+	reset_common_data();
+	mock_state->screen = &mock_screen_menu;
+	mock_state->selected_item = 4;
+	mock_ui_context.key = VB_KEY_DOWN;
+	TEST_EQ(menu_down_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"invalid menu_down_action (blocked)");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 4,
+			MOCK_IGNORE);
+
+	/* Invalid menu_down_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.key = VB_KEY_DOWN;
+	TEST_EQ(menu_down_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"invalid menu_down_action (blocked by mask)");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 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.key = VB_BUTTON_VOL_DOWN_SHORT_PRESS;
+		TEST_EQ(menu_down_action(&mock_ui_context),
+			VB2_REQUEST_UI_CONTINUE,
+			"ignore volume-down when not DETACHABLE");
+		screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 2,
+				MOCK_IGNORE);
+	}
+
+	/* menu_select_action with no item screen */
+	reset_common_data();
+	mock_state->screen = &mock_screen_base;
+	mock_ui_context.key = VB_KEY_ENTER;
+	TEST_EQ(menu_select_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"menu_select_action with no item screen");
+	screen_state_eq(mock_state, MOCK_SCREEN_BASE, MOCK_IGNORE, 0,
+			MOCK_IGNORE);
+
+	/* Try to select target 0..3 */
+	for (i = 0; i <= 3; i++) {
+		sprintf(test_name, "select target %d", i);
+		target_id = MOCK_SCREEN_TARGET0 + i;
+		reset_common_data();
+		mock_state->screen = &mock_screen_menu;
+		mock_state->selected_item = i;
+		mock_ui_context.key = VB_KEY_ENTER;
+		TEST_EQ(menu_select_action(&mock_ui_context),
+			VB2_REQUEST_UI_CONTINUE, test_name);
+		screen_state_eq(mock_state, target_id, MOCK_IGNORE, 0,
+				MOCK_IGNORE);
+	}
+
+	/* Try to select no target item */
+	reset_common_data();
+	mock_state->screen = &mock_screen_menu;
+	mock_state->selected_item = 4;
+	mock_ui_context.key = VB_KEY_ENTER;
+	TEST_EQ(menu_select_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"select no target");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 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.key = VB_BUTTON_POWER_SHORT_PRESS;
+		TEST_EQ(menu_select_action(&mock_ui_context),
+			VB2_REQUEST_UI_CONTINUE,
+			"ignore power button short press when not DETACHABLE");
+		screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 1,
+				MOCK_IGNORE);
+	}
+
+	/* menu_back_action */
+	reset_common_data();
+	mock_ui_context.key = VB_KEY_ESC;
+	TEST_EQ(menu_back_action(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"menu_back_action");
+	screen_state_eq(mock_state, VB2_SCREEN_BLANK, MOCK_IGNORE, 0,
+			MOCK_IGNORE);
+
+	VB2_DEBUG("...done.\n");
+}
+
+static void change_screen_tests(void)
+{
+	VB2_DEBUG("Testing change_screen...\n");
+
+	/* Changing screen will clear screen state except locale_id */
+	reset_common_data();
+	mock_state->screen = &mock_screen_menu;
+	mock_state->locale_id = 1;
+	mock_state->selected_item = 2;
+	mock_state->disabled_item_mask = 0x10;
+	VB2_DEBUG("change_screen will clear screen state except locale_id\n");
+	change_screen(&mock_ui_context, MOCK_SCREEN_BASE);
+	screen_state_eq(mock_state, MOCK_SCREEN_BASE, 1, 0, 0);
+
+	/* Change to screen which does not exist */
+	reset_common_data();
+	mock_state->screen = &mock_screen_menu;
+	VB2_DEBUG("change to screen which does not exist\n");
+	change_screen(&mock_ui_context, MOCK_NO_SCREEN);
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, MOCK_IGNORE,
+			MOCK_IGNORE);
+
+	VB2_DEBUG("...done.\n");
+}
+
+static void validate_selection_tests(void)
+{
+	VB2_DEBUG("Testing validate_selection...");
+
+	/* No item */
+	reset_common_data();
+	mock_state->screen = &mock_screen_base;
+	mock_state->selected_item = 2;
+	mock_state->disabled_item_mask = 0x10;
+	VB2_DEBUG("no item (fix selected_item)\n");
+	validate_selection(mock_state);
+	screen_state_eq(mock_state, MOCK_SCREEN_BASE, MOCK_IGNORE, 0,
+			MOCK_IGNORE);
+
+	/* Valid selected_item */
+	reset_common_data();
+	mock_state->screen = &mock_screen_menu;
+	mock_state->selected_item = 2;
+	mock_state->disabled_item_mask = 0x13;  /* 0b10011 */
+	VB2_DEBUG("valid selected_item\n");
+	validate_selection(mock_state);
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 2,
+			MOCK_IGNORE);
+
+	/* selected_item too large */
+	reset_common_data();
+	mock_state->screen = &mock_screen_menu;
+	mock_state->selected_item = 5;
+	mock_state->disabled_item_mask = 0x15;  /* 0b10101 */
+	VB2_DEBUG("selected_item too large\n");
+	validate_selection(mock_state);
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 1,
+			MOCK_IGNORE);
+
+	/* Select a disabled item */
+	reset_common_data();
+	mock_state->screen = &mock_screen_menu;
+	mock_state->selected_item = 4;
+	mock_state->disabled_item_mask = 0x17;  /* 0b10111 */
+	VB2_DEBUG("select a disabled item\n");
+	validate_selection(mock_state);
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 3,
+			MOCK_IGNORE);
+
+	/* No available item */
+	reset_common_data();
+	mock_state->screen = &mock_screen_menu;
+	mock_state->selected_item = 2;
+	mock_state->disabled_item_mask = 0x1f;  /* 0b11111 */
+	VB2_DEBUG("no available item\n");
+	validate_selection(mock_state);
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU, MOCK_IGNORE, 0,
+			MOCK_IGNORE);
+
+	VB2_DEBUG("...done.\n");
+}
+
+static void ui_loop_tests(void)
+{
+	VB2_DEBUG("Testing ui_loop...\n");
+
+	/* Die if no root screen */
+	reset_common_data();
+	TEST_ABORT(ui_loop(ctx, MOCK_NO_SCREEN, NULL),
+		   "die if no root screen");
+	displayed_no_extra();
+
+	/* Shutdown if requested */
+	reset_common_data();
+	mock_shutdown_request = VB_SHUTDOWN_REQUEST_OTHER;
+	TEST_EQ(ui_loop(ctx, MOCK_SCREEN_BASE, NULL),
+		VB2_REQUEST_SHUTDOWN, "shutdown if requested");
+	displayed_eq("mock_screen_base", MOCK_SCREEN_BASE, MOCK_IGNORE,
+		     MOCK_IGNORE, MOCK_IGNORE);
+	displayed_no_extra();
+
+	/* Global action */
+	reset_common_data();
+	TEST_EQ(ui_loop(ctx, VB2_SCREEN_BLANK, global_action_countdown),
+		VB2_SUCCESS, "global action");
+	TEST_EQ(global_action_called, 10, "  global action called");
+
+	/* Global action can change screen */
+	reset_common_data();
+	TEST_EQ(ui_loop(ctx, VB2_SCREEN_BLANK, global_action_change_screen),
+		VB2_SUCCESS, "global action can change screen");
+	TEST_EQ(global_action_called, 10, "  global action called");
+	displayed_eq("pass", MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE,
+		     MOCK_IGNORE);
+	displayed_eq("change to mock_screen_base", MOCK_IGNORE, MOCK_IGNORE,
+		     MOCK_IGNORE, MOCK_IGNORE);
+
+	/* KEY_UP, KEY_DOWN, and KEY_ENTER navigation */
+	reset_common_data();
+	add_mock_keypress(VB_KEY_UP);  /* (blocked) */
+	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_DOWN);  /* (blocked) */
+	add_mock_keypress(VB_KEY_UP);
+	add_mock_keypress(VB_KEY_ENTER);
+	TEST_EQ(ui_loop(ctx, MOCK_SCREEN_MENU, global_action_countdown),
+		VB2_SUCCESS, "KEY_UP, KEY_DOWN, and KEY_ENTER");
+	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 0,
+		     MOCK_IGNORE);
+	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 1,
+		     MOCK_IGNORE);
+	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 2,
+		     MOCK_IGNORE);
+	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 3,
+		     MOCK_IGNORE);
+	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 4,
+		     MOCK_IGNORE);
+	displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE, 3,
+		     MOCK_IGNORE);
+	displayed_eq("mock_screen_target_3", MOCK_SCREEN_TARGET3, MOCK_IGNORE,
+		     MOCK_IGNORE, MOCK_IGNORE);
+	displayed_no_extra();
+
+	/* For DETACHABLE */
+	if (DETACHABLE) {
+		reset_common_data();
+		add_mock_keypress(VB_BUTTON_VOL_UP_SHORT_PRESS);
+		add_mock_keypress(VB_BUTTON_VOL_DOWN_SHORT_PRESS);
+		add_mock_keypress(VB_BUTTON_VOL_DOWN_SHORT_PRESS);
+		add_mock_keypress(VB_BUTTON_VOL_DOWN_SHORT_PRESS);
+		add_mock_keypress(VB_BUTTON_VOL_DOWN_SHORT_PRESS);
+		add_mock_keypress(VB_BUTTON_VOL_DOWN_SHORT_PRESS);
+		add_mock_keypress(VB_BUTTON_VOL_UP_SHORT_PRESS);
+		add_mock_keypress(VB_BUTTON_POWER_SHORT_PRESS);
+		TEST_EQ(ui_loop(ctx, MOCK_SCREEN_MENU, global_action_countdown),
+			VB2_SUCCESS, "DETACHABLE");
+		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
+			     0, MOCK_IGNORE);
+		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
+			     1, MOCK_IGNORE);
+		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
+			     2, MOCK_IGNORE);
+		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
+			     3, MOCK_IGNORE);
+		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
+			     4, MOCK_IGNORE);
+		displayed_eq("mock_screen_menu", MOCK_SCREEN_MENU, MOCK_IGNORE,
+			     3, MOCK_IGNORE);
+		displayed_eq("mock_screen_target_3", MOCK_SCREEN_TARGET3,
+			     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
+		displayed_no_extra();
+	}
+
+	VB2_DEBUG("...done.\n");
+}
+
+int main(void)
+{
+	shutdown_required_tests();
+	menu_action_tests();
+	change_screen_tests();
+	validate_selection_tests();
+	ui_loop_tests();
+
+	return gTestSuccess ? 0 : 255;
+}