diff --git a/firmware/2lib/2ui.c b/firmware/2lib/2ui.c
index 8eea9af..3c54ecc 100644
--- a/firmware/2lib/2ui.c
+++ b/firmware/2lib/2ui.c
@@ -138,12 +138,11 @@
 	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 change_screen(ui, menu_item->target);
 	}
 
+	VB2_DEBUG("No target screen for menu item <%s>\n", menu_item->text);
+
 	return VB2_REQUEST_UI_CONTINUE;
 }
 
@@ -152,8 +151,8 @@
  */
 vb2_error_t menu_back_action(struct vb2_ui_context *ui)
 {
-	change_screen(ui, ui->root_screen->id);
-	return VB2_REQUEST_UI_CONTINUE;
+	/* TODO(kitching): Return to previous screen instead of root screen. */
+	return change_screen(ui, ui->root_screen->id);
 }
 
 /*****************************************************************************/
@@ -181,33 +180,19 @@
 /*****************************************************************************/
 /* Core UI functions */
 
-void change_screen(struct vb2_ui_context *ui, enum vb2_screen id)
+vb2_error_t change_screen(struct vb2_ui_context *ui, enum vb2_screen id)
 {
 	const struct vb2_screen_info *new_screen_info = vb2_get_screen_info(id);
+
 	if (new_screen_info == NULL) {
 		VB2_DEBUG("ERROR: Screen entry %#x not found; ignoring\n", id);
-	} else {
-		memset(&ui->state, 0, sizeof(ui->state));
-		ui->state.screen = new_screen_info;
+		return VB2_REQUEST_UI_CONTINUE;
 	}
-}
 
-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;
+	memset(&ui->state, 0, sizeof(ui->state));
+	ui->state.screen = new_screen_info;
 
-	/* 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;
+	return VB2_REQUEST_UI_CONTINUE;
 }
 
 vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
@@ -225,7 +210,9 @@
 	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);
+	rv = change_screen(&ui, ui.root_screen->id);
+	if (rv != VB2_REQUEST_UI_CONTINUE)
+		return rv;
 	memset(&prev_state, 0, sizeof(prev_state));
 
 	while (1) {
@@ -259,7 +246,6 @@
 			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));
@@ -268,7 +254,6 @@
 		/* 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;
 		}
@@ -336,10 +321,9 @@
 	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 change_screen(ui, invalid_disk ?
+				     VB2_SCREEN_RECOVERY_INVALID :
+				     VB2_SCREEN_RECOVERY_SELECT);
 	}
 
 	return VB2_REQUEST_UI_CONTINUE;
diff --git a/firmware/2lib/include/2ui_private.h b/firmware/2lib/include/2ui_private.h
index 40ee6b5..0053449 100644
--- a/firmware/2lib/include/2ui_private.h
+++ b/firmware/2lib/include/2ui_private.h
@@ -31,8 +31,7 @@
 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 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));
 
diff --git a/tests/vb2_ui_utility_tests.c b/tests/vb2_ui_utility_tests.c
index da7aaf7..e62f6c0 100644
--- a/tests/vb2_ui_utility_tests.c
+++ b/tests/vb2_ui_utility_tests.c
@@ -614,72 +614,22 @@
 	mock_state->screen = &mock_screen_menu;
 	mock_state->selected_item = 2;
 	mock_state->disabled_item_mask = 0x10;
-	VB2_DEBUG("change_screen will clear screen state\n");
-	change_screen(&mock_ui_context, MOCK_SCREEN_BASE);
+	TEST_EQ(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);
 
 	/* 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);
+	TEST_EQ(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);
 
 	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, 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, 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, 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, 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, 0, MOCK_IGNORE);
-
-	VB2_DEBUG("...done.\n");
-}
-
 static void ui_loop_tests(void)
 {
 	VB2_DEBUG("Testing ui_loop...\n");
@@ -781,7 +731,6 @@
 	shutdown_required_tests();
 	menu_action_tests();
 	change_screen_tests();
-	validate_selection_tests();
 	ui_loop_tests();
 
 	return gTestSuccess ? 0 : 255;