vboot: Replace VB2_REQUEST_UI_CONTINUE with VB2_SUCCESS

In current codebase, VB2_REQUEST_UI_CONTINUE is for successful ui
function call, and VB2_SUCCESS is for breaking the main ui loop.

Replace most of the VB2_REQUEST_UI_CONTINUE with VB2_SUCCESS, and
replace those breaking requesting VB2_SUCCESS with VB2_REQUEST_UI_EXIT.
All the "VB2_REQUEST_UI_CONTINUE" and "VB2_REQUEST_UI_EXIT" are caught
in the ui loop. VB2_REQUEST_UI_CONTINUE does nothing while
VB2_REQUEST_UI_EXIT breaks the ui loop and return VB2_SUCCESS. Returning
VB2_SUCCESS and VB2_REQUEST_UI_CONTINUE to the ui loop now should do the
same thing, that is, continue process the ui.

VB2_REQUEST_UI_CONTINUE can be used to quickly return from ui functions
to the main UI loop, because it is not VB2_SUCCESS and makes VB2_TRY
return.

BRANCH=none
BUG=b:157625765
TEST=make clean && CC=x86_64-pc-linux-gnu-clang make runtests

Signed-off-by: Chung-Sheng Wu <chungsheng@google.com>
Change-Id: If1b54657d09198196f9c646b6b907b4fd8faadce
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2714502
Tested-by: Chung-Sheng Wu <chungsheng@chromium.org>
Reviewed-by: Joel Kitching <kitching@chromium.org>
Commit-Queue: Chung-Sheng Wu <chungsheng@chromium.org>

4 files changed, 89 insertions, 77 deletions

diff --git a/firmware/2lib/2ui.c b/firmware/2lib/2ui.c
index 800f06a3..71734507 100644
--- a/firmware/2lib/2ui.c
+++ b/firmware/2lib/2ui.c
@@ -24,7 +24,7 @@
  * shutdown is required.
  *
  * @param ui		UI context pointer
- * @return VB2_REQUEST_SHUTDOWN if shutdown needed, or VB2_REQUEST_UI_CONTINUE
+ * @return VB2_REQUEST_SHUTDOWN if shutdown needed, or VB2_SUCCESS
  */
 vb2_error_t check_shutdown_request(struct vb2_ui_context *ui)
 {
@@ -64,7 +64,7 @@ vb2_error_t check_shutdown_request(struct vb2_ui_context *ui)
 	if (shutdown_request)
 		return VB2_REQUEST_SHUTDOWN;
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 /*****************************************************************************/
@@ -81,7 +81,7 @@ vb2_error_t error_exit_action(struct vb2_ui_context *ui)
 		ui->error_code = VB2_UI_ERROR_NONE;
 		ui->key = 0;
 	}
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 /*****************************************************************************/
@@ -131,7 +131,7 @@ vb2_error_t menu_navigation_action(struct vb2_ui_context *ui)
 				  ui->key, ui->key_trusted);
 	}
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 vb2_error_t vb2_ui_menu_prev(struct vb2_ui_context *ui)
@@ -139,7 +139,7 @@ vb2_error_t vb2_ui_menu_prev(struct vb2_ui_context *ui)
 	int item;
 
 	if (!DETACHABLE && ui->key == VB_BUTTON_VOL_UP_SHORT_PRESS)
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 
 	item = ui->state->selected_item - 1;
 	while (item >= 0 && VB2_GET_BIT(ui->state->hidden_item_mask, item))
@@ -148,7 +148,7 @@ vb2_error_t vb2_ui_menu_prev(struct vb2_ui_context *ui)
 	if (item >= 0)
 		ui->state->selected_item = item;
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 vb2_error_t vb2_ui_menu_next(struct vb2_ui_context *ui)
@@ -157,7 +157,7 @@ vb2_error_t vb2_ui_menu_next(struct vb2_ui_context *ui)
 	const struct vb2_menu *menu;
 
 	if (!DETACHABLE && ui->key == VB_BUTTON_VOL_DOWN_SHORT_PRESS)
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 
 	menu = get_menu(ui);
 	item = ui->state->selected_item + 1;
@@ -168,7 +168,7 @@ vb2_error_t vb2_ui_menu_next(struct vb2_ui_context *ui)
 	if (item < menu->num_items)
 		ui->state->selected_item = item;
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 vb2_error_t vb2_ui_menu_select(struct vb2_ui_context *ui)
@@ -177,11 +177,11 @@ vb2_error_t vb2_ui_menu_select(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;
+		return VB2_SUCCESS;
 
 	menu = get_menu(ui);
 	if (menu->num_items == 0)
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 
 	menu_item = &menu->items[ui->state->selected_item];
 
@@ -190,7 +190,7 @@ vb2_error_t vb2_ui_menu_select(struct vb2_ui_context *ui)
 			ui->state->selected_item)) {
 		VB2_DEBUG("Menu item <%s> disabled; ignoring\n",
 			  menu_item->text);
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 	}
 
 	if (menu_item->action) {
@@ -204,7 +204,7 @@ vb2_error_t vb2_ui_menu_select(struct vb2_ui_context *ui)
 
 	VB2_DEBUG("Menu item <%s> no action or target screen\n",
 		  menu_item->text);
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 /*****************************************************************************/
@@ -219,7 +219,7 @@ vb2_error_t vb2_ui_screen_back(struct vb2_ui_context *ui)
 		free(ui->state);
 		ui->state = tmp;
 		if (ui->state->screen->reinit)
-			return ui->state->screen->reinit(ui);
+			VB2_TRY(ui->state->screen->reinit(ui));
 	} else {
 		VB2_DEBUG("ERROR: No previous screen; ignoring\n");
 	}
@@ -233,7 +233,7 @@ static vb2_error_t default_screen_init(struct vb2_ui_context *ui)
 	ui->state->selected_item = 0;
 	if (menu->num_items > 1 && menu->items[0].is_language_select)
 		ui->state->selected_item = 1;
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 vb2_error_t vb2_ui_screen_change(struct vb2_ui_context *ui, enum vb2_screen id)
@@ -266,7 +266,7 @@ vb2_error_t vb2_ui_screen_change(struct vb2_ui_context *ui, enum vb2_screen id)
 			free(cur_state);
 		}
 		if (ui->state->screen->reinit)
-			return ui->state->screen->reinit(ui);
+			VB2_TRY(ui->state->screen->reinit(ui));
 	} else {
 		/* Allocate the requested screen on top of the stack. */
 		cur_state = malloc(sizeof(*ui->state));
@@ -279,9 +279,9 @@ vb2_error_t vb2_ui_screen_change(struct vb2_ui_context *ui, enum vb2_screen id)
 		cur_state->screen = new_screen_info;
 		ui->state = cur_state;
 		if (ui->state->screen->init)
-			return ui->state->screen->init(ui);
+			VB2_TRY(ui->state->screen->init(ui));
 		else
-			return default_screen_init(ui);
+			VB2_TRY(default_screen_init(ui));
 	}
 
 	return VB2_REQUEST_UI_CONTINUE;
@@ -290,8 +290,9 @@ vb2_error_t vb2_ui_screen_change(struct vb2_ui_context *ui, enum vb2_screen id)
 /*****************************************************************************/
 /* Core UI loop */
 
-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))
+static vb2_error_t ui_loop_impl(
+	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;
@@ -309,9 +310,11 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 	if (root_info == NULL)
 		VB2_DIE("Root screen not found.\n");
 	ui.locale_id = vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX);
+
 	rv = vb2_ui_screen_change(&ui, root_screen_id);
-	if (rv != VB2_REQUEST_UI_CONTINUE)
+	if (rv && rv != VB2_REQUEST_UI_CONTINUE)
 		return rv;
+
 	memset(&prev_state, 0, sizeof(prev_state));
 	prev_disable_timer = 0;
 	prev_error_code = VB2_UI_ERROR_NONE;
@@ -365,32 +368,32 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 
 		/* Check for shutdown request. */
 		rv = check_shutdown_request(&ui);
-		if (rv != VB2_REQUEST_UI_CONTINUE) {
+		if (rv && rv != VB2_REQUEST_UI_CONTINUE) {
 			VB2_DEBUG("Shutdown requested!\n");
 			return rv;
 		}
 
 		/* Check if we need to exit an error box. */
 		rv = error_exit_action(&ui);
-		if (rv != VB2_REQUEST_UI_CONTINUE)
+		if (rv && rv != VB2_REQUEST_UI_CONTINUE)
 			return rv;
 
 		/* Run screen action. */
 		if (ui.state->screen->action) {
 			rv = ui.state->screen->action(&ui);
-			if (rv != VB2_REQUEST_UI_CONTINUE)
+			if (rv && rv != VB2_REQUEST_UI_CONTINUE)
 				return rv;
 		}
 
 		/* Run menu navigation action. */
 		rv = menu_navigation_action(&ui);
-		if (rv != VB2_REQUEST_UI_CONTINUE)
+		if (rv && rv != VB2_REQUEST_UI_CONTINUE)
 			return rv;
 
 		/* Run global action function if available. */
 		if (global_action) {
 			rv = global_action(&ui);
-			if (rv != VB2_REQUEST_UI_CONTINUE)
+			if (rv && rv != VB2_REQUEST_UI_CONTINUE)
 				return rv;
 		}
 
@@ -403,6 +406,15 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 	return VB2_SUCCESS;
 }
 
+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 rv = ui_loop_impl(ctx, root_screen_id, global_action);
+	if (rv == VB2_REQUEST_UI_EXIT)
+		return VB2_SUCCESS;
+	return rv;
+}
+
 /*****************************************************************************/
 /* Developer mode */
 
@@ -437,7 +449,7 @@ vb2_error_t developer_action(struct vb2_ui_context *ui)
 	if (ui->key == VB_KEY_CTRL('L'))  /* L for aLtfw (formerly Legacy) */
 		return vb2_ui_developer_mode_boot_altfw_action(ui);
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 /*****************************************************************************/
@@ -454,7 +466,7 @@ vb2_error_t broken_recovery_action(struct vb2_ui_context *ui)
 	if (ui->key == '\t')
 		return vb2_ui_screen_change(ui, VB2_SCREEN_DEBUG_INFO);
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 /*****************************************************************************/
@@ -470,7 +482,7 @@ vb2_error_t manual_recovery_action(struct vb2_ui_context *ui)
 	/* See if we have a recovery kernel available yet. */
 	vb2_error_t rv = VbTryLoadKernel(ui->ctx, VB_DISK_FLAG_REMOVABLE);
 	if (rv == VB2_SUCCESS)
-		return rv;
+		return VB2_REQUEST_UI_EXIT;
 
 	/* If disk validity state changed, switch to appropriate screen. */
 	if (ui->recovery_rv != rv) {
@@ -491,7 +503,7 @@ vb2_error_t manual_recovery_action(struct vb2_ui_context *ui)
 	if (ui->key == '\t')
 		return vb2_ui_screen_change(ui, VB2_SCREEN_DEBUG_INFO);
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 /*****************************************************************************/
diff --git a/firmware/2lib/2ui_screens.c b/firmware/2lib/2ui_screens.c
index a8435304..fd1d5d2b 100644
--- a/firmware/2lib/2ui_screens.c
+++ b/firmware/2lib/2ui_screens.c
@@ -80,7 +80,7 @@ static vb2_error_t log_page_init(struct vb2_ui_context *ui)
 		ui->state->selected_item = screen->page_down_item;
 	}
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 static vb2_error_t log_page_prev_action(struct vb2_ui_context *ui)
@@ -89,7 +89,7 @@ static vb2_error_t log_page_prev_action(struct vb2_ui_context *ui)
 
 	/* Validity check. */
 	if (ui->state->current_page == 0)
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 
 	ui->state->current_page--;
 
@@ -103,7 +103,7 @@ static vb2_error_t log_page_prev_action(struct vb2_ui_context *ui)
 		VB2_SET_BIT(ui->state->disabled_item_mask,
 			    screen->page_up_item);
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 static vb2_error_t log_page_next_action(struct vb2_ui_context *ui)
@@ -112,7 +112,7 @@ static vb2_error_t log_page_next_action(struct vb2_ui_context *ui)
 
 	/* Validity check. */
 	if (ui->state->current_page == ui->state->page_count - 1)
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 
 	ui->state->current_page++;
 
@@ -126,7 +126,7 @@ static vb2_error_t log_page_next_action(struct vb2_ui_context *ui)
 		VB2_SET_BIT(ui->state->disabled_item_mask,
 			    screen->page_down_item);
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 #define PAGE_UP_ITEM ((struct vb2_menu_item){ \
@@ -214,7 +214,7 @@ static vb2_error_t language_select_init(struct vb2_ui_context *ui)
 			  "initializing selected_item to 0\n");
 		ui->state->selected_item = 0;
 	}
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 static const struct vb2_screen_info language_select_screen = {
@@ -255,7 +255,7 @@ vb2_error_t advanced_options_init(struct vb2_ui_context *ui)
 		ui->state->selected_item = ADVANCED_OPTIONS_ITEM_DEBUG_INFO;
 	}
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 static const struct vb2_menu_item advanced_options_items[] = {
@@ -325,7 +325,7 @@ static vb2_error_t debug_info_reinit(struct vb2_ui_context *ui)
 		return vb2_ui_screen_back(ui);
 	}
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 static const struct vb2_menu_item debug_info_items[] = {
@@ -389,7 +389,7 @@ static vb2_error_t firmware_log_reinit(struct vb2_ui_context *ui)
 		return vb2_ui_screen_back(ui);
 	}
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 static const struct vb2_menu_item firmware_log_items[] = {
@@ -440,7 +440,7 @@ vb2_error_t recovery_select_init(struct vb2_ui_context *ui)
 		VB2_SET_BIT(ui->state->hidden_item_mask,
 			    RECOVERY_SELECT_ITEM_DIAGNOSTICS);
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 static const struct vb2_menu_item recovery_select_items[] = {
@@ -512,7 +512,7 @@ vb2_error_t recovery_to_dev_init(struct vb2_ui_context *ui)
 
 	ui->physical_presence_button_pressed = 0;
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 static vb2_error_t recovery_to_dev_finalize(struct vb2_ui_context *ui)
@@ -524,7 +524,7 @@ static vb2_error_t recovery_to_dev_finalize(struct vb2_ui_context *ui)
 	    (vb2_get_sd(ui->ctx)->flags & VB2_SD_FLAG_DEV_MODE_ENABLED) ||
 	    !vb2_allow_recovery(ui->ctx)) {
 		VB2_DEBUG("ERROR: Dev transition validity check failed\n");
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 	}
 
 	VB2_DEBUG("Enabling dev mode and rebooting...\n");
@@ -544,7 +544,7 @@ vb2_error_t recovery_to_dev_confirm_action(struct vb2_ui_context *ui)
 		 */
 		if (PHYSICAL_PRESENCE_KEYBOARD && ui->key == VB_KEY_ENTER)
 			ui->error_code = VB2_UI_ERROR_UNTRUSTED_CONFIRMATION;
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 	}
 	return recovery_to_dev_finalize(ui);
 }
@@ -560,17 +560,17 @@ vb2_error_t recovery_to_dev_action(struct vb2_ui_context *ui)
 
 	/* Keyboard physical presence case covered by "Confirm" action. */
 	if (PHYSICAL_PRESENCE_KEYBOARD)
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 
 	pressed = vb2ex_physical_presence_pressed();
 	if (pressed) {
 		VB2_DEBUG("Physical presence button pressed, "
 			  "awaiting release\n");
 		ui->physical_presence_button_pressed = 1;
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 	}
 	if (!ui->physical_presence_button_pressed)
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 	VB2_DEBUG("Physical presence button released\n");
 
 	return recovery_to_dev_finalize(ui);
@@ -723,7 +723,7 @@ vb2_error_t developer_mode_init(struct vb2_ui_context *ui)
 
 	ui->start_time_ms = vb2ex_mtime();
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 vb2_error_t vb2_ui_developer_mode_boot_internal_action(
@@ -732,11 +732,11 @@ vb2_error_t vb2_ui_developer_mode_boot_internal_action(
 	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;
+		return VB2_SUCCESS;
 	}
 
 	VB2_TRY(VbTryLoadKernel(ui->ctx, VB_DISK_FLAG_FIXED));
-	return VB2_SUCCESS;
+	return VB2_REQUEST_UI_EXIT;
 }
 
 vb2_error_t vb2_ui_developer_mode_boot_external_action(
@@ -751,12 +751,12 @@ vb2_error_t vb2_ui_developer_mode_boot_external_action(
 		VB2_DEBUG("ERROR: Dev mode external boot not allowed\n");
 		ui->error_beep = 1;
 		ui->error_code = VB2_UI_ERROR_EXTERNAL_BOOT_NOT_ENABLED;
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 	}
 
 	rv = VbTryLoadKernel(ui->ctx, VB_DISK_FLAG_REMOVABLE);
 	if (rv == VB2_SUCCESS) {
-		return VB2_SUCCESS;
+		return VB2_REQUEST_UI_EXIT;
 	} else if (rv == VB2_ERROR_LK_NO_DISK_FOUND) {
 		if (ui->state->screen->id !=
 		    VB2_SCREEN_DEVELOPER_BOOT_EXTERNAL) {
@@ -789,7 +789,7 @@ vb2_error_t developer_mode_action(struct vb2_ui_context *ui)
 	if (ui->key)
 		ui->disable_timer = 1;
 	if (ui->disable_timer)
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 
 	elapsed_ms = vb2ex_mtime() - ui->start_time_ms;
 
@@ -814,7 +814,7 @@ vb2_error_t developer_mode_action(struct vb2_ui_context *ui)
 		return vb2_ui_menu_select(ui);
 	}
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 static const struct vb2_menu_item developer_mode_items[] = {
@@ -866,14 +866,14 @@ static vb2_error_t developer_to_norm_init(struct vb2_ui_context *ui)
 	if (!vb2_dev_boot_allowed(ui->ctx))
 		VB2_SET_BIT(ui->state->hidden_item_mask,
 			    DEVELOPER_TO_NORM_ITEM_CANCEL);
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 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;
+		return VB2_SUCCESS;
 	}
 
 	VB2_DEBUG("Leaving dev mode\n");
@@ -954,7 +954,7 @@ static vb2_error_t developer_select_bootloader_init(struct vb2_ui_context *ui)
 	/* Select the first bootloader. */
 	ui->state->selected_item =
 		ARRAY_SIZE(developer_select_bootloader_items_before);
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 vb2_error_t vb2_ui_developer_mode_boot_altfw_action(
@@ -969,13 +969,13 @@ vb2_error_t vb2_ui_developer_mode_boot_altfw_action(
 	    !vb2_dev_boot_altfw_allowed(ui->ctx)) {
 		VB2_DEBUG("ERROR: Dev mode alternate bootloader not allowed\n");
 		ui->error_code = VB2_UI_ERROR_ALTFW_DISABLED;
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 	}
 
 	if (vb2ex_get_altfw_count() == 0) {
 		VB2_DEBUG("ERROR: No alternate bootloader was found\n");
 		ui->error_code = VB2_UI_ERROR_ALTFW_EMPTY;
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 	}
 
 	if (ui->key == VB_KEY_CTRL('L')) {
@@ -991,7 +991,7 @@ vb2_error_t vb2_ui_developer_mode_boot_altfw_action(
 
 	VB2_DEBUG("ERROR: Alternate bootloader failed\n");
 	ui->error_code = VB2_UI_ERROR_ALTFW_FAILED;
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 static const struct vb2_menu *get_bootloader_menu(struct vb2_ui_context *ui)
@@ -1138,13 +1138,13 @@ static vb2_error_t diagnostics_memory_update_screen(struct vb2_ui_context *ui,
 
 	/* Early return if the memory test is done. */
 	if (ui->state->test_finished)
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 
 	vb2_error_t rv = op(reset, &log_string);
 
 	/* The test is still running but the output buffer was unchanged. */
 	if (rv == VB2_ERROR_EX_DIAG_TEST_RUNNING)
-		return VB2_REQUEST_UI_CONTINUE;
+		return VB2_SUCCESS;
 
 	if ((rv && rv != VB2_ERROR_EX_DIAG_TEST_UPDATED) || !log_string) {
 		VB2_DEBUG("ERROR: Failed to retrieve memory test status\n");
@@ -1185,28 +1185,20 @@ static vb2_error_t diagnostics_memory_update_screen(struct vb2_ui_context *ui,
 		ui->state->test_finished = 1;
 	}
 
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 static vb2_error_t diagnostics_memory_init_quick(struct vb2_ui_context *ui)
 {
-	vb2_error_t rv;
-	rv = diagnostics_memory_update_screen(
-		ui, &vb2ex_diag_memory_quick_test, 1);
-
-	if (rv != VB2_REQUEST_UI_CONTINUE)
-		return rv;
+	VB2_TRY(diagnostics_memory_update_screen(
+		ui, &vb2ex_diag_memory_quick_test, 1));
 	return log_page_init(ui);
 }
 
 static vb2_error_t diagnostics_memory_init_full(struct vb2_ui_context *ui)
 {
-	vb2_error_t rv;
-	rv = diagnostics_memory_update_screen(
-		ui, &vb2ex_diag_memory_full_test, 1);
-
-	if (rv != VB2_REQUEST_UI_CONTINUE)
-		return rv;
+	VB2_TRY(diagnostics_memory_update_screen(
+		ui, &vb2ex_diag_memory_full_test, 1));
 	return log_page_init(ui);
 }
 
diff --git a/firmware/2lib/include/2return_codes.h b/firmware/2lib/include/2return_codes.h
index 49f99358..ed940531 100644
--- a/firmware/2lib/include/2return_codes.h
+++ b/firmware/2lib/include/2return_codes.h
@@ -44,6 +44,9 @@ enum vb2_return_code {
 	/* Continue in the UI loop.  This is used in UI internal functions. */
 	VB2_REQUEST_UI_CONTINUE = 0x1005,
 
+	/* Break from the UI loop.  This is used in UI internal functions. */
+	VB2_REQUEST_UI_EXIT = 0x1006,
+
 	/* End of VB2_REQUEST_* */
 	VB2_REQUEST_END = 0x5000,
 
diff --git a/firmware/2lib/include/2ui.h b/firmware/2lib/include/2ui.h
index c8e2928b..7e815418 100644
--- a/firmware/2lib/include/2ui.h
+++ b/firmware/2lib/include/2ui.h
@@ -155,7 +155,7 @@ const struct vb2_screen_info *vb2_get_screen_info(enum vb2_screen id);
  * on 0 when we hit the start of the menu.
  *
  * @param ui		UI context pointer
- * @return VB2_REQUEST_UI_CONTINUE, or error code on error.
+ * @return VB2_SUCCESS, or error code on error.
  */
 vb2_error_t vb2_ui_menu_prev(struct vb2_ui_context *ui);
 
@@ -167,7 +167,7 @@ vb2_error_t vb2_ui_menu_prev(struct vb2_ui_context *ui);
  * on the max index when we hit the end of the menu.
  *
  * @param ui		UI context pointer
- * @return VB2_REQUEST_UI_CONTINUE, or error code on error.
+ * @return VB2_SUCCESS, or error code on error.
  */
 vb2_error_t vb2_ui_menu_next(struct vb2_ui_context *ui);
 
@@ -182,13 +182,18 @@ vb2_error_t vb2_ui_menu_next(struct vb2_ui_context *ui);
  * selecting the menu item is a no-op.
  *
  * @param ui		UI context pointer
- * @return VB2_REQUEST_UI_CONTINUE, or error code on error.
+ * @return VB2_SUCCESS, or error code on error.
  */
 vb2_error_t vb2_ui_menu_select(struct vb2_ui_context *ui);
 
 /*****************************************************************************/
 /* Screen navigation functions */
-
+/**
+ * After these functions are called, no assumptions may be made about which
+ * screen is currently displayed, and thus execution should return to ui_loop.
+ * VB2_REQUEST_UI_CONTINUE is returned rather than VB2_SUCCESS, so VB2_TRY can
+ * be used to wrapped to these functions and the callers of these functions.
+ */
 /**
  * Return back to the previous screen.
  *