Enable launching diagnostics from OS broken screen

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>
diff --git a/firmware/include/vboot_api.h b/firmware/include/vboot_api.h
index 21ac36b..66e26ef 100644
--- a/firmware/include/vboot_api.h
+++ b/firmware/include/vboot_api.h
@@ -774,6 +774,9 @@
 /* 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 6ba4668..9ccc9aa 100644
--- a/firmware/lib/vboot_ui.c
+++ b/firmware/lib/vboot_ui.c
@@ -378,6 +378,20 @@
 	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 @@
 			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 @@
 					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 fc13362..889d2cb 100644
--- a/tests/vboot_api_kernel2_tests.c
+++ b/tests/vboot_api_kernel2_tests.c
@@ -1100,8 +1100,9 @@
 	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 @@
 		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 @@
 			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");
 }