Add --serial-id option to host_forwarder.

This moves the forwarder's control channel set up logic (which can be
considered as an internal implementation detail) from the Python script to the
host_forwarder binary.

This also changes the command line interface for the host_forwarder and
device_forwarder binaries.

BUG=242846
R=bulach@chromium.org, digit@chromium.org

Review URL: https://codereview.chromium.org/18354005

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@209945 0039d316-1c4b-4281-b951-d872f2087c98

2 files changed, 127 insertions, 132 deletions

diff --git a/tools/android/forwarder2/device_forwarder_main.cc b/tools/android/forwarder2/device_forwarder_main.cc
index f822d28..26904b4 100644
--- a/tools/android/forwarder2/device_forwarder_main.cc
+++ b/tools/android/forwarder2/device_forwarder_main.cc
@@ -3,9 +3,9 @@
 // found in the LICENSE file.
 
 #include <signal.h>
-#include <stdio.h>
 #include <stdlib.h>
 
+#include <iostream>
 #include <string>
 
 #include "base/at_exit.h"
@@ -29,11 +29,9 @@ forwarder2::PipeNotifier* g_notifier = NULL;
 
 const int kBufSize = 256;
 
+const char kUnixDomainSocketPath[] = "chrome_device_forwarder";
 const char kDaemonIdentifier[] = "chrome_device_forwarder_daemon";
 
-const char kKillServerCommand[] = "kill-server";
-const char kStartCommand[] = "start";
-
 void KillHandler(int /* unused */) {
   CHECK(g_notifier);
   if (!g_notifier->Notify())
@@ -48,6 +46,8 @@ int GetExitNotifierFD() {
 
 class ServerDelegate : public Daemon::ServerDelegate {
  public:
+  ServerDelegate() : initialized_(false) {}
+
   // Daemon::ServerDelegate:
   virtual void Init() OVERRIDE {
     DCHECK(!g_notifier);
@@ -59,45 +59,28 @@ class ServerDelegate : public Daemon::ServerDelegate {
   }
 
   virtual void OnClientConnected(scoped_ptr<Socket> client_socket) OVERRIDE {
-    char buf[kBufSize];
-    const int bytes_read = client_socket->Read(buf, sizeof(buf));
-    if (bytes_read <= 0) {
-      if (client_socket->DidReceiveEvent())
-        return;
-      PError("Read()");
-      return;
-    }
-    const std::string adb_socket_path(buf, bytes_read);
-    if (adb_socket_path == adb_socket_path_) {
+    if (initialized_) {
       client_socket->WriteString("OK");
       return;
     }
-    if (!adb_socket_path_.empty()) {
-      client_socket->WriteString(
-          base::StringPrintf(
-              "ERROR: Device controller already running (adb_socket_path=%s)",
-              adb_socket_path_.c_str()));
-      return;
-    }
-    adb_socket_path_ = adb_socket_path;
     controller_thread_->message_loop()->PostTask(
         FROM_HERE,
-        base::Bind(&ServerDelegate::StartController, adb_socket_path,
-                   GetExitNotifierFD(), base::Passed(&client_socket)));
+        base::Bind(&ServerDelegate::StartController, GetExitNotifierFD(),
+                   base::Passed(&client_socket)));
+    initialized_ = true;
   }
 
   virtual void OnServerExited() OVERRIDE {}
 
  private:
-  static void StartController(const std::string& adb_socket_path,
-                              int exit_notifier_fd,
+  static void StartController(int exit_notifier_fd,
                               scoped_ptr<Socket> client_socket) {
     forwarder2::DeviceController controller(exit_notifier_fd);
-    if (!controller.Init(adb_socket_path)) {
+    if (!controller.Init(kUnixDomainSocketPath)) {
       client_socket->WriteString(
           base::StringPrintf("ERROR: Could not initialize device controller "
                              "with ADB socket path: %s",
-                             adb_socket_path.c_str()));
+                             kUnixDomainSocketPath));
       return;
     }
     client_socket->WriteString("OK");
@@ -110,23 +93,17 @@ class ServerDelegate : public Daemon::ServerDelegate {
 
   base::AtExitManager at_exit_manager_;  // Used by base::Thread.
   scoped_ptr<base::Thread> controller_thread_;
-  std::string adb_socket_path_;
+  bool initialized_;
 };
 
 class ClientDelegate : public Daemon::ClientDelegate {
  public:
-  ClientDelegate(const std::string& adb_socket)
-      : adb_socket_(adb_socket),
-        has_failed_(false) {
-  }
+  ClientDelegate() : has_failed_(false) {}
 
   bool has_failed() const { return has_failed_; }
 
   // Daemon::ClientDelegate:
   virtual void OnDaemonReady(Socket* daemon_socket) OVERRIDE {
-    // Send the adb socket path to the daemon.
-    CHECK(daemon_socket->Write(adb_socket_.c_str(),
-                               adb_socket_.length()));
     char buf[kBufSize];
     const int bytes_read = daemon_socket->Read(
         buf, sizeof(buf) - 1 /* leave space for null terminator */);
@@ -142,31 +119,26 @@ class ClientDelegate : public Daemon::ClientDelegate {
   }
 
  private:
-  const std::string adb_socket_;
   bool has_failed_;
 };
 
 int RunDeviceForwarder(int argc, char** argv) {
-  if (argc != 2) {
-    fprintf(stderr,
-            "Usage: %s kill-server|<adb_socket>\n"
-            "  <adb_socket> is the abstract Unix Domain Socket path "
-            "where Adb is configured to forward from.\n", argv[0]);
+  CommandLine::Init(argc, argv);  // Needed by logging.
+  const bool kill_server = CommandLine::ForCurrentProcess()->HasSwitch(
+      "kill-server");
+  if ((kill_server && argc != 2) || (!kill_server && argc != 1)) {
+    std::cerr << "Usage: device_forwarder [--kill-server]" << std::endl;
     return 1;
   }
-  CommandLine::Init(argc, argv);  // Needed by logging.
-  const char* const command =
-      !strcmp(argv[1], kKillServerCommand) ? kKillServerCommand : kStartCommand;
-  ClientDelegate client_delegate(argv[1]);
+  ClientDelegate client_delegate;
   ServerDelegate daemon_delegate;
   const char kLogFilePath[] = "";  // Log to logcat.
   Daemon daemon(kLogFilePath, kDaemonIdentifier, &client_delegate,
                 &daemon_delegate, &GetExitNotifierFD);
 
-  if (command == kKillServerCommand)
+  if (kill_server)
     return !daemon.Kill();
 
-  DCHECK(command == kStartCommand);
   if (!daemon.SpawnIfNeeded())
     return 1;
   return client_delegate.has_failed();
diff --git a/tools/android/forwarder2/host_forwarder_main.cc b/tools/android/forwarder2/host_forwarder_main.cc
index 0f0ca68..82bb7a8 100644
--- a/tools/android/forwarder2/host_forwarder_main.cc
+++ b/tools/android/forwarder2/host_forwarder_main.cc
@@ -4,10 +4,13 @@
 
 #include <errno.h>
 #include <signal.h>
+#include <sys/types.h>
+#include <sys/wait.h>
 #include <unistd.h>
 
 #include <cstdio>
-#include <cstring>
+#include <iostream>
+#include <limits>
 #include <string>
 #include <utility>
 #include <vector>
@@ -20,6 +23,7 @@
 #include "base/logging.h"
 #include "base/memory/linked_ptr.h"
 #include "base/memory/scoped_vector.h"
+#include "base/pickle.h"
 #include "base/posix/eintr_wrapper.h"
 #include "base/safe_strerror_posix.h"
 #include "base/strings/string_number_conversions.h"
@@ -72,40 +76,6 @@ void KillHandler(int signal_number) {
     exit(1);
 }
 
-// Format of |command|:
-//   <ADB port>:<Device port>[:<Forward to port>:<Forward to address>].
-bool ParseForwardCommand(const std::string& command,
-                         int* adb_port,
-                         int* device_port,
-                         std::string* forward_to_host,
-                         int* forward_to_port) {
-  std::vector<std::string> command_pieces;
-  base::SplitString(command, ':', &command_pieces);
-
-  if (command_pieces.size() < 2 ||
-      !base::StringToInt(command_pieces[0], adb_port) ||
-      !base::StringToInt(command_pieces[1], device_port))
-    return false;
-
-  if (command_pieces.size() > 2) {
-    if (!base::StringToInt(command_pieces[2], forward_to_port))
-      return false;
-    if (command_pieces.size() > 3)
-      *forward_to_host = command_pieces[3];
-  } else {
-    *forward_to_port = *device_port;
-  }
-  return true;
-}
-
-bool IsForwardCommandValid(const std::string& command) {
-  int adb_port, device_port, forward_to_port;
-  std::string forward_to_host;
-  std::vector<std::string> command_pieces;
-  return ParseForwardCommand(
-      command, &adb_port, &device_port, &forward_to_host, &forward_to_port);
-}
-
 class ServerDelegate : public Daemon::ServerDelegate {
  public:
   ServerDelegate() : has_failed_(false) {}
@@ -131,18 +101,21 @@ class ServerDelegate : public Daemon::ServerDelegate {
       has_failed_ = true;
       return;
     }
-    const std::string command(buf, bytes_read);
-    int adb_port = 0;
-    int device_port = 0;
-    std::string forward_to_host;
-    int forward_to_port = 0;
-    const bool succeeded = ParseForwardCommand(
-        command, &adb_port, &device_port, &forward_to_host, &forward_to_port);
-    if (!succeeded) {
-      has_failed_ = true;
-      const std::string msg = base::StringPrintf(
-          "ERROR: Could not parse forward command '%s'", command.c_str());
-      SendMessage(msg, client_socket.get());
+    const Pickle command_pickle(buf, bytes_read);
+    PickleIterator pickle_it(command_pickle);
+    std::string device_serial;
+    CHECK(pickle_it.ReadString(&device_serial));
+    int device_port;
+    if (!pickle_it.ReadInt(&device_port)) {
+      SendMessage("ERROR: missing device port", client_socket.get());
+      return;
+    }
+    const int adb_port = GetAdbPortForDevice(device_serial);
+    if (adb_port < 0) {
+      SendMessage(
+          "ERROR: could not get adb port for device. You might need to add "
+          "'adb' to your PATH or provide the device serial id.",
+          client_socket.get());
       return;
     }
     if (device_port < 0) {
@@ -156,10 +129,15 @@ class ServerDelegate : public Daemon::ServerDelegate {
           client_socket.get());
       return;
     }
+    int host_port;
+    if (!pickle_it.ReadInt(&host_port)) {
+      SendMessage("ERROR: missing host port", client_socket.get());
+      return;
+    }
     // Create a new host controller.
     scoped_ptr<HostController> host_controller(
-        new HostController(device_port, forward_to_host, forward_to_port,
-                           adb_port, GetExitNotifierFD()));
+        new HostController(device_port, "127.0.0.1", host_port, adb_port,
+                           GetExitNotifierFD()));
     if (!host_controller->Connect()) {
       has_failed_ = true;
       SendMessage("ERROR: Connection to device failed.", client_socket.get());
@@ -167,10 +145,9 @@ class ServerDelegate : public Daemon::ServerDelegate {
     }
     // Get the current allocated port.
     device_port = host_controller->device_port();
-    LOG(INFO) << "Forwarding device port " << device_port << " to host "
-              << forward_to_host << ":" << forward_to_port;
-    const std::string msg = base::StringPrintf(
-        "%d:%d", device_port, forward_to_port);
+    LOG(INFO) << "Forwarding device port " << device_port << " to host port "
+              << host_port;
+    const std::string msg = base::StringPrintf("%d:%d", device_port, host_port);
     if (!SendMessage(msg, client_socket.get()))
       return;
     host_controller->Start();
@@ -201,6 +178,29 @@ class ServerDelegate : public Daemon::ServerDelegate {
     return base::StringPrintf("%d:%d", adb_port, device_port);
   }
 
+  int GetAdbPortForDevice(const std::string& device_serial) {
+    base::hash_map<std::string, int>::const_iterator it =
+        device_serial_to_adb_port_map_.find(device_serial);
+    if (it != device_serial_to_adb_port_map_.end())
+      return it->second;
+    Socket bind_socket;
+    CHECK(bind_socket.BindTcp("127.0.0.1", 0));
+    const int port = bind_socket.GetPort();
+    bind_socket.Close();
+    const std::string serial_part = device_serial.empty() ?
+        std::string() : std::string("-s ") + device_serial;
+    const std::string command = base::StringPrintf(
+        "adb %s forward tcp:%d localabstract:chrome_device_forwarder",
+        device_serial.empty() ? "" : serial_part.c_str(),
+        port);
+    LOG(INFO) << command;
+    const int ret = system(command.c_str());
+    if (ret < 0 || !WIFEXITED(ret) || WEXITSTATUS(ret) != 0)
+      return -1;
+    device_serial_to_adb_port_map_[device_serial] = port;
+    return port;
+  }
+
   bool SendMessage(const std::string& msg, Socket* client_socket) {
     bool result = client_socket->WriteString(msg);
     DCHECK(result);
@@ -209,6 +209,7 @@ class ServerDelegate : public Daemon::ServerDelegate {
     return result;
   }
 
+  base::hash_map<std::string, int> device_serial_to_adb_port_map_;
   HostControllerMap controllers_;
   bool has_failed_;
 
@@ -217,8 +218,8 @@ class ServerDelegate : public Daemon::ServerDelegate {
 
 class ClientDelegate : public Daemon::ClientDelegate {
  public:
-  ClientDelegate(const std::string& forward_command)
-      : forward_command_(forward_command),
+  ClientDelegate(const Pickle& command_pickle)
+      : command_pickle_(command_pickle),
         has_failed_(false) {
   }
 
@@ -227,7 +228,9 @@ class ClientDelegate : public Daemon::ClientDelegate {
   // Daemon::ClientDelegate:
   virtual void OnDaemonReady(Socket* daemon_socket) OVERRIDE {
     // Send the forward command to the daemon.
-    CHECK(daemon_socket->WriteString(forward_command_));
+    CHECK_EQ(command_pickle_.size(),
+             daemon_socket->WriteNumBytes(command_pickle_.data(),
+                                          command_pickle_.size()));
     char buf[kBufSize];
     const int bytes_read = daemon_socket->Read(
         buf, sizeof(buf) - 1 /* leave space for null terminator */);
@@ -244,49 +247,69 @@ class ClientDelegate : public Daemon::ClientDelegate {
   }
 
  private:
-  const std::string forward_command_;
+  const Pickle command_pickle_;
   bool has_failed_;
 };
 
-void PrintUsage(const char* program_name) {
-  LOG(ERROR) << program_name
-             << " adb_port:from_port:to_port:to_host\n"
-                "<adb port> is the TCP port Adb is configured to forward to.\n"
-                "Note that <from_port> can be unmapped by making it negative.";
+void ExitWithUsage() {
+  std::cerr << "Usage: host_forwarder [options]\n\n"
+               "Options:\n"
+               "  --serial-id=[0-9A-Z]{16}]\n"
+               "  --map DEVICE_PORT HOST_PORT\n"
+               "  --unmap DEVICE_PORT\n"
+               "  --kill-server\n";
+  exit(1);
 }
 
-int RunHostForwarder(int argc, char** argv) {
-  if (!CommandLine::Init(argc, argv)) {
-    LOG(ERROR) << "Could not initialize command line";
-    return 1;
-  }
-  const CommandLine& command_line = *CommandLine::ForCurrentProcess();
-  const char* command = NULL;
-  int adb_port = 0;
-  if (argc != 2) {
-    PrintUsage(argv[0]);
-    return 1;
+int PortToInt(const std::string& s) {
+  int value;
+  // Note that 0 is a valid port (used for dynamic port allocation).
+  if (!base::StringToInt(s, &value) || value < 0 ||
+      value > std::numeric_limits<uint16>::max()) {
+    LOG(ERROR) << "Could not convert string " << s << " to port";
+    ExitWithUsage();
   }
-  if (!strcmp(argv[1], kKillServerCommand)) {
-    command = kKillServerCommand;
+  return value;
+}
+
+int RunHostForwarder(int argc, char** argv) {
+  CommandLine::Init(argc, argv);
+  const CommandLine& cmd_line = *CommandLine::ForCurrentProcess();
+  bool kill_server = false;
+
+  Pickle pickle;
+  pickle.WriteString(
+      cmd_line.HasSwitch("serial-id") ?
+          cmd_line.GetSwitchValueASCII("serial-id") : std::string());
+
+  const std::vector<std::string> args = cmd_line.GetArgs();
+  if (cmd_line.HasSwitch("kill-server")) {
+    kill_server = true;
+  } else if (cmd_line.HasSwitch("unmap")) {
+    if (args.size() != 1)
+      ExitWithUsage();
+    // Note the minus sign below.
+    pickle.WriteInt(-PortToInt(args[0]));
+  } else if (cmd_line.HasSwitch("map")) {
+    if (args.size() != 2)
+      ExitWithUsage();
+    pickle.WriteInt(PortToInt(args[0]));
+    pickle.WriteInt(PortToInt(args[1]));
   } else {
-    command = kForwardCommand;
-    if (!IsForwardCommandValid(argv[1])) {
-      PrintUsage(argv[0]);
-      return 1;
-    }
+    ExitWithUsage();
   }
 
-  ClientDelegate client_delegate(argv[1]);
+  if (kill_server && args.size() > 0)
+    ExitWithUsage();
+
+  ClientDelegate client_delegate(pickle);
   ServerDelegate daemon_delegate;
   Daemon daemon(
       kLogFilePath, kDaemonIdentifier, &client_delegate, &daemon_delegate,
       &GetExitNotifierFD);
 
-  if (command == kKillServerCommand)
+  if (kill_server)
     return !daemon.Kill();
-
-  DCHECK(command == kForwardCommand);
   if (!daemon.SpawnIfNeeded())
     return 1;