Ensure that in-flight message pipes are always closed and the other end is notified.

mojo/edk/system/message_pipe_dispatcher.cc

 // Copyright 2015 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 "mojo/edk/system/message_pipe_dispatcher.h"
 
 #include <stddef.h>
 #include <stdint.h>
 
 #include <utility>
@@ skipping 139 matching lines @@
 
 void MessagePipeDispatcher::InitOnIO() {
   base::AutoLock locker(lock());
   CHECK(transferable_);
   calling_init_ = true;
   if (channel_)
     channel_->Init(this);
   calling_init_ = false;
 }
 
+void MessagePipeDispatcher::CloseOnIOAndRelease() {
+  {
+    base::AutoLock locker(lock());
+    CloseOnIO();
+  }
+  Release();  // To match CloseImplNoLock.
+}
+
 void MessagePipeDispatcher::CloseOnIO() {
-  base::AutoLock locker(lock());
-  Release();  // To match CloseImplNoLock.
-  if (transferable_) {
-    if (channel_) {
-      channel_->Shutdown();
-      channel_ = nullptr;
-    }
-  } else {
-    if (non_transferable_state_ == CONNECT_CALLED ||
-        non_transferable_state_ == WAITING_FOR_READ_OR_WRITE) {
-      if (non_transferable_state_ == WAITING_FOR_READ_OR_WRITE)
-        RequestNontransferableChannel();
+  lock().AssertAcquired();
 
-      // We can't cancel the pending request yet, since the other side of the
-      // message pipe would want to get pending outgoing messages (if any) or
-      // at least know that this end was closed. So keep this object alive until
-      // then.
-      non_transferable_state_ = WAITING_FOR_CONNECT_TO_CLOSE;
-      AddRef();
-    } else if (non_transferable_state_ == CONNECTED) {
-      internal::g_broker->CloseMessagePipe(pipe_id_, this);
-      non_transferable_state_ = CLOSED;
-      channel_ = nullptr;
-    }
+  if (channel_) {
+    // If we closed the channel now then in-flight message pipes wouldn't get
+    // closed, and their other side wouldn't get a connection error notification
+    // which could lead to hangs or leaks. So we ask the other side of this
+    // message pipe to close, which ensures that we have dispatched all
+    // in-flight message pipes.
+    DCHECK(!close_requested_);
+    close_requested_ = true;
+    AddRef();
+    scoped_ptr<MessageInTransit> message(new MessageInTransit(
+        MessageInTransit::Type::QUIT_MESSAGE, 0, nullptr));
+    if (!transferable_)
+      message->set_route_id(pipe_id_);
+    channel_->WriteMessage(std::move(message));
+    return;
+  }
+
+  if (!transferable_ &&
+      (non_transferable_state_ == CONNECT_CALLED ||
+       non_transferable_state_ == WAITING_FOR_READ_OR_WRITE)) {
+    if (non_transferable_state_ == WAITING_FOR_READ_OR_WRITE)
+      RequestNontransferableChannel();
+
+    // We can't cancel the pending request yet, since the other side of the
+    // message pipe would want to get pending outgoing messages (if any) or at
+    // least know that this end was closed. So keep this object alive until
+    // then.
+    non_transferable_state_ = WAITING_FOR_CONNECT_TO_CLOSE;
+    AddRef();
   }
 }
 
 Dispatcher::Type MessagePipeDispatcher::GetType() const {
   return Type::MESSAGE_PIPE;
 }
 
 void MessagePipeDispatcher::GotNonTransferableChannel(RawChannel* channel) {
+  DCHECK(internal::g_io_thread_task_runner->RunsTasksOnCurrentThread());
   base::AutoLock locker(lock());
   channel_ = channel;
   while (!non_transferable_outgoing_message_queue_.IsEmpty()) {
     channel_->WriteMessage(
         non_transferable_outgoing_message_queue_.GetMessage());
   }
 
   if (non_transferable_state_ == WAITING_FOR_CONNECT_TO_CLOSE) {
-    // We kept this object alive until it's connected, we can release it now.
-    internal::g_broker->CloseMessagePipe(pipe_id_, this);
-    non_transferable_state_ = CLOSED;
-    channel_ = nullptr;
-    base::MessageLoop::current()->ReleaseSoon(FROM_HERE, this);
-  } else {
     non_transferable_state_ = CONNECTED;
+    // We kept this object alive until it's connected, we can close it now.
+    CloseOnIO();
+    // Balance the AddRef in CloseOnIO.
+    Release();
+    return;
   }
+
+  non_transferable_state_ = CONNECTED;
 }
 
 #if defined(OS_WIN)
 // TODO(jam): this is copied from RawChannelWin till I figure out what's the
 // best way we want to share this.
 // Since this is used for serialization of messages read/written to a MP that
 // aren't consumed by Mojo primitives yet, there could be an unbounded number of
 // them when a MP is being sent. As a result, even for POSIX we will probably
 // want to send the handles to the shell process and exchange them for tokens
 // (since we can be sure that the shell will respond to our IPCs, compared to
@@ skipping 182 matching lines @@
 MessagePipeDispatcher::MessagePipeDispatcher(bool transferable)
     : channel_(nullptr),
       serialized_read_fds_length_(0u),
       serialized_write_fds_length_(0u),
       serialized_message_fds_length_(0u),
       pipe_id_(0),
       non_transferable_state_(WAITING_FOR_READ_OR_WRITE),
       serialized_(false),
       calling_init_(false),
       write_error_(false),
-      transferable_(transferable) {
+      transferable_(transferable),
+      close_requested_(false) {
 }
 
 MessagePipeDispatcher::~MessagePipeDispatcher() {
   // |Close()|/|CloseImplNoLock()| should have taken care of the channel. The
   // exception is if they posted a task to run CloseOnIO but the IO thread shut
   // down and so when it was deleting pending tasks it caused the last reference
   // to destruct this object. In that case, safe to destroy the channel.
   if (channel_ &&
       internal::g_io_thread_task_runner->RunsTasksOnCurrentThread()) {
     if (transferable_) {
@@ skipping 21 matching lines @@
   if (!transferable_ && non_transferable_state_ == CLOSED)
     return;
 
   // We take a manual refcount because at shutdown, the task below might not get
   // a chance to execute. If that happens, the RawChannel will still call our
   // OnError method because it always runs (since it watches thread
   // destruction). So to avoid UAF, manually add a reference and only release it
   // if the task runs.
   AddRef();
   internal::g_io_thread_task_runner->PostTask(
-      FROM_HERE, base::Bind(&MessagePipeDispatcher::CloseOnIO, this));
+      FROM_HERE, base::Bind(&MessagePipeDispatcher::CloseOnIOAndRelease, this));
 }
 
 void MessagePipeDispatcher::SerializeInternal() {
   serialized_ = true;
   if (!transferable_) {
     CHECK(non_transferable_state_ == WAITING_FOR_READ_OR_WRITE)
         << "Non transferable message pipe being sent after read/write/waited. "
         << "MOJO_CREATE_MESSAGE_PIPE_OPTIONS_FLAG_TRANSFERABLE must be used if "
         << "the pipe can be sent after it's read or written. This message pipe "
         << "was previously bound at:\n"
@@ skipping 396 matching lines @@
     ScopedPlatformHandleVectorPtr platform_handles) {
   scoped_ptr<MessageInTransit> message(new MessageInTransit(message_view));
   if (message_view.transport_data_buffer_size() > 0) {
     DCHECK(message_view.transport_data_buffer());
     message->SetDispatchers(TransportData::DeserializeDispatchers(
         message_view.transport_data_buffer(),
         message_view.transport_data_buffer_size(),
         std::move(platform_handles)));
   }
 
+  bool call_release = false;
   if (started_transport_.Try()) {
-    // we're not in the middle of being sent
+    // We're not in the middle of being sent.
 
     // Can get synchronously called back in Init if there was initial data.
     scoped_ptr<base::AutoLock> locker;
-    if (!calling_init_) {
+    if (!calling_init_)
       locker.reset(new base::AutoLock(lock()));
+
+    if (message_view.type() == MessageInTransit::Type::QUIT_MESSAGE) {
+      if (transferable_) {
+        channel_->Shutdown();
+      } else {
+        internal::g_broker->CloseMessagePipe(pipe_id_, this);
+        non_transferable_state_ = CLOSED;
+      }
+      channel_ = nullptr;
+      awakable_list_.AwakeForStateChange(GetHandleSignalsStateImplNoLock());
+      if (close_requested_) {
+        // We requested the other side to close the connection while they also
+        // did the same. We must balance out the AddRef in CloseOnIO to ensure
+        // this object isn't leaked.
+        call_release = true;
+      }
+    } else {
+      bool was_empty = message_queue_.IsEmpty();
+      message_queue_.AddMessage(std::move(message));
+      if (was_empty)
+        awakable_list_.AwakeForStateChange(GetHandleSignalsStateImplNoLock());
     }
 
-    bool was_empty = message_queue_.IsEmpty();
-    message_queue_.AddMessage(std::move(message));
-    if (was_empty)
-      awakable_list_.AwakeForStateChange(GetHandleSignalsStateImplNoLock());
-
     started_transport_.Release();
   } else {
-    // If RawChannel is calling OnRead, that means it has its read_lock_
-    // acquired. That means StartSerialize can't be accessing message queue as
-    // it waits on ReleaseHandle first which acquires readlock_.
-    message_queue_.AddMessage(std::move(message));
+    if (message_view.type() == MessageInTransit::Type::QUIT_MESSAGE) {
+      // We got a request to shutdown the channel but this object is already
+      // calling into channel to serialize it. Since all the other side cares
+      // about is flushing pending messages, we bounce the quit back to it.
+      scoped_ptr<MessageInTransit> message(new MessageInTransit(
+          MessageInTransit::Type::QUIT_MESSAGE, 0, nullptr));
+      channel_->WriteMessage(std::move(message));
+    } else {
+      // If RawChannel is calling OnRead, that means it has its read_lock_
+      // acquired. That means StartSerialize can't be accessing message queue as
+      // it waits on ReleaseHandle first which acquires read_lock_.
+      message_queue_.AddMessage(std::move(message));
+    }
   }
+
+  if (call_release)
+    Release();
 }
 
 void MessagePipeDispatcher::OnError(Error error) {
   // If there's a read error, then the other side of the pipe is closed. By
   // definition, we can't write since there's no one to read it. And we can't
   // read anymore, since we just got a read erorr. So we close the pipe.
   // If there's a write error, then we stop writing. But we keep the pipe open
   // until we finish reading everything in it. This is because it's valid for
   // one endpoint to write some data and close their pipe immediately. Even
   // though the other end can't write anymore, it should still get all the data.
@@ skipping 15 matching lines @@
       break;
     case ERROR_WRITE:
       // Write errors are slightly notable: they probably shouldn't happen under
       // normal operation (but maybe the other side crashed).
       LOG(WARNING) << "MessagePipeDispatcher write error";
       DCHECK_EQ(write_error_, false) << "Should only get one write error.";
       write_error_ = true;
       break;
   }
 
+  bool call_release = false;
   if (started_transport_.Try()) {
     base::AutoLock locker(lock());
-    // We can get two OnError callbacks before the post task below completes.
-    // Although RawChannel still has a pointer to this object until Shutdown is
-    // called, that is safe since this class always does a PostTask to the IO
-    // thread to self destruct.
     if (channel_ && error != ERROR_WRITE) {
       if (transferable_) {
         channel_->Shutdown();
       } else {
         CHECK_NE(non_transferable_state_, CLOSED);
         internal::g_broker->CloseMessagePipe(pipe_id_, this);
         non_transferable_state_ = CLOSED;
       }
       channel_ = nullptr;
+      if (close_requested_) {
+        // Balance AddRef in CloseOnIO.
+        call_release = true;
+      }
     }
     awakable_list_.AwakeForStateChange(GetHandleSignalsStateImplNoLock());
     started_transport_.Release();
   } else {
     // We must be waiting to call ReleaseHandle. It will call Shutdown.
   }
+
+  if (call_release)
+    Release();
 }
 
 MojoResult MessagePipeDispatcher::AttachTransportsNoLock(
     MessageInTransit* message,
     std::vector<DispatcherTransport>* transports) {
   DCHECK(!message->has_dispatchers());
 
   // You're not allowed to send either handle to a message pipe over the message
   // pipe, so check for this. (The case of trying to write a handle to itself is
   // taken care of by |Core|. That case kind of makes sense, but leads to
@@ skipping 48 matching lines @@
   // PostTask since the broker can call us back synchronously.
   internal::g_io_thread_task_runner->PostTask(
       FROM_HERE,
       base::Bind(&Broker::ConnectMessagePipe,
                  base::Unretained(internal::g_broker), pipe_id_,
                  base::Unretained(this)));
 }
 
 }  // namespace edk
 }  // namespace mojo