Refactor the bootstrap sandbox process launching integration.

sandbox/mac/bootstrap_sandbox.cc

Index: sandbox/mac/bootstrap_sandbox.cc
diff --git a/sandbox/mac/bootstrap_sandbox.cc b/sandbox/mac/bootstrap_sandbox.cc
index a90f570eb47487524acecdd11e903c8e60eeeb40..e8997f9ef68f774a5a826312add6fee87f47c71a 100644
--- a/sandbox/mac/bootstrap_sandbox.cc
+++ b/sandbox/mac/bootstrap_sandbox.cc
@@ -10,18 +10,53 @@
 #include "base/logging.h"
 #include "base/mac/foundation_util.h"
 #include "base/mac/mach_logging.h"
+#include "base/rand_util.h"
 #include "base/strings/stringprintf.h"
 #include "sandbox/mac/launchd_interception_server.h"
+#include "sandbox/mac/pre_exec_delegate.h"
 
 namespace sandbox {
 
-const int kNotAPolicy = -1;
+namespace {
+
+struct SandboxCheckInRequest {
+  mach_msg_header_t header;
+  uint64_t token;
+};
+
+struct SandboxCheckInReply {
+  mach_msg_header_t header;
+  mach_msg_body_t body;
+  mach_msg_port_descriptor_t bootstrap_port;
+};
+
+class ScopedCallMachMsgDestroy {
+ public:
+  explicit ScopedCallMachMsgDestroy(mach_msg_header_t* message)
+      : message_(message) {}
+
+  ~ScopedCallMachMsgDestroy() {
+    if (message_)
+      mach_msg_destroy(message_);
+  }
+
+  void Disarm() {
+    message_ = nullptr;
+  }
+
+ private:
+  mach_msg_header_t* message_;
+
+  DISALLOW_COPY_AND_ASSIGN(ScopedCallMachMsgDestroy);
+};
+
+}  // namespace
 
 // static
 scoped_ptr<BootstrapSandbox> BootstrapSandbox::Create() {
   scoped_ptr<BootstrapSandbox> null;  // Used for early returns.
   scoped_ptr<BootstrapSandbox> sandbox(new BootstrapSandbox());
-  sandbox->server_.reset(new LaunchdInterceptionServer(sandbox.get()));
+  sandbox->launchd_server_.reset(new LaunchdInterceptionServer(sandbox.get()));
 
   // Check in with launchd to get the receive right for the server that is
   // published in the bootstrap namespace.
@@ -33,17 +68,67 @@ scoped_ptr<BootstrapSandbox> BootstrapSandbox::Create() {
         << "Failed to bootstrap_check_in the sandbox server.";
     return null.Pass();
   }
-  base::mac::ScopedMachReceiveRight scoped_port(port);
+  sandbox->check_in_port_.reset(port);
+
+  BootstrapSandbox* __block sandbox_ptr = sandbox.get();
+  sandbox->check_in_server_.reset(new base::DispatchSourceMach(
+      "org.chromium.sandbox.BootstrapClientManager",
+      sandbox->check_in_port_.get(),
+      ^{ sandbox_ptr->HandleChildCheckIn(); }));
+  sandbox->check_in_server_->Resume();
 
   // Start the sandbox server.
-  if (sandbox->server_->Initialize(scoped_port.get()))
-    ignore_result(scoped_port.release());  // Transferred to server_.
-  else
+  if (!sandbox->launchd_server_->Initialize(MACH_PORT_NULL))
     return null.Pass();
 
   return sandbox.Pass();
 }
 
+// Warning: This function must be safe to call in
+// PreExecDelegate::RunAsyncSafe().
+// static
+bool BootstrapSandbox::ClientCheckIn(mach_port_t sandbox_server_port,
+                                     uint64_t sandbox_token,
+                                     mach_port_t* new_bootstrap_port) {
+  // Create a reply port for the check in message.
+  mach_port_t reply_port;
+  kern_return_t kr = mach_port_allocate(mach_task_self(),
+                                        MACH_PORT_RIGHT_RECEIVE,
+                                        &reply_port);
+  if (kr != KERN_SUCCESS) {
+    RAW_LOG(ERROR, "ClientCheckIn: mach_port_allocate failed");
+    return false;
+  }
+  base::mac::ScopedMachReceiveRight scoped_reply_port(reply_port);
+
+  // Check in with the sandbox server, presenting the |sandbox_token| in
+  // exchange for a new task bootstrap port.
+  union {
+    SandboxCheckInRequest request;
+    struct {
+      SandboxCheckInReply reply;
+      mach_msg_trailer_t trailer;
+    };
+  } msg = {};
+  msg.request.header.msgh_bits = MACH_MSGH_BITS(MACH_MSG_TYPE_COPY_SEND,
+                                                MACH_MSG_TYPE_MAKE_SEND_ONCE);
+  msg.request.header.msgh_remote_port = sandbox_server_port;
+  msg.request.header.msgh_local_port = reply_port;
+  msg.request.header.msgh_size = sizeof(msg);
+  msg.request.token = sandbox_token;
+
+  kr = mach_msg(&msg.request.header, MACH_SEND_MSG | MACH_RCV_MSG,
+                sizeof(msg.request), sizeof(msg), reply_port,
+                MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
+  if (kr == KERN_SUCCESS) {
+    *new_bootstrap_port = msg.reply.bootstrap_port.name;
+    return true;
+  } else {
+    RAW_LOG(ERROR, "ClientCheckIn: mach_msg failed");
+    return false;
+  }
+}
+
 BootstrapSandbox::~BootstrapSandbox() {
 }
 
@@ -51,48 +136,35 @@ void BootstrapSandbox::RegisterSandboxPolicy(
     int sandbox_policy_id,
     const BootstrapSandboxPolicy& policy) {
   CHECK(IsPolicyValid(policy));
-  CHECK_GT(sandbox_policy_id, kNotAPolicy);
   base::AutoLock lock(lock_);
   DCHECK(policies_.find(sandbox_policy_id) == policies_.end());
   policies_.insert(std::make_pair(sandbox_policy_id, policy));
 }
 
-void BootstrapSandbox::PrepareToForkWithPolicy(int sandbox_policy_id) {
+scoped_ptr<PreExecDelegate> BootstrapSandbox::NewClient(int sandbox_policy_id) {
   base::AutoLock lock(lock_);
 
-  // Verify that this is a real policy.
-  CHECK(policies_.find(sandbox_policy_id) != policies_.end());
-  CHECK_EQ(kNotAPolicy, effective_policy_id_)
-      << "Cannot nest calls to PrepareToForkWithPolicy()";
+  DCHECK(policies_.find(sandbox_policy_id) != policies_.end());
 
-  // Store the policy for the process we're about to create.
-  effective_policy_id_ = sandbox_policy_id;
-}
+  uint64_t token;
+  while (true) {
+    token = base::RandUint64();
+    if (awaiting_processes_.find(token) == awaiting_processes_.end())
+      break;
+  }
 
-// TODO(rsesek): The |lock_| needs to be taken twice because
-// base::LaunchProcess handles both fork+exec, and holding the lock for the
-// duration would block servicing of other bootstrap messages. If a better
-// LaunchProcess existed (do arbitrary work without layering violations), this
-// could be avoided.
+  awaiting_processes_[token] = sandbox_policy_id;
+  return make_scoped_ptr(new PreExecDelegate(server_bootstrap_name_, token));
+}
 
-void BootstrapSandbox::FinishedFork(base::ProcessHandle handle) {
+void BootstrapSandbox::RevokeToken(uint64_t token) {
   base::AutoLock lock(lock_);
-
-  CHECK_NE(kNotAPolicy, effective_policy_id_)
-      << "Must PrepareToForkWithPolicy() before FinishedFork()";
-
-  // Apply the policy to the new process.
-  if (handle != base::kNullProcessHandle) {
-    const auto& existing_process = sandboxed_processes_.find(handle);
-    CHECK(existing_process == sandboxed_processes_.end());
-    sandboxed_processes_.insert(std::make_pair(handle, effective_policy_id_));
-    VLOG(3) << "Bootstrap sandbox enforced for pid " << handle;
-  }
-
-  effective_policy_id_ = kNotAPolicy;
+  const auto& it = awaiting_processes_.find(token);
+  if (it != awaiting_processes_.end())
+    awaiting_processes_.erase(it);
 }
 
-void BootstrapSandbox::ChildDied(base::ProcessHandle handle) {
+void BootstrapSandbox::InvalidateClient(base::ProcessHandle handle) {
   base::AutoLock lock(lock_);
   const auto& it = sandboxed_processes_.find(handle);
   if (it != sandboxed_processes_.end())
@@ -103,26 +175,18 @@ const BootstrapSandboxPolicy* BootstrapSandbox::PolicyForProcess(
     pid_t pid) const {
   base::AutoLock lock(lock_);
   const auto& process = sandboxed_processes_.find(pid);
-
-  // The new child could send bootstrap requests before the parent calls
-  // FinishedFork().
-  int policy_id = effective_policy_id_;
   if (process != sandboxed_processes_.end()) {
-    policy_id = process->second;
+    return &policies_.find(process->second)->second;
   }
 
-  if (policy_id == kNotAPolicy)
-    return NULL;
-
-  return &policies_.find(policy_id)->second;
+  return nullptr;
 }
 
 BootstrapSandbox::BootstrapSandbox()
     : server_bootstrap_name_(
           base::StringPrintf("%s.sandbox.%d", base::mac::BaseBundleID(),
               getpid())),
-      real_bootstrap_port_(MACH_PORT_NULL),
-      effective_policy_id_(kNotAPolicy) {
+      real_bootstrap_port_(MACH_PORT_NULL) {
   mach_port_t port = MACH_PORT_NULL;
   kern_return_t kr = task_get_special_port(
       mach_task_self(), TASK_BOOTSTRAP_PORT, &port);
@@ -130,4 +194,66 @@ BootstrapSandbox::BootstrapSandbox()
   real_bootstrap_port_.reset(port);
 }
 
+void BootstrapSandbox::HandleChildCheckIn() {
+  struct {
+    SandboxCheckInRequest request;
+    mach_msg_audit_trailer_t trailer;
+  } msg = {};
+  msg.request.header.msgh_local_port = check_in_port_.get();
+  msg.request.header.msgh_size = sizeof(msg.request);
+  const mach_msg_option_t kOptions = MACH_RCV_MSG |
+      MACH_RCV_TRAILER_TYPE(MACH_MSG_TRAILER_FORMAT_0) |
+      MACH_RCV_TRAILER_ELEMENTS(MACH_RCV_TRAILER_AUDIT);
+  kern_return_t kr = mach_msg(&msg.request.header, kOptions, 0,
+                              sizeof(msg), check_in_port_.get(),
+                              MACH_MSG_TIMEOUT_NONE, MACH_PORT_NULL);
+  if (kr != KERN_SUCCESS) {
+    MACH_LOG(ERROR, kr) << "HandleChildCheckIn mach_msg MACH_RCV_MSG";
+    return;
+  }
+
+  // Call mach_msg_destroy to clean up the reply send-once right.
+  ScopedCallMachMsgDestroy message_destroyer(&msg.request.header);
+
+  pid_t client_pid;
+  audit_token_to_au32(msg.trailer.msgh_audit, nullptr, nullptr, nullptr,
+      nullptr, nullptr, &client_pid, nullptr, nullptr);
+
+  {
+    base::AutoLock lock(lock_);
+
+    auto awaiting_it = awaiting_processes_.find(msg.request.token);
+    if (awaiting_it == awaiting_processes_.end()) {
+      LOG(ERROR) << "Received sandbox check-in message from unknown client.";
+      return;
+    }
+
+    CHECK(sandboxed_processes_.find(client_pid) == sandboxed_processes_.end());
+    sandboxed_processes_[client_pid] = awaiting_it->second;
+    awaiting_processes_.erase(awaiting_it);
+  }
+
+  SandboxCheckInReply reply = {};
+  reply.header.msgh_bits = MACH_MSGH_BITS_REMOTE(msg.request.header.msgh_bits) |
+                           MACH_MSGH_BITS_COMPLEX;
+  reply.header.msgh_remote_port = msg.request.header.msgh_remote_port;
+  reply.header.msgh_size = sizeof(reply);
+  reply.body.msgh_descriptor_count = 1;
+  reply.bootstrap_port.name = launchd_server_->server_port();
+  reply.bootstrap_port.disposition = MACH_MSG_TYPE_MAKE_SEND;
+  reply.bootstrap_port.type = MACH_MSG_PORT_DESCRIPTOR;
+
+  kr = mach_msg(&reply.header, MACH_SEND_MSG | MACH_SEND_TIMEOUT,
+                sizeof(reply), 0, MACH_PORT_NULL, 100 /*ms*/, MACH_PORT_NULL);
+  if (kr == KERN_SUCCESS) {
+    message_destroyer.Disarm();  // The send-once was consumed at mach_msg().
+  } else {
+    {
+      base::AutoLock lock(lock_);
+      sandboxed_processes_.erase(client_pid);
+    }
+    MACH_LOG(ERROR, kr) << "HandleChildCheckIn mach_msg MACH_SEND_MSG";
+  }
+}
+
 }  // namespace sandbox