vboot: Initialize USB in diagnostic mode

In addition to recovery and developer modes, diagnostic mode also has
UI. Therefore, we need to initialize USB in diagnostic mode as well.

BUG=b:229698042, b:179334129
BRANCH=none
TEST=USB KB worked in diagnostic mode

Signed-off-by: Eric Lai <eric_lai@quanta.corp-partner.google.com>
Change-Id: Ib81c63bb15c8fad441f2f4862e916ac852497f37
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/depthcharge/+/3626949
Reviewed-by: Yu-Ping Wu <yupingso@chromium.org>
diff --git a/src/vboot/stages.c b/src/vboot/stages.c
index 233f658..742e593 100644
--- a/src/vboot/stages.c
+++ b/src/vboot/stages.c
@@ -61,9 +61,17 @@
 
 int vboot_check_enable_usb(void)
 {
-	/* Initialize USB in developer or recovery mode, skip in normal mode. */
-	if (vboot_in_recovery() || vboot_in_developer())
-		dc_usb_initialize();
+	/* Initialize USB in developer, recovery mode or diagnostics mode,
+	   skip in normal mode. */
+	switch (vboot_get_context()->boot_mode) {
+		case VB2_BOOT_MODE_MANUAL_RECOVERY:
+		case VB2_BOOT_MODE_BROKEN_SCREEN:
+		case VB2_BOOT_MODE_DIAGNOSTICS:
+		case VB2_BOOT_MODE_DEVELOPER:
+			dc_usb_initialize();
+		default:
+			break;
+	}
 	return 0;
 }