Enable launching diagnostics from OS broken screenstabilize-12088.B

If firmware diagnostics is enabled allow the user to lauch diagnostics
from the OS broken screen. Currently diagnostics can only be launched
from the recovery screen.

BUG=b:130244249
TEST=Tested on arcada device
     make runmisctests
BRANCH=none

Change-Id: Ic6946338605599ea7411b5658acb2c6fc960a782
Signed-off-by: Mathew King <mathewk@chromium.org>
Reviewed-on: https://chromium-review.googlesource.com/1564448
Commit-Ready: ChromeOS CL Exonerator Bot <chromiumos-cl-exonerator@appspot.gserviceaccount.com>
Reviewed-by: Julius Werner <jwerner@chromium.org>

3 files changed, 60 insertions, 14 deletions

diff --git a/firmware/include/vboot_api.h b/firmware/include/vboot_api.h
index 21ac36b4..66e26ef0 100644
--- a/firmware/include/vboot_api.h
+++ b/firmware/include/vboot_api.h
@@ -774,6 +774,9 @@ VbError_t VbExSetVendorData(const char *vendor_data_value);
 /* Key code for CTRL + letter */
 #define VB_KEY_CTRL(letter) (letter & 0x1f)
 
+/* Key code for fn keys */
+#define VB_KEY_F(num) (num + 0x108)
+
 /* Key codes for required non-printable-ASCII characters. */
 enum VbKeyCode_t {
 	VB_KEY_ENTER = '\r',
diff --git a/firmware/lib/vboot_ui.c b/firmware/lib/vboot_ui.c
index 6ba46682..9ccc9aa8 100644
--- a/firmware/lib/vboot_ui.c
+++ b/firmware/lib/vboot_ui.c
@@ -378,6 +378,20 @@ VbError_t vb2_vendor_data_ui(struct vb2_context *ctx)
 	return VBERROR_SUCCESS;
 }
 
+static VbError_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");
+		if (vb2_get_sd(ctx)->vbsd->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;
+	}
+
+	return VBERROR_SUCCESS;
+}
+
 /*
  * User interface for confirming launch of diagnostics rom
  *
@@ -832,6 +846,10 @@ static VbError_t recovery_ui(struct vb2_context *ctx)
 			VbCheckDisplayKey(ctx, key, NULL);
 			if (VbWantShutdown(ctx, key))
 				return VBERROR_SHUTDOWN_REQUESTED;
+			else if ((retval =
+				  vb2_check_diagnostic_key(ctx, key)) !=
+				  VBERROR_SUCCESS)
+				return retval;
 			VbExSleepMs(REC_KEY_DELAY);
 		}
 	}
@@ -922,14 +940,10 @@ 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 if ((retval =
+				    vb2_check_diagnostic_key(ctx, key)) !=
+				    VBERROR_SUCCESS) {
+				return retval;
 			} else {
 				VbCheckDisplayKey(ctx, key, NULL);
 			}
diff --git a/tests/vboot_api_kernel2_tests.c b/tests/vboot_api_kernel2_tests.c
index fc13362b..889d2cba 100644
--- a/tests/vboot_api_kernel2_tests.c
+++ b/tests/vboot_api_kernel2_tests.c
@@ -1100,8 +1100,9 @@ static void VbBootRecTest(void)
 	ResetMocks();
 	shared->flags = VBSD_BOOT_REC_SWITCH_ON;
 	trust_ec = 1;
+	vbtlk_retval = VBERROR_NO_DISK_FOUND - VB_DISK_FLAG_REMOVABLE;
 	shutdown_request_calls_left = 100;
-	mock_keypress[0] = 0x03;
+	mock_keypress[0] = VB_KEY_CTRL('C');
 	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_DIAG_REQUEST), 0,
 		"todiag is zero");
 	if (DIAGNOSTIC_UI)
@@ -1112,17 +1113,21 @@ static void VbBootRecTest(void)
 		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");
+		"  todiag is updated for Ctrl-C");
 	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_OPROM_NEEDED), 0,
-		"todiag doesn't update for unneeded opom");
+		"  todiag doesn't update for unneeded opom");
+	TEST_EQ(screens_displayed[0], VB_SCREEN_RECOVERY_INSERT,
+		"  insert screen");
 
 	/* Test Diagnostic Mode via F12 - oprom needed */
 	ResetMocks();
 	shared->flags = VBSD_BOOT_REC_SWITCH_ON | VBSD_OPROM_MATTERS;
 	trust_ec = 1;
+	vbtlk_retval = VBERROR_NO_DISK_FOUND - VB_DISK_FLAG_REMOVABLE;
 	shutdown_request_calls_left = 100;
-	mock_keypress[0] = 0x114;
+	mock_keypress[0] = VB_KEY_F(12);
 	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_DIAG_REQUEST), 0,
 		"todiag is zero");
 	if (DIAGNOSTIC_UI)
@@ -1134,9 +1139,33 @@ static void VbBootRecTest(void)
 			VBERROR_SHUTDOWN_REQUESTED,
 			"F12 todiag - disabled");
 	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_DIAG_REQUEST), DIAGNOSTIC_UI,
-		"todiag is updated for F12");
+		"  todiag is updated for F12");
 	TEST_EQ(vb2_nv_get(&ctx, VB2_NV_OPROM_NEEDED), DIAGNOSTIC_UI,
-		"todiag updates opom, if need");
+		"  todiag updates opom, if need");
+	TEST_EQ(screens_displayed[0], VB_SCREEN_RECOVERY_INSERT,
+		"  insert screen");
+
+	/* Test Diagnostic Mode via Ctrl-C OS broken */
+	ResetMocks();
+	shared->flags = 0;
+	shutdown_request_calls_left = 100;
+	mock_keypress[0] = VB_KEY_CTRL('C');
+	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 os broken - enabled");
+	else
+		TEST_EQ(VbBootRecovery(&ctx),
+			VBERROR_SHUTDOWN_REQUESTED,
+			"Ctrl+C todiag os broken - 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_EQ(screens_displayed[0], VB_SCREEN_OS_BROKEN,
+		"  os broken screen");
 
 	printf("...done.\n");
 }