fastboot: Wait for reboot-fastboot to complete before signaling success.

Bug: N/A
Test: manual test
Change-Id: I07495c1befcf63607c9996d6da6b4df05c060641
diff --git a/fastboot/fastboot.cpp b/fastboot/fastboot.cpp
index fee0857..c77b3f7 100644
--- a/fastboot/fastboot.cpp
+++ b/fastboot/fastboot.cpp
@@ -1158,6 +1158,10 @@
     std::this_thread::sleep_for(std::chrono::milliseconds(1000));
 
     fb->set_transport(open_device());
+
+    if (!is_userspace_fastboot()) {
+        die("Failed to boot into userspace fastboot; one or more components might be unbootable.");
+    }
 }
 
 class ImageSource {
@@ -1314,9 +1318,6 @@
     if (!is_userspace_fastboot()) {
         reboot_to_userspace_fastboot();
     }
-    if (!is_userspace_fastboot()) {
-        die("Failed to boot into userspace; one or more components might be unbootable.");
-    }
 
     std::string super_name;
     if (fb->GetVar("super-partition-name", &super_name) != fastboot::RetCode::SUCCESS) {
@@ -1962,8 +1963,7 @@
         fb->RebootTo("recovery");
         fb->WaitForDisconnect();
     } else if (wants_reboot_fastboot) {
-        fb->RebootTo("fastboot");
-        fb->WaitForDisconnect();
+        reboot_to_userspace_fastboot();
     }
 
     fprintf(stderr, "Finished. Total time: %.3fs\n", (now() - start));