vboot/ui: add developer and to_norm screens

Developer mode can boot internal or external disk.  Support for
booting legacy firmware ("Ctrl+L") will be added in a subsequent CL.

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

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

7 files changed, 335 insertions, 95 deletions

diff --git a/firmware/2lib/2ui.c b/firmware/2lib/2ui.c
index 6f90135b..a711bde8 100644
--- a/firmware/2lib/2ui.c
+++ b/firmware/2lib/2ui.c
@@ -120,7 +120,7 @@ vb2_error_t menu_down_action(struct vb2_ui_context *ui)
 /**
  * Navigate to the target screen of the current menu item selection.
  */
-vb2_error_t menu_select_action(struct vb2_ui_context *ui)
+vb2_error_t vb2_ui_menu_select_action(struct vb2_ui_context *ui)
 {
 	const struct vb2_menu_item *menu_item;
 
@@ -138,7 +138,7 @@ vb2_error_t menu_select_action(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 change_screen(ui, menu_item->target);
+		return vb2_ui_change_screen(ui, menu_item->target);
 	}
 
 	VB2_DEBUG("Menu item <%s> no action or target screen\n",
@@ -152,19 +152,30 @@ vb2_error_t menu_select_action(struct vb2_ui_context *ui)
 vb2_error_t vb2_ui_back_action(struct vb2_ui_context *ui)
 {
 	/* TODO(kitching): Return to previous screen instead of root screen. */
-	return change_screen(ui, ui->root_screen->id);
+	return vb2_ui_change_screen(ui, ui->root_screen->id);
 }
 
 /**
  * Context-dependent keyboard shortcut Ctrl+D.
  *
  * - Manual recovery mode: Change to dev mode transition screen.
- * - Developer mode: Boot from internal disk (TODO).
+ * - Developer mode: Boot from internal disk.
  */
 vb2_error_t ctrl_d_action(struct vb2_ui_context *ui)
 {
 	if (vb2_allow_recovery(ui->ctx))
-		return change_screen(ui, VB2_SCREEN_RECOVERY_TO_DEV);
+		return change_to_dev_screen_action(ui);
+	else if (ui->ctx->flags & VB2_CONTEXT_DEVELOPER_MODE)
+		return vb2_ui_developer_mode_boot_internal_action(ui);
+
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
+vb2_error_t change_to_dev_screen_action(struct vb2_ui_context *ui)
+{
+	if (vb2_allow_recovery(ui->ctx))
+		return vb2_ui_change_screen(ui, VB2_SCREEN_RECOVERY_TO_DEV);
+
 	return VB2_REQUEST_UI_CONTINUE;
 }
 
@@ -174,13 +185,20 @@ vb2_error_t ctrl_d_action(struct vb2_ui_context *ui)
 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_KEY_ENTER,  			vb2_ui_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_BUTTON_POWER_SHORT_PRESS, 		vb2_ui_menu_select_action },
 	{ VB_KEY_ESC, 			 	vb2_ui_back_action },
 	{ VB_KEY_CTRL('D'),		 	ctrl_d_action },
-	{ ' ',				 	vb2_ui_recovery_to_dev_action },
+	{ VB_BUTTON_VOL_DOWN_LONG_PRESS,
+	  vb2_ui_developer_mode_boot_internal_action },
+	{ VB_BUTTON_VOL_UP_DOWN_COMBO_PRESS,	change_to_dev_screen_action },
+	{ ' ',					vb2_ui_recovery_to_dev_action },
+	{ VB_KEY_CTRL('U'),
+	  vb2_ui_developer_mode_boot_external_action },
+	{ VB_BUTTON_VOL_UP_LONG_PRESS,
+	  vb2_ui_developer_mode_boot_external_action },
 };
 
 vb2_error_t (*input_action_lookup(int key))(struct vb2_ui_context *ui)
@@ -195,7 +213,7 @@ vb2_error_t (*input_action_lookup(int key))(struct vb2_ui_context *ui)
 /*****************************************************************************/
 /* Core UI functions */
 
-vb2_error_t change_screen(struct vb2_ui_context *ui, enum vb2_screen id)
+vb2_error_t vb2_ui_change_screen(struct vb2_ui_context *ui, enum vb2_screen id)
 {
 	const struct vb2_screen_info *new_screen_info = vb2_get_screen_info(id);
 
@@ -227,7 +245,7 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 	ui.root_screen = vb2_get_screen_info(root_screen_id);
 	if (ui.root_screen == NULL)
 		VB2_DIE("Root screen not found.\n");
-	rv = change_screen(&ui, ui.root_screen->id);
+	rv = vb2_ui_change_screen(&ui, ui.root_screen->id);
 	if (rv != VB2_REQUEST_UI_CONTINUE)
 		return rv;
 	memset(&prev_state, 0, sizeof(prev_state));
@@ -299,27 +317,7 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 
 vb2_error_t vb2_developer_menu(struct vb2_context *ctx)
 {
-	enum vb2_dev_default_boot default_boot;
-
-	/* If dev mode was disabled, loop forever. */
-	if (!vb2_dev_boot_allowed(ctx))
-		while (1);
-
-	/* Boot from the default option. */
-	default_boot = vb2_get_dev_boot_target(ctx);
-
-	/* Boot legacy does not return on success */
-	if (default_boot == VB2_DEV_DEFAULT_BOOT_LEGACY &&
-	    vb2_dev_boot_legacy_allowed(ctx) &&
-	    VbExLegacy(VB_ALTFW_DEFAULT) == VB2_SUCCESS)
-		return VB2_SUCCESS;
-
-	if (default_boot == VB2_DEV_DEFAULT_BOOT_USB &&
-	    vb2_dev_boot_usb_allowed(ctx) &&
-	    VbTryLoadKernel(ctx, VB_DISK_FLAG_REMOVABLE) == VB2_SUCCESS)
-		return VB2_SUCCESS;
-
-	return VbTryLoadKernel(ctx, VB_DISK_FLAG_FIXED);
+	return ui_loop(ctx, VB2_SCREEN_DEVELOPER_MODE, NULL);
 }
 
 /*****************************************************************************/
@@ -350,9 +348,9 @@ vb2_error_t try_recovery_action(struct vb2_ui_context *ui)
 	invalid_disk = rv != VB2_ERROR_LK_NO_DISK_FOUND;
 	if (invalid_disk_last != invalid_disk) {
 		invalid_disk_last = invalid_disk;
-		return change_screen(ui, invalid_disk ?
-				     VB2_SCREEN_RECOVERY_INVALID :
-				     VB2_SCREEN_RECOVERY_SELECT);
+		return vb2_ui_change_screen(ui, invalid_disk ?
+					    VB2_SCREEN_RECOVERY_INVALID :
+					    VB2_SCREEN_RECOVERY_SELECT);
 	}
 
 	return VB2_REQUEST_UI_CONTINUE;
diff --git a/firmware/2lib/2ui_screens.c b/firmware/2lib/2ui_screens.c
index ca492039..356b3bfb 100644
--- a/firmware/2lib/2ui_screens.c
+++ b/firmware/2lib/2ui_screens.c
@@ -10,7 +10,8 @@
 #include "2nvstorage.h"
 #include "2ui.h"
 #include "2ui_private.h"
-#include "vboot_api.h"  /* for VB_KEY_ */
+#include "vboot_api.h"
+#include "vboot_kernel.h"
 
 #define MENU_ITEMS(a) \
 	.num_items = ARRAY_SIZE(a), \
@@ -45,12 +46,26 @@ static const struct vb2_screen_info recovery_broken_screen = {
 /******************************************************************************/
 /* VB2_SCREEN_ADVANCED_OPTIONS */
 
+#define ADVANCED_OPTIONS_ITEM_DEVELOPER_MODE 0
+#define ADVANCED_OPTIONS_ITEM_BACK 1
+
+vb2_error_t advanced_options_init(struct vb2_ui_context *ui)
+{
+	if (vb2_get_sd(ui->ctx)->flags & VB2_SD_FLAG_DEV_MODE_ENABLED) {
+		ui->state.disabled_item_mask |=
+			1 << ADVANCED_OPTIONS_ITEM_DEVELOPER_MODE;
+		ui->state.selected_item = ADVANCED_OPTIONS_ITEM_BACK;
+	}
+
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
 static const struct vb2_menu_item advanced_options_items[] = {
-	{
-		.text = "Developer mode",
+	[ADVANCED_OPTIONS_ITEM_DEVELOPER_MODE] = {
+		.text = "Enable developer mode",
 		.target = VB2_SCREEN_RECOVERY_TO_DEV,
 	},
-	{
+	[ADVANCED_OPTIONS_ITEM_BACK] = {
 		.text = "Back",
 		.action = vb2_ui_back_action,
 	},
@@ -59,6 +74,7 @@ static const struct vb2_menu_item advanced_options_items[] = {
 static const struct vb2_screen_info advanced_options_screen = {
 	.id = VB2_SCREEN_ADVANCED_OPTIONS,
 	.name = "Advanced options",
+	.init = advanced_options_init,
 	MENU_ITEMS(advanced_options_items),
 };
 
@@ -95,6 +111,7 @@ static const struct vb2_screen_info recovery_invalid_screen = {
 /* VB2_SCREEN_RECOVERY_TO_DEV */
 
 #define RECOVERY_TO_DEV_ITEM_CONFIRM 0
+#define RECOVERY_TO_DEV_ITEM_CANCEL 1
 
 vb2_error_t recovery_to_dev_init(struct vb2_ui_context *ui)
 {
@@ -109,9 +126,11 @@ vb2_error_t recovery_to_dev_init(struct vb2_ui_context *ui)
 	}
 
 	/* Disable "Confirm" button for other physical presence types. */
-	if (!PHYSICAL_PRESENCE_KEYBOARD)
-		ui->state.disabled_item_mask =
+	if (!PHYSICAL_PRESENCE_KEYBOARD) {
+		ui->state.disabled_item_mask |=
 			1 << RECOVERY_TO_DEV_ITEM_CONFIRM;
+		ui->state.selected_item = RECOVERY_TO_DEV_ITEM_CANCEL;
+	}
 
 	return VB2_REQUEST_UI_CONTINUE;
 }
@@ -158,7 +177,7 @@ vb2_error_t vb2_ui_recovery_to_dev_action(struct vb2_ui_context *ui)
 	/* Sanity check, should never happen. */
 	if ((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");
+		VB2_DEBUG("ERROR: Dev transition sanity check failed\n");
 		return VB2_REQUEST_UI_CONTINUE;
 	}
 
@@ -172,7 +191,7 @@ static const struct vb2_menu_item recovery_to_dev_items[] = {
 		.text = "Confirm",
 		.action = vb2_ui_recovery_to_dev_action,
 	},
-	{
+	[RECOVERY_TO_DEV_ITEM_CANCEL] = {
 		.text = "Cancel",
 		.action = vb2_ui_back_action,
 	},
@@ -203,6 +222,177 @@ static const struct vb2_screen_info recovery_disk_step1_screen = {
 };
 
 /******************************************************************************/
+/* VB2_SCREEN_DEVELOPER_MODE */
+
+#define DEVELOPER_MODE_ITEM_RETURN_TO_SECURE 0
+#define DEVELOPER_MODE_ITEM_BOOT_INTERNAL 1
+#define DEVELOPER_MODE_ITEM_BOOT_EXTERNAL 2
+
+vb2_error_t developer_mode_init(struct vb2_ui_context *ui)
+{
+	enum vb2_dev_default_boot default_boot =
+		vb2_get_dev_boot_target(ui->ctx);
+
+	/* Get me outta here! */
+	if (!vb2_dev_boot_allowed(ui->ctx))
+		vb2_ui_change_screen(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 |=
+			1 << DEVELOPER_MODE_ITEM_RETURN_TO_SECURE;
+
+	/* Don't show "Boot from external disk" button if not allowed. */
+	if (!vb2_dev_boot_usb_allowed(ui->ctx))
+		ui->state.disabled_item_mask |=
+			1 << DEVELOPER_MODE_ITEM_BOOT_EXTERNAL;
+
+	/* Choose the default selection. */
+	switch (default_boot) {
+	case VB2_DEV_DEFAULT_BOOT_USB:
+		ui->state.selected_item = DEVELOPER_MODE_ITEM_BOOT_EXTERNAL;
+		break;
+	default:
+		ui->state.selected_item = DEVELOPER_MODE_ITEM_BOOT_INTERNAL;
+		break;
+	}
+
+	ui->start_time = VbExGetTimer();
+
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
+vb2_error_t vb2_ui_developer_mode_boot_internal_action(
+	struct vb2_ui_context *ui)
+{
+	if (!(ui->ctx->flags & VB2_CONTEXT_DEVELOPER_MODE) ||
+	    !vb2_dev_boot_allowed(ui->ctx)) {
+		VB2_DEBUG("ERROR: Dev mode internal boot not allowed\n");
+		return VB2_REQUEST_UI_CONTINUE;
+	}
+
+	if (VbTryLoadKernel(ui->ctx, VB_DISK_FLAG_FIXED)) {
+		VB2_DEBUG("ERROR: Dev mode internal boot failed\n");
+		return VB2_REQUEST_UI_CONTINUE;
+	}
+
+	return VB2_SUCCESS;
+}
+
+vb2_error_t vb2_ui_developer_mode_boot_external_action(
+	struct vb2_ui_context *ui)
+{
+	/* Sanity check, should never happen. */
+	if (!(ui->ctx->flags & VB2_CONTEXT_DEVELOPER_MODE) ||
+	    !vb2_dev_boot_allowed(ui->ctx) ||
+	    !vb2_dev_boot_usb_allowed(ui->ctx)) {
+		VB2_DEBUG("ERROR: Dev mode external boot not allowed\n");
+		return VB2_REQUEST_UI_CONTINUE;
+	}
+
+	if (VbTryLoadKernel(ui->ctx, VB_DISK_FLAG_REMOVABLE)) {
+		VB2_DEBUG("ERROR: Dev mode external boot failed\n");
+		return VB2_REQUEST_UI_CONTINUE;
+	}
+
+	return VB2_SUCCESS;
+}
+
+vb2_error_t developer_mode_action(struct vb2_ui_context *ui)
+{
+	struct vb2_gbb_header *gbb = vb2_get_gbb(ui->ctx);
+	const int use_short = gbb->flags & VB2_GBB_FLAG_DEV_SCREEN_SHORT_DELAY;
+	uint64_t elapsed;
+
+	/* Once any user interaction occurs, stop the timer. */
+	if (ui->key)
+		ui->disable_timer = 1;
+	if (ui->disable_timer)
+		return VB2_REQUEST_UI_CONTINUE;
+
+	elapsed = VbExGetTimer() - ui->start_time;
+
+	/* If we're using short delay, wait 2 seconds and don't beep. */
+	if (use_short && elapsed > 2 * VB_USEC_PER_SEC) {
+		VB2_DEBUG("Booting default target after 2s\n");
+		ui->disable_timer = 1;
+		return vb2_ui_menu_select_action(ui);
+	}
+
+	/* Otherwise, beep at 20 and 20.5 seconds. */
+	if ((ui->beep_count == 0 && elapsed > 20 * VB_USEC_PER_SEC) ||
+	    (ui->beep_count == 1 && elapsed > 20500 * VB_USEC_PER_MSEC)) {
+		VbExBeep(250, 400);
+		ui->beep_count++;
+	}
+
+	/* Stop after 30 seconds. */
+	if (elapsed > 30 * VB_USEC_PER_SEC) {
+		VB2_DEBUG("Booting default target after 30s\n");
+		ui->disable_timer = 1;
+		return vb2_ui_menu_select_action(ui);
+	}
+
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
+static const struct vb2_menu_item developer_mode_items[] = {
+	[DEVELOPER_MODE_ITEM_RETURN_TO_SECURE] = {
+		.text = "Return to secure mode",
+		.target = VB2_SCREEN_DEVELOPER_TO_NORM,
+	},
+	[DEVELOPER_MODE_ITEM_BOOT_INTERNAL] = {
+		.text = "Boot from internal disk",
+		.action = vb2_ui_developer_mode_boot_internal_action,
+	},
+	[DEVELOPER_MODE_ITEM_BOOT_EXTERNAL] = {
+		.text = "Boot from external disk",
+		.action = vb2_ui_developer_mode_boot_external_action,
+	},
+	ADVANCED_OPTIONS_ITEM,
+};
+
+static const struct vb2_screen_info developer_mode_screen = {
+	.id = VB2_SCREEN_DEVELOPER_MODE,
+	.name = "Developer mode",
+	.init = developer_mode_init,
+	.action = developer_mode_action,
+	MENU_ITEMS(developer_mode_items),
+};
+
+/******************************************************************************/
+/* VB2_SCREEN_DEVELOPER_TO_NORM */
+
+vb2_error_t developer_to_norm_action(struct vb2_ui_context *ui)
+{
+	if (vb2_get_gbb(ui->ctx)->flags & VB2_GBB_FLAG_FORCE_DEV_SWITCH_ON) {
+		VB2_DEBUG("ERROR: dev mode forced by GBB flag\n");
+		return VB2_REQUEST_UI_CONTINUE;
+	}
+
+	VB2_DEBUG("Leaving dev mode\n");
+	vb2_nv_set(ui->ctx, VB2_NV_DISABLE_DEV_REQUEST, 1);
+	return VB2_REQUEST_REBOOT;
+}
+
+static const struct vb2_menu_item developer_to_norm_items[] = {
+	{
+		.text = "Confirm",
+		.action = developer_to_norm_action,
+	},
+	{
+		.text = "Cancel",
+		.action = vb2_ui_back_action,
+	},
+};
+
+static const struct vb2_screen_info developer_to_norm_screen = {
+	.id = VB2_SCREEN_DEVELOPER_TO_NORM,
+	.name = "Transition to normal mode",
+	MENU_ITEMS(developer_to_norm_items),
+};
+
+/******************************************************************************/
 /*
  * TODO(chromium:1035800): Refactor UI code across vboot and depthcharge.
  * Currently vboot and depthcharge maintain their own copies of menus/screens.
@@ -219,6 +409,8 @@ static const struct vb2_screen_info *screens[] = {
 	&recovery_to_dev_screen,
 	&recovery_phone_step1_screen,
 	&recovery_disk_step1_screen,
+	&developer_mode_screen,
+	&developer_to_norm_screen,
 };
 
 const struct vb2_screen_info *vb2_get_screen_info(enum vb2_screen id)
diff --git a/firmware/2lib/include/2ui.h b/firmware/2lib/include/2ui.h
index b5754020..0043963a 100644
--- a/firmware/2lib/include/2ui.h
+++ b/firmware/2lib/include/2ui.h
@@ -53,11 +53,22 @@ struct vb2_ui_context {
 	uint32_t locale_id;
 	uint32_t key;
 	int key_trusted;
+
+	/* For developer mode screen. */
+	int disable_timer;
+	uint64_t start_time;
+	int beep_count;
 };
 
 vb2_error_t vb2_ui_change_screen(struct vb2_ui_context *ui, enum vb2_screen id);
+vb2_error_t vb2_ui_menu_select_action(struct vb2_ui_context *ui);
 vb2_error_t vb2_ui_back_action(struct vb2_ui_context *ui);
+
 vb2_error_t vb2_ui_recovery_to_dev_action(struct vb2_ui_context *ui);
+vb2_error_t vb2_ui_developer_mode_boot_internal_action(
+	struct vb2_ui_context *ui);
+vb2_error_t vb2_ui_developer_mode_boot_external_action(
+	struct vb2_ui_context *ui);
 
 /**
  * Get info struct of a screen.
diff --git a/firmware/2lib/include/2ui_private.h b/firmware/2lib/include/2ui_private.h
index f7c4e247..bfc0dd0b 100644
--- a/firmware/2lib/include/2ui_private.h
+++ b/firmware/2lib/include/2ui_private.h
@@ -27,17 +27,20 @@ struct input_action {
 
 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 ctrl_d_action(struct vb2_ui_context *ui);
+vb2_error_t change_to_dev_screen_action(struct vb2_ui_context *ui);
 vb2_error_t (*input_action_lookup(int key))(struct vb2_ui_context *ui);
 
-vb2_error_t change_screen(struct vb2_ui_context *ui, enum vb2_screen id);
 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);
 
 /* From 2ui_screens.c */
+vb2_error_t advanced_options_init(struct vb2_ui_context *ui);
 vb2_error_t recovery_to_dev_init(struct vb2_ui_context *ui);
+vb2_error_t developer_mode_init(struct vb2_ui_context *ui);
+vb2_error_t developer_mode_action(struct vb2_ui_context *ui);
+vb2_error_t developer_to_norm_action(struct vb2_ui_context *ui);
 
 #endif  /* VBOOT_REFERENCE_2UI_PRIVATE_H_ */
diff --git a/tests/vb2_ui_action_tests.c b/tests/vb2_ui_action_tests.c
index dd87bc33..03b4eed7 100644
--- a/tests/vb2_ui_action_tests.c
+++ b/tests/vb2_ui_action_tests.c
@@ -142,7 +142,7 @@ static vb2_error_t global_action_countdown(struct vb2_ui_context *ui)
 
 static vb2_error_t global_action_change_screen(struct vb2_ui_context *ui)
 {
-	change_screen(ui, MOCK_SCREEN_BASE);
+	vb2_ui_change_screen(ui, MOCK_SCREEN_BASE);
 	return VB2_REQUEST_UI_CONTINUE;
 }
 
@@ -500,7 +500,8 @@ static void menu_action_tests(void)
 	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,
+	TEST_EQ(vb2_ui_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, 0, MOCK_IGNORE);
 
@@ -512,7 +513,7 @@ static void menu_action_tests(void)
 		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),
+		TEST_EQ(vb2_ui_menu_select_action(&mock_ui_context),
 			VB2_REQUEST_UI_CONTINUE, test_name);
 		screen_state_eq(mock_state, target_id, 0, MOCK_IGNORE);
 	}
@@ -522,7 +523,8 @@ static void menu_action_tests(void)
 	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,
+	TEST_EQ(vb2_ui_menu_select_action(&mock_ui_context),
+		VB2_REQUEST_UI_CONTINUE,
 		"select no target");
 	screen_state_eq(mock_state, MOCK_SCREEN_MENU, 4, MOCK_IGNORE);
 
@@ -532,7 +534,7 @@ static void menu_action_tests(void)
 		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),
+		TEST_EQ(vb2_ui_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, 1, MOCK_IGNORE);
diff --git a/tests/vb2_ui_tests.c b/tests/vb2_ui_tests.c
index 2f652880..1bd9e8b9 100644
--- a/tests/vb2_ui_tests.c
+++ b/tests/vb2_ui_tests.c
@@ -45,6 +45,11 @@ static int mock_key_trusted[64];
 static int mock_key_count;
 static int mock_key_total;
 
+static uint64_t mock_get_timer_last;
+static uint64_t mock_time;
+static const uint64_t mock_time_start = 31ULL * VB_USEC_PER_SEC;
+static int mock_vbexbeep_called;
+
 static enum vb2_dev_default_boot mock_default_boot;
 static int mock_dev_boot_allowed;
 static int mock_dev_boot_legacy_allowed;
@@ -135,8 +140,15 @@ static void displayed_no_extra(void)
 			"  no extra screens");
 }
 
+/* Type of test to reset for */
+enum reset_type {
+	FOR_DEVELOPER,
+	FOR_BROKEN_RECOVERY,
+	FOR_MANUAL_RECOVERY,
+};
+
 /* Reset mock data (for use before each test) */
-static void reset_common_data(void)
+static void reset_common_data(enum reset_type t)
 {
 	TEST_SUCC(vb2api_init(workbuf, sizeof(workbuf), &ctx),
 		  "vb2api_init failed");
@@ -145,6 +157,9 @@ static void reset_common_data(void)
 
 	vb2_nv_init(ctx);
 
+	if (t == FOR_DEVELOPER)
+		ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
+
 	sd = vb2_get_sd(ctx);
 
 	/* For try_recovery_action */
@@ -171,7 +186,10 @@ static void reset_common_data(void)
 	mock_displayed_i = 0;
 
 	/* For shutdown_required */
-	mock_calls_until_shutdown = 10;
+	if (t == FOR_DEVELOPER)
+		mock_calls_until_shutdown = 2000;  /* Larger than 30s */
+	else
+		mock_calls_until_shutdown = 10;
 
 	/* For VbExKeyboardRead */
 	memset(mock_key, 0, sizeof(mock_key));
@@ -181,6 +199,11 @@ static void reset_common_data(void)
 	/* Avoid iteration #0 which has a screen change by global action */
 	add_mock_keypress(0);
 
+	/* For vboot_audio.h */
+	mock_get_timer_last = 0;
+	mock_time = mock_time_start;
+	mock_vbexbeep_called = 0;
+
 	/* For dev_boot* in 2misc.h */
 	mock_default_boot = VB2_DEV_DEFAULT_BOOT_DISK;
 	mock_dev_boot_allowed = 1;
@@ -261,6 +284,23 @@ uint32_t VbExKeyboardReadWithFlags(uint32_t *key_flags)
 	return 0;
 }
 
+uint64_t VbExGetTimer(void)
+{
+	mock_get_timer_last = mock_time;
+	return mock_time;
+}
+
+void VbExSleepMs(uint32_t msec)
+{
+	mock_time += msec * VB_USEC_PER_MSEC;
+}
+
+vb2_error_t VbExBeep(uint32_t msec, uint32_t frequency)
+{
+	mock_vbexbeep_called++;
+	return VB2_SUCCESS;
+}
+
 enum vb2_dev_default_boot vb2_get_dev_boot_target(struct vb2_context *c)
 {
 	return mock_default_boot;
@@ -311,51 +351,45 @@ static void developer_tests(void)
 {
 	VB2_DEBUG("Testing developer mode...\n");
 
-	/* Proceed */
-	reset_common_data();
+	/* Proceed to internal disk after timeout */
+	reset_common_data(FOR_DEVELOPER);
 	add_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_FIXED);
-	TEST_EQ(vb2_developer_menu(ctx), VB2_SUCCESS, "proceed");
-	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");
-
-	/* Proceed to legacy */
-	reset_common_data();
-	mock_default_boot = VB2_DEV_DEFAULT_BOOT_LEGACY;
-	mock_dev_boot_legacy_allowed = 1;
-	TEST_EQ(vb2_developer_menu(ctx), VB2_SUCCESS, "proceed to legacy");
-	TEST_EQ(mock_vbexlegacy_called, 1, "  try legacy");
-	TEST_EQ(mock_altfw_num_last, 0, "  check altfw_num");
-	displayed_no_extra();
-	TEST_EQ(mock_vbtlk_count, mock_vbtlk_total, "  used up mock_vbtlk");
-
-	/* Proceed to legacy only if enabled */
-	reset_common_data();
-	add_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_FIXED);
-	mock_default_boot = VB2_DEV_DEFAULT_BOOT_LEGACY;
 	TEST_EQ(vb2_developer_menu(ctx), VB2_SUCCESS,
-		"default legacy not enabled");
-	TEST_EQ(mock_vbexlegacy_called, 0, "  not legacy");
+		"proceed to internal disk after timeout");
+	displayed_eq("dev mode", VB2_SCREEN_DEVELOPER_MODE, MOCK_IGNORE,
+		     MOCK_IGNORE, MOCK_IGNORE);
 	displayed_no_extra();
+	TEST_TRUE(mock_get_timer_last - mock_time_start >=
+		  30 * VB_USEC_PER_SEC, "  finished delay");
+	TEST_EQ(mock_vbexbeep_called, 2, "  beeped twice");
 	TEST_EQ(mock_vbtlk_count, mock_vbtlk_total, "  used up mock_vbtlk");
 
-	/* Proceed to USB */
-	reset_common_data();
+	/* Proceed to USB after timeout */
+	reset_common_data(FOR_DEVELOPER);
 	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(vb2_developer_menu(ctx), VB2_SUCCESS,
+		"proceed to USB after timeout");
+	displayed_eq("dev mode", VB2_SCREEN_DEVELOPER_MODE, MOCK_IGNORE,
+		     MOCK_IGNORE, MOCK_IGNORE);
 	displayed_no_extra();
+	TEST_TRUE(mock_get_timer_last - mock_time_start >=
+		  30 * VB_USEC_PER_SEC, "  finished delay");
+	TEST_EQ(mock_vbexbeep_called, 2, "  beeped twice");
 	TEST_EQ(mock_vbtlk_count, mock_vbtlk_total, "  used up mock_vbtlk");
 
-	/* Proceed to USB only if enabled */
-	reset_common_data();
-	add_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_FIXED);
+	/* Default boot USB not allowed, don't boot */
+	reset_common_data(FOR_DEVELOPER);
 	mock_default_boot = VB2_DEV_DEFAULT_BOOT_USB;
-	TEST_EQ(vb2_developer_menu(ctx), VB2_SUCCESS,
-		"default USB not enabled");
+	TEST_EQ(vb2_developer_menu(ctx), VB2_REQUEST_SHUTDOWN,
+		"default USB not allowed, don't boot");
+	displayed_eq("dev mode", VB2_SCREEN_DEVELOPER_MODE, MOCK_IGNORE,
+		     MOCK_IGNORE, MOCK_IGNORE);
 	displayed_no_extra();
+	TEST_TRUE(mock_get_timer_last - mock_time_start >=
+		  30 * VB_USEC_PER_SEC, "  finished delay");
+	TEST_EQ(mock_vbexbeep_called, 2, "  beeped twice");
 	TEST_EQ(mock_vbtlk_count, mock_vbtlk_total, "  used up mock_vbtlk");
 
 	VB2_DEBUG("...done.\n");
@@ -367,7 +401,7 @@ static void broken_recovery_tests(void)
 
 	/* BROKEN screen shutdown request */
 	if (!DETACHABLE) {
-		reset_common_data();
+		reset_common_data(FOR_BROKEN_RECOVERY);
 		add_mock_keypress(VB_BUTTON_POWER_SHORT_PRESS);
 		mock_calls_until_shutdown = -1;
 		TEST_EQ(vb2_broken_recovery_menu(ctx),
@@ -379,7 +413,7 @@ static void broken_recovery_tests(void)
 	}
 
 	/* Shortcuts that are always ignored in BROKEN */
-	reset_common_data();
+	reset_common_data(FOR_BROKEN_RECOVERY);
 	add_mock_key(VB_KEY_CTRL('D'), 1);
 	add_mock_key(VB_KEY_CTRL('U'), 1);
 	add_mock_key(VB_KEY_CTRL('L'), 1);
@@ -388,7 +422,7 @@ static void broken_recovery_tests(void)
 	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");
+	TEST_EQ(mock_calls_until_shutdown, 0, "  loop forever");
 	displayed_eq("broken screen", VB2_SCREEN_RECOVERY_BROKEN,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	displayed_no_extra();
@@ -401,7 +435,7 @@ static void manual_recovery_tests(void)
 	VB2_DEBUG("Testing manual recovery mode...\n");
 
 	/* Timeout, shutdown */
-	reset_common_data();
+	reset_common_data(FOR_MANUAL_RECOVERY);
 	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");
@@ -411,7 +445,7 @@ static void manual_recovery_tests(void)
 
 	/* Power button short pressed = shutdown request */
 	if (!DETACHABLE) {
-		reset_common_data();
+		reset_common_data(FOR_MANUAL_RECOVERY);
 		add_mock_keypress(VB_BUTTON_POWER_SHORT_PRESS);
 		add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND,
 			       VB_DISK_FLAG_REMOVABLE);
@@ -424,7 +458,7 @@ static void manual_recovery_tests(void)
 	}
 
 	/* Item 1 = phone recovery */
-	reset_common_data();
+	reset_common_data(FOR_MANUAL_RECOVERY);
 	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,
@@ -436,7 +470,7 @@ static void manual_recovery_tests(void)
 	displayed_no_extra();
 
 	/* Item 2 = external disk recovery */
-	reset_common_data();
+	reset_common_data(FOR_MANUAL_RECOVERY);
 	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);
@@ -451,7 +485,7 @@ static void manual_recovery_tests(void)
 	displayed_no_extra();
 
 	/* Boots if we have a valid image on first try */
-	reset_common_data();
+	reset_common_data(FOR_MANUAL_RECOVERY);
 	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,
@@ -461,7 +495,7 @@ static void manual_recovery_tests(void)
 	displayed_no_extra();
 
 	/* Boots eventually if we get a valid image later */
-	reset_common_data();
+	reset_common_data(FOR_MANUAL_RECOVERY);
 	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);
@@ -473,7 +507,7 @@ static void manual_recovery_tests(void)
 	displayed_no_extra();
 
 	/* Invalid image, then remove, then valid image */
-	reset_common_data();
+	reset_common_data(FOR_MANUAL_RECOVERY);
 	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);
diff --git a/tests/vb2_ui_utility_tests.c b/tests/vb2_ui_utility_tests.c
index 5350e1c6..b616bcb8 100644
--- a/tests/vb2_ui_utility_tests.c
+++ b/tests/vb2_ui_utility_tests.c
@@ -240,7 +240,7 @@ static void change_screen_tests(void)
 	mock_state->screen = &mock_screen_menu;
 	mock_state->selected_item = 2;
 	mock_state->disabled_item_mask = 0x10;
-	TEST_EQ(change_screen(&mock_ui_context, MOCK_SCREEN_BASE),
+	TEST_EQ(vb2_ui_change_screen(&mock_ui_context, MOCK_SCREEN_BASE),
 		VB2_REQUEST_UI_CONTINUE,
 		"change_screen will clear screen state");
 	screen_state_eq(mock_state, MOCK_SCREEN_BASE, 0, 0);
@@ -248,7 +248,7 @@ static void change_screen_tests(void)
 	/* Change to screen which does not exist */
 	reset_common_data();
 	mock_state->screen = &mock_screen_menu;
-	TEST_EQ(change_screen(&mock_ui_context, MOCK_NO_SCREEN),
+	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);