Fix memory leak in unittests with new Mojo EDK.

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 "base/bind.h"
 #include "base/debug/stack_trace.h"
 #include "base/logging.h"
 #include "base/message_loop/message_loop.h"
@@ skipping 136 matching lines @@
   base::AutoLock locker(lock());
   CHECK(transferable_);
   calling_init_ = true;
   if (channel_)
     channel_->Init(this);
   calling_init_ = false;
 }
 
 void MessagePipeDispatcher::CloseOnIO() {
   base::AutoLock locker(lock());
-  Release();  // To match CloseImplNoLock.
+  if (release_reference_in_close_)
+    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();
@@ skipping 231 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),
+      release_reference_in_close_(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 15 matching lines @@
 }
 
 void MessagePipeDispatcher::CloseImplNoLock() {
   lock().AssertAcquired();
   // We take a manual refcount because at shutdown, the task below might not get
   // a chance to execute. If that happens, the RawChannel's 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();
+  release_reference_in_close_ = true;
   internal::g_io_thread_task_runner->PostTask(
       FROM_HERE, base::Bind(&MessagePipeDispatcher::CloseOnIO, 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 "
@@ skipping 453 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 (release_reference_in_close_) {
+        // CloseOnIO might never get called on process shutdown. Release the
+        // extra reference here to avoid leaks in tests.
+        call_release = true;

[msw, 2015/12/18 01:54:26]

nit: could you actually just call Release() here and remove |call_release|?

[jam, 2015/12/18 02:03:47]

that's what I did in patchset 1 but that failed since if this was the last
reference, the autolock's destructor above would cause a UAF

+        release_reference_in_close_ = false;
+      }
     }
     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