adb: support wait-for- with multiple states.

Test: manual, adb wait-for-device-recovery
Change-Id: I92034c477d28f847e24ec72bbb87b252fba875e1
diff --git a/adb/client/commandline.cpp b/adb/client/commandline.cpp
index 1186060..7c7da08 100644
--- a/adb/client/commandline.cpp
+++ b/adb/client/commandline.cpp
@@ -194,8 +194,8 @@
         "     generate adb public/private key; private key stored in FILE,\n"
         "\n"
         "scripting:\n"
-        " wait-for[-TRANSPORT]-STATE\n"
-        "     wait for device to be in the given state\n"
+        " wait-for[-TRANSPORT]-STATE...\n"
+        "     wait for device to be in a given state\n"
         "     STATE: device, recovery, rescue, sideload, bootloader, or disconnect\n"
         "     TRANSPORT: usb, local, or any [default=any]\n"
         " get-state                print offline | bootloader | device\n"
@@ -1057,17 +1057,16 @@
 static bool wait_for_device(const char* service,
                             std::optional<std::chrono::milliseconds> timeout = std::nullopt) {
     std::vector<std::string> components = android::base::Split(service, "-");
-    if (components.size() < 3 || components.size() > 4) {
+    if (components.size() < 3) {
         fprintf(stderr, "adb: couldn't parse 'wait-for' command: %s\n", service);
         return false;
     }
 
-    TransportType t;
-    adb_get_transport(&t, nullptr, nullptr);
-
-    // Was the caller vague about what they'd like us to wait for?
-    // If so, check they weren't more specific in their choice of transport type.
-    if (components.size() == 3) {
+    // If the first thing after "wait-for-" wasn't a TRANSPORT, insert whatever
+    // the current transport implies.
+    if (components[2] != "usb" && components[2] != "local" && components[2] != "any") {
+        TransportType t;
+        adb_get_transport(&t, nullptr, nullptr);
         auto it = components.begin() + 2;
         if (t == kTransportUsb) {
             components.insert(it, "usb");
@@ -1076,23 +1075,9 @@
         } else {
             components.insert(it, "any");
         }
-    } else if (components[2] != "any" && components[2] != "local" && components[2] != "usb") {
-        fprintf(stderr, "adb: unknown type %s; expected 'any', 'local', or 'usb'\n",
-                components[2].c_str());
-        return false;
     }
 
-    if (components[3] != "any" && components[3] != "bootloader" && components[3] != "device" &&
-        components[3] != "recovery" && components[3] != "rescue" && components[3] != "sideload" &&
-        components[3] != "disconnect") {
-        fprintf(stderr,
-                "adb: unknown state %s; "
-                "expected 'any', 'bootloader', 'device', 'recovery', 'rescue', 'sideload', or "
-                "'disconnect'\n",
-                components[3].c_str());
-        return false;
-    }
-
+    // Stitch it back together and send it over...
     std::string cmd = format_host_command(android::base::Join(components, "-").c_str());
     if (timeout) {
         std::thread([timeout]() {
diff --git a/adb/services.cpp b/adb/services.cpp
index 853d658..d87948c 100644
--- a/adb/services.cpp
+++ b/adb/services.cpp
@@ -95,56 +95,6 @@
 }
 
 #if ADB_HOST
-struct state_info {
-    TransportType transport_type;
-    std::string serial;
-    TransportId transport_id;
-    ConnectionState state;
-};
-
-static void wait_for_state(unique_fd fd, state_info* sinfo) {
-    D("wait_for_state %d", sinfo->state);
-
-    while (true) {
-        bool is_ambiguous = false;
-        std::string error = "unknown error";
-        const char* serial = sinfo->serial.length() ? sinfo->serial.c_str() : nullptr;
-        atransport* t = acquire_one_transport(sinfo->transport_type, serial, sinfo->transport_id,
-                                              &is_ambiguous, &error);
-        if (sinfo->state == kCsOffline) {
-            // wait-for-disconnect uses kCsOffline, we don't actually want to wait for 'offline'.
-            if (t == nullptr) {
-                SendOkay(fd);
-                break;
-            }
-        } else if (t != nullptr &&
-                   (sinfo->state == kCsAny || sinfo->state == t->GetConnectionState())) {
-            SendOkay(fd);
-            break;
-        }
-
-        if (!is_ambiguous) {
-            adb_pollfd pfd = {.fd = fd.get(), .events = POLLIN};
-            int rc = adb_poll(&pfd, 1, 100);
-            if (rc < 0) {
-                SendFail(fd, error);
-                break;
-            } else if (rc > 0 && (pfd.revents & POLLHUP) != 0) {
-                // The other end of the socket is closed, probably because the other side was
-                // terminated, bail out.
-                break;
-            }
-
-            // Try again...
-        } else {
-            SendFail(fd, error);
-            break;
-        }
-    }
-
-    D("wait_for_state is done");
-}
-
 void connect_emulator(const std::string& port_spec, std::string* response) {
     std::vector<std::string> pieces = android::base::Split(port_spec, ",");
     if (pieces.size() != 2) {
@@ -201,6 +151,80 @@
     adb_wifi_pair_device(host, password, response);
     SendProtocolString(fd.get(), response);
 }
+
+static void wait_service(unique_fd fd, std::string serial, TransportId transport_id,
+                         std::string spec) {
+    std::vector<std::string> components = android::base::Split(spec, "-");
+    if (components.size() < 2) {
+        SendFail(fd, "short wait-for-: " + spec);
+        return;
+    }
+
+    TransportType transport_type;
+    if (components[0] == "local") {
+        transport_type = kTransportLocal;
+    } else if (components[0] == "usb") {
+        transport_type = kTransportUsb;
+    } else if (components[0] == "any") {
+        transport_type = kTransportAny;
+    } else {
+        SendFail(fd, "bad wait-for- transport: " + spec);
+        return;
+    }
+
+    std::vector<ConnectionState> states;
+    for (size_t i = 1; i < components.size(); ++i) {
+        if (components[i] == "device") {
+            states.push_back(kCsDevice);
+        } else if (components[i] == "recovery") {
+            states.push_back(kCsRecovery);
+        } else if (components[i] == "rescue") {
+            states.push_back(kCsRescue);
+        } else if (components[i] == "sideload") {
+            states.push_back(kCsSideload);
+        } else if (components[i] == "bootloader") {
+            states.push_back(kCsBootloader);
+        } else if (components[i] == "any") {
+            states.push_back(kCsAny);
+        } else if (components[i] == "disconnect") {
+            states.push_back(kCsOffline);
+        } else {
+            SendFail(fd, "bad wait-for- state: " + spec);
+            return;
+        }
+    }
+
+    while (true) {
+        bool is_ambiguous = false;
+        std::string error = "unknown error";
+        atransport* t =
+                acquire_one_transport(transport_type, !serial.empty() ? serial.c_str() : nullptr,
+                                      transport_id, &is_ambiguous, &error);
+
+        for (const auto& state : states) {
+            // wait-for-disconnect uses kCsOffline, we don't actually want to wait for 'offline'.
+            if ((t == nullptr && state == kCsOffline) || (t != nullptr && state == kCsAny) ||
+                (t != nullptr && state == t->GetConnectionState())) {
+                SendOkay(fd);
+                return;
+            }
+        }
+
+        if (is_ambiguous) {
+            SendFail(fd, error);
+            return;
+        }
+
+        // Sleep before retrying.
+        adb_pollfd pfd = {.fd = fd.get(), .events = POLLIN};
+        if (adb_poll(&pfd, 1, 100) != 0) {
+            // The other end of the socket is closed, probably because the
+            // client terminated. Bail out.
+            SendFail(fd, error);
+            return;
+        }
+    }
+}
 #endif
 
 #if ADB_HOST
@@ -211,45 +235,10 @@
     } else if (name == "track-devices-l") {
         return create_device_tracker(true);
     } else if (android::base::ConsumePrefix(&name, "wait-for-")) {
-        std::shared_ptr<state_info> sinfo = std::make_shared<state_info>();
-        if (sinfo == nullptr) {
-            fprintf(stderr, "couldn't allocate state_info: %s", strerror(errno));
-            return nullptr;
-        }
-
-        sinfo->serial = serial;
-        sinfo->transport_id = transport_id;
-
-        if (android::base::ConsumePrefix(&name, "local")) {
-            sinfo->transport_type = kTransportLocal;
-        } else if (android::base::ConsumePrefix(&name, "usb")) {
-            sinfo->transport_type = kTransportUsb;
-        } else if (android::base::ConsumePrefix(&name, "any")) {
-            sinfo->transport_type = kTransportAny;
-        } else {
-            return nullptr;
-        }
-
-        if (name == "-device") {
-            sinfo->state = kCsDevice;
-        } else if (name == "-recovery") {
-            sinfo->state = kCsRecovery;
-        } else if (name == "-rescue") {
-            sinfo->state = kCsRescue;
-        } else if (name == "-sideload") {
-            sinfo->state = kCsSideload;
-        } else if (name == "-bootloader") {
-            sinfo->state = kCsBootloader;
-        } else if (name == "-any") {
-            sinfo->state = kCsAny;
-        } else if (name == "-disconnect") {
-            sinfo->state = kCsOffline;
-        } else {
-            return nullptr;
-        }
-
-        unique_fd fd = create_service_thread(
-                "wait", [sinfo](unique_fd fd) { wait_for_state(std::move(fd), sinfo.get()); });
+        std::string spec(name);
+        unique_fd fd =
+                create_service_thread("wait", std::bind(wait_service, std::placeholders::_1,
+                                                        std::string(serial), transport_id, spec));
         return create_local_socket(std::move(fd));
     } else if (android::base::ConsumePrefix(&name, "connect:")) {
         std::string host(name);