vboot/ui: implement language selection screen

Implement language selection screen, and add language item to all the
other screens. Add a default screen init function default_screen_init()
to initialize the default selection to the second item if the first item
is the language selection.

BRANCH=none
BUG=b:146399181, b:144968920
TEST=make runtests
TEST=USE="menu_ui" emerge-nami depthcharge

Cq-Depend: chromium:2193151, chromium:2192508
Change-Id: I3251b0095ec29ec26cc27745b1089e60894c892c
Signed-off-by: Yu-Ping Wu <yupingso@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2196095
Reviewed-by: Joel Kitching <kitching@chromium.org>

7 files changed, 328 insertions, 71 deletions

diff --git a/firmware/2lib/2ui.c b/firmware/2lib/2ui.c
index eaa2bc4f..ee0da1f1 100644
--- a/firmware/2lib/2ui.c
+++ b/firmware/2lib/2ui.c
@@ -73,6 +73,21 @@ vb2_error_t check_shutdown_request(struct vb2_ui_context *ui)
 /*****************************************************************************/
 /* Menu navigation functions */
 
+const struct vb2_menu *get_menu(struct vb2_ui_context *ui)
+{
+	const struct vb2_menu *menu;
+	static const struct vb2_menu empty_menu = {
+		.num_items = 0,
+		.items = NULL,
+	};
+	if (ui->state.screen->get_menu) {
+		menu = ui->state.screen->get_menu(ui);
+		return menu ? menu : &empty_menu;
+	} else {
+		return &ui->state.screen->menu;
+	}
+}
+
 vb2_error_t menu_navigation_action(struct vb2_ui_context *ui)
 {
 	uint32_t key = ui->key;
@@ -126,16 +141,18 @@ vb2_error_t vb2_ui_menu_prev(struct vb2_ui_context *ui)
 vb2_error_t vb2_ui_menu_next(struct vb2_ui_context *ui)
 {
 	int item;
+	const struct vb2_menu *menu;
 
 	if (!DETACHABLE && ui->key == VB_BUTTON_VOL_DOWN_SHORT_PRESS)
 		return VB2_REQUEST_UI_CONTINUE;
 
+	menu = get_menu(ui);
 	item = ui->state.selected_item + 1;
-	while (item < ui->state.screen->num_items &&
+	while (item < menu->num_items &&
 	       ((1 << item) & ui->state.disabled_item_mask))
 		item++;
 	/* Only update if item is valid */
-	if (item < ui->state.screen->num_items)
+	if (item < menu->num_items)
 		ui->state.selected_item = item;
 
 	return VB2_REQUEST_UI_CONTINUE;
@@ -143,15 +160,17 @@ vb2_error_t vb2_ui_menu_next(struct vb2_ui_context *ui)
 
 vb2_error_t vb2_ui_menu_select(struct vb2_ui_context *ui)
 {
+	const struct vb2_menu *menu;
 	const struct vb2_menu_item *menu_item;
 
 	if (!DETACHABLE && ui->key == VB_BUTTON_POWER_SHORT_PRESS)
 		return VB2_REQUEST_UI_CONTINUE;
 
-	if (ui->state.screen->num_items == 0)
+	menu = get_menu(ui);
+	if (menu->num_items == 0)
 		return VB2_REQUEST_UI_CONTINUE;
 
-	menu_item = &ui->state.screen->items[ui->state.selected_item];
+	menu_item = &menu->items[ui->state.selected_item];
 
 	if (menu_item->action) {
 		VB2_DEBUG("Menu item <%s> run action\n", menu_item->text);
@@ -175,6 +194,15 @@ vb2_error_t vb2_ui_change_root(struct vb2_ui_context *ui)
 	return vb2_ui_change_screen(ui, ui->root_screen->id);
 }
 
+static vb2_error_t default_screen_init(struct vb2_ui_context *ui)
+{
+	const struct vb2_menu *menu = get_menu(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;
+}
+
 vb2_error_t vb2_ui_change_screen(struct vb2_ui_context *ui, enum vb2_screen id)
 {
 	const struct vb2_screen_info *new_screen_info;
@@ -195,8 +223,8 @@ vb2_error_t vb2_ui_change_screen(struct vb2_ui_context *ui, enum vb2_screen id)
 
 	if (ui->state.screen->init)
 		return ui->state.screen->init(ui);
-
-	return VB2_REQUEST_UI_CONTINUE;
+	else
+		return default_screen_init(ui);
 }
 
 /*****************************************************************************/
@@ -207,6 +235,7 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 {
 	struct vb2_ui_context ui;
 	struct vb2_screen_state prev_state;
+	const struct vb2_menu *menu;
 	uint32_t key_flags;
 	vb2_error_t rv;
 
@@ -215,6 +244,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");
+	ui.locale_id = vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX);
 	rv = vb2_ui_change_screen(&ui, ui.root_screen->id);
 	if (rv != VB2_REQUEST_UI_CONTINUE)
 		return rv;
@@ -225,11 +255,12 @@ vb2_error_t ui_loop(struct vb2_context *ctx, enum vb2_screen root_screen_id,
 		if (memcmp(&prev_state, &ui.state, sizeof(ui.state))) {
 			memcpy(&prev_state, &ui.state, sizeof(ui.state));
 
+			menu = get_menu(&ui);
 			VB2_DEBUG("<%s> menu item <%s>\n",
 				  ui.state.screen->name,
-				  ui.state.screen->num_items ?
-				  ui.state.screen->items[
-				  ui.state.selected_item].text : "null");
+				  menu->num_items ?
+				  menu->items[ui.state.selected_item].text :
+				  "null");
 
 			vb2ex_display_ui(ui.state.screen->id, ui.locale_id,
 					 ui.state.selected_item,
diff --git a/firmware/2lib/2ui_screens.c b/firmware/2lib/2ui_screens.c
index 50a61423..612a08c8 100644
--- a/firmware/2lib/2ui_screens.c
+++ b/firmware/2lib/2ui_screens.c
@@ -5,6 +5,7 @@
  * Firmware screen definitions.
  */
 
+#include "2api.h"
 #include "2common.h"
 #include "2misc.h"
 #include "2nvstorage.h"
@@ -13,9 +14,16 @@
 #include "vboot_api.h"
 #include "vboot_kernel.h"
 
-#define MENU_ITEMS(a) \
+#define MENU_ITEMS(a) ((struct vb2_menu){ \
 	.num_items = ARRAY_SIZE(a), \
-	.items = a
+	.items = a, \
+})
+
+#define LANGUAGE_SELECT_ITEM { \
+	.text = "Language selection", \
+	.target = VB2_SCREEN_LANGUAGE_SELECT, \
+	.is_language_select = 1, \
+}
 
 #define ADVANCED_OPTIONS_ITEM { \
 	.text = "Advanced options", \
@@ -31,26 +39,105 @@ static const struct vb2_screen_info blank_screen = {
 };
 
 /******************************************************************************/
+/* VB2_SCREEN_LANGUAGE_SELECT */
+
+static vb2_error_t language_select_action(struct vb2_ui_context *ui)
+{
+	vb2_error_t rv;
+	ui->locale_id = ui->state.selected_item;
+	VB2_DEBUG("Locale changed to %u\n", ui->locale_id);
+
+	/* Write locale id back to nvdata. */
+	vb2_nv_set(ui->ctx, VB2_NV_LOCALIZATION_INDEX, ui->locale_id);
+
+	/* Commit nvdata changes immediately, in case of three-finger salute
+	   reboot.  Ignore commit errors in recovery mode. */
+	rv = vb2ex_commit_data(ui->ctx);
+	if (rv && !(ui->ctx->flags & VB2_CONTEXT_RECOVERY_MODE))
+		return rv;
+
+	return vb2_ui_change_root(ui);
+}
+
+const struct vb2_menu *get_language_menu(struct vb2_ui_context *ui)
+{
+	int i;
+	uint32_t num_locales;
+	struct vb2_menu_item *items;
+
+	if (ui->language_menu.num_items > 0)
+		return &ui->language_menu;
+
+	num_locales = vb2ex_get_locale_count();
+	if (num_locales == 0) {
+		VB2_DEBUG("WARNING: No locales available; assuming 1 locale\n");
+		num_locales = 1;
+	}
+
+	items = malloc(num_locales * sizeof(struct vb2_menu_item));
+	if (!items) {
+		VB2_DEBUG("ERROR: malloc failed for language items\n");
+		return NULL;
+	}
+
+	for (i = 0; i < num_locales; i++) {
+		items[i].text = "Some language";
+		items[i].action = language_select_action;
+	}
+
+	ui->language_menu.num_items = num_locales;
+	ui->language_menu.items = items;
+	return &ui->language_menu;
+}
+
+static vb2_error_t language_select_init(struct vb2_ui_context *ui)
+{
+	const struct vb2_menu *menu = get_menu(ui);
+	if (menu->num_items == 0) {
+		VB2_DEBUG("ERROR: No menu items found; "
+			  "rejecting entering language selection screen\n");
+		return vb2_ui_change_root(ui);
+	}
+	if (ui->locale_id < menu->num_items) {
+		ui->state.selected_item = ui->locale_id;
+	} else {
+		VB2_DEBUG("WARNING: Current locale not found in menu items; "
+			  "initializing selected_item to 0\n");
+		ui->state.selected_item = 0;
+	}
+	return VB2_REQUEST_UI_CONTINUE;
+}
+
+static const struct vb2_screen_info language_select_screen = {
+	.id = VB2_SCREEN_LANGUAGE_SELECT,
+	.name = "Language selection screen",
+	.init = language_select_init,
+	.get_menu = get_language_menu,
+};
+
+/******************************************************************************/
 /* VB2_SCREEN_RECOVERY_BROKEN */
 
 static const struct vb2_menu_item recovery_broken_items[] = {
+	LANGUAGE_SELECT_ITEM,
 	ADVANCED_OPTIONS_ITEM,
 };
 
 static const struct vb2_screen_info recovery_broken_screen = {
 	.id = VB2_SCREEN_RECOVERY_BROKEN,
 	.name = "Recover broken device",
-	MENU_ITEMS(recovery_broken_items),
+	.menu = MENU_ITEMS(recovery_broken_items),
 };
 
 /******************************************************************************/
 /* VB2_SCREEN_ADVANCED_OPTIONS */
 
-#define ADVANCED_OPTIONS_ITEM_DEVELOPER_MODE 0
-#define ADVANCED_OPTIONS_ITEM_BACK 1
+#define ADVANCED_OPTIONS_ITEM_DEVELOPER_MODE 1
+#define ADVANCED_OPTIONS_ITEM_BACK 2
 
 vb2_error_t advanced_options_init(struct vb2_ui_context *ui)
 {
+	ui->state.selected_item = ADVANCED_OPTIONS_ITEM_DEVELOPER_MODE;
 	if (vb2_get_sd(ui->ctx)->flags & VB2_SD_FLAG_DEV_MODE_ENABLED) {
 		ui->state.disabled_item_mask |=
 			1 << ADVANCED_OPTIONS_ITEM_DEVELOPER_MODE;
@@ -61,6 +148,7 @@ vb2_error_t advanced_options_init(struct vb2_ui_context *ui)
 }
 
 static const struct vb2_menu_item advanced_options_items[] = {
+	LANGUAGE_SELECT_ITEM,
 	[ADVANCED_OPTIONS_ITEM_DEVELOPER_MODE] = {
 		.text = "Enable developer mode",
 		.target = VB2_SCREEN_RECOVERY_TO_DEV,
@@ -75,17 +163,18 @@ 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),
+	.menu = MENU_ITEMS(advanced_options_items),
 };
 
 /******************************************************************************/
 /* VB2_SCREEN_RECOVERY_SELECT */
 
-#define RECOVERY_SELECT_ITEM_PHONE 0
-#define RECOVERY_SELECT_ITEM_EXTERNAL_DISK 1
+#define RECOVERY_SELECT_ITEM_PHONE 1
+#define RECOVERY_SELECT_ITEM_EXTERNAL_DISK 2
 
 vb2_error_t recovery_select_init(struct vb2_ui_context *ui)
 {
+	ui->state.selected_item = RECOVERY_SELECT_ITEM_PHONE;
 	if (!vb2api_phone_recovery_enabled(ui->ctx)) {
 		VB2_DEBUG("WARNING: Phone recovery not available\n");
 		ui->state.disabled_item_mask |=
@@ -96,6 +185,7 @@ vb2_error_t recovery_select_init(struct vb2_ui_context *ui)
 }
 
 static const struct vb2_menu_item recovery_select_items[] = {
+	LANGUAGE_SELECT_ITEM,
 	[RECOVERY_SELECT_ITEM_PHONE] = {
 		.text = "Recovery using phone",
 		.target = VB2_SCREEN_RECOVERY_PHONE_STEP1,
@@ -111,22 +201,27 @@ static const struct vb2_screen_info recovery_select_screen = {
 	.id = VB2_SCREEN_RECOVERY_SELECT,
 	.name = "Recovery method selection",
 	.init = recovery_select_init,
-	MENU_ITEMS(recovery_select_items),
+	.menu = MENU_ITEMS(recovery_select_items),
 };
 
 /******************************************************************************/
 /* VB2_SCREEN_RECOVERY_INVALID */
 
+static const struct vb2_menu_item recovery_invalid_items[] = {
+	LANGUAGE_SELECT_ITEM,
+};
+
 static const struct vb2_screen_info recovery_invalid_screen = {
 	.id = VB2_SCREEN_RECOVERY_INVALID,
 	.name = "Invalid recovery inserted",
+	.menu = MENU_ITEMS(recovery_invalid_items),
 };
 
 /******************************************************************************/
 /* VB2_SCREEN_RECOVERY_TO_DEV */
 
-#define RECOVERY_TO_DEV_ITEM_CONFIRM 0
-#define RECOVERY_TO_DEV_ITEM_CANCEL 1
+#define RECOVERY_TO_DEV_ITEM_CONFIRM 1
+#define RECOVERY_TO_DEV_ITEM_CANCEL 2
 
 vb2_error_t recovery_to_dev_init(struct vb2_ui_context *ui)
 {
@@ -140,6 +235,8 @@ vb2_error_t recovery_to_dev_init(struct vb2_ui_context *ui)
 		return vb2_ui_change_root(ui);
 	}
 
+	ui->state.selected_item = RECOVERY_TO_DEV_ITEM_CONFIRM;
+
 	/* Disable "Confirm" button for other physical presence types. */
 	if (!PHYSICAL_PRESENCE_KEYBOARD) {
 		ui->state.disabled_item_mask |=
@@ -207,6 +304,7 @@ vb2_error_t recovery_to_dev_action(struct vb2_ui_context *ui)
 }
 
 static const struct vb2_menu_item recovery_to_dev_items[] = {
+	LANGUAGE_SELECT_ITEM,
 	[RECOVERY_TO_DEV_ITEM_CONFIRM] = {
 		.text = "Confirm",
 		.action = recovery_to_dev_confirm_action,
@@ -222,31 +320,41 @@ static const struct vb2_screen_info recovery_to_dev_screen = {
 	.name = "Transition to developer mode",
 	.init = recovery_to_dev_init,
 	.action = recovery_to_dev_action,
-	MENU_ITEMS(recovery_to_dev_items),
+	.menu = MENU_ITEMS(recovery_to_dev_items),
 };
 
 /******************************************************************************/
 /* VB2_SCREEN_RECOVERY_PHONE_STEP1 */
 
+static const struct vb2_menu_item recovery_phone_step1_items[] = {
+	LANGUAGE_SELECT_ITEM,
+};
+
 static const struct vb2_screen_info recovery_phone_step1_screen = {
 	.id = VB2_SCREEN_RECOVERY_PHONE_STEP1,
 	.name = "Phone recovery step 1",
+	.menu = MENU_ITEMS(recovery_phone_step1_items),
 };
 
 /******************************************************************************/
 /* VB2_SCREEN_RECOVERY_DISK_STEP1 */
 
+static const struct vb2_menu_item recovery_disk_step1_items[] = {
+	LANGUAGE_SELECT_ITEM,
+};
+
 static const struct vb2_screen_info recovery_disk_step1_screen = {
 	.id = VB2_SCREEN_RECOVERY_DISK_STEP1,
 	.name = "Disk recovery step 1",
+	.menu = MENU_ITEMS(recovery_disk_step1_items),
 };
 
 /******************************************************************************/
 /* 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
+#define DEVELOPER_MODE_ITEM_RETURN_TO_SECURE 1
+#define DEVELOPER_MODE_ITEM_BOOT_INTERNAL 2
+#define DEVELOPER_MODE_ITEM_BOOT_EXTERNAL 3
 
 vb2_error_t developer_mode_init(struct vb2_ui_context *ui)
 {
@@ -353,6 +461,7 @@ vb2_error_t developer_mode_action(struct vb2_ui_context *ui)
 }
 
 static const struct vb2_menu_item developer_mode_items[] = {
+	LANGUAGE_SELECT_ITEM,
 	[DEVELOPER_MODE_ITEM_RETURN_TO_SECURE] = {
 		.text = "Return to secure mode",
 		.target = VB2_SCREEN_DEVELOPER_TO_NORM,
@@ -373,7 +482,7 @@ static const struct vb2_screen_info developer_mode_screen = {
 	.name = "Developer mode",
 	.init = developer_mode_init,
 	.action = developer_mode_action,
-	MENU_ITEMS(developer_mode_items),
+	.menu = MENU_ITEMS(developer_mode_items),
 };
 
 /******************************************************************************/
@@ -392,6 +501,7 @@ vb2_error_t developer_to_norm_action(struct vb2_ui_context *ui)
 }
 
 static const struct vb2_menu_item developer_to_norm_items[] = {
+	LANGUAGE_SELECT_ITEM,
 	{
 		.text = "Confirm",
 		.action = developer_to_norm_action,
@@ -405,7 +515,7 @@ static const struct vb2_menu_item developer_to_norm_items[] = {
 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),
+	.menu = MENU_ITEMS(developer_to_norm_items),
 };
 
 /******************************************************************************/
@@ -418,6 +528,7 @@ static const struct vb2_screen_info developer_to_norm_screen = {
  */
 static const struct vb2_screen_info *screens[] = {
 	&blank_screen,
+	&language_select_screen,
 	&recovery_broken_screen,
 	&advanced_options_screen,
 	&recovery_select_screen,
diff --git a/firmware/2lib/include/2ui.h b/firmware/2lib/include/2ui.h
index eee1565f..49f51b32 100644
--- a/firmware/2lib/include/2ui.h
+++ b/firmware/2lib/include/2ui.h
@@ -16,6 +16,24 @@
 
 struct vb2_ui_context;  /* Forward declaration */
 
+struct vb2_menu_item {
+	/* Text description */
+	const char *text;
+	/* Target screen */
+	enum vb2_screen target;
+	/* Action function takes precedence over target screen if non-NULL. */
+	vb2_error_t (*action)(struct vb2_ui_context *ui);
+	/* Whether the item is language selection */
+	int is_language_select;
+};
+
+struct vb2_menu {
+	/* Number of menu items */
+	uint16_t num_items;
+	/* List of menu items */
+	const struct vb2_menu_item *items;
+};
+
 struct vb2_screen_info {
 	/* Screen id */
 	enum vb2_screen id;
@@ -25,19 +43,13 @@ struct vb2_screen_info {
 	vb2_error_t (*init)(struct vb2_ui_context *ui);
 	/* Action function runs repeatedly while on the screen. */
 	vb2_error_t (*action)(struct vb2_ui_context *ui);
-	/* Number of menu items */
-	uint16_t num_items;
-	/* List of menu items */
-	const struct vb2_menu_item *items;
-};
-
-struct vb2_menu_item {
-	/* Text description */
-	const char *text;
-	/* Target screen */
-	enum vb2_screen target;
-	/* Action function takes precedence over target screen if non-NULL. */
-	vb2_error_t (*action)(struct vb2_ui_context *ui);
+	/* Menu items. */
+	struct vb2_menu menu;
+	/*
+	 * Custom function for getting menu items. If non-null, field 'menu'
+	 * will be ignored.
+	 */
+	const struct vb2_menu *(*get_menu)(struct vb2_ui_context *ui);
 };
 
 struct vb2_screen_state {
@@ -73,6 +85,9 @@ struct vb2_ui_context {
 
 	/* For to_dev transition flow. */
 	int physical_presence_button_pressed;
+
+	/* For language selection screen. */
+	struct vb2_menu language_menu;
 };
 
 vb2_error_t vb2_ui_developer_mode_boot_internal_action(
diff --git a/firmware/2lib/include/2ui_private.h b/firmware/2lib/include/2ui_private.h
index c1ecd8be..e1e47af3 100644
--- a/firmware/2lib/include/2ui_private.h
+++ b/firmware/2lib/include/2ui_private.h
@@ -12,6 +12,7 @@
 
 /* From 2ui.c */
 vb2_error_t check_shutdown_request(struct vb2_ui_context *ui);
+const struct vb2_menu *get_menu(struct vb2_ui_context *ui);
 vb2_error_t menu_navigation_action(struct vb2_ui_context *ui);
 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));
@@ -19,6 +20,7 @@ vb2_error_t developer_action(struct vb2_ui_context *ui);
 vb2_error_t manual_recovery_action(struct vb2_ui_context *ui);
 
 /* From 2ui_screens.c */
+const struct vb2_menu *get_language_menu(struct vb2_ui_context *ui);
 vb2_error_t advanced_options_init(struct vb2_ui_context *ui);
 vb2_error_t recovery_select_init(struct vb2_ui_context *ui);
 vb2_error_t recovery_to_dev_init(struct vb2_ui_context *ui);
diff --git a/tests/vb2_ui_action_tests.c b/tests/vb2_ui_action_tests.c
index ddb9ee38..e778f6ee 100644
--- a/tests/vb2_ui_action_tests.c
+++ b/tests/vb2_ui_action_tests.c
@@ -105,18 +105,13 @@ static vb2_error_t mock_action_flag2(struct vb2_ui_context *ui)
 
 /* Mock screens */
 struct vb2_screen_info mock_screen_temp;
-const struct vb2_menu_item mock_empty_menu[] = {};
 const struct vb2_screen_info mock_screen_blank = {
 	.id = VB2_SCREEN_BLANK,
 	.name = "mock_screen_blank",
-	.num_items = ARRAY_SIZE(mock_empty_menu),
-	.items = mock_empty_menu,
 };
 const struct vb2_screen_info mock_screen_base = {
 	.id = MOCK_SCREEN_BASE,
 	.name = "mock_screen_base: menuless screen",
-	.num_items = ARRAY_SIZE(mock_empty_menu),
-	.items = mock_empty_menu,
 };
 const struct vb2_menu_item mock_screen_menu_items[] = {
 	{
@@ -142,33 +137,27 @@ const struct vb2_menu_item mock_screen_menu_items[] = {
 const struct vb2_screen_info mock_screen_menu = {
 	.id = MOCK_SCREEN_MENU,
 	.name = "mock_screen_menu: screen with 5 items",
-	.num_items = ARRAY_SIZE(mock_screen_menu_items),
-	.items = mock_screen_menu_items,
+	.menu = {
+		.num_items = ARRAY_SIZE(mock_screen_menu_items),
+		.items = mock_screen_menu_items,
+	},
 };
 const struct vb2_screen_info mock_screen_target0 = {
 	.id = MOCK_SCREEN_TARGET0,
 	.name = "mock_screen_target0",
-	.num_items = ARRAY_SIZE(mock_empty_menu),
-	.items = mock_empty_menu,
 };
 const struct vb2_screen_info mock_screen_target1 = {
 	.id = MOCK_SCREEN_TARGET1,
 	.name = "mock_screen_target1",
-	.num_items = ARRAY_SIZE(mock_empty_menu),
-	.items = mock_empty_menu,
 };
 const struct vb2_screen_info mock_screen_target2 = {
 	.id = MOCK_SCREEN_TARGET2,
 	.name = "mock_screen_target2",
-	.num_items = ARRAY_SIZE(mock_empty_menu),
-	.items = mock_empty_menu,
 };
 const struct vb2_screen_info mock_screen_action = {
 	.id = MOCK_SCREEN_ACTION,
 	.name = "mock_screen_action",
 	.action = mock_action_countdown,
-	.num_items = ARRAY_SIZE(mock_empty_menu),
-	.items = mock_empty_menu,
 };
 const struct vb2_menu_item mock_screen_all_action_items[] = {
 	{
@@ -180,8 +169,10 @@ const struct vb2_screen_info mock_screen_all_action = {
 	.id = MOCK_SCREEN_ALL_ACTION,
 	.name = "mock_screen_all_action",
 	.action = mock_action_flag0,
-	.num_items = ARRAY_SIZE(mock_screen_all_action_items),
-	.items = mock_screen_all_action_items,
+	.menu = {
+		.num_items = ARRAY_SIZE(mock_screen_all_action_items),
+		.items = mock_screen_all_action_items,
+	},
 };
 
 static void screen_state_eq(const struct vb2_screen_state *state,
@@ -293,8 +284,6 @@ static void reset_common_data(void)
 	mock_screen_temp = (struct vb2_screen_info){
 	      .id = MOCK_NO_SCREEN,
 	      .name = "mock_screen_temp",
-	      .num_items = ARRAY_SIZE(mock_empty_menu),
-	      .items = mock_empty_menu,
 	};
 
 	/* Mock ui_context based on mock screens */
diff --git a/tests/vb2_ui_tests.c b/tests/vb2_ui_tests.c
index 4066e5b8..aef133eb 100644
--- a/tests/vb2_ui_tests.c
+++ b/tests/vb2_ui_tests.c
@@ -39,6 +39,8 @@ static struct display_call mock_displayed[64];
 static int mock_displayed_count;
 static int mock_displayed_i;
 
+static uint32_t mock_locale_count;
+
 static int mock_calls_until_shutdown;
 
 /* Iteration counter starts from 0
@@ -194,6 +196,9 @@ static void reset_common_data(enum reset_type t)
 	mock_displayed_count = 0;
 	mock_displayed_i = 0;
 
+	/* For vb2ex_get_locale_count */
+	mock_locale_count = 1;
+
 	/* For check_shutdown_request */
 	if (t == FOR_DEVELOPER)
 		mock_calls_until_shutdown = 2000;  /* Larger than 30s */
@@ -280,6 +285,11 @@ vb2_error_t vb2ex_display_ui(enum vb2_screen screen,
 	return VB2_SUCCESS;
 }
 
+uint32_t vb2ex_get_locale_count(void)
+{
+	return mock_locale_count;
+}
+
 uint32_t VbExIsShutdownRequested(void)
 {
 	if (mock_calls_until_shutdown < 0)  /* Never request shutdown */
@@ -502,7 +512,7 @@ static void manual_recovery_tests(void)
 	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
 		"phone recovery");
 	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
-		     MOCK_IGNORE, 0, MOCK_IGNORE);
+		     MOCK_IGNORE, 1, MOCK_IGNORE);
 	displayed_eq("phone recovery", VB2_SCREEN_RECOVERY_PHONE_STEP1,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	displayed_no_extra();
@@ -514,9 +524,9 @@ static void manual_recovery_tests(void)
 	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
 		"external disk recovery");
 	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
-		     MOCK_IGNORE, 0, MOCK_IGNORE);
-	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
 		     MOCK_IGNORE, 1, MOCK_IGNORE);
+	displayed_eq("recovery select", VB2_SCREEN_RECOVERY_SELECT,
+		     MOCK_IGNORE, 2, MOCK_IGNORE);
 	displayed_eq("disk recovery", VB2_SCREEN_RECOVERY_DISK_STEP1,
 		     MOCK_IGNORE, MOCK_IGNORE, MOCK_IGNORE);
 	displayed_no_extra();
@@ -703,11 +713,64 @@ static void manual_recovery_tests(void)
 	VB2_DEBUG("...done.\n");
 }
 
+static void language_selection_tests(void)
+{
+	VB2_DEBUG("Testing language selection...\n");
+
+	/* Enter language menu and change language */
+	reset_common_data(FOR_MANUAL_RECOVERY);
+	mock_locale_count = 100;
+	vb2_nv_set(ctx, VB2_NV_LOCALIZATION_INDEX, 23);
+	add_mock_keypress(VB_KEY_UP);
+	add_mock_keypress(VB_KEY_ENTER);	/* select language */
+	add_mock_keypress(VB_KEY_DOWN);
+	add_mock_keypress(VB_KEY_ENTER);	/* select locale 24 */
+	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
+		"change language");
+	displayed_eq("RECOVERY_SELECT default", VB2_SCREEN_RECOVERY_SELECT,
+		     23, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_eq("RECOVERY_SELECT lang", VB2_SCREEN_RECOVERY_SELECT,
+		     23, 0, MOCK_IGNORE);
+	displayed_eq("LANGUAGE_SELECT 23", VB2_SCREEN_LANGUAGE_SELECT,
+		     23, 23, MOCK_IGNORE);
+	displayed_eq("LANGUAGE_SELECT 24", VB2_SCREEN_LANGUAGE_SELECT,
+		     23, 24, MOCK_IGNORE);
+	displayed_eq("RECOVERY_SELECT new locale", VB2_SCREEN_RECOVERY_SELECT,
+		     24, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_no_extra();
+	TEST_EQ(vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX), 24,
+		"  locale 24 saved to nvdata");
+
+	/* Locale count = 0 */
+	reset_common_data(FOR_MANUAL_RECOVERY);
+	mock_locale_count = 0;
+	vb2_nv_set(ctx, VB2_NV_LOCALIZATION_INDEX, 23);
+	add_mock_keypress(VB_KEY_UP);
+	add_mock_keypress(VB_KEY_ENTER);	/* select language */
+	add_mock_keypress(VB_KEY_ENTER);	/* select locale 0 */
+	add_mock_vbtlk(VB2_ERROR_LK_NO_DISK_FOUND, VB_DISK_FLAG_REMOVABLE);
+	TEST_EQ(vb2_manual_recovery_menu(ctx), VB2_REQUEST_SHUTDOWN,
+		"enter language menu");
+	displayed_eq("RECOVERY_SELECT default", VB2_SCREEN_RECOVERY_SELECT,
+		     23, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_eq("RECOVERY_SELECT lang", VB2_SCREEN_RECOVERY_SELECT,
+		     23, 0, MOCK_IGNORE);
+	displayed_eq("LANGUAGE_SELECT index 0", VB2_SCREEN_LANGUAGE_SELECT,
+		     23, 0, MOCK_IGNORE);
+	displayed_eq("RECOVERY_SELECT locale 0", VB2_SCREEN_RECOVERY_SELECT,
+		     0, MOCK_IGNORE, MOCK_IGNORE);
+	displayed_no_extra();
+
+	VB2_DEBUG("...done.\n");
+}
+
 int main(void)
 {
 	developer_tests();
 	broken_recovery_tests();
 	manual_recovery_tests();
+	language_selection_tests();
 
 	return gTestSuccess ? 0 : 255;
 }
diff --git a/tests/vb2_ui_utility_tests.c b/tests/vb2_ui_utility_tests.c
index 257e69d9..533aa514 100644
--- a/tests/vb2_ui_utility_tests.c
+++ b/tests/vb2_ui_utility_tests.c
@@ -29,6 +29,7 @@ static uint8_t workbuf[VB2_KERNEL_WORKBUF_RECOMMENDED_SIZE]
 static struct vb2_context *ctx;
 static struct vb2_gbb_header gbb;
 
+static uint32_t mock_locale_count;
 static int mock_shutdown_request;
 
 static struct vb2_ui_context mock_ui_context;
@@ -43,22 +44,18 @@ static vb2_error_t mock_action_base(struct vb2_ui_context *ui)
 }
 
 /* Mock screens */
-const struct vb2_menu_item mock_empty_menu[] = {};
 struct vb2_screen_info mock_screen_blank = {
 	.id = VB2_SCREEN_BLANK,
 	.name = "mock_screen_blank",
-	.num_items = ARRAY_SIZE(mock_empty_menu),
-	.items = mock_empty_menu,
 };
 struct vb2_screen_info mock_screen_base = {
 	.id = MOCK_SCREEN_BASE,
 	.name = "mock_screen_base: menuless screen",
-	.num_items = ARRAY_SIZE(mock_empty_menu),
-	.items = mock_empty_menu,
 };
 struct vb2_menu_item mock_screen_menu_items[] = {
 	{
-		.text = "option 0",
+		.text = "option 0: language selection",
+		.is_language_select = 1,
 	},
 	{
 		.text = "option 1",
@@ -76,14 +73,14 @@ struct vb2_menu_item mock_screen_menu_items[] = {
 const struct vb2_screen_info mock_screen_menu = {
 	.id = MOCK_SCREEN_MENU,
 	.name = "mock_screen_menu: screen with 5 options",
-	.num_items = ARRAY_SIZE(mock_screen_menu_items),
-	.items = mock_screen_menu_items,
+	.menu = {
+		.num_items = ARRAY_SIZE(mock_screen_menu_items),
+		.items = mock_screen_menu_items,
+	},
 };
 const struct vb2_screen_info mock_screen_root = {
 	.id = MOCK_SCREEN_ROOT,
 	.name = "mock_screen_root",
-	.num_items = ARRAY_SIZE(mock_empty_menu),
-	.items = mock_empty_menu,
 };
 
 static void screen_state_eq(const struct vb2_screen_state *state,
@@ -115,6 +112,9 @@ static void reset_common_data(void)
 
 	vb2_nv_init(ctx);
 
+	/* For vb2ex_get_locale_count */
+	mock_locale_count = 1;
+
 	/* For check_shutdown_request */
 	mock_shutdown_request = MOCK_IGNORE;
 
@@ -133,6 +133,11 @@ struct vb2_gbb_header *vb2_get_gbb(struct vb2_context *c)
 	return &gbb;
 }
 
+uint32_t vb2ex_get_locale_count(void)
+{
+	return mock_locale_count;
+}
+
 uint32_t VbExIsShutdownRequested(void)
 {
 	if (mock_shutdown_request != MOCK_IGNORE)
@@ -299,6 +304,46 @@ static void change_screen_tests(void)
 		VB2_SUCCESS, "change to screen with init");
 	TEST_EQ(mock_action_called, 1, "  action called once");
 
+	/* Change to screen without init; using default init() */
+	reset_common_data();
+	mock_state->screen = &mock_screen_base;
+	TEST_EQ(vb2_ui_change_screen(&mock_ui_context, MOCK_SCREEN_MENU),
+		VB2_REQUEST_UI_CONTINUE,
+		"change to screen with language selection");
+	screen_state_eq(mock_state, MOCK_SCREEN_MENU,
+			1,	/* Since index 0 is the language selection */
+			0);
+
+	VB2_DEBUG("...done.\n");
+}
+
+static void get_language_menu_tests(void)
+{
+	const struct vb2_menu *menu;
+	const struct vb2_menu_item *items;
+	VB2_DEBUG("Testing get_language_menu...\n");
+
+	/* Only allocate menu items once */
+	reset_common_data();
+	mock_locale_count = 7;
+	menu = get_language_menu(&mock_ui_context);
+	TEST_PTR_NEQ(menu, NULL, "get language menu");
+	TEST_EQ(menu->num_items, 7, "  correct locale count");
+	TEST_PTR_NEQ(menu->items, NULL, "  items not null");
+	items = menu->items;
+
+	menu = get_language_menu(&mock_ui_context);
+	TEST_PTR_NEQ(menu, NULL, "get language menu again");
+	TEST_EQ(menu->num_items, 7, "  correct locale count again");
+	TEST_PTR_EQ(menu->items, items, "  same pointer of items");
+
+	/* Locale count = 0 */
+	reset_common_data();
+	mock_locale_count = 0;
+	menu = get_language_menu(&mock_ui_context);
+	TEST_PTR_NEQ(menu, NULL, "menu not null");
+	TEST_EQ(menu->num_items, 1, "  locale count 1");
+
 	VB2_DEBUG("...done.\n");
 }
 
@@ -307,6 +352,7 @@ int main(void)
 	check_shutdown_request_tests();
 	vb2_ui_change_root_tests();
 	change_screen_tests();
+	get_language_menu_tests();
 
 	return gTestSuccess ? 0 : 255;
 }