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>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2848044
Reviewed-by: Nick Vaccaro <nvaccaro@google.com>
Tested-by: Nick Vaccaro <nvaccaro@google.com>

6 files changed, 150 insertions, 152 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 b4907d60..8579c8f0 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 5d5f73d7..c67415c4 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.
  *
diff --git a/tests/vb2_ui_action_tests.c b/tests/vb2_ui_action_tests.c
index dc4ebfa8..8e27548e 100644
--- a/tests/vb2_ui_action_tests.c
+++ b/tests/vb2_ui_action_tests.c
@@ -80,8 +80,8 @@ static uint32_t mock_action_countdown_limit;
 static vb2_error_t mock_action_countdown(struct vb2_ui_context *ui)
 {
 	if (++mock_action_called >= mock_action_countdown_limit)
-		return VB2_SUCCESS;
-	return VB2_REQUEST_UI_CONTINUE;
+		return VB2_REQUEST_UI_EXIT;
+	return VB2_SUCCESS;
 }
 
 static vb2_error_t mock_action_screen_change(struct vb2_ui_context *ui)
@@ -99,29 +99,29 @@ static int mock_action_flags;
 static vb2_error_t mock_action_flag0(struct vb2_ui_context *ui)
 {
 	if ((1 << 0) & mock_action_flags)
-		return VB2_SUCCESS;
-	return VB2_REQUEST_UI_CONTINUE;
+		return VB2_REQUEST_UI_EXIT;
+	return VB2_SUCCESS;
 }
 
 static vb2_error_t mock_action_flag1(struct vb2_ui_context *ui)
 {
 	if ((1 << 1) & mock_action_flags)
-		return VB2_SUCCESS;
-	return VB2_REQUEST_UI_CONTINUE;
+		return VB2_REQUEST_UI_EXIT;
+	return VB2_SUCCESS;
 }
 
 static vb2_error_t mock_action_flag2(struct vb2_ui_context *ui)
 {
 	if ((1 << 2) & mock_action_flags)
-		return VB2_SUCCESS;
-	return VB2_REQUEST_UI_CONTINUE;
+		return VB2_REQUEST_UI_EXIT;
+	return VB2_SUCCESS;
 }
 
 static uint32_t mock_action_delay_ms;
 static vb2_error_t mock_action_msleep(struct vb2_ui_context *ui)
 {
 	vb2ex_msleep(mock_action_delay_ms);
-	return VB2_REQUEST_UI_CONTINUE;
+	return VB2_SUCCESS;
 }
 
 /* Mock screens */
@@ -533,7 +533,7 @@ static void menu_prev_tests(void)
 	mock_ui_context.state->screen = &mock_screen_menu;
 	mock_ui_context.state->selected_item = 2;
 	mock_ui_context.key = VB_KEY_UP;
-	TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+	TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_SUCCESS,
 		"valid action");
 	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 1,
 			MOCK_IGNORE);
@@ -544,7 +544,7 @@ static void menu_prev_tests(void)
 	mock_ui_context.state->selected_item = 2;
 	mock_ui_context.state->hidden_item_mask = 0x0a;  /* 0b01010 */
 	mock_ui_context.key = VB_KEY_UP;
-	TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+	TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_SUCCESS,
 		"valid action with hidden mask");
 	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 0,
 			MOCK_IGNORE);
@@ -555,7 +555,7 @@ static void menu_prev_tests(void)
 	mock_ui_context.state->selected_item = 2;
 	mock_ui_context.state->disabled_item_mask = 0x0a;  /* 0b01010 */
 	mock_ui_context.key = VB_KEY_UP;
-	TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+	TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_SUCCESS,
 		"valid action with disabled mask");
 	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 1,
 			MOCK_IGNORE);
@@ -565,7 +565,7 @@ static void menu_prev_tests(void)
 	mock_ui_context.state->screen = &mock_screen_menu;
 	mock_ui_context.state->selected_item = 0;
 	mock_ui_context.key = VB_KEY_UP;
-	TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+	TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_SUCCESS,
 		"invalid action (blocked)");
 	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 0,
 			MOCK_IGNORE);
@@ -576,7 +576,7 @@ static void menu_prev_tests(void)
 	mock_ui_context.state->selected_item = 2;
 	mock_ui_context.state->hidden_item_mask = 0x0b;  /* 0b01011 */
 	mock_ui_context.key = VB_KEY_UP;
-	TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+	TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_SUCCESS,
 		"invalid action (blocked by mask)");
 	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 2,
 			MOCK_IGNORE);
@@ -587,8 +587,7 @@ static void menu_prev_tests(void)
 		mock_ui_context.state->screen = &mock_screen_menu;
 		mock_ui_context.state->selected_item = 2;
 		mock_ui_context.key = VB_BUTTON_VOL_UP_SHORT_PRESS;
-		TEST_EQ(vb2_ui_menu_prev(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE,
+		TEST_EQ(vb2_ui_menu_prev(&mock_ui_context), VB2_SUCCESS,
 			"ignore volume-up when not DETACHABLE");
 		screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 2,
 				MOCK_IGNORE);
@@ -606,7 +605,7 @@ static void menu_next_tests(void)
 	mock_ui_context.state->screen = &mock_screen_menu;
 	mock_ui_context.state->selected_item = 2;
 	mock_ui_context.key = VB_KEY_DOWN;
-	TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+	TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_SUCCESS,
 		"valid action");
 	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 3,
 			MOCK_IGNORE);
@@ -617,7 +616,7 @@ static void menu_next_tests(void)
 	mock_ui_context.state->selected_item = 2;
 	mock_ui_context.state->hidden_item_mask = 0x0a;  /* 0b01010 */
 	mock_ui_context.key = VB_KEY_DOWN;
-	TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+	TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_SUCCESS,
 		"valid action with hidden mask");
 	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 4,
 			MOCK_IGNORE);
@@ -628,7 +627,7 @@ static void menu_next_tests(void)
 	mock_ui_context.state->selected_item = 2;
 	mock_ui_context.state->disabled_item_mask = 0x0a;  /* 0b01010 */
 	mock_ui_context.key = VB_KEY_DOWN;
-	TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+	TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_SUCCESS,
 		"valid action with disabled mask");
 	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 3,
 			MOCK_IGNORE);
@@ -638,7 +637,7 @@ static void menu_next_tests(void)
 	mock_ui_context.state->screen = &mock_screen_menu;
 	mock_ui_context.state->selected_item = 4;
 	mock_ui_context.key = VB_KEY_DOWN;
-	TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+	TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_SUCCESS,
 		"invalid action (blocked)");
 	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 4,
 			MOCK_IGNORE);
@@ -649,7 +648,7 @@ static void menu_next_tests(void)
 	mock_ui_context.state->selected_item = 2;
 	mock_ui_context.state->hidden_item_mask = 0x1a;  /* 0b11010 */
 	mock_ui_context.key = VB_KEY_DOWN;
-	TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+	TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_SUCCESS,
 		"invalid action (blocked by mask)");
 	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 2,
 			MOCK_IGNORE);
@@ -660,8 +659,7 @@ static void menu_next_tests(void)
 		mock_ui_context.state->screen = &mock_screen_menu;
 		mock_ui_context.state->selected_item = 2;
 		mock_ui_context.key = VB_BUTTON_VOL_DOWN_SHORT_PRESS;
-		TEST_EQ(vb2_ui_menu_next(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE,
+		TEST_EQ(vb2_ui_menu_next(&mock_ui_context), VB2_SUCCESS,
 			"ignore volume-down when not DETACHABLE");
 		screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 2,
 				MOCK_IGNORE);
@@ -684,8 +682,7 @@ static void menu_select_tests(void)
 	reset_common_data();
 	mock_ui_context.state->screen = &mock_screen_base;
 	mock_ui_context.key = VB_KEY_ENTER;
-	TEST_EQ(vb2_ui_menu_select(&mock_ui_context),
-		VB2_REQUEST_UI_CONTINUE,
+	TEST_EQ(vb2_ui_menu_select(&mock_ui_context), VB2_SUCCESS,
 		"vb2_ui_menu_select with no item screen");
 	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_BASE, 0,
 			MOCK_IGNORE);
@@ -702,8 +699,8 @@ static void menu_select_tests(void)
 	mock_ui_context.state->screen = &mock_screen_menu;
 	mock_ui_context.state->selected_item = 2;
 	mock_ui_context.key = VB_KEY_ENTER;
-	TEST_EQ(vb2_ui_menu_select(&mock_ui_context),
-		VB2_REQUEST_UI_CONTINUE, "select an item with a target");
+	TEST_EQ(vb2_ui_menu_select(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
+		"select an item with a target");
 	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_TARGET2, 0,
 			MOCK_IGNORE);
 
@@ -721,8 +718,7 @@ static void menu_select_tests(void)
 	mock_ui_context.state->screen = &mock_screen_menu;
 	mock_ui_context.state->selected_item = 4;
 	mock_ui_context.key = VB_KEY_ENTER;
-	TEST_EQ(vb2_ui_menu_select(&mock_ui_context),
-		VB2_REQUEST_UI_CONTINUE,
+	TEST_EQ(vb2_ui_menu_select(&mock_ui_context), VB2_SUCCESS,
 		"select an item with neither targets nor actions");
 	screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 4,
 			MOCK_IGNORE);
@@ -733,8 +729,8 @@ static void menu_select_tests(void)
 	mock_ui_context.state->selected_item = 3;
 	mock_ui_context.state->disabled_item_mask = 0x08;  /* 0b01000 */
 	mock_ui_context.key = VB_KEY_ENTER;
-	TEST_EQ(vb2_ui_menu_select(&mock_ui_context),
-		VB2_REQUEST_UI_CONTINUE, "cannot select a disabled item");
+	TEST_EQ(vb2_ui_menu_select(&mock_ui_context), VB2_SUCCESS,
+		"cannot select a disabled item");
 	TEST_EQ(mock_action_called, 0, "  no action called");
 
 	/* Ignore power button short press when not DETACHABLE */
@@ -743,8 +739,7 @@ static void menu_select_tests(void)
 		mock_ui_context.state->screen = &mock_screen_menu;
 		mock_ui_context.state->selected_item = 1;
 		mock_ui_context.key = VB_BUTTON_POWER_SHORT_PRESS;
-		TEST_EQ(vb2_ui_menu_select(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE,
+		TEST_EQ(vb2_ui_menu_select(&mock_ui_context), VB2_SUCCESS,
 			"ignore power button short press when not DETACHABLE");
 		screen_state_eq(mock_ui_context.state, MOCK_SCREEN_MENU, 1,
 				MOCK_IGNORE);
@@ -761,9 +756,8 @@ static void vb2_ui_developer_mode_boot_altfw_action_tests(void)
 	reset_common_data();
 	mock_dev_boot_altfw_allowed = 1;
 	TEST_EQ(vb2_ui_developer_mode_boot_altfw_action(&mock_ui_context),
-		VB2_REQUEST_UI_CONTINUE, "not allowed: not in dev mode");
-	TEST_EQ(mock_run_altfw_called, 0,
-		"  vb2ex_run_altfw not called");
+		VB2_SUCCESS, "not allowed: not in dev mode");
+	TEST_EQ(mock_run_altfw_called, 0, "  vb2ex_run_altfw not called");
 
 	/* Not allowed: dev boot not allowed */
 	reset_common_data();
@@ -771,18 +765,15 @@ static void vb2_ui_developer_mode_boot_altfw_action_tests(void)
 	mock_dev_boot_allowed = 0;
 	mock_dev_boot_altfw_allowed = 1;
 	TEST_EQ(vb2_ui_developer_mode_boot_altfw_action(&mock_ui_context),
-		VB2_REQUEST_UI_CONTINUE, "not allowed: dev boot not allowed");
-	TEST_EQ(mock_run_altfw_called, 0,
-		"  vb2ex_run_altfw not called");
+		VB2_SUCCESS, "not allowed: dev boot not allowed");
+	TEST_EQ(mock_run_altfw_called, 0, "  vb2ex_run_altfw not called");
 
 	/* Not allowed: boot altfw not allowed */
 	reset_common_data();
 	ctx->flags |= VB2_CONTEXT_DEVELOPER_MODE;
 	TEST_EQ(vb2_ui_developer_mode_boot_altfw_action(&mock_ui_context),
-		VB2_REQUEST_UI_CONTINUE,
-		"not allowed: boot altfw not allowed");
-	TEST_EQ(mock_run_altfw_called, 0,
-		"  vb2ex_run_altfw not called");
+		VB2_SUCCESS, "not allowed: boot altfw not allowed");
+	TEST_EQ(mock_run_altfw_called, 0, "  vb2ex_run_altfw not called");
 
 	/* Allowed */
 	reset_common_data();
@@ -790,9 +781,8 @@ static void vb2_ui_developer_mode_boot_altfw_action_tests(void)
 	mock_dev_boot_altfw_allowed = 1;
 	mock_ui_context.state->selected_item = 2;
 	TEST_EQ(vb2_ui_developer_mode_boot_altfw_action(&mock_ui_context),
-		VB2_REQUEST_UI_CONTINUE, "allowed");
-	TEST_EQ(mock_run_altfw_called, 1,
-		"  vb2ex_run_altfw called once");
+		VB2_SUCCESS, "allowed");
+	TEST_EQ(mock_run_altfw_called, 1, "  vb2ex_run_altfw called once");
 	TEST_EQ(mock_altfw_last, 2, "  select bootloader #2");
 
 	/* CTRL+L = default bootloader */
@@ -802,9 +792,8 @@ static void vb2_ui_developer_mode_boot_altfw_action_tests(void)
 	mock_ui_context.key = VB_KEY_CTRL('L');
 	mock_ui_context.state->selected_item = 4;  /* Ignored */
 	TEST_EQ(vb2_ui_developer_mode_boot_altfw_action(&mock_ui_context),
-		VB2_REQUEST_UI_CONTINUE, "allowed: ctrl+l");
-	TEST_EQ(mock_run_altfw_called, 1,
-		"  vb2ex_run_altfw called once");
+		VB2_SUCCESS, "allowed: ctrl+l");
+	TEST_EQ(mock_run_altfw_called, 1, "  vb2ex_run_altfw called once");
 	TEST_EQ(mock_altfw_last, 0, "  select bootloader #0");
 
 	VB2_DEBUG("...done.\n");
@@ -817,8 +806,8 @@ static void manual_recovery_action_tests(void)
 	/* SUCCESS */
 	reset_common_data();
 	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
-	TEST_EQ(manual_recovery_action(&mock_ui_context),
-		VB2_SUCCESS, "SUCCESS");
+	TEST_EQ(manual_recovery_action(&mock_ui_context), VB2_REQUEST_UI_EXIT,
+		"EXIT");
 	TEST_EQ(mock_get_screen_info_called, 0, "  no change_screen");
 
 	/* NO_DISK_FOUND */
@@ -839,8 +828,8 @@ static void manual_recovery_action_tests(void)
 	TEST_EQ(manual_recovery_action(&mock_ui_context),
 		VB2_REQUEST_UI_CONTINUE, "INVALID_KERNEL");
 	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
-	TEST_EQ(manual_recovery_action(&mock_ui_context),
-		VB2_SUCCESS, "SUCCESS");
+	TEST_EQ(manual_recovery_action(&mock_ui_context), VB2_REQUEST_UI_EXIT,
+		"EXIT");
 	screen_state_eq(mock_ui_context.state, VB2_SCREEN_RECOVERY_INVALID,
 			MOCK_IGNORE, MOCK_IGNORE);
 
@@ -863,8 +852,8 @@ static void manual_recovery_action_tests(void)
 	TEST_EQ(manual_recovery_action(&mock_ui_context),
 		VB2_REQUEST_UI_CONTINUE, "NO_DISK_FOUND");
 	set_mock_vbtlk(VB2_SUCCESS, VB_DISK_FLAG_REMOVABLE);
-	TEST_EQ(manual_recovery_action(&mock_ui_context),
-		VB2_SUCCESS, "SUCCESS");
+	TEST_EQ(manual_recovery_action(&mock_ui_context), VB2_REQUEST_UI_EXIT,
+		"EXIT");
 	screen_state_eq(mock_ui_context.state, VB2_SCREEN_RECOVERY_SELECT,
 			MOCK_IGNORE, MOCK_IGNORE);
 
diff --git a/tests/vb2_ui_utility_tests.c b/tests/vb2_ui_utility_tests.c
index fd2eda8c..85a514cc 100644
--- a/tests/vb2_ui_utility_tests.c
+++ b/tests/vb2_ui_utility_tests.c
@@ -179,14 +179,13 @@ static void check_shutdown_request_tests(void)
 	if (!DETACHABLE) {
 		reset_common_data();
 		mock_shutdown_request = 0;
-		TEST_EQ(check_shutdown_request(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE,
+		TEST_EQ(check_shutdown_request(&mock_ui_context), VB2_SUCCESS,
 			"release, press, hold, and release");
 		mock_shutdown_request = VB_SHUTDOWN_REQUEST_POWER_BUTTON;
-		TEST_EQ(check_shutdown_request(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE, "  press");
-		TEST_EQ(check_shutdown_request(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE, "  hold");
+		TEST_EQ(check_shutdown_request(&mock_ui_context), VB2_SUCCESS,
+			"  press");
+		TEST_EQ(check_shutdown_request(&mock_ui_context), VB2_SUCCESS,
+			"  hold");
 		mock_shutdown_request = 0;
 		TEST_EQ(check_shutdown_request(&mock_ui_context),
 			VB2_REQUEST_SHUTDOWN, "  release");
@@ -196,8 +195,8 @@ static void check_shutdown_request_tests(void)
 	if (!DETACHABLE) {
 		reset_common_data();
 		mock_shutdown_request = VB_SHUTDOWN_REQUEST_POWER_BUTTON;
-		TEST_EQ(check_shutdown_request(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE, "press is ignored");
+		TEST_EQ(check_shutdown_request(&mock_ui_context), VB2_SUCCESS,
+			"press is ignored");
 	}
 
 	/* Power button short press from key */
@@ -222,13 +221,13 @@ static void check_shutdown_request_tests(void)
 	reset_common_data();
 	gbb.flags |= VB2_GBB_FLAG_DISABLE_LID_SHUTDOWN;
 	mock_shutdown_request = VB_SHUTDOWN_REQUEST_LID_CLOSED;
-	TEST_EQ(check_shutdown_request(&mock_ui_context),
-		VB2_REQUEST_UI_CONTINUE, "lid ignored");
+	TEST_EQ(check_shutdown_request(&mock_ui_context), VB2_SUCCESS,
+		"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(check_shutdown_request(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE, "  lidsw + pwdsw");
+		TEST_EQ(check_shutdown_request(&mock_ui_context), VB2_SUCCESS,
+			"  lidsw + pwdsw");
 		mock_shutdown_request = 0;
 		TEST_EQ(check_shutdown_request(&mock_ui_context),
 			VB2_REQUEST_SHUTDOWN, "  pwdsw release");
@@ -250,18 +249,17 @@ static void check_shutdown_request_tests(void)
 		/* Flag pwdsw */
 		reset_common_data();
 		mock_shutdown_request = VB_SHUTDOWN_REQUEST_POWER_BUTTON;
-		TEST_EQ(check_shutdown_request(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE, "DETACHABLE: ignore pwdsw");
+		TEST_EQ(check_shutdown_request(&mock_ui_context), VB2_SUCCESS,
+			"DETACHABLE: ignore pwdsw");
 		mock_shutdown_request = 0;
-		TEST_EQ(check_shutdown_request(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE, "  ignore on release");
+		TEST_EQ(check_shutdown_request(&mock_ui_context), VB2_SUCCESS,
+			"  ignore on release");
 
 		/* Power button short press */
 		reset_common_data();
 		mock_shutdown_request = 0;
 		mock_ui_context.key = VB_BUTTON_POWER_SHORT_PRESS;
-		TEST_EQ(check_shutdown_request(&mock_ui_context),
-			VB2_REQUEST_UI_CONTINUE,
+		TEST_EQ(check_shutdown_request(&mock_ui_context), VB2_SUCCESS,
 			"DETACHABLE: ignore power button short press");
 	}
 
@@ -293,8 +291,7 @@ static void screen_stack_tests(void)
 
 	/* Screen back with empty stack */
 	reset_common_data();
-	TEST_EQ(vb2_ui_screen_back(&mock_ui_context),
-		VB2_REQUEST_UI_CONTINUE,
+	TEST_EQ(vb2_ui_screen_back(&mock_ui_context), VB2_REQUEST_UI_CONTINUE,
 		"screen back with empty stack");
 	TEST_PTR_EQ(mock_ui_context.state, NULL, "  stack is empty");