VbBootRecovery: Make second check for 'remove' devices if none found

There is some inherent latency between the time the USB root hub is
initialized and the time USB devices are detected. This can lead to a
situation where USB media is attached, yet not found when we do our
initial device poll. The device may be detected in subsequent polls, so
the media can be booted and no 'remove' screen will be displayed.

With this change, if no media to remove is initially found, a second
poll will be made after a 500ms delay. This will be enough time for USB
devices to be correctly detected in our test cases.

Also, it is necessary to change the unit test due to the fact that we
now call VbExDiskGetInfo twice before actually displaying any screen.

TEST=Manual on Monroe. Insert USB media and trigger recovery boot.
Verify 'remove' screen is seen, 'insert' screen is seen after removing
media, and system boots after re-inserting media. Also passes
vboot_reference unit tests.
BUG=chrome-os-partner:23840
BRANCH=Panther, Monroe

Signed-off-by: Shawn Nematbakhsh <shawnn@chromium.org>
Change-Id: Ia902c3a126588cd7ea618f2dbbca6b38d35d6ea0
Reviewed-on: https://chromium-review.googlesource.com/179757
Reviewed-by: Randall Spangler <rspangler@chromium.org>
diff --git a/firmware/lib/vboot_api_kernel.c b/firmware/lib/vboot_api_kernel.c
index 5c498ee..4cf2a94 100644
--- a/firmware/lib/vboot_api_kernel.c
+++ b/firmware/lib/vboot_api_kernel.c
@@ -392,8 +392,9 @@
 }
 
 /* Delay in recovery mode */
-#define REC_DISK_DELAY 1000     /* Check disks every 1s */
-#define REC_KEY_DELAY  20       /* Check keys every 20ms */
+#define REC_DISK_DELAY       1000     /* Check disks every 1s */
+#define REC_KEY_DELAY        20       /* Check keys every 20ms */
+#define REC_MEDIA_INIT_DELAY 500      /* Check removable media every 500ms */
 
 VbError_t VbBootRecovery(VbCommonParams *cparams, LoadKernelParams *p)
 {
@@ -416,7 +417,19 @@
 
 		VBDEBUG(("VbBootRecovery() forcing device removal\n"));
 
+		/* If no media is detected initially, delay and make one extra
+		 * attempt, in case devices appear later than expected. */
+		if (VBERROR_SUCCESS != VbExDiskGetInfo(&disk_info, &disk_count,
+						       VB_DISK_FLAG_REMOVABLE))
+			disk_count = 0;
+
+		VbExDiskFreeInfo(disk_info, NULL);
+		if (0 == disk_count)
+			VbExSleepMs(REC_MEDIA_INIT_DELAY);
+
 		while (1) {
+			disk_info = NULL;
+			disk_count = 0;
 			if (VBERROR_SUCCESS !=
 			    VbExDiskGetInfo(&disk_info, &disk_count,
 					    VB_DISK_FLAG_REMOVABLE))
diff --git a/tests/vboot_api_kernel2_tests.c b/tests/vboot_api_kernel2_tests.c
index 1cb2d83..21ea306 100644
--- a/tests/vboot_api_kernel2_tests.c
+++ b/tests/vboot_api_kernel2_tests.c
@@ -419,6 +419,7 @@
 	shutdown_request_calls_left = 100;
 	mock_num_disks[0] = 1;
 	mock_num_disks[1] = 1;
+	mock_num_disks[2] = 1;
 	vbtlk_retval = VBERROR_NO_DISK_FOUND - VB_DISK_FLAG_REMOVABLE;
 	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
 		"Remove");
@@ -455,6 +456,21 @@
 	TEST_EQ(screens_displayed[0], VB_SCREEN_RECOVERY_INSERT,
 		"  insert screen");
 
+	/* Removal if no disk initially found, but found on second attempt */
+	ResetMocks();
+	shutdown_request_calls_left = 100;
+	mock_num_disks[0] = 0;
+	mock_num_disks[1] = 1;
+	vbtlk_retval = VBERROR_NO_DISK_FOUND - VB_DISK_FLAG_REMOVABLE;
+	TEST_EQ(VbBootRecovery(&cparams, &lkp), VBERROR_SHUTDOWN_REQUESTED,
+		"Remove");
+	TEST_EQ(screens_displayed[0], VB_SCREEN_RECOVERY_REMOVE,
+		"  remove screen");
+	TEST_EQ(screens_displayed[1], VB_SCREEN_BLANK,
+		"  blank screen");
+	TEST_EQ(screens_displayed[2], VB_SCREEN_RECOVERY_INSERT,
+		"  insert screen");
+
 	/* Bad disk count doesn't require removal */
 	ResetMocks();
 	mock_num_disks[0] = -1;