[Android] Add --unmap-all to forwarder2.

tools/android/forwarder2/host_forwarder_main.cc

 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include <errno.h>
 #include <signal.h>
 #include <stddef.h>
 #include <stdint.h>
 #include <sys/types.h>
 #include <sys/wait.h>
@@ skipping 33 matching lines @@
 #include "tools/android/forwarder2/util.h"
 
 namespace forwarder2 {
 namespace {
 
 const char kLogFilePath[] = "/tmp/host_forwarder_log";
 const char kDaemonIdentifier[] = "chrome_host_forwarder_daemon";
 
 const int kBufSize = 256;
 
+enum : int {
+  MAP = 0,
+  UNMAP = 1,
+  UNMAP_ALL = 2,
+};
+
 // Needs to be global to be able to be accessed from the signal handler.
 PipeNotifier* g_notifier = NULL;
 
 // Lets the daemon fetch the exit notifier file descriptor.
 int GetExitNotifierFD() {
   DCHECK(g_notifier);
   return g_notifier->receiver_fd();
 }
 
 void KillHandler(int signal_number) {
@@ skipping 29 matching lines @@
   ~HostControllersManager() {
     if (!thread_.get())
       return;
     // Delete the controllers on the thread they were created on.
     thread_->task_runner()->DeleteSoon(
         FROM_HERE, controllers_.release());
   }
 
   void HandleRequest(const std::string& adb_path,
                      const std::string& device_serial,
+                     int command,
                      int device_port,
                      int host_port,
                      std::unique_ptr<Socket> client_socket) {
     // Lazy initialize so that the CLI process doesn't get this thread created.
     InitOnce();
     thread_->task_runner()->PostTask(
         FROM_HERE,
         base::Bind(&HostControllersManager::HandleRequestOnInternalThread,
-                   base::Unretained(this), adb_path, device_serial, device_port,
-                   host_port, base::Passed(&client_socket)));
+                   base::Unretained(this), adb_path, device_serial, command,
+                   device_port, host_port, base::Passed(&client_socket)));
   }
 
   bool has_failed() const { return has_failed_; }
 
  private:
   typedef base::hash_map<
       std::string, linked_ptr<HostController> > HostControllerMap;
 
   static std::string MakeHostControllerMapKey(int adb_port, int device_port) {
     return base::StringPrintf("%d:%d", adb_port, device_port);
@@ skipping 22 matching lines @@
       return;
     }
     DCHECK(manager->thread_->task_runner()->RunsTasksOnCurrentThread());
     // Note that this will delete |controller| which is owned by the map.
     DeleteRefCountedValueInMap(
         MakeHostControllerMapKey(
             controller->adb_port(), controller->device_port()),
         manager->controllers_.get());
   }
 
-  void HandleRequestOnInternalThread(const std::string& adb_path,
-                                     const std::string& device_serial,
-                                     int device_port,
-                                     int host_port,
-                                     std::unique_ptr<Socket> client_socket) {
-    const int adb_port = GetAdbPortForDevice(adb_path, 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.\n",
-          client_socket.get());
-      return;
-    }
-    if (device_port < 0) {
-      // Remove the previously created host controller.
-      const std::string controller_key = MakeHostControllerMapKey(
-          adb_port, -device_port);
-      const bool controller_did_exist = DeleteRefCountedValueInMap(
-          controller_key, controllers_.get());
-      if (!controller_did_exist) {
-        SendMessage("ERROR: could not unmap port.\n", client_socket.get());
-        LogExistingControllers(client_socket);
-      } else {
-        SendMessage("OK", client_socket.get());
-      }
-
-      RemoveAdbPortForDeviceIfNeeded(adb_path, device_serial);
-      return;
-    }
+  void Map(const std::string& adb_path,
+           const std::string& device_serial,
+           int adb_port,
+           int device_port,
+           int host_port,
+           Socket* client_socket) {
     if (host_port < 0) {
-      SendMessage("ERROR: missing host port\n", client_socket.get());
+      SendMessage("ERROR: missing host port\n", client_socket);
       return;
     }
     const bool use_dynamic_port_allocation = device_port == 0;
     if (!use_dynamic_port_allocation) {
       const std::string controller_key = MakeHostControllerMapKey(
           adb_port, device_port);
       if (controllers_->find(controller_key) != controllers_->end()) {
         LOG(INFO) << "Already forwarding device port " << device_port
                   << " to host port " << host_port;
         SendMessage(base::StringPrintf("%d:%d", device_port, host_port),
-                    client_socket.get());
+                    client_socket);
         return;
       }
     }
     // Create a new host controller.
     std::unique_ptr<HostController> host_controller(HostController::Create(
         device_serial, device_port, host_port, adb_port, GetExitNotifierFD(),
         base::Bind(&HostControllersManager::DeleteHostController,
                    weak_ptr_factory_.GetWeakPtr())));
     if (!host_controller.get()) {
       has_failed_ = true;
-      SendMessage("ERROR: Connection to device failed.\n", client_socket.get());
+      SendMessage("ERROR: Connection to device failed.\n", client_socket);
       LogExistingControllers(client_socket);
       return;
     }
     // Get the current allocated port.
     device_port = host_controller->device_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()))
+    if (!SendMessage(msg, client_socket))
       return;
     host_controller->Start();
     controllers_->insert(
         std::make_pair(MakeHostControllerMapKey(adb_port, device_port),
                        linked_ptr<HostController>(host_controller.release())));
   }
 
-  void LogExistingControllers(const std::unique_ptr<Socket>& client_socket) {
-    SendMessage("ERROR: Existing controllers:\n", client_socket.get());
-    for (const auto& controller : *controllers_) {
-      SendMessage(base::StringPrintf("ERROR: %s\n", controller.first.c_str()),
+  void Unmap(const std::string& adb_path,
+             const std::string& device_serial,
+             int adb_port,
+             int device_port,
+             Socket* client_socket) {
+    // Remove the previously created host controller.
+    const std::string controller_key =
+        MakeHostControllerMapKey(adb_port, device_port);
+    const bool controller_did_exist =
+        DeleteRefCountedValueInMap(controller_key, controllers_.get());
+    if (!controller_did_exist) {
+      SendMessage("ERROR: could not unmap port.\n", client_socket);
+      LogExistingControllers(client_socket);
+    } else {
+      SendMessage("OK", client_socket);
+    }
+
+    RemoveAdbPortForDeviceIfNeeded(adb_path, device_serial);
+  }
+
+  void UnmapAll(const std::string& adb_path,
+                const std::string& device_serial,
+                int adb_port,
+                Socket* client_socket) {
+    const std::string adb_port_str = base::StringPrintf("%d", adb_port);
+    for (HostControllerMap::const_iterator controller_key =
+             controllers_->cbegin();
+         controller_key != controllers_->cend(); ++controller_key) {
+      std::vector<std::string> pieces =
+          base::SplitString(controller_key->first, ":", base::KEEP_WHITESPACE,
+                            base::SPLIT_WANT_ALL);
+      if (pieces.size() == 2) {
+        if (pieces[0] == adb_port_str) {
+          DeleteRefCountedValueInMapFromIterator(controller_key,
+                                                 controllers_.get());
+        }
+      } else {
+        LOG(ERROR) << "Unexpected controller key: " << controller_key->first;
+      }
+    }
+
+    RemoveAdbPortForDeviceIfNeeded(adb_path, device_serial);
+    SendMessage("OK", client_socket);
+  }
+
+  void HandleRequestOnInternalThread(const std::string& adb_path,
+                                     const std::string& device_serial,
+                                     int command,
+                                     int device_port,
+                                     int host_port,
+                                     std::unique_ptr<Socket> client_socket) {
+    const int adb_port = GetAdbPortForDevice(adb_path, 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.\n",
           client_socket.get());
+      return;
+    }
+    switch (command) {
+      case MAP:
+        Map(adb_path, device_serial, adb_port, device_port, host_port,
+            client_socket.get());
+        break;
+      case UNMAP:
+        Unmap(adb_path, device_serial, adb_port, device_port,
+              client_socket.get());
+        break;
+      case UNMAP_ALL:
+        UnmapAll(adb_path, device_serial, adb_port, client_socket.get());
+        break;
+      default:
+        SendMessage(
+            base::StringPrintf("ERROR: unrecognized command %d\n", command),
+            client_socket.get());
+        break;
     }
   }
 
+  void LogExistingControllers(Socket* client_socket) {
+    SendMessage("ERROR: Existing controllers:\n", client_socket);
+    for (const auto& controller : *controllers_) {
+      SendMessage(base::StringPrintf("ERROR: %s\n", controller.first.c_str()),
+                  client_socket);
+    }
+  }
+
   void RemoveAdbPortForDeviceIfNeeded(const std::string& adb_path,
                                       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;
 
     int port = it->second;
     const std::string prefix = base::StringPrintf("%d:", port);
     for (HostControllerMap::const_iterator others = controllers_->begin();
@@ skipping 97 matching lines @@
     const int bytes_read = client_socket->Read(buf, sizeof(buf));
     if (bytes_read <= 0) {
       if (client_socket->DidReceiveEvent())
         return;
       PError("Read()");
       has_failed_ = true;
       return;
     }
     const base::Pickle command_pickle(buf, bytes_read);
     base::PickleIterator pickle_it(command_pickle);
+
     std::string device_serial;
     CHECK(pickle_it.ReadString(&device_serial));
-    int device_port;
-    if (!pickle_it.ReadInt(&device_port)) {
-      client_socket->WriteString("ERROR: missing device port\n");
+
+    int command;
+    if (!pickle_it.ReadInt(&command)) {
+      client_socket->WriteString("Error: missing command\n");
       return;
     }
+
+    int device_port;
+    if (!pickle_it.ReadInt(&device_port))
+      device_port = -1;
+
     int host_port;
     if (!pickle_it.ReadInt(&host_port))
       host_port = -1;
-    controllers_manager_.HandleRequest(adb_path_, device_serial, device_port,
-                                       host_port, std::move(client_socket));
+    controllers_manager_.HandleRequest(adb_path_, device_serial, command,
+                                       device_port, host_port,
+                                       std::move(client_socket));
   }
 
  private:
   std::string adb_path_;
   bool has_failed_;
   HostControllersManager controllers_manager_;
 
   DISALLOW_COPY_AND_ASSIGN(ServerDelegate);
 };
 
@@ skipping 29 matching lines @@
   const base::Pickle command_pickle_;
   bool has_failed_;
 };
 
 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"
+               "  --unmap-all\n"
                "  --adb PATH_TO_ADB\n"
                "  --kill-server\n";
   exit(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_t>::max()) {
@@ skipping 13 matching lines @@
   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]));
+    pickle.WriteInt(UNMAP);
+    pickle.WriteInt(PortToInt(args[0]));
+  } else if (cmd_line.HasSwitch("unmap-all")) {
+    pickle.WriteInt(UNMAP_ALL);
   } else if (cmd_line.HasSwitch("map")) {
     if (args.size() != 2)
       ExitWithUsage();
+    pickle.WriteInt(MAP);
     pickle.WriteInt(PortToInt(args[0]));
     pickle.WriteInt(PortToInt(args[1]));
   } else {
     ExitWithUsage();
   }
 
   if (cmd_line.HasSwitch("adb")) {
     adb_path = cmd_line.GetSwitchValueASCII("adb");
   }
 
@@ skipping 13 matching lines @@
 
   return client_delegate.has_failed() || daemon_delegate.has_failed();
 }
 
 }  // namespace
 }  // namespace forwarder2
 
 int main(int argc, char** argv) {
   return forwarder2::RunHostForwarder(argc, argv);
 }