vboot: add diagnostic mode

This change adds diagnostic mode.  When enabled for a board (based on
defconfig in depthcharge) the user can press Ctrl-C or F12 at a recovery
mode screen, at which point an nv bit is set and the system reboots.
Upon reboot, if the nv bit is set then the user is prompted to confirm
launch of the diagnostic rom via the power button.  If user confirms
then the diagnostic payload is verified and run (if verify fails or
payload doesn't run then a recovery reason is recorded and system
reboots to recovery mode).  If the user does not confirm then the system
reboots.

BUG=b:124358784
BRANCH=None
TEST=Locally built and flashed using change that enabled feature for atlas
and set to use payload 2 (tianocore) rather than 5 (diagnostic).  Confirmed
that Ctrl-C is functional or not based on defconfig and that Ctrl-C
sets NV bit and reboots.  Confirmed that NV bit can be set and
queried via crossystem. Confirmed that during boot confirmation screen
appears or not based on NV bit.  Confirmed that pressing power button
caused payload to be verified and run.  Confirmed that non-matching hash
(build configured to use sha1 rather than sha256) caused payload to not
be run and system reboot to recovery.  Confirmed that Esc or timeout
caused system to reboot.
CQ-DEPEND=CL:1471056

Change-Id: I8979d4eeb443bf64b727ee86a814c46d1d27ff37
Signed-off-by: Matt Delco <delco@google.com>
Reviewed-on: https://chromium-review.googlesource.com/1470723
Reviewed-by: Julius Werner <jwerner@chromium.org>

8 files changed, 429 insertions, 5 deletions

diff --git a/Makefile b/Makefile
index 5b6ea76f..47918c95 100644
--- a/Makefile
+++ b/Makefile
@@ -191,6 +191,18 @@ ifneq (${TPM2_MODE},)
 CFLAGS += -DTPM2_MODE
 endif
 
+# enable all features during local compile (permits testing)
+ifeq (${CHROMEOS_ENVIRONMENT},1)
+DIAGNOSTIC_UI := 1
+endif
+
+# pass DIAGNOSTIC_UI= (or =0) to make to disable feature
+ifneq ($(filter-out 0,${DIAGNOSTIC_UI}),)
+CFLAGS += -DDIAGNOSTIC_UI=1
+else
+CFLAGS += -DDIAGNOSTIC_UI=0
+endif
+
 # NOTE: We don't use these files but they are useful for other packages to
 # query about required compiling/linking flags.
 PC_IN_FILES = vboot_host.pc.in
diff --git a/firmware/2lib/include/2recovery_reasons.h b/firmware/2lib/include/2recovery_reasons.h
index a35892c2..f8e6dbe9 100644
--- a/firmware/2lib/include/2recovery_reasons.h
+++ b/firmware/2lib/include/2recovery_reasons.h
@@ -219,6 +219,12 @@ enum vb2_nv_recovery {
 	/* Recovery hash space lock error in RO firmware */
 	VB2_RECOVERY_RO_TPM_REC_HASH_L_ERROR = 0x5f,
 
+	/* Failed to disable the TPM [prior to running untrusted code] */
+	VB2_RECOVERY_TPM_DISABLE_FAILED = 0x60,
+
+	/* Alt FW Failed hash verification */
+	VB2_RECOVERY_ALTFW_HASH_FAILED = 0x61,
+
 	/* Unspecified/unknown error in rewritable firmware */
 	VB2_RECOVERY_RW_UNSPECIFIED  = 0x7f,
 
diff --git a/firmware/include/vboot_api.h b/firmware/include/vboot_api.h
index 9d8c8a0d..6202fe91 100644
--- a/firmware/include/vboot_api.h
+++ b/firmware/include/vboot_api.h
@@ -687,6 +687,8 @@ enum VbScreenType_t {
 	VB_SCREEN_SET_VENDOR_DATA = 0x214,
 	/* Confirm vendor data menu screen */
 	VB_SCREEN_CONFIRM_VENDOR_DATA = 0x215,
+	/* Confirm reboot for running diagnostics rom */
+	VB_SCREEN_CONFIRM_DIAG = 0x216,
 };
 
 /**
@@ -1017,6 +1019,7 @@ enum {
 };
 
 enum VbAltFwIndex_t {
+	VB_ALTFW_DIAGNOSTIC = -1,
 	VB_ALTFW_DEFAULT = 0,
 	VB_ALTFW_FIRST = 1,
 	VB_ALTFW_SECOND,
@@ -1037,7 +1040,8 @@ enum VbAltFwIndex_t {
  *     >0 (i.e., positive #) run a payload by # based in altfw/list file
  *     <0 (i.e., negative #) run a specific payload by name without using
  *        the altfw/list file.  Typically payloads in this category will be
- *        verified before they are run. Currently no #s are defined.
+ *        verified before they are run. Currently these #s are defined:
+ *          -1 diagnostic payload
  */
 int VbExLegacy(enum VbAltFwIndex_t altfw_num);
 
diff --git a/firmware/lib/include/vboot_kernel.h b/firmware/lib/include/vboot_kernel.h
index c3c4585b..7a8087e4 100644
--- a/firmware/lib/include/vboot_kernel.h
+++ b/firmware/lib/include/vboot_kernel.h
@@ -60,6 +60,11 @@ VbError_t VbBootNormal(struct vb2_context *ctx);
 VbError_t VbBootDeveloper(struct vb2_context *ctx);
 
 /**
+ * Handle a diagnostic-mode boot.
+ */
+VbError_t VbBootDiagnostic(struct vb2_context *ctx);
+
+/**
  * Handle a recovery-mode boot.
  */
 VbError_t VbBootRecovery(struct vb2_context *ctx);
diff --git a/firmware/lib/vboot_api_kernel.c b/firmware/lib/vboot_api_kernel.c
index 17b3b7b2..78f35744 100644
--- a/firmware/lib/vboot_api_kernel.c
+++ b/firmware/lib/vboot_api_kernel.c
@@ -431,6 +431,27 @@ VbError_t VbSelectAndLoadKernel(VbCommonParams *cparams,
 		else
 			retval = VbBootRecovery(&ctx);
 		VbExEcEnteringMode(0, VB_EC_RECOVERY);
+	} else if (DIAGNOSTIC_UI && vb2_nv_get(&ctx, VB2_NV_DIAG_REQUEST)) {
+		struct vb2_shared_data *sd = vb2_get_sd(&ctx);
+		if (sd->vbsd->flags & VBSD_OPROM_MATTERS)
+			vb2_nv_set(&ctx, VB2_NV_OPROM_NEEDED, 0);
+		vb2_nv_set(&ctx, VB2_NV_DIAG_REQUEST, 0);
+
+		/*
+		 * Diagnostic boot. This has a UI but only power button
+		 * is used for input so no detachable-specific UI is needed.
+		 * This mode is also 1-shot so it's placed before developer
+		 * mode.
+		 */
+		retval = VbBootDiagnostic(&ctx);
+		/*
+		 * The diagnostic menu should either boot a rom, or
+		 * return either of reboot or shutdown.  The following
+		 * check is a safety precaution.
+		 */
+		if (!retval) {
+			retval = VBERROR_REBOOT_REQUIRED;
+		}
 	} else if (ctx.flags & VB2_CONTEXT_DEVELOPER_MODE) {
 		if (kparams->inflags & VB_SALK_INFLAGS_VENDOR_DATA_SETTABLE)
 			ctx.flags |= VB2_CONTEXT_VENDOR_DATA_SETTABLE;
diff --git a/firmware/lib/vboot_ui.c b/firmware/lib/vboot_ui.c
index 6d19a164..2e3bb270 100644
--- a/firmware/lib/vboot_ui.c
+++ b/firmware/lib/vboot_ui.c
@@ -17,6 +17,7 @@
 #include "gbb_header.h"
 #include "load_kernel_fw.h"
 #include "rollback_index.h"
+#include "tlcl.h"
 #include "utility.h"
 #include "vb2_common.h"
 #include "vboot_api.h"
@@ -44,7 +45,10 @@ static void VbAllowUsbBoot(struct vb2_context *ctx)
  * Checks GBB flags against VbExIsShutdownRequested() shutdown request to
  * determine if a shutdown is required.
  *
- * Returns true if a shutdown is required and false if no shutdown is required.
+ * Returns zero or more of the following flags (if any are set then typically
+ * shutdown is required):
+ * VB_SHUTDOWN_REQUEST_LID_CLOSED
+ * VB_SHUTDOWN_REQUEST_POWER_BUTTON
  */
 static int VbWantShutdown(struct vb2_context *ctx, uint32_t key)
 {
@@ -70,7 +74,7 @@ static int VbWantShutdown(struct vb2_context *ctx, uint32_t key)
 	if (sd->gbb_flags & GBB_FLAG_DISABLE_LID_SHUTDOWN)
 		shutdown_request &= ~VB_SHUTDOWN_REQUEST_LID_CLOSED;
 
-	return !!shutdown_request;
+	return shutdown_request;
 }
 
 uint32_t VbTryUsb(struct vb2_context *ctx)
@@ -366,6 +370,124 @@ VbError_t vb2_vendor_data_ui(struct vb2_context *ctx)
 	return VBERROR_SUCCESS;
 }
 
+/*
+ * User interface for confirming launch of diagnostics rom
+ *
+ * This asks the user to confirm the launch of the diagnostics rom. The user
+ * can press the power button to confirm or press escape. There is a 30-second
+ * timeout which acts the same as escape.
+ */
+VbError_t vb2_diagnostics_ui(struct vb2_context *ctx)
+{
+	int active = 1;
+	int power_button_was_pressed = 0;
+	VbError_t result = VBERROR_REBOOT_REQUIRED;
+	int action_confirmed = 0;
+	uint64_t start_time_us;
+
+	VbDisplayScreen(ctx, VB_SCREEN_CONFIRM_DIAG, 0, NULL);
+
+	start_time_us = VbExGetTimer();
+
+	/* We'll loop until the user decides what to do */
+	do {
+		uint32_t key = VbExKeyboardRead();
+		/*
+		 * VbExIsShutdownRequested() is almost an adequate substitute
+		 * for adding a new flag to VbExGetSwitches().  The main
+		 * issue is that the former doesn't consult the power button
+		 * on detachables, and this function wants to see for itself
+		 * that the power button isn't currently pressed.
+		 */
+		uint32_t power_pressed =
+			VbExGetSwitches(VB_SWITCH_FLAG_PHYS_PRESENCE_PRESSED);
+		/*
+		 * TODO(delco): Remove this workaround.  On Wilco a button
+		 * press is only reported a single time regardless of the
+		 * duration of the press. Until it's changed to report the
+		 * live/current status of the button we can't ignore when
+		 * VbWantShutdown() reports a button press (well, we can
+		 * ignore it but the user might have to press the power button
+		 * more than once for this code to react).
+		 */
+		int shutdown = VbWantShutdown(ctx, 0);
+		if (shutdown & VB_SHUTDOWN_REQUEST_POWER_BUTTON) {
+			power_pressed = 1;
+		}
+
+		if (power_pressed) {
+			power_button_was_pressed = 1;
+		} else if (power_button_was_pressed) {
+			VB2_DEBUG("vb2_diagnostics_ui() - power released\n");
+			action_confirmed = 1;
+			active = 0;
+			break;
+		}
+
+		/* Check the lid and ignore the power button. */
+		if (shutdown & VB_SHUTDOWN_REQUEST_LID_CLOSED) {
+			VB2_DEBUG("vb2_diagnostics_ui() - shutdown request\n");
+			result = VBERROR_SHUTDOWN_REQUESTED;
+			active = 0;
+			break;
+		}
+
+		switch (key) {
+		case 0:
+			/* nothing pressed */
+			break;
+		case VB_KEY_ESC:
+			/* Escape pressed - reboot */
+			VB2_DEBUG("vb2_diagnostics_ui() - user pressed Esc\n");
+			active = 0;
+			break;
+		default:
+			VB2_DEBUG("vb2_diagnostics_ui() - pressed key %d\n",
+				  key);
+			VbCheckDisplayKey(ctx, key);
+			break;
+		}
+		if (VbExGetTimer() - start_time_us >= 30 * VB_USEC_PER_SEC) {
+			VB2_DEBUG("vb2_diagnostics_ui() - timeout\n");
+			break;
+		}
+		if (active) {
+			VbExSleepMs(DEV_KEY_DELAY);
+		}
+	} while (active);
+
+	VbDisplayScreen(ctx, VB_SCREEN_BLANK, 0, NULL);
+
+	if (action_confirmed) {
+		VB2_DEBUG("Diagnostic requested, running\n");
+
+		/*
+		 * The following helps avoid use of the TPM after
+		 * it's disabled (e.g., when vb2_run_altfw() calls
+		 * RollbackKernelLock() ).
+		 */
+
+		if (RollbackKernelLock(0)) {
+			VB2_DEBUG("Failed to lock TPM PP\n");
+			vb2_fail(ctx, VB2_RECOVERY_TPM_DISABLE_FAILED, 0);
+		} else if (vb2ex_tpm_set_mode(VB2_TPM_MODE_DISABLED) !=
+			   VB2_SUCCESS) {
+			VB2_DEBUG("Failed to disable TPM\n");
+			vb2_fail(ctx, VB2_RECOVERY_TPM_DISABLE_FAILED, 0);
+		} else {
+			vb2_run_altfw(ctx, VB_ALTFW_DIAGNOSTIC);
+			VB2_DEBUG("Diagnostic failed to run\n");
+			/*
+			 * Assuming failure was due to bad hash, though
+			 * the rom could just be missing or invalid.
+			 */
+			vb2_fail(ctx, VB2_RECOVERY_ALTFW_HASH_FAILED, 0);
+		}
+	}
+
+	return result;
+}
+
 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"
@@ -650,6 +772,14 @@ VbError_t VbBootDeveloper(struct vb2_context *ctx)
 	return retval;
 }
 
+VbError_t VbBootDiagnostic(struct vb2_context *ctx)
+{
+	vb2_init_ui();
+	VbError_t retval = vb2_diagnostics_ui(ctx);
+	VbDisplayScreen(ctx, VB_SCREEN_BLANK, 0, NULL);
+	return retval;
+}
+
 /* Delay in recovery mode */
 #define REC_DISK_DELAY       1000     /* Check disks every 1s */
 #define REC_KEY_DELAY        20       /* Check keys every 20ms */
@@ -786,6 +916,14 @@ static VbError_t recovery_ui(struct vb2_context *ctx)
 					i = 4;
 					break;
 				}
+			} else if (DIAGNOSTIC_UI &&
+				   (key == VB_KEY_CTRL('C') ||
+				    key == 0x114)) {       /* F12 */
+				VB2_DEBUG("Diagnostic requested, rebooting\n");
+                                if (shared->flags & VBSD_OPROM_MATTERS)
+					vb2_nv_set(ctx, VB2_NV_OPROM_NEEDED, 1);
+				vb2_nv_set(ctx, VB2_NV_DIAG_REQUEST, 1);
+				return VBERROR_REBOOT_REQUIRED;
 			} else {
 				VbCheckDisplayKey(ctx, key);
 			}
diff --git a/tests/vboot_api_kernel2_tests.c b/tests/vboot_api_kernel2_tests.c
index 006e8812..67414ab5 100644
--- a/tests/vboot_api_kernel2_tests.c
+++ b/tests/vboot_api_kernel2_tests.c
@@ -34,10 +34,12 @@ static struct vb2_shared_data *sd;
 
 static int shutdown_request_calls_left;
 static int shutdown_request_power_held;
+static int shutdown_via_lid_close;
 static int audio_looping_calls_left;
 static uint32_t vbtlk_retval;
 static int vbexlegacy_called;
 static enum VbAltFwIndex_t altfw_num;
+static uint64_t current_ticks;
 static int trust_ec;
 static int virtdev_set;
 static uint32_t virtdev_retval;
@@ -51,6 +53,8 @@ static uint32_t screens_displayed[8];
 static uint32_t screens_count = 0;
 static uint32_t mock_num_disks[8];
 static uint32_t mock_num_disks_count;
+static int tpm_set_mode_called;
+static enum vb2_tpm_mode tpm_mode;
 
 static char set_vendor_data[32];
 static int set_vendor_data_called;
@@ -79,10 +83,12 @@ static void ResetMocks(void)
 
 	shutdown_request_calls_left = -1;
 	shutdown_request_power_held = -1;
+	shutdown_via_lid_close = 0;
 	audio_looping_calls_left = 30;
 	vbtlk_retval = 1000;
 	vbexlegacy_called = 0;
 	altfw_num = -100;
+	current_ticks = 0;
 	trust_ec = 0;
 	virtdev_set = 0;
 	virtdev_retval = 0;
@@ -101,6 +107,9 @@ static void ResetMocks(void)
 
 	memset(mock_num_disks, 0, sizeof(mock_num_disks));
 	mock_num_disks_count = 0;
+
+	tpm_set_mode_called = 0;
+	tpm_mode = VB2_TPM_MODE_ENABLED_TENTATIVE;
 }
 
 /* Mock functions */
@@ -108,7 +117,9 @@ static void ResetMocks(void)
 uint32_t VbExIsShutdownRequested(void)
 {
 	if (shutdown_request_calls_left == 0)
-		return 1;
+		return shutdown_via_lid_close ?
+			VB_SHUTDOWN_REQUEST_LID_CLOSED :
+			VB_SHUTDOWN_REQUEST_POWER_BUTTON;
 	else if (shutdown_request_calls_left > 0)
 		shutdown_request_calls_left--;
 
@@ -154,7 +165,18 @@ int VbExLegacy(enum VbAltFwIndex_t _altfw_num)
 	vbexlegacy_called++;
 	altfw_num = _altfw_num;
 
-	return 0;
+	/* VbExLegacy() can only return failure, or not return at all. */
+	return VBERROR_UNKNOWN;
+}
+
+void VbExSleepMs(uint32_t msec)
+{
+	current_ticks += (uint64_t)msec * VB_USEC_PER_MSEC;
+}
+
+uint64_t VbExGetTimer(void)
+{
+	return current_ticks;
 }
 
 VbError_t VbExDiskGetInfo(VbDiskInfo **infos_ptr, uint32_t *count,
@@ -220,6 +242,23 @@ VbError_t VbExSetVendorData(const char *vendor_data_value)
 	return VBERROR_SUCCESS;
 }
 
+int vb2ex_tpm_set_mode(enum vb2_tpm_mode mode_val)
+{
+	tpm_set_mode_called = 1;
+	/*
+	 * This mock will pretend that any call will fail if the tpm is
+	 * already disabled (e.g., as if the code always tries to contact the
+	 * tpm to issue a command).  The real version may eventually be changed
+	 * to return success if the incoming request is also to disable, but
+	 * the point here is to have a way to simulate failure.
+	 */
+	if (tpm_mode == VB2_TPM_MODE_DISABLED) {
+		return VB2_ERROR_UNKNOWN;
+	}
+	tpm_mode = mode_val;
+	return VB2_SUCCESS;
+}
+
 /* Tests */
 
 static void VbUserConfirmsTest(void)
@@ -1051,6 +1090,145 @@ static void VbBootRecTest(void)
 		VBERROR_TPM_SET_BOOT_MODE_STATE,
 		"Ctrl+D todev failure");
 
+	/* Test Diagnostic Mode via Ctrl-C when no oprom needed */
+	ResetMocks();
+	shared->flags = VBSD_BOOT_REC_SWITCH_ON;
+	trust_ec = 1;
+	shutdown_request_calls_left = 100;
+	mock_keypress[0] = 0x03;
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_DIAG_REQUEST), 0,
+		"todiag is zero");
+	if (DIAGNOSTIC_UI)
+		TEST_EQ(VbBootRecovery(&ctx),
+			VBERROR_REBOOT_REQUIRED,
+			"Ctrl+C todiag - enabled");
+	else
+		TEST_EQ(VbBootRecovery(&ctx),
+			VBERROR_SHUTDOWN_REQUESTED,
+			"Ctrl+C todiag - disabled");
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_DIAG_REQUEST), DIAGNOSTIC_UI,
+		"todiag is updated for Ctrl-C");
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_OPROM_NEEDED), 0,
+		"todiag doesn't update for unneeded opom");
+
+	/* Test Diagnostic Mode via F12 - oprom needed */
+	ResetMocks();
+	shared->flags = VBSD_BOOT_REC_SWITCH_ON | VBSD_OPROM_MATTERS;
+	trust_ec = 1;
+	shutdown_request_calls_left = 100;
+	mock_keypress[0] = 0x114;
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_DIAG_REQUEST), 0,
+		"todiag is zero");
+	if (DIAGNOSTIC_UI)
+		TEST_EQ(VbBootRecovery(&ctx),
+			VBERROR_REBOOT_REQUIRED,
+			"F12 todiag - enabled");
+	else
+		TEST_EQ(VbBootRecovery(&ctx),
+			VBERROR_SHUTDOWN_REQUESTED,
+			"F12 todiag - disabled");
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_DIAG_REQUEST), DIAGNOSTIC_UI,
+		"todiag is updated for F12");
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_OPROM_NEEDED), DIAGNOSTIC_UI,
+		"todiag updates opom, if need");
+
+	printf("...done.\n");
+}
+
+static void VbBootDiagTest(void)
+{
+	printf("Testing VbBootDiagnostic()...\n");
+
+	/* No key pressed - timeout. */
+	ResetMocks();
+	TEST_EQ(VbBootDiagnostic(&ctx), VBERROR_REBOOT_REQUIRED, "Timeout");
+	TEST_EQ(screens_displayed[0], VB_SCREEN_CONFIRM_DIAG,
+		"  confirm screen");
+	TEST_EQ(screens_displayed[1], VB_SCREEN_BLANK,
+		"  blank screen");
+	TEST_EQ(tpm_set_mode_called, 0, "  no tpm call");
+	TEST_EQ(vbexlegacy_called, 0, "  not legacy");
+	TEST_EQ(current_ticks, 30 * VB_USEC_PER_SEC,
+		"  waited for 30 seconds");
+
+	/* Esc key pressed. */
+	ResetMocks();
+	mock_keypress[0] = VB_KEY_ESC;
+	TEST_EQ(VbBootDiagnostic(&ctx), VBERROR_REBOOT_REQUIRED, "Esc key");
+	TEST_EQ(screens_displayed[0], VB_SCREEN_CONFIRM_DIAG,
+		"  confirm screen");
+	TEST_EQ(screens_displayed[1], VB_SCREEN_BLANK,
+		"  blank screen");
+	TEST_EQ(tpm_set_mode_called, 0, "  no tpm call");
+	TEST_EQ(vbexlegacy_called, 0, "  not legacy");
+	TEST_EQ(current_ticks, 0, "  didn't wait at all");
+
+	/* Shutdown requested via lid close */
+	ResetMocks();
+	shutdown_via_lid_close = 1;
+	shutdown_request_calls_left = 10;
+	TEST_EQ(VbBootDiagnostic(&ctx), VBERROR_SHUTDOWN_REQUESTED, "Shutdown");
+	TEST_EQ(screens_displayed[0], VB_SCREEN_CONFIRM_DIAG,
+		"  confirm screen");
+	TEST_EQ(screens_displayed[1], VB_SCREEN_BLANK,
+		"  blank screen");
+	TEST_EQ(tpm_set_mode_called, 0, "  no tpm call");
+	TEST_EQ(vbexlegacy_called, 0, "  not legacy");
+	TEST_TRUE(current_ticks < VB_USEC_PER_SEC, "  didn't wait long");
+
+	/* Power button pressed but not released. */
+	ResetMocks();
+	mock_switches_are_stuck = 1;
+	mock_switches[0] = VB_SWITCH_FLAG_PHYS_PRESENCE_PRESSED;
+	TEST_EQ(VbBootDiagnostic(&ctx), VBERROR_REBOOT_REQUIRED, "Power held");
+	TEST_EQ(screens_displayed[0], VB_SCREEN_CONFIRM_DIAG,
+		"  confirm screen");
+	TEST_EQ(screens_displayed[1], VB_SCREEN_BLANK,
+		"  blank screen");
+	TEST_EQ(tpm_set_mode_called, 0, "  no tpm call");
+	TEST_EQ(vbexlegacy_called, 0, "  not legacy");
+
+	/* Power button is pressed and released. */
+	ResetMocks();
+	mock_switches[0] = 0;
+	mock_switches[1] = VB_SWITCH_FLAG_PHYS_PRESENCE_PRESSED;
+	mock_switches[2] = 0;
+	TEST_EQ(VbBootDiagnostic(&ctx), VBERROR_REBOOT_REQUIRED, "Confirm");
+	TEST_EQ(screens_displayed[0], VB_SCREEN_CONFIRM_DIAG,
+		"  confirm screen");
+	TEST_EQ(screens_displayed[1], VB_SCREEN_BLANK,
+		"  blank screen");
+	TEST_EQ(tpm_set_mode_called, 1, "  tpm call");
+	TEST_EQ(tpm_mode, VB2_TPM_MODE_DISABLED, "  tpm disabled");
+	TEST_EQ(vbexlegacy_called, 1, "  legacy");
+	TEST_EQ(altfw_num, VB_ALTFW_DIAGNOSTIC, "  check altfw_num");
+	/*
+	 * Ideally we'd that no recovery request was recorded, but
+	 * VbExLegacy() can only fail or crash the tests.
+	 */
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_RECOVERY_REQUEST),
+		VB2_RECOVERY_ALTFW_HASH_FAILED,
+		"  recovery request");
+
+        /* Power button confirm, but now with a tpm failure. */
+	ResetMocks();
+	tpm_mode = VB2_TPM_MODE_DISABLED;
+	mock_switches[0] = 0;
+	mock_switches[1] = VB_SWITCH_FLAG_PHYS_PRESENCE_PRESSED;
+	mock_switches[2] = 0;
+	TEST_EQ(VbBootDiagnostic(&ctx), VBERROR_REBOOT_REQUIRED,
+		"Confirm but tpm fail");
+	TEST_EQ(screens_displayed[0], VB_SCREEN_CONFIRM_DIAG,
+		"  confirm screen");
+	TEST_EQ(screens_displayed[1], VB_SCREEN_BLANK,
+		"  blank screen");
+	TEST_EQ(tpm_set_mode_called, 1, "  tpm call");
+	TEST_EQ(tpm_mode, VB2_TPM_MODE_DISABLED, "  tpm disabled");
+	TEST_EQ(vbexlegacy_called, 0, "  legacy not called");
+	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_RECOVERY_REQUEST),
+		VB2_RECOVERY_TPM_DISABLE_FAILED,
+		"  recovery request");
+
 	printf("...done.\n");
 }
 
@@ -1061,6 +1239,8 @@ int main(void)
 	VbBootTest();
 	VbBootDevTest();
 	VbBootRecTest();
+	if (DIAGNOSTIC_UI)
+		VbBootDiagTest();
 
 	return gTestSuccess ? 0 : 255;
 }
diff --git a/tests/vboot_api_kernel4_tests.c b/tests/vboot_api_kernel4_tests.c
index 82d6d5db..2d7e04e1 100644
--- a/tests/vboot_api_kernel4_tests.c
+++ b/tests/vboot_api_kernel4_tests.c
@@ -40,6 +40,10 @@ static struct RollbackSpaceFwmp rfr_fwmp;
 static int rkr_retval, rkw_retval, rkl_retval, rfr_retval;
 static VbError_t vbboot_retval;
 
+static uint32_t mock_switches[8];
+static uint32_t mock_switches_count;
+static int mock_switches_are_stuck;
+
 /* Reset mock data (for use before each test) */
 static void ResetMocks(void)
 {
@@ -75,6 +79,10 @@ static void ResetMocks(void)
 	rkr_version = new_version = 0x10002;
 	rkr_retval = rkw_retval = rkl_retval = VBERROR_SUCCESS;
 	vbboot_retval = VBERROR_SUCCESS;
+
+	memset(mock_switches, 0, sizeof(mock_switches));
+	mock_switches_count = 0;
+	mock_switches_are_stuck = 0;
 }
 
 /* Mock functions */
@@ -144,6 +152,14 @@ VbError_t VbBootRecovery(struct vb2_context *ctx)
 	return vbboot_retval;
 }
 
+VbError_t VbBootDiagnostic(struct vb2_context *ctx)
+{
+	if (vbboot_retval == -4)
+		return VBERROR_SIMULATED;
+
+	return vbboot_retval;
+}
+
 static void test_slk(VbError_t retval, int recovery_reason, const char *desc)
 {
 	TEST_EQ(VbSelectAndLoadKernel(&cparams, &kparams), retval, desc);
@@ -151,6 +167,21 @@ static void test_slk(VbError_t retval, int recovery_reason, const char *desc)
 		recovery_reason, "  recovery reason");
 }
 
+uint32_t VbExGetSwitches(uint32_t request_mask)
+{
+	if (mock_switches_are_stuck)
+		return mock_switches[0] & request_mask;
+	if (mock_switches_count < ARRAY_SIZE(mock_switches))
+		return mock_switches[mock_switches_count++] & request_mask;
+	else
+		return 0;
+}
+
+int vb2ex_tpm_set_mode(enum vb2_tpm_mode mode_val)
+{
+	return VB2_SUCCESS;
+}
+
 /* Tests */
 
 static void VbSlkTest(void)
@@ -226,6 +257,33 @@ static void VbSlkTest(void)
 	vbboot_retval = -1;
 	test_slk(VBERROR_SIMULATED, 0, "Normal boot bad");
 
+	/* Check that NV_DIAG_REQUEST triggers diagnostic UI */
+	if (DIAGNOSTIC_UI) {
+		ResetMocks();
+		mock_switches[1] = VB_SWITCH_FLAG_PHYS_PRESENCE_PRESSED;
+		vb2_nv_set(&ctx, VB2_NV_DIAG_REQUEST, 1);
+		vb2_nv_set(&ctx, VB2_NV_OPROM_NEEDED, 1);
+		vbboot_retval = -4;
+		test_slk(VBERROR_SIMULATED, 0, "Normal boot with diag");
+		TEST_EQ(vb2_nv_get(&ctx, VB2_NV_DIAG_REQUEST), 0,
+			"  diag not requested");
+		TEST_EQ(vb2_nv_get(&ctx, VB2_NV_OPROM_NEEDED), 1,
+			"  oprom still needed");
+
+		ResetMocks();
+		mock_switches[1] = VB_SWITCH_FLAG_PHYS_PRESENCE_PRESSED;
+		vb2_nv_set(&ctx, VB2_NV_DIAG_REQUEST, 1);
+		vb2_nv_set(&ctx, VB2_NV_OPROM_NEEDED, 1);
+		shared->flags |= VBSD_OPROM_MATTERS;
+		vbboot_retval = -4;
+		test_slk(VBERROR_SIMULATED, 0,
+			 "Normal boot with diag and oprom");
+		TEST_EQ(vb2_nv_get(&ctx, VB2_NV_DIAG_REQUEST), 0,
+			"  diag not requested");
+		TEST_EQ(vb2_nv_get(&ctx, VB2_NV_OPROM_NEEDED), 0,
+			"  oprom not needed");
+	}
+
 	/* Boot dev */
 	ResetMocks();
 	shared->flags |= VBSD_BOOT_DEV_SWITCH_ON;