From 504d393176a56763585490258aee70d3490e7a71 Mon Sep 17 00:00:00 2001
From: Joshua Duong <joshuaduong@google.com>
Date: Tue, 31 Mar 2020 10:58:50 -0700
Subject: [PATCH] [adb client] Add "adb mdns check" command.

This command will check if the mdns daemon is available on the host
machine.

Bug: 152510294

Test: pkill -9 mdnsd; adb mdns check; mdnsd; adb mdns check;
Test: test_adb.py
Change-Id: If644678a339763817a8a7adcbdc545626d161aba
---
 adb/adb.cpp                   | 18 ++++++++++++++++++
 adb/adb_wifi.h                |  2 ++
 adb/client/commandline.cpp    | 20 ++++++++++++++++++++
 adb/client/transport_mdns.cpp | 14 ++++++++++++++
 adb/test_adb.py               | 11 +++++++++++
 5 files changed, 65 insertions(+)

diff --git a/adb/adb.cpp b/adb/adb.cpp
index 06fdb69a1..aced079ae 100644
--- a/adb/adb.cpp
+++ b/adb/adb.cpp
@@ -1076,6 +1076,20 @@ void adb_set_reject_kill_server(bool value) {
     g_reject_kill_server = value;
 }
 
+static bool handle_mdns_request(std::string_view service, int reply_fd) {
+    if (!android::base::ConsumePrefix(&service, "mdns:")) {
+        return false;
+    }
+
+    if (service == "check") {
+        std::string check = mdns_check();
+        SendOkay(reply_fd, check);
+        return true;
+    }
+
+    return false;
+}
+
 HostRequestResult handle_host_request(std::string_view service, TransportType type,
                                       const char* serial, TransportId transport_id, int reply_fd,
                                       asocket* s) {
@@ -1320,6 +1334,10 @@ HostRequestResult handle_host_request(std::string_view service, TransportType ty
         return HostRequestResult::Handled;
     }
 
+    if (handle_mdns_request(service, reply_fd)) {
+        return HostRequestResult::Handled;
+    }
+
     return HostRequestResult::Unhandled;
 }
 
diff --git a/adb/adb_wifi.h b/adb/adb_wifi.h
index 585748c91..49d3a9134 100644
--- a/adb/adb_wifi.h
+++ b/adb/adb_wifi.h
@@ -27,6 +27,8 @@ void adb_wifi_pair_device(const std::string& host, const std::string& password,
                           std::string& response);
 bool adb_wifi_is_known_host(const std::string& host);
 
+std::string mdns_check();
+
 #else  // !ADB_HOST
 
 struct AdbdAuthContext;
diff --git a/adb/client/commandline.cpp b/adb/client/commandline.cpp
index 29f9dc14c..7b8b2d770 100644
--- a/adb/client/commandline.cpp
+++ b/adb/client/commandline.cpp
@@ -127,6 +127,7 @@ static void help() {
         "       localfilesystem:<unix domain socket name>\n"
         " reverse --remove REMOTE  remove specific reverse socket connection\n"
         " reverse --remove-all     remove all reverse socket connections from device\n"
+        " mdns check               check if mdns discovery is available\n"
         "\n"
         "file transfer:\n"
         " push [--sync] [-z ALGORITHM] [-Z] LOCAL... REMOTE\n"
@@ -1910,6 +1911,25 @@ int adb_commandline(int argc, const char** argv) {
 
         ReadOrderlyShutdown(fd);
         return 0;
+    } else if (!strcmp(argv[0], "mdns")) {
+        --argc;
+        if (argc < 1) error_exit("mdns requires an argument");
+        ++argv;
+
+        std::string error;
+        if (!adb_check_server_version(&error)) {
+            error_exit("failed to check server version: %s", error.c_str());
+        }
+
+        std::string query = "host:mdns:";
+        if (!strcmp(argv[0], "check")) {
+            if (argc != 1) error_exit("mdns %s doesn't take any arguments", argv[0]);
+            query += "check";
+        } else {
+            error_exit("unknown mdns command [%s]", argv[0]);
+        }
+
+        return adb_query_command(query);
     }
     /* do_sync_*() commands */
     else if (!strcmp(argv[0], "ls")) {
diff --git a/adb/client/transport_mdns.cpp b/adb/client/transport_mdns.cpp
index 22b9b1808..38a760fcb 100644
--- a/adb/client/transport_mdns.cpp
+++ b/adb/client/transport_mdns.cpp
@@ -545,3 +545,17 @@ void init_mdns_transport_discovery(void) {
     ResolvedService::initAdbSecure();
     std::thread(init_mdns_transport_discovery_thread).detach();
 }
+
+std::string mdns_check() {
+    uint32_t daemon_version;
+    uint32_t sz = sizeof(daemon_version);
+
+    auto dnserr = DNSServiceGetProperty(kDNSServiceProperty_DaemonVersion, &daemon_version, &sz);
+    std::string result = "ERROR: mdns daemon unavailable";
+    if (dnserr != kDNSServiceErr_NoError) {
+        return result;
+    }
+
+    result = android::base::StringPrintf("mdns daemon version [%u]", daemon_version);
+    return result;
+}
diff --git a/adb/test_adb.py b/adb/test_adb.py
index c872fb0f7..6989e3ba9 100755
--- a/adb/test_adb.py
+++ b/adb/test_adb.py
@@ -576,6 +576,17 @@ class PowerTest(unittest.TestCase):
         # If the power event was detected, the adb shell command should be broken very quickly.
         self.assertLess(end - start, 2)
 
+"""Use 'adb mdns check' to see if mdns discovery is available."""
+def is_adb_mdns_available():
+    with adb_server() as server_port:
+        output = subprocess.check_output(["adb", "-P", str(server_port),
+                                          "mdns", "check"]).strip()
+        return output.startswith(b"mdns daemon version")
+
+@unittest.skipIf(not is_adb_mdns_available(), "mdns feature not available")
+class MdnsTest(unittest.TestCase):
+    """Tests for adb mdns."""
+    pass
 
 def main():
     """Main entrypoint."""