[remoting host] Fix bugs in MojoIpcServer

* Fix a bug such that MojoIpcServer doesn't remove a connection from
  active_connections_ when the IPC is closed remotely.
* Keep sending invitations when OnMessagePipeReady() is called with a
  failing result, which can happen if the remote process has
  exited/crashed before the pending remote is bound.

Change-Id: Ib359b62f7f89acea0994bf1d8020e3e3b773bbc3
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/3230238
Reviewed-by: Joe Downing <joedow@chromium.org>
Commit-Queue: Yuwei Huang <yuweih@chromium.org>
Cr-Commit-Position: refs/heads/main@{#933133}

remoting/host/mojo_ipc/mojo_ipc_server.cc

@@ -110,7 +110,10 @@ void MojoIpcServerBase::StopServer() {
 
 void MojoIpcServerBase::Close(mojo::ReceiverId id) {
   UntrackMessagePipe(id);
-  active_connections_.erase(id);
+  auto it = active_connections_.find(id);
+  if (it != active_connections_.end()) {
+    active_connections_.erase(it);
+  }
 }
 
 void MojoIpcServerBase::SendInvitation() {
@@ -123,6 +126,13 @@ void MojoIpcServerBase::SendInvitation() {
                      weak_factory_.GetWeakPtr()));
 }
 
+void MojoIpcServerBase::OnIpcDisconnected() {
+  if (disconnect_handler_) {
+    disconnect_handler_.Run();
+  }
+  Close(current_receiver());
+}
+
 void MojoIpcServerBase::OnInvitationSent(
     std::unique_ptr<MojoIpcServerBase::PendingConnection> pending_connection) {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
@@ -153,17 +163,24 @@ void MojoIpcServerBase::OnMessagePipeReady(
     const mojo::HandleSignalsState& state) {
   DCHECK_CALLED_ON_VALID_SEQUENCE(sequence_checker_);
 
+  pending_message_pipe_watcher_.Cancel();
+  if (result == MOJO_RESULT_FAILED_PRECONDITION) {
+    LOG(WARNING) << "The message pipe will never become ready.";
+    pending_connection_.reset();
+    SendInvitation();
+    return;
+  }
   if (result != MOJO_RESULT_OK) {
-    LOG(ERROR) << "Message pipe is not ready. Result: " << result;
+    LOG(ERROR) << "Unexpected message pipe result: " << result;
     pending_connection_.reset();
     return;
   }
   if (state.peer_closed()) {
     LOG(ERROR) << "Message pipe is closed.";
     pending_connection_.reset();
+    SendInvitation();
     return;
   }
-  pending_message_pipe_watcher_.Cancel();
 
   DCHECK(pending_connection_->message_pipe.is_valid());
   auto receiver_id =

remoting/host/mojo_ipc/mojo_ipc_server.h

@@ -7,6 +7,7 @@
 
 #include <memory>
 
+#include "base/bind.h"
 #include "base/callback.h"
 #include "base/containers/flat_map.h"
 #include "base/memory/ref_counted.h"
@@ -44,13 +45,19 @@ class MojoIpcServerBase : public IpcServer {
     on_invitation_sent_callback_for_testing_ = callback;
   }
 
+  size_t GetNumberOfActiveConnectionsForTesting() const {
+    return active_connections_.size();
+  }
+
  protected:
   explicit MojoIpcServerBase(
       const mojo::NamedPlatformChannel::ServerName& server_name);
   ~MojoIpcServerBase() override;
 
   void SendInvitation();
 
+  void OnIpcDisconnected();
+
   virtual mojo::ReceiverId TrackMessagePipe(
       mojo::ScopedMessagePipeHandle message_pipe) = 0;
 
@@ -60,6 +67,8 @@ class MojoIpcServerBase : public IpcServer {
 
   SEQUENCE_CHECKER(sequence_checker_);
 
+  base::RepeatingClosure disconnect_handler_;
+
  private:
   using ActiveConnectionMap =
       base::flat_map<mojo::ReceiverId,
@@ -123,15 +132,18 @@ class MojoIpcServer final : public MojoIpcServerBase {
   // ExtractMessagePipe() with the same ID.
   MojoIpcServer(const mojo::NamedPlatformChannel::ServerName& server_name,
                 Interface* interface_impl)
-      : MojoIpcServerBase(server_name), interface_impl_(interface_impl) {}
+      : MojoIpcServerBase(server_name), interface_impl_(interface_impl) {
+    receiver_set_.set_disconnect_handler(base::BindRepeating(
+        &MojoIpcServer::OnIpcDisconnected, base::Unretained(this)));
+  }
 
   ~MojoIpcServer() override = default;
 
   MojoIpcServer(const MojoIpcServer&) = delete;
   MojoIpcServer& operator=(const MojoIpcServer&) = delete;
 
   void set_disconnect_handler(base::RepeatingClosure handler) override {
-    receiver_set_.set_disconnect_handler(handler);
+    disconnect_handler_ = handler;
   }
 
   mojo::ReceiverId current_receiver() const override {

remoting/host/mojo_ipc/mojo_ipc_server_unittest.cc

@@ -17,6 +17,7 @@
 #include "base/test/mock_callback.h"
 #include "base/test/task_environment.h"
 #include "base/threading/platform_thread.h"
+#include "base/threading/sequenced_task_runner_handle.h"
 #include "base/time/time.h"
 #include "mojo/public/cpp/bindings/pending_remote.h"
 #include "mojo/public/cpp/bindings/receiver_set.h"
@@ -198,11 +199,40 @@ TEST_F(MojoIpcServerTest, CloseReceiver_RemoteDisconnected) {
   mojo::IsolatedConnection client_connection;
   auto echo_remote = ConnectAndCreateEchoRemote(client_connection);
   SendEchoAndVerifyResponse(echo_remote);
+  ASSERT_EQ(1u, ipc_server_->GetNumberOfActiveConnectionsForTesting());
 
   base::RunLoop disconnect_run_loop;
   echo_remote.set_disconnect_handler(disconnect_run_loop.QuitClosure());
   ipc_server_->Close(last_echo_string_receiver_id_);
   disconnect_run_loop.Run();
+  ASSERT_EQ(0u, ipc_server_->GetNumberOfActiveConnectionsForTesting());
+}
+
+TEST_F(MojoIpcServerTest, CloseNonexistentReceiver_NoCrash) {
+  ASSERT_EQ(0u, ipc_server_->GetNumberOfActiveConnectionsForTesting());
+  ipc_server_->Close(1u);
+  ASSERT_EQ(0u, ipc_server_->GetNumberOfActiveConnectionsForTesting());
+}
+
+TEST_F(MojoIpcServerTest, RemoteDisconnected_ConnectionRemoved) {
+  mojo::IsolatedConnection client_connection;
+  auto echo_remote = ConnectAndCreateEchoRemote(client_connection);
+  SendEchoAndVerifyResponse(echo_remote);
+  ASSERT_EQ(1u, ipc_server_->GetNumberOfActiveConnectionsForTesting());
+
+  base::RunLoop disconnect_run_loop;
+  ipc_server_->set_disconnect_handler(disconnect_run_loop.QuitClosure());
+  echo_remote.reset();
+  disconnect_run_loop.Run();
+  ASSERT_EQ(0u, ipc_server_->GetNumberOfActiveConnectionsForTesting());
+}
+
+TEST_F(MojoIpcServerTest, RemoteDisconnectedBeforeBound_NewInvitationIsSent) {
+  mojo::IsolatedConnection client_connection;
+  auto handle = client_connection.Connect(ConnectToTestServer());
+  base::SequencedTaskRunnerHandle::Get()->PostTask(
+      FROM_HERE, base::BindLambdaForTesting([&]() { handle.reset(); }));
+  WaitForInvitationSent();
 }
 
 TEST_F(MojoIpcServerTest, ParallelIpcs) {