adb: implement shell protocol.

Adds functionality for handling stdin/stdout/stderr streams and exit
codes using the shell protocol.

This CL just contains implementation for adbd which will not yet be
enabled. Once we have the ability to query transport features from the
adb client, another CL will add the implementation for the client side
and update the feature list to turn this on.

Note: this CL must be submitted together with a minadbd CL to update
the service_to_fd() function signature.

Bug: http://b/23030641

Change-Id: Ibed55e9c1946d8a35190696163ff63e8fb880238
diff --git a/adb/shell_service.cpp b/adb/shell_service.cpp
index 5f80a59..0274ae3 100644
--- a/adb/shell_service.cpp
+++ b/adb/shell_service.cpp
@@ -14,6 +14,67 @@
  * limitations under the License.
  */
 
+// Functionality for launching and managing shell subprocesses.
+//
+// There are two types of subprocesses, PTY or raw. PTY is typically used for
+// an interactive session, raw for non-interactive. There are also two methods
+// of communication with the subprocess, passing raw data or using a simple
+// protocol to wrap packets. The protocol allows separating stdout/stderr and
+// passing the exit code back, but is not backwards compatible.
+//   ----------------+--------------------------------------
+//   Type  Protocol  |   Exit code?  Separate stdout/stderr?
+//   ----------------+--------------------------------------
+//   PTY   No        |   No          No
+//   Raw   No        |   No          No
+//   PTY   Yes       |   Yes         No
+//   Raw   Yes       |   Yes         Yes
+//   ----------------+--------------------------------------
+//
+// Non-protocol subprocesses work by passing subprocess stdin/out/err through
+// a single pipe which is registered with a local socket in adbd. The local
+// socket uses the fdevent loop to pass raw data between this pipe and the
+// transport, which then passes data back to the adb client. Cleanup is done by
+// waiting in a separate thread for the subprocesses to exit and then signaling
+// a separate fdevent to close out the local socket from the main loop.
+//
+// ------------------+-------------------------+------------------------------
+//   Subprocess      |  adbd subprocess thread |   adbd main fdevent loop
+// ------------------+-------------------------+------------------------------
+//                   |                         |
+//   stdin/out/err <----------------------------->       LocalSocket
+//      |            |                         |
+//      |            |      Block on exit      |
+//      |            |           *             |
+//      v            |           *             |
+//     Exit         --->      Unblock          |
+//                   |           |             |
+//                   |           v             |
+//                   |   Notify shell exit FD --->    Close LocalSocket
+// ------------------+-------------------------+------------------------------
+//
+// The protocol requires the thread to intercept stdin/out/err in order to
+// wrap/unwrap data with shell protocol packets.
+//
+// ------------------+-------------------------+------------------------------
+//   Subprocess      |  adbd subprocess thread |   adbd main fdevent loop
+// ------------------+-------------------------+------------------------------
+//                   |                         |
+//     stdin/out   <--->      Protocol       <--->       LocalSocket
+//     stderr       --->      Protocol        --->       LocalSocket
+//       |           |                         |
+//       v           |                         |
+//      Exit        --->  Exit code protocol  --->       LocalSocket
+//                   |           |             |
+//                   |           v             |
+//                   |   Notify shell exit FD --->    Close LocalSocket
+// ------------------+-------------------------+------------------------------
+//
+// An alternate approach is to put the protocol wrapping/unwrapping in the main
+// fdevent loop, which has the advantage of being able to re-use the existing
+// select() code for handling data streams. However, implementation turned out
+// to be more complex due to partial reads and non-blocking I/O so this model
+// was chosen instead.
+
 #define TRACE_TAG TRACE_SHELL
 
 #include "shell_service.h"
@@ -22,8 +83,11 @@
 
 #include <errno.h>
 #include <pty.h>
+#include <sys/select.h>
 #include <termios.h>
 
+#include <memory>
+
 #include <base/logging.h>
 #include <base/stringprintf.h>
 #include <paths.h>
@@ -110,7 +174,8 @@
 
 class Subprocess {
   public:
-    Subprocess(const std::string& command, SubprocessType type);
+    Subprocess(const std::string& command, SubprocessType type,
+               SubprocessProtocol protocol);
     ~Subprocess();
 
     const std::string& command() const { return command_; }
@@ -129,26 +194,42 @@
     int OpenPtyChildFd(const char* pts_name, ScopedFd* error_sfd);
 
     static void* ThreadHandler(void* userdata);
+    void PassDataStreams();
     void WaitForExit();
 
+    ScopedFd* SelectLoop(fd_set* master_read_set_ptr,
+                         fd_set* master_write_set_ptr);
+
+    // Input/output stream handlers. Success returns nullptr, failure returns
+    // a pointer to the failed FD.
+    ScopedFd* PassInput();
+    ScopedFd* PassOutput(ScopedFd* sfd, ShellProtocol::Id id);
+
     const std::string command_;
     SubprocessType type_;
-
+    SubprocessProtocol protocol_;
     pid_t pid_ = -1;
     ScopedFd local_socket_sfd_;
 
+    // Shell protocol variables.
+    ScopedFd stdinout_sfd_, stderr_sfd_, protocol_sfd_;
+    std::unique_ptr<ShellProtocol> input_, output_;
+    size_t input_bytes_left_ = 0;
+
     DISALLOW_COPY_AND_ASSIGN(Subprocess);
 };
 
-Subprocess::Subprocess(const std::string& command, SubprocessType type)
-        : command_(command), type_(type) {
+Subprocess::Subprocess(const std::string& command, SubprocessType type,
+                       SubprocessProtocol protocol)
+        : command_(command), type_(type), protocol_(protocol) {
 }
 
 Subprocess::~Subprocess() {
 }
 
 bool Subprocess::ForkAndExec() {
-    ScopedFd parent_sfd, child_sfd, parent_error_sfd, child_error_sfd;
+    ScopedFd child_stdinout_sfd, child_stderr_sfd;
+    ScopedFd parent_error_sfd, child_error_sfd;
     char pts_name[PATH_MAX];
 
     // Create a socketpair for the fork() child to report any errors back to
@@ -161,9 +242,14 @@
     if (type_ == SubprocessType::kPty) {
         int fd;
         pid_ = forkpty(&fd, pts_name, nullptr, nullptr);
-        parent_sfd.Reset(fd);
+        stdinout_sfd_.Reset(fd);
     } else {
-        if (!CreateSocketpair(&parent_sfd, &child_sfd)) {
+        if (!CreateSocketpair(&stdinout_sfd_, &child_stdinout_sfd)) {
+            return false;
+        }
+        // Raw subprocess + shell protocol allows for splitting stderr.
+        if (protocol_ == SubprocessProtocol::kShell &&
+                !CreateSocketpair(&stderr_sfd_, &child_stderr_sfd)) {
             return false;
         }
         pid_ = fork();
@@ -179,16 +265,19 @@
         init_subproc_child();
 
         if (type_ == SubprocessType::kPty) {
-            child_sfd.Reset(OpenPtyChildFd(pts_name, &child_error_sfd));
+            child_stdinout_sfd.Reset(OpenPtyChildFd(pts_name, &child_error_sfd));
         }
 
-        dup2(child_sfd.fd(), STDIN_FILENO);
-        dup2(child_sfd.fd(), STDOUT_FILENO);
-        dup2(child_sfd.fd(), STDERR_FILENO);
+        dup2(child_stdinout_sfd.fd(), STDIN_FILENO);
+        dup2(child_stdinout_sfd.fd(), STDOUT_FILENO);
+        dup2(child_stderr_sfd.valid() ? child_stderr_sfd.fd() : child_stdinout_sfd.fd(),
+             STDERR_FILENO);
 
         // exec doesn't trigger destructors, close the FDs manually.
-        parent_sfd.Reset();
-        child_sfd.Reset();
+        stdinout_sfd_.Reset();
+        stderr_sfd_.Reset();
+        child_stdinout_sfd.Reset();
+        child_stderr_sfd.Reset();
         parent_error_sfd.Reset();
         close_on_exec(child_error_sfd.fd());
 
@@ -203,7 +292,8 @@
     }
 
     // Subprocess parent.
-    D("subprocess parent: subprocess FD = %d", parent_sfd.fd());
+    D("subprocess parent: stdin/stdout FD = %d, stderr FD = %d",
+      stdinout_sfd_.fd(), stderr_sfd_.fd());
 
     // Wait to make sure the subprocess exec'd without error.
     child_error_sfd.Reset();
@@ -213,7 +303,38 @@
         return false;
     }
 
-    local_socket_sfd_.Reset(parent_sfd.Release());
+    if (protocol_ == SubprocessProtocol::kNone) {
+        // No protocol: all streams pass through the stdinout FD and hook
+        // directly into the local socket for raw data transfer.
+        local_socket_sfd_.Reset(stdinout_sfd_.Release());
+    } else {
+        // Shell protocol: create another socketpair to intercept data.
+        if (!CreateSocketpair(&protocol_sfd_, &local_socket_sfd_)) {
+            return false;
+        }
+        D("protocol FD = %d", protocol_sfd_.fd());
+
+        input_.reset(new ShellProtocol(protocol_sfd_.fd()));
+        output_.reset(new ShellProtocol(protocol_sfd_.fd()));
+        if (!input_ || !output_) {
+            LOG(ERROR) << "failed to allocate shell protocol objects";
+            return false;
+        }
+
+        // Don't let reads/writes to the subprocess block our thread. This isn't
+        // likely but could happen under unusual circumstances, such as if we
+        // write a ton of data to stdin but the subprocess never reads it and
+        // the pipe fills up.
+        for (int fd : {stdinout_sfd_.fd(), stderr_sfd_.fd()}) {
+            if (fd >= 0) {
+                int flags = fcntl(fd, F_GETFL, 0);
+                if (flags < 0 || fcntl(fd, F_SETFL, flags | O_NONBLOCK) < 0) {
+                    PLOG(ERROR) << "error making FD " << fd << " non-blocking";
+                    return false;
+                }
+            }
+        }
+    }
 
     if (!adb_thread_create(ThreadHandler, this)) {
         PLOG(ERROR) << "failed to create subprocess thread";
@@ -259,6 +380,7 @@
     adb_thread_setname(android::base::StringPrintf(
             "shell srvc %d", subprocess->local_socket_fd()));
 
+    subprocess->PassDataStreams();
     subprocess->WaitForExit();
 
     D("deleting Subprocess");
@@ -267,25 +389,192 @@
     return nullptr;
 }
 
+void Subprocess::PassDataStreams() {
+    if (!protocol_sfd_.valid()) {
+        return;
+    }
+
+    // Start by trying to read from the protocol FD, stdout, and stderr.
+    fd_set master_read_set, master_write_set;
+    FD_ZERO(&master_read_set);
+    FD_ZERO(&master_write_set);
+    for (ScopedFd* sfd : {&protocol_sfd_, &stdinout_sfd_, &stderr_sfd_}) {
+        if (sfd->valid()) {
+            FD_SET(sfd->fd(), &master_read_set);
+        }
+    }
+
+    // Pass data until the protocol FD or both the subprocess pipes die, at
+    // which point we can't pass any more data.
+    while (protocol_sfd_.valid() &&
+            (stdinout_sfd_.valid() || stderr_sfd_.valid())) {
+        ScopedFd* dead_sfd = SelectLoop(&master_read_set, &master_write_set);
+        if (dead_sfd) {
+            D("closing FD %d", dead_sfd->fd());
+            FD_CLR(dead_sfd->fd(), &master_read_set);
+            FD_CLR(dead_sfd->fd(), &master_write_set);
+            dead_sfd->Reset();
+        }
+    }
+}
+
+namespace {
+
+inline bool ValidAndInSet(const ScopedFd& sfd, fd_set* set) {
+    return sfd.valid() && FD_ISSET(sfd.fd(), set);
+}
+
+}   // namespace
+
+ScopedFd* Subprocess::SelectLoop(fd_set* master_read_set_ptr,
+                                 fd_set* master_write_set_ptr) {
+    fd_set read_set, write_set;
+    int select_n = std::max(std::max(protocol_sfd_.fd(), stdinout_sfd_.fd()),
+                            stderr_sfd_.fd()) + 1;
+    ScopedFd* dead_sfd = nullptr;
+
+    // Keep calling select() and passing data until an FD closes/errors.
+    while (!dead_sfd) {
+        memcpy(&read_set, master_read_set_ptr, sizeof(read_set));
+        memcpy(&write_set, master_write_set_ptr, sizeof(write_set));
+        if (select(select_n, &read_set, &write_set, nullptr, nullptr) < 0) {
+            if (errno == EINTR) {
+                continue;
+            } else {
+                PLOG(ERROR) << "select failed, closing subprocess pipes";
+                stdinout_sfd_.Reset();
+                stderr_sfd_.Reset();
+                return nullptr;
+            }
+        }
+
+        // Read stdout, write to protocol FD.
+        if (ValidAndInSet(stdinout_sfd_, &read_set)) {
+            dead_sfd = PassOutput(&stdinout_sfd_, ShellProtocol::kIdStdout);
+        }
+
+        // Read stderr, write to protocol FD.
+        if (!dead_sfd && ValidAndInSet(stderr_sfd_, &read_set)) {
+            dead_sfd = PassOutput(&stderr_sfd_, ShellProtocol::kIdStderr);
+        }
+
+        // Read protocol FD, write to stdin.
+        if (!dead_sfd && ValidAndInSet(protocol_sfd_, &read_set)) {
+            dead_sfd = PassInput();
+            // If we didn't finish writing, block on stdin write.
+            if (input_bytes_left_) {
+                FD_CLR(protocol_sfd_.fd(), master_read_set_ptr);
+                FD_SET(stdinout_sfd_.fd(), master_write_set_ptr);
+            }
+        }
+
+        // Continue writing to stdin; only happens if a previous write blocked.
+        if (!dead_sfd && ValidAndInSet(stdinout_sfd_, &write_set)) {
+            dead_sfd = PassInput();
+            // If we finished writing, go back to blocking on protocol read.
+            if (!input_bytes_left_) {
+                FD_SET(protocol_sfd_.fd(), master_read_set_ptr);
+                FD_CLR(stdinout_sfd_.fd(), master_write_set_ptr);
+            }
+        }
+    }  // while (!dead_sfd)
+
+    return dead_sfd;
+}
+
+ScopedFd* Subprocess::PassInput() {
+    // Only read a new packet if we've finished writing the last one.
+    if (!input_bytes_left_) {
+        if (!input_->Read()) {
+            // Read() uses ReadFdExactly() which sets errno to 0 on EOF.
+            if (errno != 0) {
+                PLOG(ERROR) << "error reading protocol FD "
+                            << protocol_sfd_.fd();
+            }
+            return &protocol_sfd_;
+        }
+
+        // We only care about stdin packets.
+        if (stdinout_sfd_.valid() && input_->id() == ShellProtocol::kIdStdin) {
+            input_bytes_left_ = input_->data_length();
+        } else {
+            input_bytes_left_ = 0;
+        }
+    }
+
+    if (input_bytes_left_ > 0) {
+        int index = input_->data_length() - input_bytes_left_;
+        int bytes = adb_write(stdinout_sfd_.fd(), input_->data() + index,
+                              input_bytes_left_);
+        if (bytes == 0 || (bytes < 0 && errno != EAGAIN)) {
+            if (bytes < 0) {
+                PLOG(ERROR) << "error reading stdin FD " << stdinout_sfd_.fd();
+            }
+            // stdin is done, mark this packet as finished and we'll just start
+            // dumping any further data received from the protocol FD.
+            input_bytes_left_ = 0;
+            return &stdinout_sfd_;
+        } else if (bytes > 0) {
+            input_bytes_left_ -= bytes;
+        }
+    }
+
+    return nullptr;
+}
+
+ScopedFd* Subprocess::PassOutput(ScopedFd* sfd, ShellProtocol::Id id) {
+    int bytes = adb_read(sfd->fd(), output_->data(), output_->data_capacity());
+    if (bytes == 0 || (bytes < 0 && errno != EAGAIN)) {
+        if (bytes < 0) {
+            PLOG(ERROR) << "error reading output FD " << sfd->fd();
+        }
+        return sfd;
+    }
+
+    if (bytes > 0 && !output_->Write(id, bytes)) {
+        if (errno != 0) {
+            PLOG(ERROR) << "error reading protocol FD " << protocol_sfd_.fd();
+        }
+        return &protocol_sfd_;
+    }
+
+    return nullptr;
+}
+
 void Subprocess::WaitForExit() {
+    int exit_code = 1;
+
     D("waiting for pid %d", pid_);
     while (true) {
         int status;
         if (pid_ == waitpid(pid_, &status, 0)) {
             D("post waitpid (pid=%d) status=%04x", pid_, status);
             if (WIFSIGNALED(status)) {
+                exit_code = 0x80 | WTERMSIG(status);
                 D("subprocess killed by signal %d", WTERMSIG(status));
                 break;
             } else if (!WIFEXITED(status)) {
                 D("subprocess didn't exit");
                 break;
             } else if (WEXITSTATUS(status) >= 0) {
+                exit_code = WEXITSTATUS(status);
                 D("subprocess exit code = %d", WEXITSTATUS(status));
                 break;
             }
         }
     }
 
+    // If we have an open protocol FD send an exit packet.
+    if (protocol_sfd_.valid()) {
+        output_->data()[0] = exit_code;
+        if (output_->Write(ShellProtocol::kIdExit, 1)) {
+            D("wrote the exit code packet: %d", exit_code);
+        } else {
+            PLOG(ERROR) << "failed to write the exit code packet";
+        }
+        protocol_sfd_.Reset();
+    }
+
     // Pass the local socket FD to the shell cleanup fdevent.
     if (SHELL_EXIT_NOTIFY_FD >= 0) {
         int fd = local_socket_sfd_.fd();
@@ -305,11 +594,13 @@
 
 }  // namespace
 
-int StartSubprocess(const char *name, SubprocessType type) {
-    D("starting %s subprocess: '%s'",
-      type == SubprocessType::kRaw ? "raw" : "PTY", name);
+int StartSubprocess(const char *name, SubprocessType type,
+                    SubprocessProtocol protocol) {
+    D("starting %s subprocess (protocol=%s): '%s'",
+      type == SubprocessType::kRaw ? "raw" : "PTY",
+      protocol == SubprocessProtocol::kNone ? "none" : "shell", name);
 
-    Subprocess* subprocess = new Subprocess(name, type);
+    Subprocess* subprocess = new Subprocess(name, type, protocol);
     if (!subprocess) {
         LOG(ERROR) << "failed to allocate new subprocess";
         return -1;