Fix crashes in ChromotingHost

Capturer should be owned by ScreenRecoder. This patch also addes test to test
reconnection of ChromotingHost.

BUG=69969
TEST=ChromotingHostTest.*

Review URL: http://codereview.chromium.org/6266010

git-svn-id: svn://svn.chromium.org/chrome/trunk/src@72543 0039d316-1c4b-4281-b951-d872f2087c98

10 files changed, 255 insertions, 46 deletions

diff --git a/remoting/host/chromoting_host.cc b/remoting/host/chromoting_host.cc
index c769d39..21cc126 100644
--- a/remoting/host/chromoting_host.cc
+++ b/remoting/host/chromoting_host.cc
@@ -24,30 +24,35 @@
 #include "remoting/protocol/session_config.h"
 
 using remoting::protocol::ConnectionToClient;
+using remoting::protocol::InputStub;
 
 namespace remoting {
 
 // static
 ChromotingHost* ChromotingHost::Create(ChromotingHostContext* context,
                                        MutableHostConfig* config) {
-  return Create(context, config,
-                Capturer::Create(context->main_message_loop()));
+  Capturer* capturer = Capturer::Create(context->main_message_loop());
+  InputStub* input_stub = CreateEventExecutor(context->main_message_loop(),
+                                              capturer);
+  return Create(context, config, capturer, input_stub);
 }
 
 // static
 ChromotingHost* ChromotingHost::Create(ChromotingHostContext* context,
                                        MutableHostConfig* config,
-                                       Capturer* capturer) {
-  return new ChromotingHost(context, config, capturer);
+                                       Capturer* capturer,
+                                       InputStub* input_stub) {
+  return new ChromotingHost(context, config, capturer, input_stub);
 }
 
 ChromotingHost::ChromotingHost(ChromotingHostContext* context,
-                               MutableHostConfig* config, Capturer* capturer)
+                               MutableHostConfig* config,
+                               Capturer* capturer,
+                               InputStub* input_stub)
     : context_(context),
       config_(config),
       capturer_(capturer),
-      input_stub_(CreateEventExecutor(
-          context->main_message_loop(), capturer)),
+      input_stub_(input_stub),
       host_stub_(new HostStubFake()),
       state_(kInitial),
       protocol_config_(protocol::CandidateSessionConfig::CreateDefault()) {
@@ -121,12 +126,6 @@ void ChromotingHost::Shutdown() {
     state_ = kStopped;
   }
 
-  // Tell the session to stop and then disconnect all clients.
-  if (recorder_.get()) {
-    recorder_->Stop(NULL);
-    recorder_->RemoveAllConnections();
-  }
-
   // Disconnect the client.
   if (connection_) {
     connection_->Disconnect();
@@ -148,9 +147,13 @@ void ChromotingHost::Shutdown() {
     jingle_client_->Close();
   }
 
-  // Lastly call the shutdown task.
-  if (shutdown_task_.get()) {
+  // Tell the recorder to stop and then disconnect all clients.
+  if (recorder_.get()) {
+    recorder_->RemoveAllConnections();
+    recorder_->Stop(shutdown_task_.release());
+  } else {
     shutdown_task_->Run();
+    shutdown_task_.reset();
   }
 }
 
@@ -169,7 +172,7 @@ void ChromotingHost::OnClientConnected(ConnectionToClient* connection) {
     recorder_ = new ScreenRecorder(context_->main_message_loop(),
                                    context_->encode_message_loop(),
                                    context_->network_message_loop(),
-                                   capturer_.release(),
+                                   capturer_.get(),
                                    encoder);
   }
 
@@ -187,6 +190,7 @@ void ChromotingHost::OnClientDisconnected(ConnectionToClient* connection) {
   if (recorder_.get()) {
     recorder_->RemoveConnection(connection);
     recorder_->Stop(NULL);
+    recorder_ = NULL;
   }
 
   // Close the connection to connection just to be safe.
@@ -206,7 +210,7 @@ void ChromotingHost::OnConnectionOpened(ConnectionToClient* connection) {
   context_->main_message_loop()->PostTask(
       FROM_HERE,
       NewRunnableMethod(this, &ChromotingHost::OnClientConnected,
-                        connection_));
+                        make_scoped_refptr(connection)));
 }
 
 void ChromotingHost::OnConnectionClosed(ConnectionToClient* connection) {
@@ -216,7 +220,7 @@ void ChromotingHost::OnConnectionClosed(ConnectionToClient* connection) {
   context_->main_message_loop()->PostTask(
       FROM_HERE,
       NewRunnableMethod(this, &ChromotingHost::OnClientDisconnected,
-                        connection_));
+                        make_scoped_refptr(connection)));
 }
 
 void ChromotingHost::OnConnectionFailed(ConnectionToClient* connection) {
@@ -226,7 +230,7 @@ void ChromotingHost::OnConnectionFailed(ConnectionToClient* connection) {
   context_->main_message_loop()->PostTask(
       FROM_HERE,
       NewRunnableMethod(this, &ChromotingHost::OnClientDisconnected,
-                        connection_));
+                        make_scoped_refptr(connection)));
 }
 
 ////////////////////////////////////////////////////////////////////////////
@@ -299,8 +303,7 @@ void ChromotingHost::OnNewClientSession(
 
   VLOG(1) << "Client connected: " << session->jid();
 
-  // If we accept the connected then create a client object and set the
-  // callback.
+  // If we accept the connected then create a connection object.
   connection_ = new ConnectionToClient(context_->network_message_loop(),
                                        this, host_stub_.get(),
                                        input_stub_.get());
diff --git a/remoting/host/chromoting_host.h b/remoting/host/chromoting_host.h
index 33680a6..c99dac2 100644
--- a/remoting/host/chromoting_host.h
+++ b/remoting/host/chromoting_host.h
@@ -64,13 +64,14 @@ class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
                        public protocol::ConnectionToClient::EventHandler,
                        public JingleClient::Callback {
  public:
-  // Factory methods that must be used to create ChromotingHost
-  // instances.  Default capturer is used if it is not specified.
+  // Factory methods that must be used to create ChromotingHost instances.
+  // Default capturer and input stub are used if it is not specified.
   static ChromotingHost* Create(ChromotingHostContext* context,
                                 MutableHostConfig* config);
   static ChromotingHost* Create(ChromotingHostContext* context,
                                 MutableHostConfig* config,
-                                Capturer* capturer);
+                                Capturer* capturer,
+                                protocol::InputStub* input_stub);
 
   // Asynchronously start the host process.
   //
@@ -114,7 +115,7 @@ class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
  private:
   friend class base::RefCountedThreadSafe<ChromotingHost>;
   ChromotingHost(ChromotingHostContext* context, MutableHostConfig* config,
-                 Capturer* capturer);
+                 Capturer* capturer, protocol::InputStub* input_stub);
   virtual ~ChromotingHost();
 
   enum State {
diff --git a/remoting/host/chromoting_host_unittest.cc b/remoting/host/chromoting_host_unittest.cc
new file mode 100644
index 0000000..2a90307
--- /dev/null
+++ b/remoting/host/chromoting_host_unittest.cc
@@ -0,0 +1,172 @@
+// Copyright (c) 2010 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 "base/task.h"
+#include "remoting/host/capturer_fake.h"
+#include "remoting/host/chromoting_host.h"
+#include "remoting/host/chromoting_host_context.h"
+#include "remoting/host/in_memory_host_config.h"
+#include "remoting/proto/video.pb.h"
+#include "remoting/protocol/mock_objects.h"
+#include "remoting/protocol/session_config.h"
+#include "testing/gmock/include/gmock/gmock.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+using testing::_;
+using testing::AnyNumber;
+using testing::DeleteArg;
+using testing::DoAll;
+using testing::InSequence;
+using testing::InvokeWithoutArgs;
+using testing::Return;
+
+namespace remoting {
+
+namespace {
+
+void PostQuitTask(MessageLoop* message_loop) {
+  message_loop->PostTask(FROM_HERE, new MessageLoop::QuitTask());
+}
+
+// Run the task and delete it afterwards. This action is used to deal with
+// done callbacks.
+ACTION(RunDoneTask) {
+  arg1->Run();
+  delete arg1;
+}
+
+ACTION_P(QuitMainMessageLoop, message_loop) {
+  PostQuitTask(message_loop);
+}
+
+}  // namepace
+
+class ChromotingHostTest : public testing::Test {
+ public:
+  ChromotingHostTest() {
+  }
+
+  virtual void SetUp() {
+    config_ = new InMemoryHostConfig();
+    context_.Start();
+    Capturer* capturer = new CapturerFake(context_.main_message_loop());
+    input_stub_ = new protocol::MockInputStub();
+    host_ = ChromotingHost::Create(&context_, config_, capturer, input_stub_);
+    connection_ = new protocol::MockConnectionToClient();
+    session_ = new protocol::MockSession();
+    session_config_.reset(protocol::SessionConfig::CreateDefault());
+
+    ON_CALL(*connection_.get(), video_stub())
+        .WillByDefault(Return(&video_stub_));
+    ON_CALL(*connection_.get(), session())
+        .WillByDefault(Return(session_));
+    ON_CALL(*session_.get(), config())
+        .WillByDefault(Return(session_config_.get()));
+  }
+
+  virtual void TearDown() {
+    context_.Stop();
+  }
+
+  // Helper metjod to pretend a client is connected to ChromotingHost.
+  void InjectClientConnection() {
+    context_.network_message_loop()->PostTask(
+        FROM_HERE,
+        NewRunnableMethod(host_.get(),
+                          &ChromotingHost::OnConnectionOpened,
+                          connection_));
+  }
+
+  // Helper method to remove a client connection from ChromotongHost.
+  void RemoveClientConnection() {
+    context_.network_message_loop()->PostTask(
+        FROM_HERE,
+        NewRunnableMethod(host_.get(),
+                          &ChromotingHost::OnConnectionClosed,
+                          connection_));
+  }
+
+ protected:
+  MessageLoop message_loop_;
+  scoped_refptr<ChromotingHost> host_;
+  scoped_refptr<InMemoryHostConfig> config_;
+  ChromotingHostContext context_;
+  scoped_refptr<protocol::MockConnectionToClient> connection_;
+  scoped_refptr<protocol::MockSession> session_;
+  scoped_ptr<protocol::SessionConfig> session_config_;
+  protocol::MockVideoStub video_stub_;
+  protocol::MockInputStub* input_stub_;
+};
+
+TEST_F(ChromotingHostTest, StartAndShutdown) {
+  host_->Start(NewRunnableFunction(&PostQuitTask, &message_loop_));
+
+  message_loop_.PostTask(FROM_HERE,
+                         NewRunnableMethod(host_.get(),
+                                           &ChromotingHost::Shutdown));
+  message_loop_.Run();
+}
+
+TEST_F(ChromotingHostTest, Connect) {
+  host_->Start(NewRunnableFunction(&PostQuitTask, &message_loop_));
+
+  ON_CALL(video_stub_, ProcessVideoPacket(_, _))
+      .WillByDefault(
+          DoAll(DeleteArg<0>(), DeleteArg<1>()));
+
+  // When the video packet is received we first shutdown ChromotingHost
+  // then execute the done task.
+  InSequence s;
+  EXPECT_CALL(video_stub_, ProcessVideoPacket(_, _))
+      .WillOnce(DoAll(
+          InvokeWithoutArgs(host_.get(), &ChromotingHost::Shutdown),
+          RunDoneTask()))
+      .RetiresOnSaturation();
+  EXPECT_CALL(video_stub_, ProcessVideoPacket(_, _))
+      .Times(AnyNumber());
+
+  InjectClientConnection();
+  message_loop_.Run();
+}
+
+TEST_F(ChromotingHostTest, Reconnect) {
+  host_->Start(NewRunnableFunction(&PostQuitTask, &message_loop_));
+
+  ON_CALL(video_stub_, ProcessVideoPacket(_, _))
+      .WillByDefault(
+          DoAll(DeleteArg<0>(), DeleteArg<1>()));
+
+  // When the video packet is received we first disconnect the mock
+  // connection.
+  InSequence s;
+  EXPECT_CALL(video_stub_, ProcessVideoPacket(_, _))
+      .WillOnce(DoAll(
+          InvokeWithoutArgs(this, &ChromotingHostTest::RemoveClientConnection),
+          RunDoneTask()))
+      .RetiresOnSaturation();
+  EXPECT_CALL(video_stub_, ProcessVideoPacket(_, _))
+      .Times(AnyNumber());
+
+  // If Disconnect() is called we can break the main message loop.
+  EXPECT_CALL(*connection_.get(), Disconnect())
+      .WillOnce(QuitMainMessageLoop(&message_loop_))
+      .RetiresOnSaturation();
+
+  InjectClientConnection();
+  message_loop_.Run();
+
+  // Connect the client again.
+  EXPECT_CALL(video_stub_, ProcessVideoPacket(_, _))
+      .WillOnce(DoAll(
+          InvokeWithoutArgs(host_.get(), &ChromotingHost::Shutdown),
+          RunDoneTask()))
+      .RetiresOnSaturation();
+  EXPECT_CALL(video_stub_, ProcessVideoPacket(_, _))
+      .Times(AnyNumber());
+
+  InjectClientConnection();
+  message_loop_.Run();
+}
+
+}  // namespace remoting
diff --git a/remoting/host/screen_recorder.cc b/remoting/host/screen_recorder.cc
index 76ebcc8..c154197 100644
--- a/remoting/host/screen_recorder.cc
+++ b/remoting/host/screen_recorder.cc
@@ -102,8 +102,8 @@ void ScreenRecorder::RemoveAllConnections() {
 
 Capturer* ScreenRecorder::capturer() {
   DCHECK_EQ(capture_loop_, MessageLoop::current());
-  DCHECK(capturer_.get());
-  return capturer_.get();
+  DCHECK(capturer_);
+  return capturer_;
 }
 
 Encoder* ScreenRecorder::encoder() {
diff --git a/remoting/host/screen_recorder.h b/remoting/host/screen_recorder.h
index 1115e8f..cddeb56 100644
--- a/remoting/host/screen_recorder.h
+++ b/remoting/host/screen_recorder.h
@@ -70,7 +70,7 @@ class ScreenRecorder : public base::RefCountedThreadSafe<ScreenRecorder> {
  public:
 
   // Construct a ScreenRecorder. Message loops and threads are provided.
-  // This object does not own capturer and encoder.
+  // This object does not own capturer but owns encoder.
   ScreenRecorder(MessageLoop* capture_loop,
                  MessageLoop* encode_loop,
                  MessageLoop* network_loop,
@@ -160,7 +160,7 @@ class ScreenRecorder : public base::RefCountedThreadSafe<ScreenRecorder> {
 
   // Reference to the capturer. This member is always accessed on the capture
   // thread.
-  scoped_ptr<Capturer> capturer_;
+  Capturer* capturer_;
 
   // Reference to the encoder. This member is always accessed on the encode
   // thread.
diff --git a/remoting/host/screen_recorder_unittest.cc b/remoting/host/screen_recorder_unittest.cc
index c9f4123..85ba54a 100644
--- a/remoting/host/screen_recorder_unittest.cc
+++ b/remoting/host/screen_recorder_unittest.cc
@@ -76,12 +76,11 @@ class ScreenRecorderTest : public testing::Test {
 
   virtual void SetUp() {
     // Capturer and Encoder are owned by ScreenRecorder.
-    capturer_ = new MockCapturer();
     encoder_ = new MockEncoder();
     connection_ = new MockConnectionToClient();
     record_ = new ScreenRecorder(
         &message_loop_, &message_loop_, &message_loop_,
-        capturer_, encoder_);
+        &capturer_, encoder_);
   }
 
  protected:
@@ -89,7 +88,7 @@ class ScreenRecorderTest : public testing::Test {
   scoped_refptr<MockConnectionToClient> connection_;
 
   // The following mock objects are owned by ScreenRecorder.
-  MockCapturer* capturer_;
+  MockCapturer capturer_;
   MockEncoder* encoder_;
   MessageLoop message_loop_;
  private:
@@ -108,11 +107,11 @@ TEST_F(ScreenRecorderTest, OneRecordCycle) {
   }
   scoped_refptr<CaptureData> data(new CaptureData(planes, kWidth,
                                                   kHeight, kFormat));
-  EXPECT_CALL(*capturer_, width()).WillRepeatedly(Return(kWidth));
-  EXPECT_CALL(*capturer_, height()).WillRepeatedly(Return(kHeight));
+  EXPECT_CALL(capturer_, width()).WillRepeatedly(Return(kWidth));
+  EXPECT_CALL(capturer_, height()).WillRepeatedly(Return(kHeight));
 
   // First the capturer is called.
-  EXPECT_CALL(*capturer_, CaptureInvalidRects(NotNull()))
+  EXPECT_CALL(capturer_, CaptureInvalidRects(NotNull()))
       .WillOnce(RunCallback(update_rects, data));
 
   // Expect the encoder be called.
@@ -158,11 +157,11 @@ TEST_F(ScreenRecorderTest, StartAndStop) {
   }
   scoped_refptr<CaptureData> data(new CaptureData(planes, kWidth,
                                                   kHeight, kFormat));
-  EXPECT_CALL(*capturer_, width()).WillRepeatedly(Return(kWidth));
-  EXPECT_CALL(*capturer_, height()).WillRepeatedly(Return(kHeight));
+  EXPECT_CALL(capturer_, width()).WillRepeatedly(Return(kWidth));
+  EXPECT_CALL(capturer_, height()).WillRepeatedly(Return(kHeight));
 
   // First the capturer is called.
-  EXPECT_CALL(*capturer_, CaptureInvalidRects(NotNull()))
+  EXPECT_CALL(capturer_, CaptureInvalidRects(NotNull()))
       .WillRepeatedly(RunCallback(update_rects, data));
 
   // Expect the encoder be called.
diff --git a/remoting/host/simple_host_process.cc b/remoting/host/simple_host_process.cc
index 56ad378..289da0f 100644
--- a/remoting/host/simple_host_process.cc
+++ b/remoting/host/simple_host_process.cc
@@ -32,6 +32,7 @@
 #include "remoting/host/capturer_fake.h"
 #include "remoting/host/chromoting_host.h"
 #include "remoting/host/chromoting_host_context.h"
+#include "remoting/host/event_executor.h"
 #include "remoting/host/json_host_config.h"
 #include "remoting/proto/video.pb.h"
 
@@ -116,9 +117,12 @@ int main(int argc, char** argv) {
 
   bool fake = cmd_line->HasSwitch(kFakeSwitchName);
   if (fake) {
+    remoting::Capturer* capturer =
+        new remoting::CapturerFake(context.main_message_loop());
+    remoting::protocol::InputStub* input_stub =
+        CreateEventExecutor(context.main_message_loop(), capturer);
     host = ChromotingHost::Create(
-        &context, config,
-        new remoting::CapturerFake(context.main_message_loop()));
+        &context, config, capturer, input_stub);
   } else {
     host = ChromotingHost::Create(&context, config);
   }
diff --git a/remoting/protocol/connection_to_client.h b/remoting/protocol/connection_to_client.h
index 749ef5c..b7aa052 100644
--- a/remoting/protocol/connection_to_client.h
+++ b/remoting/protocol/connection_to_client.h
@@ -55,10 +55,10 @@ class ConnectionToClient :
 
   virtual ~ConnectionToClient();
 
-  virtual void Init(protocol::Session* session);
+  virtual void Init(Session* session);
 
   // Returns the connection in use.
-  virtual protocol::Session* session();
+  virtual Session* session();
 
   // Disconnect the client connection. This method is allowed to be called
   // more than once and calls after the first one will be ignored.
@@ -78,15 +78,15 @@ class ConnectionToClient :
 
  private:
   // Callback for protocol Session.
-  void OnSessionStateChange(protocol::Session::State state);
+  void OnSessionStateChange(Session::State state);
 
   // Process a libjingle state change event on the |loop_|.
-  void StateChangeTask(protocol::Session::State state);
+  void StateChangeTask(Session::State state);
 
   void OnClosed();
 
   // The libjingle channel used to send and receive data from the remote client.
-  scoped_refptr<protocol::Session> session_;
+  scoped_refptr<Session> session_;
 
   scoped_ptr<VideoWriter> video_writer_;
 
diff --git a/remoting/protocol/mock_objects.h b/remoting/protocol/mock_objects.h
index 045a414..99826c2 100644
--- a/remoting/protocol/mock_objects.h
+++ b/remoting/protocol/mock_objects.h
@@ -6,9 +6,11 @@
 #define REMOTING_PROTOCOL_MOCK_OBJECTS_H_
 
 #include "remoting/proto/internal.pb.h"
+#include "remoting/protocol/client_stub.h"
 #include "remoting/protocol/connection_to_client.h"
 #include "remoting/protocol/host_stub.h"
 #include "remoting/protocol/input_stub.h"
+#include "remoting/protocol/session.h"
 #include "remoting/protocol/video_stub.h"
 #include "testing/gmock/include/gmock/gmock.h"
 
@@ -21,8 +23,10 @@ class MockConnectionToClient : public ConnectionToClient {
  public:
   MockConnectionToClient() {}
 
-  MOCK_METHOD1(Init, void(ChromotocolConnection* connection));
+  MOCK_METHOD1(Init, void(Session* session));
   MOCK_METHOD0(video_stub, VideoStub*());
+  MOCK_METHOD0(client_stub, ClientStub*());
+  MOCK_METHOD0(session, Session*());
   MOCK_METHOD0(Disconnect, void());
 
  private:
@@ -78,6 +82,31 @@ class MockVideoStub : public VideoStub {
   DISALLOW_COPY_AND_ASSIGN(MockVideoStub);
 };
 
+class MockSession : public Session {
+ public:
+  MockSession() {}
+
+  MOCK_METHOD1(SetStateChangeCallback, void(StateChangeCallback* callback));
+  MOCK_METHOD0(control_channel, net::Socket*());
+  MOCK_METHOD0(event_channel, net::Socket*());
+  MOCK_METHOD0(video_channel, net::Socket*());
+  MOCK_METHOD0(video_rtp_channel, net::Socket*());
+  MOCK_METHOD0(video_rtcp_channel, net::Socket*());
+  MOCK_METHOD0(jid, const std::string&());
+  MOCK_METHOD0(message_loop, MessageLoop*());
+  MOCK_METHOD0(candidate_config, const CandidateSessionConfig*());
+  MOCK_METHOD0(config, const SessionConfig*());
+  MOCK_METHOD1(set_config, void(const SessionConfig* config));
+  MOCK_METHOD0(initiator_token, const std::string&());
+  MOCK_METHOD1(set_initiator_token, void(const std::string& initiator_token));
+  MOCK_METHOD0(receiver_token, const std::string&());
+  MOCK_METHOD1(set_receiver_token, void(const std::string& receiver_token));
+  MOCK_METHOD1(Close, void(Task* closed_task));
+
+ private:
+  DISALLOW_COPY_AND_ASSIGN(MockSession);
+};
+
 }  // namespace protocol
 }  // namespace remoting
 
diff --git a/remoting/remoting.gyp b/remoting/remoting.gyp
index 70dc0fd..c112bab 100644
--- a/remoting/remoting.gyp
+++ b/remoting/remoting.gyp
@@ -467,6 +467,7 @@
         'client/mock_objects.h',
         'host/access_verifier_unittest.cc',
         'host/chromoting_host_context_unittest.cc',
+        'host/chromoting_host_unittest.cc',
         'host/differ_unittest.cc',
         'host/differ_block_unittest.cc',
         'host/heartbeat_sender_unittest.cc',