Move ConnectionToClient::EventHandler from ChromotingHost to ClientSession

BUG=None
TEST=None

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

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

8 files changed, 171 insertions, 179 deletions

diff --git a/remoting/host/chromoting_host.cc b/remoting/host/chromoting_host.cc
index 434552f..de8c78d 100644
--- a/remoting/host/chromoting_host.cc
+++ b/remoting/host/chromoting_host.cc
@@ -130,7 +130,9 @@ void ChromotingHost::Shutdown(const base::Closure& shutdown_task) {
 
   // Disconnect all of the clients, implicitly stopping the ScreenRecorder.
   while (!clients_.empty()) {
-    OnClientDisconnected(clients_.front()->connection());
+    scoped_refptr<ClientSession> client = clients_.front();
+    client->Disconnect();
+    OnClientDisconnected(client);
   }
 
   ShutdownNetwork();
@@ -142,40 +144,32 @@ void ChromotingHost::AddStatusObserver(HostStatusObserver* observer) {
 }
 
 ////////////////////////////////////////////////////////////////////////////
-// protocol::ConnectionToClient::EventHandler implementations
-void ChromotingHost::OnConnectionOpened(ConnectionToClient* connection) {
+// protocol::ClientSession::EventHandler implementation.
+void ChromotingHost::OnSessionAuthenticated(ClientSession* client) {
   DCHECK(context_->network_message_loop()->BelongsToCurrentThread());
-  VLOG(1) << "Connection to client established.";
+  protocol::Session* session = client->connection()->session();
   context_->main_message_loop()->PostTask(
-      FROM_HERE, base::Bind(&ChromotingHost::ProcessPreAuthentication, this,
-                              make_scoped_refptr(connection)));
+      FROM_HERE, base::Bind(&ChromotingHost::AddAuthenticatedClient,
+                            this, make_scoped_refptr(client),
+                            session->config(), session->jid()));
 }
 
-void ChromotingHost::OnConnectionClosed(ConnectionToClient* connection) {
+void ChromotingHost::OnSessionClosed(ClientSession* client) {
   DCHECK(context_->network_message_loop()->BelongsToCurrentThread());
 
   VLOG(1) << "Connection to client closed.";
   context_->main_message_loop()->PostTask(
       FROM_HERE, base::Bind(&ChromotingHost::OnClientDisconnected, this,
-                            make_scoped_refptr(connection)));
-}
-
-void ChromotingHost::OnConnectionFailed(ConnectionToClient* connection) {
-  DCHECK(context_->network_message_loop()->BelongsToCurrentThread());
-
-  LOG(ERROR) << "Connection failed unexpectedly.";
-  context_->main_message_loop()->PostTask(
-      FROM_HERE, base::Bind(&ChromotingHost::OnClientDisconnected, this,
-                            make_scoped_refptr(connection)));
+                            make_scoped_refptr(client)));
 }
 
-void ChromotingHost::OnSequenceNumberUpdated(ConnectionToClient* connection,
+void ChromotingHost::OnSessionSequenceNumber(ClientSession* session,
                                              int64 sequence_number) {
   // Update the sequence number in ScreenRecorder.
   if (MessageLoop::current() != context_->main_message_loop()) {
     context_->main_message_loop()->PostTask(
-        FROM_HERE, base::Bind(&ChromotingHost::OnSequenceNumberUpdated, this,
-                              make_scoped_refptr(connection), sequence_number));
+        FROM_HERE, base::Bind(&ChromotingHost::OnSessionSequenceNumber, this,
+                              make_scoped_refptr(session), sequence_number));
     return;
   }
 
@@ -292,19 +286,14 @@ void ChromotingHost::OnIncomingSession(
 
   LOG(INFO) << "Client connected: " << session->jid();
 
-  // We accept the connection, so create a connection object.
-  ConnectionToClient* connection = new ConnectionToClient(
-      context_->network_message_loop(), this);
-  connection->Init(session);
-
   // Create a client object.
+  scoped_refptr<protocol::ConnectionToClient> connection =
+      new protocol::ConnectionToClient(context_->network_message_loop(),
+                                       session);
   ClientSession* client = new ClientSession(
-      this,
-      connection,
+      this, connection,
       desktop_environment_->event_executor(),
       desktop_environment_->capturer());
-  connection->set_host_stub(client);
-  connection->set_input_stub(client);
 
   clients_.push_back(client);
 }
@@ -348,32 +337,22 @@ void ChromotingHost::SetUiStrings(const UiStrings& ui_strings) {
   ui_strings_ = ui_strings;
 }
 
-void ChromotingHost::OnClientDisconnected(ConnectionToClient* connection) {
+void ChromotingHost::OnClientDisconnected(ClientSession* client) {
   DCHECK_EQ(context_->main_message_loop(), MessageLoop::current());
 
-  // Find the client session corresponding to the given connection.
+  scoped_refptr<ClientSession> client_ref = client;
+
   ClientList::iterator it;
   for (it = clients_.begin(); it != clients_.end(); ++it) {
-    if (it->get()->connection() == connection)
+    if (it->get() == client)
       break;
   }
-  if (it == clients_.end())
-    return;
-
-  scoped_refptr<ClientSession> client = *it;
-
   clients_.erase(it);
 
   if (recorder_.get()) {
-    recorder_->RemoveConnection(connection);
+    recorder_->RemoveConnection(client->connection());
   }
 
-  // Close the connection to client just to be safe.
-  // TODO(garykac): This should be removed when we revisit our shutdown and
-  // disconnect code. This should only need to be done in
-  // ClientSession::Disconnect().
-  connection->Disconnect();
-
   for (StatusObserverList::iterator it = status_observers_.begin();
        it != status_observers_.end(); ++it) {
     (*it)->OnClientDisconnected(client->client_jid());
@@ -387,12 +366,8 @@ void ChromotingHost::OnClientDisconnected(ConnectionToClient* connection) {
 
     // Disable the "curtain" if there are no more active clients.
     EnableCurtainMode(false);
-    if (is_it2me_) {
-      desktop_environment_->OnLastDisconnect();
-    }
+    desktop_environment_->OnLastDisconnect();
   }
-
-  client->OnDisconnected();
 }
 
 // TODO(sergeyu): Move this to SessionManager?
@@ -436,18 +411,8 @@ void ChromotingHost::EnableCurtainMode(bool enable) {
   is_curtained_ = enable;
 }
 
-void ChromotingHost::OnAuthenticationComplete(
-    scoped_refptr<ConnectionToClient> connection) {
-  DCHECK(context_->network_message_loop()->BelongsToCurrentThread());
-
-  context_->main_message_loop()->PostTask(
-      FROM_HERE, base::Bind(&ChromotingHost::AddAuthenticatedClient,
-                            this, connection, connection->session()->config(),
-                            connection->session()->jid()));
-}
-
 void ChromotingHost::AddAuthenticatedClient(
-    scoped_refptr<ConnectionToClient> connection,
+    ClientSession* client,
     const protocol::SessionConfig& config,
     const std::string& jid) {
   DCHECK_EQ(context_->main_message_loop(), MessageLoop::current());
@@ -456,13 +421,14 @@ void ChromotingHost::AddAuthenticatedClient(
   // Iterate over a copy of the list of clients, to avoid mutating the list
   // while iterating over it.
   ClientList clients_copy(clients_);
-  for (ClientList::const_iterator client = clients_copy.begin();
-       client != clients_copy.end(); client++) {
-    ConnectionToClient* connection_other = client->get()->connection();
-    if (connection_other != connection) {
-      OnClientDisconnected(connection_other);
+  for (ClientList::const_iterator other_client = clients_copy.begin();
+       other_client != clients_copy.end(); ++other_client) {
+    if ((*other_client) != client) {
+      (*other_client)->Disconnect();
+      OnClientDisconnected(*other_client);
     }
   }
+
   // Those disconnections should have killed the screen recorder.
   CHECK(recorder_.get() == NULL);
 
@@ -480,7 +446,7 @@ void ChromotingHost::AddAuthenticatedClient(
   }
 
   // Immediately add the connection and start the session.
-  recorder_->AddConnection(connection);
+  recorder_->AddConnection(client->connection());
   recorder_->Start();
   // Notify observers that there is at least one authenticated client.
   for (StatusObserverList::iterator it = status_observers_.begin();
@@ -499,22 +465,6 @@ void ChromotingHost::AddAuthenticatedClient(
   }
 }
 
-void ChromotingHost::ProcessPreAuthentication(
-    const scoped_refptr<ConnectionToClient>& connection) {
-  DCHECK_EQ(context_->main_message_loop(), MessageLoop::current());
-  // Find the client session corresponding to the given connection.
-  ClientList::iterator client;
-  for (client = clients_.begin(); client != clients_.end(); ++client) {
-    if (client->get()->connection() == connection)
-      break;
-  }
-  CHECK(client != clients_.end());
-
-  context_->network_message_loop()->PostTask(
-      FROM_HERE, base::Bind(&ClientSession::OnAuthenticationComplete,
-                            client->get()));
-}
-
 void ChromotingHost::StopScreenRecorder() {
   DCHECK(MessageLoop::current() == context_->main_message_loop());
   DCHECK(recorder_.get());
diff --git a/remoting/host/chromoting_host.h b/remoting/host/chromoting_host.h
index 2596afb..160f5dc 100644
--- a/remoting/host/chromoting_host.h
+++ b/remoting/host/chromoting_host.h
@@ -26,7 +26,6 @@ class Task;
 namespace remoting {
 
 namespace protocol {
-class ConnectionToClient;
 class HostStub;
 class InputStub;
 class SessionConfig;
@@ -64,7 +63,6 @@ class ScreenRecorder;
 //    return to the idle state. We then go to step (2) if there a new
 //    incoming connection.
 class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
-                       public protocol::ConnectionToClient::EventHandler,
                        public ClientSession::EventHandler,
                        public SignalStrategy::StatusObserver,
                        public protocol::SessionManager::Listener {
@@ -99,15 +97,6 @@ class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
   void AddStatusObserver(HostStatusObserver* observer);
 
   ////////////////////////////////////////////////////////////////////////////
-  // protocol::ConnectionToClient::EventHandler implementation.
-  // TODO(sergeyu): Move this to ClientSession.
-  virtual void OnConnectionOpened(protocol::ConnectionToClient* client);
-  virtual void OnConnectionClosed(protocol::ConnectionToClient* client);
-  virtual void OnConnectionFailed(protocol::ConnectionToClient* client);
-  virtual void OnSequenceNumberUpdated(protocol::ConnectionToClient* client,
-                                       int64 sequence_number);
-
-  ////////////////////////////////////////////////////////////////////////////
   // SignalStrategy::StatusObserver implementation.
   virtual void OnStateChange(
       SignalStrategy::StatusObserver::State state) OVERRIDE;
@@ -115,8 +104,10 @@ class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
 
   ////////////////////////////////////////////////////////////////////////////
   // ClientSession::EventHandler implementation.
-  virtual void OnAuthenticationComplete(
-      scoped_refptr<protocol::ConnectionToClient> client);
+  virtual void OnSessionAuthenticated(ClientSession* client) OVERRIDE;
+  virtual void OnSessionClosed(ClientSession* session) OVERRIDE;
+  virtual void OnSessionSequenceNumber(ClientSession* session,
+                                       int64 sequence_number) OVERRIDE;
 
   // SessionManager::Listener implementation.
   virtual void OnSessionManagerInitialized() OVERRIDE;
@@ -124,11 +115,6 @@ class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
       protocol::Session* session,
       protocol::SessionManager::IncomingSessionResponse* response) OVERRIDE;
 
-  void AddAuthenticatedClient(
-      scoped_refptr<protocol::ConnectionToClient> connection,
-      const protocol::SessionConfig& config,
-      const std::string& jid);
-
   // Sets desired configuration for the protocol. Ownership of the
   // |config| is transferred to the object. Must be called before Start().
   void set_protocol_config(protocol::CandidateSessionConfig* config);
@@ -178,20 +164,21 @@ class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
   virtual ~ChromotingHost();
 
   // This method is called if a client is disconnected from the host.
-  void OnClientDisconnected(protocol::ConnectionToClient* client);
+  void OnClientDisconnected(ClientSession* client);
 
   // Creates encoder for the specified configuration.
   Encoder* CreateEncoder(const protocol::SessionConfig& config);
 
   std::string GenerateHostAuthToken(const std::string& encoded_client_token);
 
+  void AddAuthenticatedClient(ClientSession* client,
+                              const protocol::SessionConfig& config,
+                              const std::string& jid);
+
   int AuthenticatedClientsCount() const;
 
   void EnableCurtainMode(bool enable);
 
-  void ProcessPreAuthentication(
-      const scoped_refptr<protocol::ConnectionToClient>& connection);
-
   void StopScreenRecorder();
   void OnScreenRecorderStopped();
 
diff --git a/remoting/host/chromoting_host_unittest.cc b/remoting/host/chromoting_host_unittest.cc
index d5e4df6..45c33e9 100644
--- a/remoting/host/chromoting_host_unittest.cc
+++ b/remoting/host/chromoting_host_unittest.cc
@@ -101,16 +101,29 @@ class ChromotingHostTest : public testing::Test {
     host_ = ChromotingHost::Create(&context_, config_,
                                    desktop_environment_.get(),
                                    access_verifier, false);
-    connection_ = new MockConnectionToClient(
-        &handler_, &host_stub_, event_executor_);
-    connection2_ = new MockConnectionToClient(
-        &handler_, &host_stub2_, &event_executor2_);
-    session_.reset(new MockSession());
-    session2_.reset(new MockSession());
+    session_ = new MockSession();
+    session2_ = new MockSession();
     session_config_ = SessionConfig::GetDefault();
     session_jid_ = "user@domain/rest-of-jid";
     session2_jid_ = "user2@domain/rest-of-jid";
     session_config2_ = SessionConfig::GetDefault();
+    EXPECT_CALL(*session_, jid())
+        .WillRepeatedly(ReturnRef(session_jid_));
+    EXPECT_CALL(*session2_, jid())
+        .WillRepeatedly(ReturnRef(session2_jid_));
+    EXPECT_CALL(*session_, SetStateChangeCallback(_))
+        .Times(AnyNumber());
+    EXPECT_CALL(*session2_, SetStateChangeCallback(_))
+        .Times(AnyNumber());
+    EXPECT_CALL(*session_, config())
+        .WillRepeatedly(ReturnRef(session_config_));
+    EXPECT_CALL(*session2_, config())
+        .WillRepeatedly(ReturnRef(session_config2_));
+
+    connection_ = new MockConnectionToClient(
+        session_, &host_stub_, event_executor_);
+    connection2_ = new MockConnectionToClient(
+        session2_, &host_stub2_, &event_executor2_);
 
     ON_CALL(video_stub_, ProcessVideoPacket(_, _))
         .WillByDefault(DeleteArg<0>());
@@ -121,21 +134,13 @@ class ChromotingHostTest : public testing::Test {
     ON_CALL(*connection_.get(), client_stub())
         .WillByDefault(Return(&client_stub_));
     ON_CALL(*connection_.get(), session())
-        .WillByDefault(Return(session_.get()));
+        .WillByDefault(Return(session_));
     ON_CALL(*connection2_.get(), video_stub())
         .WillByDefault(Return(&video_stub2_));
     ON_CALL(*connection2_.get(), client_stub())
         .WillByDefault(Return(&client_stub2_));
     ON_CALL(*connection2_.get(), session())
-        .WillByDefault(Return(session2_.get()));
-    ON_CALL(*session_.get(), config())
-        .WillByDefault(ReturnRef(session_config_));
-    ON_CALL(*session2_.get(), config())
-        .WillByDefault(ReturnRef(session_config2_));
-    EXPECT_CALL(*session_, jid())
-        .WillRepeatedly(ReturnRef(session_jid_));
-    EXPECT_CALL(*session2_, jid())
-        .WillRepeatedly(ReturnRef(session2_jid_));
+        .WillByDefault(Return(session2_));
     EXPECT_CALL(*connection_.get(), video_stub())
         .Times(AnyNumber());
     EXPECT_CALL(*connection_.get(), client_stub())
@@ -148,10 +153,6 @@ class ChromotingHostTest : public testing::Test {
         .Times(AnyNumber());
     EXPECT_CALL(*connection2_.get(), session())
         .Times(AnyNumber());
-    EXPECT_CALL(*session_.get(), config())
-        .Times(AnyNumber());
-    EXPECT_CALL(*session2_.get(), config())
-        .Times(AnyNumber());
   }
 
   virtual void TearDown() OVERRIDE {
@@ -160,9 +161,8 @@ class ChromotingHostTest : public testing::Test {
 
   // Helper method to pretend a client is connected to ChromotingHost.
   void SimulateClientConnection(int connection_index, bool authenticate) {
-    scoped_refptr<MockConnectionToClient> connection =
+    scoped_refptr<protocol::ConnectionToClient> connection =
         (connection_index == 0) ? connection_ : connection2_;
-
     scoped_refptr<ClientSession> client = new ClientSession(
         host_.get(),
         connection,
@@ -175,16 +175,22 @@ class ChromotingHostTest : public testing::Test {
                               host_, client));
     if (authenticate) {
       context_.network_message_loop()->PostTask(
-          FROM_HERE, base::Bind(&ClientSession::OnAuthenticationComplete,
-                                client.get()));
+          FROM_HERE, base::Bind(&ClientSession::OnConnectionOpened,
+                                client.get(), connection));
+    }
+
+    if (connection_index == 0) {
+      client_ = client;
+    } else {
+      client2_ = client;
     }
   }
 
   // Helper method to remove a client connection from ChromotingHost.
-  void RemoveClientConnection() {
+  void RemoveClientSession() {
     context_.network_message_loop()->PostTask(
-        FROM_HERE, base::Bind(&ChromotingHost::OnClientDisconnected,
-                              host_.get(), connection_));
+        FROM_HERE, base::Bind(
+            &ClientSession::OnConnectionClosed, client_, connection_));
   }
 
   static void AddClientToHost(scoped_refptr<ChromotingHost> host,
@@ -205,15 +211,17 @@ class ChromotingHostTest : public testing::Test {
   scoped_refptr<InMemoryHostConfig> config_;
   MockChromotingHostContext context_;
   scoped_refptr<MockConnectionToClient> connection_;
+  scoped_refptr<ClientSession> client_;
   std::string session_jid_;
-  scoped_ptr<MockSession> session_;
+  MockSession* session_;  // Owned by |connection_|.
   SessionConfig session_config_;
   MockVideoStub video_stub_;
   MockClientStub client_stub_;
   MockHostStub host_stub_;
   scoped_refptr<MockConnectionToClient> connection2_;
+  scoped_refptr<ClientSession> client2_;
   std::string session2_jid_;
-  scoped_ptr<MockSession> session2_;
+  MockSession* session2_;  // Owned by |connection2_|.
   SessionConfig session_config2_;
   MockVideoStub video_stub2_;
   MockClientStub client_stub2_;
@@ -277,8 +285,7 @@ TEST_F(ChromotingHostTest, DISABLED_Reconnect) {
         .Times(0);
     EXPECT_CALL(video_stub_, ProcessVideoPacket(_, _))
         .WillOnce(DoAll(
-            InvokeWithoutArgs(this,
-                              &ChromotingHostTest::RemoveClientConnection),
+            InvokeWithoutArgs(this, &ChromotingHostTest::RemoveClientSession),
             RunDoneTask()))
         .RetiresOnSaturation();
     EXPECT_CALL(video_stub_, ProcessVideoPacket(_, _))
@@ -377,7 +384,7 @@ TEST_F(ChromotingHostTest, CurtainModeFail) {
   EXPECT_CALL(*connection_.get(), Disconnect())
       .WillOnce(QuitMainMessageLoop(&message_loop_));
   SimulateClientConnection(0, false);
-  RemoveClientConnection();
+  RemoveClientSession();
   message_loop_.Run();
 }
 
@@ -443,9 +450,6 @@ TEST_F(ChromotingHostTest, CurtainModeIT2Me) {
             InvokeWithoutArgs(this, &ChromotingHostTest::ShutdownHost),
             RunDoneTask()))
         .RetiresOnSaturation();
-    EXPECT_CALL(video_stub_, ProcessVideoPacket(_, _))
-        .Times(AnyNumber())
-        .InSequence(s1, s2);
     EXPECT_CALL(*connection_.get(), Disconnect())
         .InSequence(s1, s2)
         .RetiresOnSaturation();
diff --git a/remoting/host/client_session.cc b/remoting/host/client_session.cc
index 6cf1c7f..d59c3a6 100644
--- a/remoting/host/client_session.cc
+++ b/remoting/host/client_session.cc
@@ -6,6 +6,7 @@
 
 #include <algorithm>
 
+#include "base/message_loop_proxy.h"
 #include "base/task.h"
 #include "remoting/host/capturer.h"
 #include "remoting/proto/event.pb.h"
@@ -38,14 +39,16 @@ ClientSession::ClientSession(
       authenticated_(false),
       awaiting_continue_approval_(false),
       remote_mouse_button_state_(0) {
-}
+  connection_->SetEventHandler(this);
 
-ClientSession::~ClientSession() {
+  // TODO(sergeyu): Currently ConnectionToClient expects stubs to be
+  // set before channels are connected. Make it possible to set stubs
+  // later and set them only when connection is authenticated.
+  connection_->set_host_stub(this);
+  connection_->set_input_stub(this);
 }
 
-void ClientSession::OnAuthenticationComplete() {
-  authenticated_ = true;
-  event_handler_->OnAuthenticationComplete(connection_.get());
+ClientSession::~ClientSession() {
 }
 
 void ClientSession::InjectKeyEvent(const KeyEvent& event) {
@@ -86,9 +89,40 @@ void ClientSession::InjectMouseEvent(const MouseEvent& event) {
   }
 }
 
-void ClientSession::OnDisconnected() {
-  RestoreEventState();
+void ClientSession::OnConnectionOpened(
+    protocol::ConnectionToClient* connection) {
+  DCHECK_EQ(connection_.get(), connection);
+  authenticated_ = true;
+  event_handler_->OnSessionAuthenticated(this);
+}
+
+void ClientSession::OnConnectionClosed(
+    protocol::ConnectionToClient* connection) {
+  DCHECK_EQ(connection_.get(), connection);
+  scoped_refptr<ClientSession> self = this;
+  event_handler_->OnSessionClosed(this);
+  Disconnect();
+}
+
+void ClientSession::OnConnectionFailed(
+    protocol::ConnectionToClient* connection) {
+  DCHECK_EQ(connection_.get(), connection);
+  // TODO(sergeyu): Log failure reason?
+  scoped_refptr<ClientSession> self = this;
+  event_handler_->OnSessionClosed(this);
+  Disconnect();
+}
+
+void ClientSession::OnSequenceNumberUpdated(
+    protocol::ConnectionToClient* connection, int64 sequence_number) {
+  DCHECK_EQ(connection_.get(), connection);
+  event_handler_->OnSessionSequenceNumber(this, sequence_number);
+}
+
+void ClientSession::Disconnect() {
+  connection_->Disconnect();
   authenticated_ = false;
+  RestoreEventState();
 }
 
 void ClientSession::LocalMouseMoved(const SkIPoint& mouse_pos) {
diff --git a/remoting/host/client_session.h b/remoting/host/client_session.h
index cfc1fe6..07fd387 100644
--- a/remoting/host/client_session.h
+++ b/remoting/host/client_session.h
@@ -22,16 +22,20 @@ class Capturer;
 // per-client state.
 class ClientSession : public protocol::HostStub,
                       public protocol::InputStub,
+                      public protocol::ConnectionToClient::EventHandler,
                       public base::RefCountedThreadSafe<ClientSession> {
  public:
   // Callback interface for passing events to the ChromotingHost.
+  // Called only for external events, i.e. OnSessionClosed() is not
+  // called when Disconnect() is called.
   class EventHandler {
    public:
     virtual ~EventHandler() {}
 
-    // Called to signal that authentication has succeeded.
-    virtual void OnAuthenticationComplete(
-        scoped_refptr<protocol::ConnectionToClient> client) = 0;
+    virtual void OnSessionAuthenticated(ClientSession* client) = 0;
+    virtual void OnSessionClosed(ClientSession* client) = 0;
+    virtual void OnSessionSequenceNumber(ClientSession* client,
+                                         int64 sequence_number) = 0;
   };
 
   // Takes ownership of |user_authenticator|. Does not take ownership of
@@ -45,12 +49,21 @@ class ClientSession : public protocol::HostStub,
   virtual void InjectKeyEvent(const protocol::KeyEvent& event);
   virtual void InjectMouseEvent(const protocol::MouseEvent& event);
 
-  // Notifier called when the client is being disconnected.
-  // This should only be called by ChromotingHost.
-  void OnDisconnected();
-
-  // Set the authenticated flag.
-  void OnAuthenticationComplete();
+  // protocol::ConnectionToClient::EventHandler interface.
+  virtual void OnConnectionOpened(
+      protocol::ConnectionToClient* connection) OVERRIDE;
+  virtual void OnConnectionClosed(
+      protocol::ConnectionToClient* connection) OVERRIDE;
+  virtual void OnConnectionFailed(
+      protocol::ConnectionToClient* connection) OVERRIDE;
+  virtual void OnSequenceNumberUpdated(
+      protocol::ConnectionToClient* connection, int64 sequence_number) OVERRIDE;
+
+  // Disconnects the session and destroys the transport. Event handler
+  // is guaranteed not to be called after this method is called. Can
+  // be called multiple times. The object should not be used after
+  // this method returns.
+  void Disconnect();
 
   protocol::ConnectionToClient* connection() const {
     return connection_.get();
diff --git a/remoting/host/client_session_unittest.cc b/remoting/host/client_session_unittest.cc
index 38bd4be..548b49f 100644
--- a/remoting/host/client_session_unittest.cc
+++ b/remoting/host/client_session_unittest.cc
@@ -27,36 +27,31 @@ class ClientSessionTest : public testing::Test {
 
   virtual void SetUp() {
     client_jid_ = "user@domain/rest-of-jid";
-    EXPECT_CALL(session_, jid()).WillRepeatedly(ReturnRef(client_jid_));
-
-    connection_ = new MockConnectionToClient(
-        &connection_event_handler_, &host_stub_, &input_stub_);
-
-    EXPECT_CALL(*connection_, session()).WillRepeatedly(Return(&session_));
 
     // Set up a large default screen size that won't affect most tests.
     default_screen_size_.set(1000, 1000);
     EXPECT_CALL(capturer_, size_most_recent())
         .WillRepeatedly(ReturnRef(default_screen_size_));
 
+    protocol::MockSession* session = new MockSession();
+    EXPECT_CALL(*session, jid()).WillRepeatedly(ReturnRef(client_jid_));
+    EXPECT_CALL(*session, SetStateChangeCallback(_));
+
     client_session_ = new ClientSession(
         &session_event_handler_,
-        connection_,
-        &input_stub_,
-        &capturer_);
+        new protocol::ConnectionToClient(
+            base::MessageLoopProxy::current(), session),
+        &input_stub_, &capturer_);
   }
 
  protected:
   SkISize default_screen_size_;
   MessageLoop message_loop_;
   std::string client_jid_;
-  MockSession session_;
-  MockConnectionToClientEventHandler connection_event_handler_;
   MockHostStub host_stub_;
   MockInputStub input_stub_;
   MockCapturer capturer_;
   MockClientSessionEventHandler session_event_handler_;
-  scoped_refptr<MockConnectionToClient> connection_;
   scoped_refptr<ClientSession> client_session_;
 };
 
@@ -102,7 +97,7 @@ TEST_F(ClientSessionTest, InputStubFilter) {
   mouse_event3.set_y(301);
 
   InSequence s;
-  EXPECT_CALL(session_event_handler_, OnAuthenticationComplete(_));
+  EXPECT_CALL(session_event_handler_, OnSessionAuthenticated(_));
   EXPECT_CALL(input_stub_, InjectKeyEvent(EqualsKeyEvent(2, true)));
   EXPECT_CALL(input_stub_, InjectKeyEvent(EqualsKeyEvent(2, false)));
   EXPECT_CALL(input_stub_, InjectMouseEvent(EqualsMouseEvent(200, 201)));
@@ -111,12 +106,12 @@ TEST_F(ClientSessionTest, InputStubFilter) {
   // because the client isn't authenticated yet.
   client_session_->InjectKeyEvent(key_event1);
   client_session_->InjectMouseEvent(mouse_event1);
-  client_session_->OnAuthenticationComplete();
+  client_session_->OnConnectionOpened(client_session_->connection());
   // These events should get through to the input stub.
   client_session_->InjectKeyEvent(key_event2_down);
   client_session_->InjectKeyEvent(key_event2_up);
   client_session_->InjectMouseEvent(mouse_event2);
-  client_session_->OnDisconnected();
+  client_session_->Disconnect();
   // These events should not get through to the input stub,
   // because the client has disconnected.
   client_session_->InjectKeyEvent(key_event3);
@@ -135,11 +130,12 @@ TEST_F(ClientSessionTest, LocalInputTest) {
   mouse_event3.set_y(301);
 
   InSequence s;
-  EXPECT_CALL(session_event_handler_, OnAuthenticationComplete(_));
+  EXPECT_CALL(session_event_handler_,
+              OnSessionAuthenticated(client_session_.get()));
   EXPECT_CALL(input_stub_, InjectMouseEvent(EqualsMouseEvent(100, 101)));
   EXPECT_CALL(input_stub_, InjectMouseEvent(EqualsMouseEvent(200, 201)));
 
-  client_session_->OnAuthenticationComplete();
+  client_session_->OnConnectionOpened(client_session_->connection());
   // This event should get through to the input stub.
   client_session_->InjectMouseEvent(mouse_event1);
   // This one should too because the local event echoes the remote one.
@@ -152,7 +148,7 @@ TEST_F(ClientSessionTest, LocalInputTest) {
   client_session_->InjectMouseEvent(mouse_event3);
   // TODO(jamiewalch): Verify that remote inputs are re-enabled eventually
   // (via dependency injection, not sleep!)
-  client_session_->OnDisconnected();
+  client_session_->Disconnect();
 }
 
 TEST_F(ClientSessionTest, RestoreEventState) {
@@ -185,8 +181,9 @@ TEST_F(ClientSessionTest, ClampMouseEvents) {
   EXPECT_CALL(capturer_, size_most_recent())
       .WillRepeatedly(ReturnRef(screen));
 
-  EXPECT_CALL(session_event_handler_, OnAuthenticationComplete(_));
-  client_session_->OnAuthenticationComplete();
+  EXPECT_CALL(session_event_handler_,
+              OnSessionAuthenticated(client_session_.get()));
+  client_session_->OnConnectionOpened(client_session_->connection());
 
   int input_x[3] = { -999, 100, 999 };
   int expected_x[3] = { 0, 100, 199 };
diff --git a/remoting/host/host_mock_objects.h b/remoting/host/host_mock_objects.h
index 5dbdd7b..0613499 100644
--- a/remoting/host/host_mock_objects.h
+++ b/remoting/host/host_mock_objects.h
@@ -96,9 +96,11 @@ class MockClientSessionEventHandler : public ClientSession::EventHandler {
   MockClientSessionEventHandler();
   virtual ~MockClientSessionEventHandler();
 
-  MOCK_METHOD1(OnAuthenticationComplete,
-               void(scoped_refptr<protocol::ConnectionToClient>));
-
+  MOCK_METHOD1(OnSessionAuthenticated, void(ClientSession* client));
+  MOCK_METHOD1(OnSessionClosed, void(ClientSession* client));
+  MOCK_METHOD1(OnSessionFailed, void(ClientSession* client));
+  MOCK_METHOD2(OnSessionSequenceNumber, void(ClientSession* client,
+                                             int64 sequence_number));
  private:
   DISALLOW_COPY_AND_ASSIGN(MockClientSessionEventHandler);
 };
diff --git a/remoting/host/screen_recorder_unittest.cc b/remoting/host/screen_recorder_unittest.cc
index 51996dce..9540997 100644
--- a/remoting/host/screen_recorder_unittest.cc
+++ b/remoting/host/screen_recorder_unittest.cc
@@ -17,6 +17,7 @@
 using ::remoting::protocol::MockConnectionToClient;
 using ::remoting::protocol::MockConnectionToClientEventHandler;
 using ::remoting::protocol::MockHostStub;
+using ::remoting::protocol::MockSession;
 using ::remoting::protocol::MockVideoStub;
 
 using ::testing::_;
@@ -74,8 +75,11 @@ class ScreenRecorderTest : public testing::Test {
     // Capturer and Encoder are owned by ScreenRecorder.
     encoder_ = new MockEncoder();
 
+    session_ = new MockSession();
+    EXPECT_CALL(*session_, SetStateChangeCallback(_));
     connection_ = new MockConnectionToClient(
-        &handler_, &host_stub_, &event_executor_);
+        session_, &host_stub_, &event_executor_);
+    connection_->SetEventHandler(&handler_);
 
     record_ = new ScreenRecorder(
         &message_loop_, &message_loop_,
@@ -89,6 +93,7 @@ class ScreenRecorderTest : public testing::Test {
   MockConnectionToClientEventHandler handler_;
   MockHostStub host_stub_;
   MockEventExecutor event_executor_;
+  MockSession* session_;  // Owned by |connection_|.
   scoped_refptr<MockConnectionToClient> connection_;
 
   // The following mock objects are owned by ScreenRecorder.