vboot: reshuffle legacy UI files

- Use "vboot_ui_legacy" prefix for all legacy UI-related files.

- Merge vboot_display.{c,h} and vboot_ui_legacy_common.c into
  vboot_ui_legacy.{c,h}.

- Move VbDisplayScreen and VbDisplayMenu implementation into
  their respective vboot_ui_legacy_*.c files.

- Update VbCheckDisplayKey to take |screen| argument to avoid
  reading disp_current_screen global variable.

BUG=b:124141368, chromium:968464
TEST=make clean && make runtests
BRANCH=none

Change-Id: I777551e4968ca22282901d22a262a8f2ec849702
Signed-off-by: Joel Kitching <kitching@google.com>
Cq-Depend: chromium:2214615
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/vboot_reference/+/2112322
Tested-by: Joel Kitching <kitching@chromium.org>
Commit-Queue: Joel Kitching <kitching@chromium.org>
Reviewed-by: Joel Kitching <kitching@chromium.org>

16 files changed, 255 insertions, 293 deletions

diff --git a/Makefile b/Makefile
index b35dc913..d69a63a3 100644
--- a/Makefile
+++ b/Makefile
@@ -208,11 +208,6 @@ else
 CFLAGS += -DLEGACY_MENU_UI=0
 endif
 
-# enable all features during local compile (permits testing)
-ifeq (${FIRMWARE_ARCH},)
-DIAGNOSTIC_UI := 1
-endif
-
 # pass DIAGNOSTIC_UI= (or =0) to make to disable feature
 ifneq ($(filter-out 0,${DIAGNOSTIC_UI}),)
 CFLAGS += -DDIAGNOSTIC_UI=1
@@ -386,10 +381,9 @@ FWLIB_SRCS = \
 	firmware/lib/gpt_misc.c \
 	firmware/lib/vboot_api_kernel.c \
 	firmware/lib/vboot_audio.c \
-	firmware/lib/vboot_display.c \
 	firmware/lib/vboot_kernel.c \
+	firmware/lib/vboot_ui_legacy.c \
 	firmware/lib/vboot_ui_legacy_clamshell.c \
-	firmware/lib/vboot_ui_legacy_common.c \
 	firmware/lib/vboot_ui_legacy_menu.c \
 	firmware/lib/vboot_ui_legacy_wilco.c \
 	firmware/lib20/api_kernel.c \
@@ -676,17 +670,17 @@ TEST_OBJS += ${TESTLIB_OBJS}
 
 # And some compiled tests.
 TEST_NAMES = \
-	tests/chromeos_config_tests \
 	tests/cgptlib_test \
+	tests/chromeos_config_tests \
 	tests/sha_benchmark \
 	tests/subprocess_tests \
 	tests/vboot_api_kernel4_tests \
 	tests/vboot_api_kernel_tests \
-	tests/vboot_display_tests \
 	tests/vboot_kernel_tests \
-	tests/vboot_legacy_clamshell_beep_tests \
-	tests/vboot_legacy_clamshell_tests \
-	tests/vboot_legacy_menu_tests \
+	tests/vboot_ui_legacy_clamshell_beep_tests \
+	tests/vboot_ui_legacy_clamshell_tests \
+	tests/vboot_ui_legacy_menu_tests \
+	tests/vboot_ui_legacy_tests \
 	tests/verify_kernel
 
 ifeq (${MOCK_TPM}${TPM2_MODE},)
@@ -1223,11 +1217,11 @@ ifeq (${MOCK_TPM}${TPM2_MODE},)
 endif
 	${RUNTEST} ${BUILD_RUN}/tests/vboot_api_kernel4_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vboot_api_kernel_tests
-	${RUNTEST} ${BUILD_RUN}/tests/vboot_display_tests
 	${RUNTEST} ${BUILD_RUN}/tests/vboot_kernel_tests
-	${RUNTEST} ${BUILD_RUN}/tests/vboot_legacy_clamshell_beep_tests
-	${RUNTEST} ${BUILD_RUN}/tests/vboot_legacy_clamshell_tests
-	${RUNTEST} ${BUILD_RUN}/tests/vboot_legacy_menu_tests
+	${RUNTEST} ${BUILD_RUN}/tests/vboot_ui_legacy_clamshell_beep_tests
+	${RUNTEST} ${BUILD_RUN}/tests/vboot_ui_legacy_clamshell_tests
+	${RUNTEST} ${BUILD_RUN}/tests/vboot_ui_legacy_menu_tests
+	${RUNTEST} ${BUILD_RUN}/tests/vboot_ui_legacy_tests
 
 .PHONY: run2tests
 run2tests: install_for_test
diff --git a/firmware/lib/include/vboot_display.h b/firmware/lib/include/vboot_display.h
deleted file mode 100644
index e4d11872..00000000
--- a/firmware/lib/include/vboot_display.h
+++ /dev/null
@@ -1,22 +0,0 @@
-/* Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
- * Use of this source code is governed by a BSD-style license that can be
- * found in the LICENSE file.
- *
- * Display functions used in kernel selection.
- */
-
-#ifndef VBOOT_REFERENCE_VBOOT_DISPLAY_H_
-#define VBOOT_REFERENCE_VBOOT_DISPLAY_H_
-
-#include "vboot_api.h"
-
-vb2_error_t VbDisplayScreen(struct vb2_context *ctx, uint32_t screen, int force,
-			    const VbScreenData *data);
-vb2_error_t VbDisplayMenu(struct vb2_context *ctx,
-			  uint32_t screen, int force, uint32_t selected_index,
-			  uint32_t disabled_idx_mask);
-vb2_error_t VbDisplayDebugInfo(struct vb2_context *ctx);
-vb2_error_t VbCheckDisplayKey(struct vb2_context *ctx, uint32_t key,
-			      const VbScreenData *data);
-
-#endif  /* VBOOT_REFERENCE_VBOOT_DISPLAY_H_ */
diff --git a/firmware/lib/include/vboot_ui_legacy_common.h b/firmware/lib/include/vboot_ui_legacy.h
index ae081ae1..a0f6f124 100644
--- a/firmware/lib/include/vboot_ui_legacy_common.h
+++ b/firmware/lib/include/vboot_ui_legacy.h
@@ -5,8 +5,19 @@
  * Common code used by both legacy_clamshell_ui and legacy_menu_ui.
  */
 
-#ifndef VBOOT_REFERENCE_VBOOT_UI_COMMON_H_
-#define VBOOT_REFERENCE_VBOOT_UI_COMMON_H_
+#ifndef VBOOT_REFERENCE_VBOOT_UI_LEGACY_H_
+#define VBOOT_REFERENCE_VBOOT_UI_LEGACY_H_
+
+extern const char dev_disable_msg[];
+
+vb2_error_t VbDisplayScreen(struct vb2_context *ctx, uint32_t screen, int force,
+			    const VbScreenData *data);
+vb2_error_t VbDisplayMenu(struct vb2_context *ctx,
+			uint32_t screen, int force, uint32_t selected_index,
+			uint32_t disabled_idx_mask);
+vb2_error_t VbDisplayDebugInfo(struct vb2_context *ctx);
+vb2_error_t VbCheckDisplayKey(struct vb2_context *ctx, uint32_t key,
+			      uint32_t screen, const VbScreenData *data);
 
 #define KEY_DELAY_MS	20	/* Delay between key scans in UI loops */
 
@@ -72,4 +83,4 @@ void vb2_try_altfw(struct vb2_context *ctx, int allowed,
  */
 int vb2_want_shutdown(struct vb2_context *ctx, uint32_t key);
 
-#endif  /* VBOOT_REFERENCE_VBOOT_UI_COMMON_H_ */
+#endif  /* VBOOT_REFERENCE_VBOOT_UI_LEGACY_H_ */
diff --git a/firmware/lib/include/vboot_ui_legacy_menu_private.h b/firmware/lib/include/vboot_ui_legacy_menu_private.h
index 76edd655..9c1a6b48 100644
--- a/firmware/lib/include/vboot_ui_legacy_menu_private.h
+++ b/firmware/lib/include/vboot_ui_legacy_menu_private.h
@@ -10,7 +10,6 @@
 #define VBOOT_REFERENCE_VBOOT_UI_LEGACY_MENU_PRIVATE_H_
 
 #include "2api.h"
-#include "vboot_api.h"
 
 struct vb2_menu_item {
 	const char *text;
diff --git a/firmware/lib/vboot_display.c b/firmware/lib/vboot_ui_legacy.c
index 9461d5d9..f1288418 100644
--- a/firmware/lib/vboot_display.c
+++ b/firmware/lib/vboot_ui_legacy.c
@@ -1,90 +1,24 @@
-/* Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
+/* Copyright 2018 The Chromium OS Authors. All rights reserved.
  * Use of this source code is governed by a BSD-style license that can be
  * found in the LICENSE file.
  *
- * Display functions used in kernel selection.
+ * High-level firmware wrapper API - user interface for RW firmware
  */
 
 #include "2api.h"
 #include "2common.h"
 #include "2misc.h"
 #include "2nvstorage.h"
-#include "2recovery_reasons.h"
-#include "2sha.h"
 #include "2sysincludes.h"
 #include "vboot_api.h"
-#include "vboot_display.h"
 #include "vboot_kernel.h"
-#include "vboot_struct.h"
+#include "vboot_test.h"
+#include "vboot_ui_legacy.h"
 
-static uint32_t disp_current_screen = VB_SCREEN_BLANK;
-static uint32_t disp_current_index = 0;
-static uint32_t disp_disabled_idx_mask = 0;
-
-__attribute__((weak))
-uint32_t vb2ex_get_locale_count(void) {
-	return 0;
-}
-
-__attribute__((weak))
-vb2_error_t VbExGetAltFwIdxMask(void) {
-	return 0;
-}
-
-vb2_error_t VbDisplayScreen(struct vb2_context *ctx, uint32_t screen, int force,
-			    const VbScreenData *data)
-{
-	uint32_t locale;
-
-	/* If requested screen is the same as the current one, we're done. */
-	if (disp_current_screen == screen && !force)
-		return VB2_SUCCESS;
-
-	/* Keep track of the currently displayed screen */
-	disp_current_screen = screen;
-
-	/* Read the locale last saved */
-	locale = vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX);
-
-	return VbExDisplayScreen(screen, locale, data);
-}
-
-vb2_error_t VbDisplayMenu(struct vb2_context *ctx, uint32_t screen, int force,
-			  uint32_t selected_index, uint32_t disabled_idx_mask)
-{
-	uint32_t locale;
-	uint32_t redraw_base_screen = 0;
-
-	/*
-	 * If requested screen/selected_index is the same as the current one,
-	 * we're done.
-	 */
-	if (disp_current_screen == screen &&
-	    disp_current_index == selected_index &&
-	    !force)
-		return VB2_SUCCESS;
-
-	/*
-	 * If current screen is not the same, make sure we redraw the base
-	 * screen as well to avoid having artifacts from the menu.
-	 */
-	if (disp_current_screen != screen || force)
-		redraw_base_screen = 1;
-
-	/*
-	 * Keep track of the currently displayed screen and
-	 * selected_index
-	 */
-	disp_current_screen = screen;
-	disp_current_index = selected_index;
-	disp_disabled_idx_mask = disabled_idx_mask;
-
-	/* Read the locale last saved */
-	locale = vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX);
-
-	return VbExDisplayMenu(screen, locale, selected_index,
-			       disabled_idx_mask, redraw_base_screen);
-}
+const char dev_disable_msg[] =
+	"Developer mode is disabled on this device by system policy.\n"
+	"For more information, see http://dev.chromium.org/chromium-os/fwmp\n"
+	"\n";
 
 static void Uint8ToString(char *buf, uint8_t val)
 {
@@ -225,7 +159,7 @@ vb2_error_t VbDisplayDebugInfo(struct vb2_context *ctx)
 }
 
 vb2_error_t VbCheckDisplayKey(struct vb2_context *ctx, uint32_t key,
-			    const VbScreenData *data)
+			      uint32_t screen, const VbScreenData *data)
 {
 	uint32_t loc = 0;
 	uint32_t count = 0;
@@ -236,7 +170,7 @@ vb2_error_t VbCheckDisplayKey(struct vb2_context *ctx, uint32_t key,
 		return VbDisplayDebugInfo(ctx);
 	case VB_KEY_ESC:
 		/* Force redraw current screen (to clear Tab debug output) */
-		return VbDisplayScreen(ctx, disp_current_screen, 1, data);
+		return VbDisplayScreen(ctx, screen, 1, data);
 	case VB_KEY_LEFT:
 	case VB_KEY_RIGHT:
 	case VB_KEY_UP:
@@ -265,8 +199,107 @@ vb2_error_t VbCheckDisplayKey(struct vb2_context *ctx, uint32_t key,
 			 vb2ex_commit_data(ctx);
 
 		/* Force redraw of current screen */
-		return VbDisplayScreen(ctx, disp_current_screen, 1, data);
+		return VbDisplayScreen(ctx, screen, 1, data);
 	}
 
 	return VB2_SUCCESS;
 }
+
+static enum {
+	POWER_BUTTON_HELD_SINCE_BOOT = 0,
+	POWER_BUTTON_RELEASED,
+	POWER_BUTTON_PRESSED,  /* Must have been previously released */
+} power_button_state;
+
+void vb2_reset_power_button(void) {
+	power_button_state = POWER_BUTTON_HELD_SINCE_BOOT;
+}
+
+void vb2_error_beep(enum vb2_beep_type beep)
+{
+	switch (beep) {
+	case VB_BEEP_FAILED:
+		VbExBeep(250, 200);
+		break;
+	default:
+	case VB_BEEP_NOT_ALLOWED:
+		VbExBeep(120, 400);
+		VbExSleepMs(120);
+		VbExBeep(120, 400);
+		break;
+	}
+}
+
+void vb2_error_notify(const char *print_msg,
+		      const char *log_msg,
+		      enum vb2_beep_type beep)
+{
+	if (print_msg)
+		VbExDisplayDebugInfo(print_msg, 0);
+	if (!log_msg)
+		log_msg = print_msg;
+	if (log_msg)
+		VB2_DEBUG(log_msg);
+	vb2_error_beep(beep);
+}
+
+void vb2_error_no_altfw(void)
+{
+	VB2_DEBUG("Legacy boot is disabled\n");
+	VbExDisplayDebugInfo("WARNING: Booting legacy BIOS has not been "
+			     "enabled. Refer to the developer-mode "
+			     "documentation for details.\n", 0);
+	vb2_error_beep(VB_BEEP_NOT_ALLOWED);
+}
+
+void vb2_try_altfw(struct vb2_context *ctx, int allowed,
+		   enum VbAltFwIndex_t altfw_num)
+{
+	if (!allowed) {
+		vb2_error_no_altfw();
+		return;
+	}
+
+	if (vb2ex_commit_data(ctx)) {
+		vb2_error_notify("Error committing data on legacy boot.\n",
+				 NULL, VB_BEEP_FAILED);
+		return;
+	}
+
+	/* Will not return if successful */
+	VbExLegacy(altfw_num);
+
+	vb2_error_notify("Legacy boot failed. Missing BIOS?\n", NULL,
+			 VB_BEEP_FAILED);
+}
+
+int vb2_want_shutdown(struct vb2_context *ctx, uint32_t key)
+{
+	struct vb2_gbb_header *gbb = vb2_get_gbb(ctx);
+	uint32_t shutdown_request = VbExIsShutdownRequested();
+
+	/*
+	 * Ignore power button push until after we have seen it released.
+	 * This avoids shutting down immediately if the power button is still
+	 * being held on startup. After we've recognized a valid power button
+	 * push then don't report the event until after the button is released.
+	 */
+	if (shutdown_request & VB_SHUTDOWN_REQUEST_POWER_BUTTON) {
+		shutdown_request &= ~VB_SHUTDOWN_REQUEST_POWER_BUTTON;
+		if (power_button_state == POWER_BUTTON_RELEASED)
+			power_button_state = POWER_BUTTON_PRESSED;
+	} else {
+		if (power_button_state == POWER_BUTTON_PRESSED)
+			shutdown_request |= VB_SHUTDOWN_REQUEST_POWER_BUTTON;
+		power_button_state = POWER_BUTTON_RELEASED;
+	}
+
+	if (key == VB_BUTTON_POWER_SHORT_PRESS)
+		shutdown_request |= VB_SHUTDOWN_REQUEST_POWER_BUTTON;
+
+	/* If desired, ignore shutdown request due to lid closure. */
+	if (gbb->flags & VB2_GBB_FLAG_DISABLE_LID_SHUTDOWN)
+		shutdown_request &= ~VB_SHUTDOWN_REQUEST_LID_CLOSED;
+
+	return shutdown_request;
+}
diff --git a/firmware/lib/vboot_ui_legacy_clamshell.c b/firmware/lib/vboot_ui_legacy_clamshell.c
index 3b2c91dc..e91e090c 100644
--- a/firmware/lib/vboot_ui_legacy_clamshell.c
+++ b/firmware/lib/vboot_ui_legacy_clamshell.c
@@ -16,13 +16,33 @@
 #include "vb2_common.h"
 #include "vboot_api.h"
 #include "vboot_audio.h"
-#include "vboot_display.h"
 #include "vboot_kernel.h"
 #include "vboot_struct.h"
 #include "vboot_test.h"
-#include "vboot_ui_legacy_common.h"
+#include "vboot_ui_legacy.h"
 #include "vboot_ui_legacy_wilco.h"
 
+static uint32_t disp_current_screen = VB_SCREEN_BLANK;
+
+test_mockable
+vb2_error_t VbDisplayScreen(struct vb2_context *ctx, uint32_t screen, int force,
+			    const VbScreenData *data)
+{
+	uint32_t locale;
+
+	/* If requested screen is the same as the current one, we're done. */
+	if (disp_current_screen == screen && !force)
+		return VB2_SUCCESS;
+
+	/* Keep track of the currently displayed screen */
+	disp_current_screen = screen;
+
+	/* Read the locale last saved */
+	locale = vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX);
+
+	return VbExDisplayScreen(screen, locale, data);
+}
+
 static vb2_error_t VbTryUsb(struct vb2_context *ctx)
 {
 	int retval = VbTryLoadKernel(ctx, VB_DISK_FLAG_REMOVABLE);
@@ -108,7 +128,7 @@ int VbUserConfirms(struct vb2_context *ctx, uint32_t confirm_flags)
 					return 1;
 				}
 			}
-			VbCheckDisplayKey(ctx, key, NULL);
+			VbCheckDisplayKey(ctx, key, disp_current_screen, NULL);
 		}
 		VbExSleepMs(KEY_DELAY_MS);
 	} while (!shutdown_requested);
@@ -159,7 +179,7 @@ static vb2_error_t vb2_altfw_ui(struct vb2_context *ctx)
 			break;
 		default:
 			VB2_DEBUG("developer UI - pressed key %#x\n", key);
-			VbCheckDisplayKey(ctx, key, NULL);
+			VbCheckDisplayKey(ctx, key, disp_current_screen, NULL);
 			break;
 		}
 		VbExSleepMs(KEY_DELAY_MS);
@@ -171,11 +191,6 @@ static vb2_error_t vb2_altfw_ui(struct vb2_context *ctx)
 	return 0;
 }
 
-static const char dev_disable_msg[] =
-	"Developer mode is disabled on this device by system policy.\n"
-	"For more information, see http://dev.chromium.org/chromium-os/fwmp\n"
-	"\n";
-
 static vb2_error_t vb2_developer_ui(struct vb2_context *ctx)
 {
 	struct vb2_shared_data *sd = vb2_get_sd(ctx);
@@ -384,7 +399,7 @@ static vb2_error_t vb2_developer_ui(struct vb2_context *ctx)
 			break;
 		default:
 			VB2_DEBUG("developer UI - pressed key %#x\n", key);
-			VbCheckDisplayKey(ctx, key, NULL);
+			VbCheckDisplayKey(ctx, key, disp_current_screen, NULL);
 			break;
 		}
 
@@ -443,7 +458,7 @@ static vb2_error_t recovery_ui(struct vb2_context *ctx)
 		VB2_DEBUG("recovery UI - waiting for manual recovery\n");
 		while (1) {
 			key = VbExKeyboardRead();
-			VbCheckDisplayKey(ctx, key, NULL);
+			VbCheckDisplayKey(ctx, key, disp_current_screen, NULL);
 			if (vb2_want_shutdown(ctx, key))
 				return VB2_REQUEST_SHUTDOWN;
 			else if ((retval =
@@ -512,7 +527,7 @@ static vb2_error_t recovery_ui(struct vb2_context *ctx)
 			   VB2_SUCCESS) {
 			return retval;
 		} else {
-			VbCheckDisplayKey(ctx, key, NULL);
+			VbCheckDisplayKey(ctx, key, disp_current_screen, NULL);
 		}
 		if (vb2_want_shutdown(ctx, key))
 			return VB2_REQUEST_SHUTDOWN;
diff --git a/firmware/lib/vboot_ui_legacy_common.c b/firmware/lib/vboot_ui_legacy_common.c
deleted file mode 100644
index 948b63e2..00000000
--- a/firmware/lib/vboot_ui_legacy_common.c
+++ /dev/null
@@ -1,113 +0,0 @@
-/* Copyright 2018 The Chromium OS Authors. All rights reserved.
- * Use of this source code is governed by a BSD-style license that can be
- * found in the LICENSE file.
- *
- * High-level firmware wrapper API - user interface for RW firmware
- */
-
-#include "2common.h"
-#include "2misc.h"
-#include "2sysincludes.h"
-#include "vboot_api.h"
-#include "vboot_kernel.h"
-#include "vboot_test.h"
-#include "vboot_ui_legacy_common.h"
-
-static enum {
-	POWER_BUTTON_HELD_SINCE_BOOT = 0,
-	POWER_BUTTON_RELEASED,
-	POWER_BUTTON_PRESSED,  /* Must have been previously released */
-} power_button_state;
-
-void vb2_reset_power_button(void) {
-	power_button_state = POWER_BUTTON_HELD_SINCE_BOOT;
-}
-
-void vb2_error_beep(enum vb2_beep_type beep)
-{
-	switch (beep) {
-	case VB_BEEP_FAILED:
-		VbExBeep(250, 200);
-		break;
-	default:
-	case VB_BEEP_NOT_ALLOWED:
-		VbExBeep(120, 400);
-		VbExSleepMs(120);
-		VbExBeep(120, 400);
-		break;
-	}
-}
-
-void vb2_error_notify(const char *print_msg,
-		      const char *log_msg,
-		      enum vb2_beep_type beep)
-{
-	if (print_msg)
-		VbExDisplayDebugInfo(print_msg, 0);
-	if (!log_msg)
-		log_msg = print_msg;
-	if (log_msg)
-		VB2_DEBUG(log_msg);
-	vb2_error_beep(beep);
-}
-
-void vb2_error_no_altfw(void)
-{
-	VB2_DEBUG("Legacy boot is disabled\n");
-	VbExDisplayDebugInfo("WARNING: Booting legacy BIOS has not been "
-			     "enabled. Refer to the developer-mode "
-			     "documentation for details.\n", 0);
-	vb2_error_beep(VB_BEEP_NOT_ALLOWED);
-}
-
-void vb2_try_altfw(struct vb2_context *ctx, int allowed,
-		   enum VbAltFwIndex_t altfw_num)
-{
-	if (!allowed) {
-		vb2_error_no_altfw();
-		return;
-	}
-
-	if (vb2ex_commit_data(ctx)) {
-		vb2_error_notify("Error committing data on legacy boot.\n",
-				 NULL, VB_BEEP_FAILED);
-		return;
-	}
-
-	/* Will not return if successful */
-	VbExLegacy(altfw_num);
-
-	vb2_error_notify("Legacy boot failed. Missing BIOS?\n", NULL,
-			 VB_BEEP_FAILED);
-}
-
-int vb2_want_shutdown(struct vb2_context *ctx, uint32_t key)
-{
-	struct vb2_gbb_header *gbb = vb2_get_gbb(ctx);
-	uint32_t shutdown_request = VbExIsShutdownRequested();
-
-	/*
-	 * Ignore power button push until after we have seen it released.
-	 * This avoids shutting down immediately if the power button is still
-	 * being held on startup. After we've recognized a valid power button
-	 * push then don't report the event until after the button is released.
-	 */
-	if (shutdown_request & VB_SHUTDOWN_REQUEST_POWER_BUTTON) {
-		shutdown_request &= ~VB_SHUTDOWN_REQUEST_POWER_BUTTON;
-		if (power_button_state == POWER_BUTTON_RELEASED)
-			power_button_state = POWER_BUTTON_PRESSED;
-	} else {
-		if (power_button_state == POWER_BUTTON_PRESSED)
-			shutdown_request |= VB_SHUTDOWN_REQUEST_POWER_BUTTON;
-		power_button_state = POWER_BUTTON_RELEASED;
-	}
-
-	if (key == VB_BUTTON_POWER_SHORT_PRESS)
-		shutdown_request |= VB_SHUTDOWN_REQUEST_POWER_BUTTON;
-
-	/* If desired, ignore shutdown request due to lid closure. */
-	if (gbb->flags & VB2_GBB_FLAG_DISABLE_LID_SHUTDOWN)
-		shutdown_request &= ~VB_SHUTDOWN_REQUEST_LID_CLOSED;
-
-	return shutdown_request;
-}
diff --git a/firmware/lib/vboot_ui_legacy_menu.c b/firmware/lib/vboot_ui_legacy_menu.c
index 26de4871..bcdd41ce 100644
--- a/firmware/lib/vboot_ui_legacy_menu.c
+++ b/firmware/lib/vboot_ui_legacy_menu.c
@@ -15,25 +15,61 @@
 #include "vb2_common.h"
 #include "vboot_api.h"
 #include "vboot_audio.h"
-#include "vboot_display.h"
 #include "vboot_kernel.h"
 #include "vboot_struct.h"
-#include "vboot_ui_legacy_common.h"
+#include "vboot_ui_legacy.h"
 #include "vboot_ui_legacy_menu_private.h"
 
-static const char dev_disable_msg[] =
-	"Developer mode is disabled on this device by system policy.\n"
-	"For more information, see http://dev.chromium.org/chromium-os/fwmp\n"
-	"\n";
-
 static VB_MENU current_menu, prev_menu;
-static int current_menu_idx, disabled_idx_mask, usb_nogood, force_redraw;
+static int current_menu_idx, cur_disabled_idx_mask, usb_nogood, force_redraw;
 static uint32_t default_boot;
 static uint32_t disable_dev_boot;
 static uint32_t altfw_allowed;
 static struct vb2_menu menus[];
 static const char no_legacy[] = "Legacy boot failed. Missing BIOS?\n";
 
+static uint32_t disp_current_screen = VB_SCREEN_BLANK;
+static uint32_t disp_current_index = 0;
+static uint32_t disp_disabled_idx_mask = 0;
+
+test_mockable
+vb2_error_t VbDisplayMenu(struct vb2_context *ctx, uint32_t screen, int force,
+			  uint32_t selected_index, uint32_t disabled_idx_mask)
+{
+	uint32_t locale;
+	uint32_t redraw_base_screen = 0;
+
+	/*
+	 * If requested screen/selected_index is the same as the current one,
+	 * we're done.
+	 */
+	if (disp_current_screen == screen &&
+	    disp_current_index == selected_index &&
+	    !force)
+		return VB2_SUCCESS;
+
+	/*
+	 * If current screen is not the same, make sure we redraw the base
+	 * screen as well to avoid having artifacts from the menu.
+	 */
+	if (disp_current_screen != screen || force)
+		redraw_base_screen = 1;
+
+	/*
+	 * Keep track of the currently displayed screen and
+	 * selected_index
+	 */
+	disp_current_screen = screen;
+	disp_current_index = selected_index;
+	disp_disabled_idx_mask = disabled_idx_mask;
+
+	/* Read the locale last saved */
+	locale = vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX);
+
+	return VbExDisplayMenu(screen, locale, selected_index,
+			       disabled_idx_mask, redraw_base_screen);
+}
+
 /**
  * Checks GBB flags against VbExIsShutdownRequested() shutdown request to
  * determine if a shutdown is required.
@@ -61,7 +97,7 @@ static int VbWantShutdownMenu(struct vb2_context *ctx)
 /* (Re-)Draw the menu identified by current_menu[_idx] to the screen. */
 static vb2_error_t vb2_draw_current_screen(struct vb2_context *ctx) {
 	vb2_error_t ret = VbDisplayMenu(ctx, menus[current_menu].screen,
-			force_redraw, current_menu_idx, disabled_idx_mask);
+			force_redraw, current_menu_idx, cur_disabled_idx_mask);
 	force_redraw = 0;
 	return ret;
 }
@@ -69,7 +105,7 @@ static vb2_error_t vb2_draw_current_screen(struct vb2_context *ctx) {
 /* Flash the screen to black to catch user awareness, then redraw menu. */
 static void vb2_flash_screen(struct vb2_context *ctx)
 {
-	VbDisplayScreen(ctx, VB_SCREEN_BLANK, 0, NULL);
+	VbDisplayMenu(ctx, VB_SCREEN_BLANK, 0, 0, 0);
 	VbExSleepMs(50);
 	vb2_draw_current_screen(ctx);
 }
@@ -97,25 +133,25 @@ static void vb2_change_menu(VB_MENU new_current_menu,
 	prev_menu = current_menu;
 	current_menu = new_current_menu;
 
-	/* Reconfigure disabled_idx_mask for the new menu */
-	disabled_idx_mask = 0;
+	/* Reconfigure cur_disabled_idx_mask for the new menu */
+	cur_disabled_idx_mask = 0;
 	/* Disable Network Boot Option */
 	if (current_menu == VB_MENU_DEV)
-		disabled_idx_mask |= 1 << VB_DEV_NETWORK;
+		cur_disabled_idx_mask |= 1 << VB_DEV_NETWORK;
 	/* Disable cancel option if enterprise disabled dev mode */
 	if (current_menu == VB_MENU_TO_NORM &&
 	    disable_dev_boot == 1)
-		disabled_idx_mask |= 1 << VB_TO_NORM_CANCEL;
+		cur_disabled_idx_mask |= 1 << VB_TO_NORM_CANCEL;
 
 	/* Enable menu items for the selected bootloaders */
 	if (current_menu == VB_MENU_ALT_FW) {
-		disabled_idx_mask = ~(VbExGetAltFwIdxMask() >> 1);
+		cur_disabled_idx_mask = ~(VbExGetAltFwIdxMask() >> 1);
 
 		/* Make sure 'cancel' is shown even with an invalid mask */
-		disabled_idx_mask &= (1 << VB_ALTFW_COUNT) - 1;
+		cur_disabled_idx_mask &= (1 << VB_ALTFW_COUNT) - 1;
 	}
 	/* We assume that there is at least one enabled item */
-	while ((1 << new_current_menu_idx) & disabled_idx_mask)
+	while ((1 << new_current_menu_idx) & cur_disabled_idx_mask)
 		new_current_menu_idx++;
 	if (new_current_menu_idx < menus[current_menu].size)
 		current_menu_idx = new_current_menu_idx;
@@ -392,7 +428,7 @@ static vb2_error_t power_off_action(struct vb2_context *ctx)
 
 /**
  * Updates current_menu_idx upon an up/down key press, taking into
- * account disabled indices (from disabled_idx_mask).  The cursor
+ * account disabled indices (from cur_disabled_idx_mask).  The cursor
  * will not wrap, meaning that we block on the 0 or max index when
  * we hit the ends of the menu.
  *
@@ -408,7 +444,7 @@ static void vb2_update_selection(uint32_t key) {
 	case VB_KEY_UP:
 		idx = current_menu_idx - 1;
 		while (idx >= 0 &&
-		       ((1 << idx) & disabled_idx_mask))
+		       ((1 << idx) & cur_disabled_idx_mask))
 		  idx--;
 		/* Only update if idx is valid */
 		if (idx >= 0)
@@ -418,7 +454,7 @@ static void vb2_update_selection(uint32_t key) {
 	case VB_KEY_DOWN:
 		idx = current_menu_idx + 1;
 		while (idx < menus[current_menu].size &&
-		       ((1 << idx) & disabled_idx_mask))
+		       ((1 << idx) & cur_disabled_idx_mask))
 		  idx++;
 		/* Only update if idx is valid */
 		if (idx < menus[current_menu].size)
@@ -823,7 +859,7 @@ vb2_error_t VbBootDeveloperLegacyMenu(struct vb2_context *ctx)
 	if (VB2_SUCCESS != retval)
 		return retval;
 	retval = vb2_developer_legacy_menu(ctx);
-	VbDisplayScreen(ctx, VB_SCREEN_BLANK, 0, NULL);
+	VbDisplayMenu(ctx, VB_SCREEN_BLANK, 0, 0, 0);
 	return retval;
 }
 
diff --git a/firmware/lib/vboot_ui_legacy_wilco.c b/firmware/lib/vboot_ui_legacy_wilco.c
index 4cba2857..fc379ee4 100644
--- a/firmware/lib/vboot_ui_legacy_wilco.c
+++ b/firmware/lib/vboot_ui_legacy_wilco.c
@@ -2,16 +2,14 @@
  * Use of this source code is governed by a BSD-style license that can be
  * found in the LICENSE file.
  *
- * Common code used by both vboot_ui and Wilco-specific
- * feature support for vboot_ui
+ * Code used by vboot_ui_legacy_clamshell for Wilco-specific features.
  */
 
 #include "2common.h"
 #include "2nvstorage.h"
 #include "2sysincludes.h"
 #include "vboot_api.h"
-#include "vboot_display.h"
-#include "vboot_ui_legacy_common.h"
+#include "vboot_ui_legacy.h"
 #include "vboot_ui_legacy_wilco.h"
 
 static inline int is_vowel(uint32_t key)
@@ -107,7 +105,8 @@ static vb2_error_t vb2_enter_vendor_data_ui(struct vb2_context *ctx,
 			break;
 		default:
 			VB2_DEBUG("Vendor Data UI - pressed key %#x\n", key);
-			VbCheckDisplayKey(ctx, key, &data);
+			VbCheckDisplayKey(ctx, key, VB_SCREEN_SET_VENDOR_DATA,
+					  &data);
 			break;
 		}
 		VbExSleepMs(KEY_DELAY_MS);
@@ -195,7 +194,8 @@ static vb2_error_t vb2_confirm_vendor_data_ui(struct vb2_context *ctx,
 		default:
 			VB2_DEBUG("Confirm Vendor Data UI - pressed "
 				  "key %#x\n", key_confirm);
-			VbCheckDisplayKey(ctx, key_confirm, data);
+			VbCheckDisplayKey(ctx, key_confirm,
+					  VB_SCREEN_CONFIRM_VENDOR_DATA, data);
 			break;
 		}
 		VbExSleepMs(KEY_DELAY_MS);
@@ -261,13 +261,13 @@ vb2_error_t vb2_vendor_data_ui(struct vb2_context *ctx)
 			break;
 		default:
 			break;
-        }
+		}
 	} while (1);
-    return VB2_SUCCESS;
+	return VB2_SUCCESS;
 }
 
-vb2_error_t vb2_check_diagnostic_key(struct vb2_context *ctx,
-				     uint32_t key) {
+vb2_error_t vb2_check_diagnostic_key(struct vb2_context *ctx, uint32_t key)
+{
 	if (DIAGNOSTIC_UI && (key == VB_KEY_CTRL('C') || key == VB_KEY_F(12))) {
 		VB2_DEBUG("Diagnostic mode requested, rebooting\n");
 		vb2_nv_set(ctx, VB2_NV_DIAG_REQUEST, 1);
@@ -330,7 +330,8 @@ vb2_error_t vb2_diagnostics_ui(struct vb2_context *ctx)
 		default:
 			VB2_DEBUG("vb2_diagnostics_ui() - pressed key %#x\n",
 				  key);
-			VbCheckDisplayKey(ctx, key, NULL);
+			VbCheckDisplayKey(ctx, key, VB_SCREEN_CONFIRM_DIAG,
+					  NULL);
 			break;
 		}
 		if (VbExGetTimer() - start_time_us >= 30 * VB_USEC_PER_SEC) {
diff --git a/firmware/stub/vboot_api_stub.c b/firmware/stub/vboot_api_stub.c
index 6d854557..f59c4009 100644
--- a/firmware/stub/vboot_api_stub.c
+++ b/firmware/stub/vboot_api_stub.c
@@ -44,6 +44,16 @@ vb2_error_t VbExDisplayDebugInfo(const char *info_str, int full_info)
 	return VB2_SUCCESS;
 }
 
+uint32_t vb2ex_get_locale_count(void)
+{
+	return 0;
+}
+
+vb2_error_t VbExGetAltFwIdxMask(void)
+{
+	return 0;
+}
+
 uint32_t VbExKeyboardRead(void)
 {
 	return 0;
diff --git a/tests/vb2_auxfw_sync_tests.c b/tests/vb2_auxfw_sync_tests.c
index 75b139be..358f706b 100644
--- a/tests/vb2_auxfw_sync_tests.c
+++ b/tests/vb2_auxfw_sync_tests.c
@@ -17,9 +17,9 @@
 #include "load_kernel_fw.h"
 #include "test_common.h"
 #include "vboot_audio.h"
-#include "vboot_display.h"
 #include "vboot_kernel.h"
 #include "vboot_struct.h"
+#include "vboot_ui_legacy.h"
 
 /* Mock data */
 static struct vb2_context *ctx;
diff --git a/tests/vb2_ec_sync_tests.c b/tests/vb2_ec_sync_tests.c
index 25f047d1..51e3dfb9 100644
--- a/tests/vb2_ec_sync_tests.c
+++ b/tests/vb2_ec_sync_tests.c
@@ -14,9 +14,9 @@
 #include "load_kernel_fw.h"
 #include "test_common.h"
 #include "vboot_audio.h"
-#include "vboot_display.h"
 #include "vboot_kernel.h"
 #include "vboot_struct.h"
+#include "vboot_ui_legacy.h"
 
 /* Mock data */
 static int ec_ro_updated;
diff --git a/tests/vboot_legacy_clamshell_beep_tests.c b/tests/vboot_ui_legacy_clamshell_beep_tests.c
index a8c00d78..9d0d12e7 100644
--- a/tests/vboot_legacy_clamshell_beep_tests.c
+++ b/tests/vboot_ui_legacy_clamshell_beep_tests.c
@@ -19,7 +19,6 @@
 #include "host_common.h"
 #include "load_kernel_fw.h"
 #include "test_common.h"
-#include "vboot_display.h"
 #include "vboot_kernel.h"
 #include "vboot_struct.h"
 
diff --git a/tests/vboot_legacy_clamshell_tests.c b/tests/vboot_ui_legacy_clamshell_tests.c
index 77944ba8..5d973e66 100644
--- a/tests/vboot_legacy_clamshell_tests.c
+++ b/tests/vboot_ui_legacy_clamshell_tests.c
@@ -15,11 +15,10 @@
 #include "test_common.h"
 #include "tss_constants.h"
 #include "vboot_audio.h"
-#include "vboot_display.h"
 #include "vboot_kernel.h"
 #include "vboot_struct.h"
 #include "vboot_test.h"
-#include "vboot_ui_legacy_common.h"
+#include "vboot_ui_legacy.h"
 
 /* Mock data */
 static LoadKernelParams lkp;
@@ -565,7 +564,7 @@ static void VbBootDevTest(void)
 	mock_keypress[0] = VB_KEY_UP;
 	vbtlk_expect_fixed = 1;
 	TEST_EQ(VbBootDeveloperLegacyClamshell(ctx), VB2_ERROR_MOCK,
-                "Up arrow");
+		"Up arrow");
 
 	/* Shutdown requested in loop */
 	ResetMocks();
@@ -624,7 +623,7 @@ static void VbBootDevTest(void)
 	mock_keypress[1] = VB_KEY_ENTER;
 	vbtlk_expect_fixed = 1;
 	TEST_EQ(VbBootDeveloperLegacyClamshell(ctx), VB2_ERROR_MOCK,
-                "Enter ignored");
+		"Enter ignored");
 
 	/* Enter does if GBB flag set */
 	ResetMocks();
@@ -786,7 +785,7 @@ static void VbBootDevTest(void)
 	mock_keypress[0] = VB_KEY_CTRL('U');
 	vbtlk_expect_fixed = 1;
 	TEST_EQ(VbBootDeveloperLegacyClamshell(ctx), VB2_ERROR_MOCK,
-                "Ctrl+U normal");
+		"Ctrl+U normal");
 
 	/* Ctrl+U enabled, with good USB boot */
 	ResetMocks();
diff --git a/tests/vboot_legacy_menu_tests.c b/tests/vboot_ui_legacy_menu_tests.c
index 4a0dfaa0..6fa2d404 100644
--- a/tests/vboot_legacy_menu_tests.c
+++ b/tests/vboot_ui_legacy_menu_tests.c
@@ -16,10 +16,10 @@
 #include "tss_constants.h"
 #include "vboot_api.h"
 #include "vboot_audio.h"
-#include "vboot_display.h"
 #include "vboot_kernel.h"
 #include "vboot_struct.h"
 #include "vboot_test.h"
+#include "vboot_ui_legacy.h"
 #include "vboot_ui_legacy_menu_private.h"
 
 /* Mock data */
diff --git a/tests/vboot_display_tests.c b/tests/vboot_ui_legacy_tests.c
index c902614a..04da02e5 100644
--- a/tests/vboot_display_tests.c
+++ b/tests/vboot_ui_legacy_tests.c
@@ -17,8 +17,8 @@
 #include "2sysincludes.h"
 #include "host_common.h"
 #include "test_common.h"
-#include "vboot_display.h"
 #include "vboot_kernel.h"
+#include "vboot_ui_legacy.h"
 
 /* Mock data */
 static char debug_info[4096];
@@ -80,28 +80,28 @@ static void DebugInfoTest(void)
 static void DisplayKeyTest(void)
 {
 	ResetMocks();
-	VbCheckDisplayKey(ctx, 'q', NULL);
+	VbCheckDisplayKey(ctx, 'q', 0, NULL);
 	TEST_EQ(*debug_info, '\0', "DisplayKey q = does nothing");
 
 	ResetMocks();
-	VbCheckDisplayKey(ctx, '\t', NULL);
+	VbCheckDisplayKey(ctx, '\t', 0, NULL);
 	TEST_NEQ(*debug_info, '\0', "DisplayKey tab = display");
 
 	/* Toggle localization */
 	ResetMocks();
 	vb2_nv_set(ctx, VB2_NV_LOCALIZATION_INDEX, 0);
-	VbCheckDisplayKey(ctx, VB_KEY_DOWN, NULL);
+	VbCheckDisplayKey(ctx, VB_KEY_DOWN, 0, NULL);
 	TEST_EQ(vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX), 2,
 		"DisplayKey up");
-	VbCheckDisplayKey(ctx, VB_KEY_LEFT, NULL);
+	VbCheckDisplayKey(ctx, VB_KEY_LEFT, 0, NULL);
 	vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX);
 	TEST_EQ(vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX), 1,
 		"DisplayKey left");
-	VbCheckDisplayKey(ctx, VB_KEY_RIGHT, NULL);
+	VbCheckDisplayKey(ctx, VB_KEY_RIGHT, 0, NULL);
 	vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX);
 	TEST_EQ(vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX), 2,
 		"DisplayKey right");
-	VbCheckDisplayKey(ctx, VB_KEY_UP, NULL);
+	VbCheckDisplayKey(ctx, VB_KEY_UP, 0, NULL);
 	vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX);
 	TEST_EQ(vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX), 0,
 		"DisplayKey up");
@@ -110,7 +110,7 @@ static void DisplayKeyTest(void)
 	ResetMocks();
 	vb2_nv_set(ctx, VB2_NV_LOCALIZATION_INDEX, 1);
 	mock_localization_count = 0;
-	VbCheckDisplayKey(ctx, VB_KEY_UP, NULL);
+	VbCheckDisplayKey(ctx, VB_KEY_UP, 0, NULL);
 	TEST_EQ(vb2_nv_get(ctx, VB2_NV_LOCALIZATION_INDEX), 0,
 		"DisplayKey invalid");
 }