fastboot: Automatically reboot to userspace fastboot.

In order to flash logical partitions, "flashall" must be run against
userspace fastboot. When the bootloader supports the "reboot-fastboot"
command, we now attempt to automatically reboot into userspace fastboot.

We do this by closing the transport, sleeping for one second, and then
polling for a new connection until one is available. FastBootDriver is
then assigned the new transport.

Bug: 78793464
Test: fastboot flashall on device with logical partitions, while booted
      into bootloader fastboot

Change-Id: I6949871b93ab5e352784cabe0963c6ecfc0af36d

fastboot/engine.cpp

@@ -79,12 +79,18 @@ struct Action {
 static std::vector<std::unique_ptr<Action>> action_list;
 static fastboot::FastBootDriver* fb = nullptr;
 
+static constexpr char kStatusFormat[] = "%-50s ";
+
 void fb_init(fastboot::FastBootDriver& fbi) {
     fb = &fbi;
     auto cb = [](std::string& info) { fprintf(stderr, "(bootloader) %s\n", info.c_str()); };
     fb->SetInfoCallback(cb);
 }
 
+void fb_reinit(Transport* transport) {
+    fb->set_transport(transport);
+}
+
 const std::string fb_get_error() {
     return fb->Error();
 }
@@ -332,7 +338,7 @@ int64_t fb_execute_queue() {
     for (auto& a : action_list) {
         a->start = now();
         if (!a->msg.empty()) {
-            fprintf(stderr, "%-50s ", a->msg.c_str());
+            fprintf(stderr, kStatusFormat, a->msg.c_str());
             verbose("\n");
         }
         if (a->op == OP_DOWNLOAD) {
@@ -372,3 +378,20 @@ int64_t fb_execute_queue() {
     action_list.clear();
     return status;
 }
+
+bool fb_reboot_to_userspace() {
+    // First ensure that the queue is flushed.
+    fb_execute_queue();
+
+    fprintf(stderr, kStatusFormat, "Rebooting to userspace fastboot");
+    verbose("\n");
+
+    if (fb->RebootTo("fastboot") != fastboot::RetCode::SUCCESS) {
+        fprintf(stderr, "FAILED (%s)\n", fb->Error().c_str());
+        return false;
+    }
+    fprintf(stderr, "OKAY\n");
+
+    fb->set_transport(nullptr);
+    return true;
+}

fastboot/engine.h

@@ -50,6 +50,7 @@ const std::string fb_get_error();
 /* engine.c - high level command queue engine */
 
 void fb_init(fastboot::FastBootDriver& fbi);
+void fb_reinit(Transport* transport);
 
 bool fb_getvar(const std::string& key, std::string* value);
 void fb_queue_flash(const std::string& partition, void* data, uint32_t sz);
@@ -74,6 +75,7 @@ void fb_queue_delete_partition(const std::string& partition);
 void fb_queue_resize_partition(const std::string& partition, const std::string& size);
 int64_t fb_execute_queue();
 void fb_set_active(const std::string& slot);
+bool fb_reboot_to_userspace();
 
 /* Current product */
 extern char cur_product[FB_RESPONSE_SZ + 1];

fastboot/fastboot.cpp

@@ -246,13 +246,8 @@ static int list_devices_callback(usb_ifc_info* info) {
 // The returned Transport is a singleton, so multiple calls to this function will return the same
 // object, and the caller should not attempt to delete the returned Transport.
 static Transport* open_device() {
-    static Transport* transport = nullptr;
     bool announce = true;
 
-    if (transport != nullptr) {
-        return transport;
-    }
-
     Socket::Protocol protocol = Socket::Protocol::kTcp;
     std::string host;
     int port = 0;
@@ -277,6 +272,7 @@ static Transport* open_device() {
         }
     }
 
+    Transport* transport = nullptr;
     while (true) {
         if (!host.empty()) {
             std::string error;
@@ -1179,6 +1175,17 @@ static bool is_logical(const std::string& partition) {
     return fb_getvar("is-logical:" + partition, &value) && value == "yes";
 }
 
+static void reboot_to_userspace_fastboot() {
+    if (!fb_reboot_to_userspace()) {
+        die("Must reboot to userspace fastboot to flash logical partitions");
+    }
+
+    // Give the current connection time to close.
+    std::this_thread::sleep_for(std::chrono::milliseconds(1000));
+
+    fb_reinit(open_device());
+}
+
 static void update_super_partition(bool force_wipe) {
     if (!if_partition_exists("super", "")) {
         return;
@@ -1189,7 +1196,7 @@ static void update_super_partition(bool force_wipe) {
     }
 
     if (!is_userspace_fastboot()) {
-        die("Must have userspace fastboot to flash logical partitions");
+        reboot_to_userspace_fastboot();
     }
 
     int fd = open(image.c_str(), O_RDONLY);

fastboot/fastboot_driver.cpp

@@ -52,11 +52,15 @@ namespace fastboot {
 /*************************** PUBLIC *******************************/
 FastBootDriver::FastBootDriver(Transport* transport, std::function<void(std::string&)> info,
                                bool no_checks)
-    : transport(transport) {
+    : transport_(transport) {
     info_cb_ = info;
     disable_checks_ = no_checks;
 }
 
+FastBootDriver::~FastBootDriver() {
+    set_transport(nullptr);
+}
+
 RetCode FastBootDriver::Boot(std::string* response, std::vector<std::string>* info) {
     return RawCommand(Commands::BOOT, response, info);
 }
@@ -93,6 +97,11 @@ RetCode FastBootDriver::Reboot(std::string* response, std::vector<std::string>*
     return RawCommand(Commands::REBOOT, response, info);
 }
 
+RetCode FastBootDriver::RebootTo(std::string target, std::string* response,
+                                 std::vector<std::string>* info) {
+    return RawCommand("reboot-" + target, response, info);
+}
+
 RetCode FastBootDriver::SetActive(const std::string& part, std::string* response,
                                   std::vector<std::string>* info) {
     return RawCommand(Commands::SET_ACTIVE + part, response, info);
@@ -332,7 +341,7 @@ std::string FastBootDriver::Error() {
 }
 
 RetCode FastBootDriver::WaitForDisconnect() {
-    return transport->WaitForDisconnect() ? IO_ERROR : SUCCESS;
+    return transport_->WaitForDisconnect() ? IO_ERROR : SUCCESS;
 }
 
 /****************************** PROTECTED *************************************/
@@ -344,7 +353,7 @@ RetCode FastBootDriver::RawCommand(const std::string& cmd, std::string* response
         return BAD_ARG;
     }
 
-    if (transport->Write(cmd.c_str(), cmd.size()) != static_cast<int>(cmd.size())) {
+    if (transport_->Write(cmd.c_str(), cmd.size()) != static_cast<int>(cmd.size())) {
         error_ = ErrnoStr("Write to device failed");
         return IO_ERROR;
     }
@@ -378,7 +387,7 @@ RetCode FastBootDriver::HandleResponse(std::string* response, std::vector<std::s
     // erase response
     set_response("");
     while ((std::chrono::system_clock::now() - start) < std::chrono::seconds(RESP_TIMEOUT)) {
-        int r = transport->Read(status, FB_RESPONSE_SZ);
+        int r = transport_->Read(status, FB_RESPONSE_SZ);
         if (r < 0) {
             error_ = ErrnoStr("Status read failed");
             return IO_ERROR;
@@ -472,7 +481,7 @@ RetCode FastBootDriver::SendBuffer(const void* buf, size_t size) {
         return BAD_ARG;
     }
     // Write the buffer
-    ssize_t tmp = transport->Write(buf, size);
+    ssize_t tmp = transport_->Write(buf, size);
 
     if (tmp < 0) {
         error_ = ErrnoStr("Write to device failed in SendBuffer()");
@@ -493,7 +502,7 @@ RetCode FastBootDriver::ReadBuffer(std::vector<char>& buf) {
 
 RetCode FastBootDriver::ReadBuffer(void* buf, size_t size) {
     // Read the buffer
-    ssize_t tmp = transport->Read(buf, size);
+    ssize_t tmp = transport_->Read(buf, size);
 
     if (tmp < 0) {
         error_ = ErrnoStr("Read from device failed in ReadBuffer()");
@@ -539,4 +548,12 @@ int FastBootDriver::SparseWriteCallback(std::vector<char>& tpbuf, const char* da
     return 0;
 }
 
+void FastBootDriver::set_transport(Transport* transport) {
+    if (transport_) {
+        transport_->Close();
+        delete transport_;
+    }
+    transport_ = transport;
+}
+
 }  // End namespace fastboot

fastboot/fastboot_driver.h

@@ -66,6 +66,7 @@ class FastBootDriver {
     FastBootDriver(Transport* transport,
                    std::function<void(std::string&)> info = [](std::string&) {},
                    bool no_checks = false);
+    ~FastBootDriver();
 
     RetCode Boot(std::string* response = nullptr, std::vector<std::string>* info = nullptr);
     RetCode Continue(std::string* response = nullptr, std::vector<std::string>* info = nullptr);
@@ -87,6 +88,8 @@ class FastBootDriver {
     RetCode GetVarAll(std::vector<std::string>* response);
     RetCode Powerdown(std::string* response = nullptr, std::vector<std::string>* info = nullptr);
     RetCode Reboot(std::string* response = nullptr, std::vector<std::string>* info = nullptr);
+    RetCode RebootTo(std::string target, std::string* response = nullptr,
+                     std::vector<std::string>* info = nullptr);
     RetCode SetActive(const std::string& part, std::string* response = nullptr,
                       std::vector<std::string>* info = nullptr);
     RetCode Upload(const std::string& outfile, std::string* response = nullptr,
@@ -109,6 +112,10 @@ class FastBootDriver {
     std::string Error();
     RetCode WaitForDisconnect();
 
+    // Note: changing the transport will close and delete the existing one.
+    void set_transport(Transport* transport);
+    Transport* transport() const { return transport_; }
+
     // This is temporarily public for engine.cpp
     RetCode RawCommand(const std::string& cmd, std::string* response = nullptr,
                        std::vector<std::string>* info = nullptr, int* dsize = nullptr);
@@ -136,7 +143,7 @@ class FastBootDriver {
         static const std::string VERIFY;
     };
 
-    Transport* const transport;
+    Transport* transport_;
 
   private:
     RetCode SendBuffer(int fd, size_t size);