16 files changed, 356 insertions, 242 deletions

diff --git a/remoting/host/chromoting_host.cc b/remoting/host/chromoting_host.cc
index df77848..6b5803e 100644
--- a/remoting/host/chromoting_host.cc
+++ b/remoting/host/chromoting_host.cc
@@ -57,7 +57,6 @@ ChromotingHost::ChromotingHost(ChromotingHostContext* context,
       state_(kInitial),
       protocol_config_(protocol::CandidateSessionConfig::CreateDefault()) {
   DCHECK(desktop_environment_.get());
-  desktop_environment_->set_event_handler(this);
 }
 
 ChromotingHost::~ChromotingHost() {
@@ -140,10 +139,11 @@ void ChromotingHost::Shutdown() {
     recorder_->RemoveAllConnections();
   }
 
-  // Disconnect the client.
-  if (connection_) {
-    connection_->Disconnect();
+  // Disconnect the clients.
+  for (size_t i = 0; i < clients_.size(); i++) {
+    clients_[i]->Disconnect();
   }
+  clients_.clear();
 
   // Stop the heartbeat sender.
   if (heartbeat_sender_) {
@@ -172,24 +172,6 @@ void ChromotingHost::Shutdown() {
 // This method is called when a client connects.
 void ChromotingHost::OnClientConnected(ConnectionToClient* connection) {
   DCHECK_EQ(context_->main_message_loop(), MessageLoop::current());
-
-  // Create a new RecordSession if there was none.
-  if (!recorder_.get()) {
-    // Then we create a ScreenRecorder passing the message loops that
-    // it should run on.
-    DCHECK(desktop_environment_->capturer());
-
-    Encoder* encoder = CreateEncoder(connection->session()->config());
-
-    recorder_ = new ScreenRecorder(context_->main_message_loop(),
-                                   context_->encode_message_loop(),
-                                   context_->network_message_loop(),
-                                   desktop_environment_->capturer(),
-                                   encoder);
-  }
-
-  // Immediately add the connection and start the session.
-  recorder_->AddConnection(connection);
 }
 
 void ChromotingHost::OnClientDisconnected(ConnectionToClient* connection) {
@@ -199,15 +181,25 @@ void ChromotingHost::OnClientDisconnected(ConnectionToClient* connection) {
   // TODO(hclam): Stop only if the last connection disconnected.
   if (recorder_.get()) {
     recorder_->RemoveConnection(connection);
-    recorder_->Stop(NULL);
-    recorder_ = NULL;
+    // The recorder only exists to serve the unique authenticated client.
+    // If that client has disconnected, then we can kill the recorder.
+    if (connection->client_authenticated()) {
+      recorder_->Stop(NULL);
+      recorder_ = NULL;
+    }
   }
 
   // Close the connection to connection just to be safe.
   connection->Disconnect();
 
   // Also remove reference to ConnectionToClient from this object.
-  connection_ = NULL;
+  std::vector<scoped_refptr<ClientSession> >::iterator it;
+  for (it = clients_.begin(); it != clients_.end(); ++it) {
+    if (it->get()->connection() == connection) {
+      clients_.erase(it);
+      break;
+    }
+  }
 }
 
 ////////////////////////////////////////////////////////////////////////////
@@ -287,8 +279,7 @@ void ChromotingHost::OnNewClientSession(
     protocol::Session* session,
     protocol::SessionManager::IncomingSessionResponse* response) {
   base::AutoLock auto_lock(lock_);
-  // TODO(hclam): Allow multiple clients to connect to the host.
-  if (connection_.get() || state_ != kStarted) {
+  if (state_ != kStarted) {
     *response = protocol::SessionManager::DECLINE;
     return;
   }
@@ -323,12 +314,19 @@ void ChromotingHost::OnNewClientSession(
 
   VLOG(1) << "Client connected: " << session->jid();
 
-  // If we accept the connected then create a connection object.
-  connection_ = new ConnectionToClient(context_->network_message_loop(),
-                                       this,
-                                       desktop_environment_.get(),
-                                       desktop_environment_->input_stub());
-  connection_->Init(session);
+  // We accept the connection, so create a connection object.
+  ConnectionToClient* connection = new ConnectionToClient(
+      context_->network_message_loop(),
+      this,
+      desktop_environment_->input_stub());
+
+  // Create a client object.
+  ClientSession* client = new ClientSession(this, connection);
+  connection->set_host_stub(client);
+
+  connection->Init(session);
+
+  clients_.push_back(client);
 }
 
 void ChromotingHost::set_protocol_config(
@@ -338,8 +336,8 @@ void ChromotingHost::set_protocol_config(
   protocol_config_.reset(config);
 }
 
-protocol::HostStub* ChromotingHost::host_stub() const {
-  return desktop_environment_.get();
+void ChromotingHost::AddClient(ClientSession* client) {
+  clients_.push_back(client);
 }
 
 void ChromotingHost::OnServerClosed() {
@@ -371,34 +369,68 @@ std::string ChromotingHost::GenerateHostAuthToken(
   return encoded_client_token;
 }
 
-void ChromotingHost::LocalLoginSucceeded() {
+void ChromotingHost::LocalLoginSucceeded(
+    scoped_refptr<ConnectionToClient> connection) {
   if (MessageLoop::current() != context_->main_message_loop()) {
     context_->main_message_loop()->PostTask(
         FROM_HERE,
-        NewRunnableMethod(this, &ChromotingHost::LocalLoginSucceeded));
+        NewRunnableMethod(this,
+                          &ChromotingHost::LocalLoginSucceeded,
+                          connection));
     return;
   }
 
   protocol::LocalLoginStatus* status = new protocol::LocalLoginStatus();
   status->set_success(true);
-  connection_->client_stub()->BeginSessionResponse(
+  connection->client_stub()->BeginSessionResponse(
       status, new DeleteTask<protocol::LocalLoginStatus>(status));
 
-  connection_->OnClientAuthenticated();
+  connection->OnClientAuthenticated();
+
+  // Disconnect all other clients.
+  // Iterate over a copy of the list of clients, to avoid mutating the list
+  // while iterating over it.
+  std::vector<scoped_refptr<ClientSession> > clients_copy(clients_);
+  std::vector<scoped_refptr<ClientSession> >::const_iterator client;
+  for (client = clients_copy.begin(); client != clients_copy.end(); client++) {
+    ConnectionToClient* connection_other = client->get()->connection();
+    if (connection_other != connection) {
+      OnClientDisconnected(connection_other);
+    }
+  }
+  // Those disconnections should have killed the screen recorder.
+  CHECK(recorder_.get() == NULL);
+
+  // Create a new RecordSession if there was none.
+  if (!recorder_.get()) {
+    // Then we create a ScreenRecorder passing the message loops that
+    // it should run on.
+    Encoder* encoder = CreateEncoder(connection->session()->config());
+
+    recorder_ = new ScreenRecorder(context_->main_message_loop(),
+                                   context_->encode_message_loop(),
+                                   context_->network_message_loop(),
+                                   desktop_environment_->capturer(),
+                                   encoder);
+  }
+
+  // Immediately add the connection and start the session.
+  recorder_->AddConnection(connection);
   recorder_->Start();
 }
 
-void ChromotingHost::LocalLoginFailed() {
+void ChromotingHost::LocalLoginFailed(
+    scoped_refptr<ConnectionToClient> connection) {
   if (MessageLoop::current() != context_->main_message_loop()) {
     context_->main_message_loop()->PostTask(
         FROM_HERE,
-        NewRunnableMethod(this, &ChromotingHost::LocalLoginFailed));
+        NewRunnableMethod(this, &ChromotingHost::LocalLoginFailed, connection));
     return;
   }
 
   protocol::LocalLoginStatus* status = new protocol::LocalLoginStatus();
   status->set_success(false);
-  connection_->client_stub()->BeginSessionResponse(
+  connection->client_stub()->BeginSessionResponse(
       status, new DeleteTask<protocol::LocalLoginStatus>(status));
 }
 
diff --git a/remoting/host/chromoting_host.h b/remoting/host/chromoting_host.h
index fc6daaf..f754c27 100644
--- a/remoting/host/chromoting_host.h
+++ b/remoting/host/chromoting_host.h
@@ -12,6 +12,7 @@
 #include "remoting/base/encoder.h"
 #include "remoting/host/access_verifier.h"
 #include "remoting/host/capturer.h"
+#include "remoting/host/client_session.h"
 #include "remoting/host/desktop_environment.h"
 #include "remoting/host/heartbeat_sender.h"
 #include "remoting/jingle_glue/jingle_client.h"
@@ -63,7 +64,7 @@ class ScreenRecorder;
 //    incoming connection.
 class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
                        public protocol::ConnectionToClient::EventHandler,
-                       public DesktopEnvironment::EventHandler,
+                       public ClientSession::EventHandler,
                        public JingleClient::Callback {
  public:
   // Factory methods that must be used to create ChromotingHost instances.
@@ -105,9 +106,11 @@ class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
   virtual void OnStateChange(JingleClient* client, JingleClient::State state);
 
   ////////////////////////////////////////////////////////////////////////////
-  // DesktopEnvironment::EventHandler implementations
-  virtual void LocalLoginSucceeded();
-  virtual void LocalLoginFailed();
+  // ClientSession::EventHandler implementations
+  virtual void LocalLoginSucceeded(
+      scoped_refptr<protocol::ConnectionToClient> client);
+  virtual void LocalLoginFailed(
+      scoped_refptr<protocol::ConnectionToClient> client);
 
   // Callback for ChromotingServer.
   void OnNewClientSession(
@@ -118,13 +121,8 @@ class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
   // |config| is transferred to the object. Must be called before Start().
   void set_protocol_config(protocol::CandidateSessionConfig* config);
 
-  // This getter is only used in unit test.
-  protocol::HostStub* host_stub() const;
-
   // This setter is only used in unit test to simulate client connection.
-  void set_connection(protocol::ConnectionToClient* conn) {
-    connection_ = conn;
-  }
+  void AddClient(ClientSession* client);
 
  private:
   friend class base::RefCountedThreadSafe<ChromotingHost>;
@@ -167,9 +165,8 @@ class ChromotingHost : public base::RefCountedThreadSafe<ChromotingHost>,
 
   AccessVerifier access_verifier_;
 
-  // A ConnectionToClient manages the connectino to a remote client.
-  // TODO(hclam): Expand this to a list of clients.
-  scoped_refptr<protocol::ConnectionToClient> connection_;
+  // The connections to remote clients.
+  std::vector<scoped_refptr<ClientSession> > clients_;
 
   // Session manager for the host process.
   scoped_refptr<ScreenRecorder> recorder_;
diff --git a/remoting/host/chromoting_host_unittest.cc b/remoting/host/chromoting_host_unittest.cc
index 23295ae..5338c08 100644
--- a/remoting/host/chromoting_host_unittest.cc
+++ b/remoting/host/chromoting_host_unittest.cc
@@ -6,12 +6,12 @@
 #include "remoting/host/capturer_fake.h"
 #include "remoting/host/chromoting_host.h"
 #include "remoting/host/chromoting_host_context.h"
-#include "remoting/host/desktop_environment_fake.h"
 #include "remoting/host/host_mock_objects.h"
 #include "remoting/host/in_memory_host_config.h"
 #include "remoting/proto/video.pb.h"
 #include "remoting/protocol/protocol_mock_objects.h"
 #include "remoting/protocol/session_config.h"
+#include "testing/gmock_mutant.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
@@ -27,6 +27,7 @@ using ::remoting::protocol::SessionConfig;
 
 using testing::_;
 using testing::AnyNumber;
+using testing::CreateFunctor;
 using testing::DeleteArg;
 using testing::DoAll;
 using testing::InSequence;
@@ -41,20 +42,6 @@ void PostQuitTask(MessageLoop* message_loop) {
   message_loop->PostTask(FROM_HERE, new MessageLoop::QuitTask());
 }
 
-void BeginSessionRequest(protocol::HostStub* host_stub) {
-  LocalLoginCredentials* credentials =
-      new LocalLoginCredentials();
-  credentials->set_type(protocol::PASSWORD);
-  credentials->set_username("hello");
-
-  const std::string password = "world!";
-  credentials->set_credential(password.data(), password.length());
-
-  host_stub->BeginSessionRequest(
-      credentials,
-      new DeleteTask<LocalLoginCredentials>(credentials));
-}
-
 // Run the task and delete it afterwards. This action is used to deal with
 // done callbacks.
 ACTION(RunDoneTask) {
@@ -66,7 +53,7 @@ ACTION_P(QuitMainMessageLoop, message_loop) {
   PostQuitTask(message_loop);
 }
 
-}  // namepace
+}  // namespace
 
 class ChromotingHostTest : public testing::Test {
  public:
@@ -90,57 +77,91 @@ class ChromotingHostTest : public testing::Test {
 
     Capturer* capturer = new CapturerFake(context_.main_message_loop());
     host_stub_ = new MockHostStub();
+    host_stub2_ = new MockHostStub();
     input_stub_ = new MockInputStub();
+    input_stub2_ = new MockInputStub();
     DesktopEnvironment* desktop =
-        new DesktopEnvironmentFake(capturer, input_stub_);
+        new DesktopEnvironment(capturer, input_stub_);
     host_ = ChromotingHost::Create(&context_, config_, desktop);
     connection_ = new MockConnectionToClient(
         &message_loop_, &handler_, host_stub_, input_stub_);
+    connection2_ = new MockConnectionToClient(
+        &message_loop_, &handler_, host_stub2_, input_stub2_);
     session_ = new MockSession();
+    session2_ = new MockSession();
     session_config_.reset(SessionConfig::CreateDefault());
+    session_config2_.reset(SessionConfig::CreateDefault());
 
     ON_CALL(video_stub_, ProcessVideoPacket(_, _))
         .WillByDefault(
             DoAll(DeleteArg<0>(), DeleteArg<1>()));
+    ON_CALL(video_stub2_, ProcessVideoPacket(_, _))
+        .WillByDefault(
+            DoAll(DeleteArg<0>(), DeleteArg<1>()));
     ON_CALL(*connection_.get(), video_stub())
         .WillByDefault(Return(&video_stub_));
     ON_CALL(*connection_.get(), client_stub())
         .WillByDefault(Return(&client_stub_));
     ON_CALL(*connection_.get(), session())
         .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_));
     ON_CALL(*session_.get(), config())
         .WillByDefault(Return(session_config_.get()));
+    ON_CALL(*session2_.get(), config())
+        .WillByDefault(Return(session_config2_.get()));
     EXPECT_CALL(*connection_.get(), video_stub())
         .Times(AnyNumber());
     EXPECT_CALL(*connection_.get(), client_stub())
         .Times(AnyNumber());
     EXPECT_CALL(*connection_.get(), session())
         .Times(AnyNumber());
+    EXPECT_CALL(*connection2_.get(), video_stub())
+        .Times(AnyNumber());
+    EXPECT_CALL(*connection2_.get(), client_stub())
+        .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() {
   }
 
-  // Helper metjod to pretend a client is connected to ChromotingHost.
-  void SimulateClientConnection() {
+  // Helper method to pretend a client is connected to ChromotingHost.
+  void SimulateClientConnection(int connection_index) {
+    scoped_refptr<MockConnectionToClient> connection =
+        (connection_index == 0) ? connection_ : connection2_;
+    scoped_refptr<ClientSession> client = new ClientSession(host_.get(),
+                                                            connection);
+    connection->set_host_stub(client.get());
+
     context_.network_message_loop()->PostTask(
         FROM_HERE,
         NewRunnableMethod(host_.get(),
-                          &ChromotingHost::set_connection,
-                          connection_));
+                          &ChromotingHost::AddClient,
+                          client));
     context_.network_message_loop()->PostTask(
         FROM_HERE,
         NewRunnableMethod(host_.get(),
                           &ChromotingHost::OnClientConnected,
-                          connection_));
+                          connection));
     context_.network_message_loop()->PostTask(
         FROM_HERE,
-        NewRunnableFunction(&BeginSessionRequest, host_->host_stub()));
+        NewRunnableMethod(host_.get(),
+                          &ChromotingHost::LocalLoginSucceeded,
+                          connection));
   }
 
-  // Helper method to remove a client connection from ChromotongHost.
+  // Helper method to remove a client connection from ChromotingHost.
   void RemoveClientConnection() {
     context_.network_message_loop()->PostTask(
         FROM_HERE,
@@ -162,6 +183,13 @@ class ChromotingHostTest : public testing::Test {
   MockClientStub client_stub_;
   MockHostStub* host_stub_;
   MockInputStub* input_stub_;
+  scoped_refptr<MockConnectionToClient> connection2_;
+  scoped_refptr<MockSession> session2_;
+  scoped_ptr<SessionConfig> session_config2_;
+  MockVideoStub video_stub2_;
+  MockClientStub client_stub2_;
+  MockHostStub* host_stub2_;
+  MockInputStub* input_stub2_;
 };
 
 TEST_F(ChromotingHostTest, StartAndShutdown) {
@@ -192,7 +220,7 @@ TEST_F(ChromotingHostTest, Connect) {
   EXPECT_CALL(*connection_.get(), Disconnect())
       .RetiresOnSaturation();
 
-  SimulateClientConnection();
+  SimulateClientConnection(0);
   message_loop_.Run();
 }
 
@@ -222,7 +250,7 @@ TEST_F(ChromotingHostTest, Reconnect) {
       .WillOnce(QuitMainMessageLoop(&message_loop_))
       .RetiresOnSaturation();
 
-  SimulateClientConnection();
+  SimulateClientConnection(0);
   message_loop_.Run();
 
   // Connect the client again.
@@ -239,7 +267,49 @@ TEST_F(ChromotingHostTest, Reconnect) {
   EXPECT_CALL(*connection_.get(), Disconnect())
       .RetiresOnSaturation();
 
-  SimulateClientConnection();
+  SimulateClientConnection(0);
+  message_loop_.Run();
+}
+
+TEST_F(ChromotingHostTest, ConnectTwice) {
+  host_->Start(NewRunnableFunction(&PostQuitTask, &message_loop_));
+
+  EXPECT_CALL(client_stub_, BeginSessionResponse(_, _))
+      .Times(1)
+      .WillRepeatedly(RunDoneTask());
+
+  EXPECT_CALL(client_stub2_, BeginSessionResponse(_, _))
+      .Times(1)
+      .WillRepeatedly(RunDoneTask());
+
+  // When a video packet is received we connect the second mock
+  // connection.
+  {
+    InSequence s;
+    EXPECT_CALL(video_stub_, ProcessVideoPacket(_, _))
+        .WillOnce(DoAll(
+            InvokeWithoutArgs(
+                CreateFunctor(
+                    this, &ChromotingHostTest::SimulateClientConnection, 1)),
+            RunDoneTask()))
+        .RetiresOnSaturation();
+    EXPECT_CALL(video_stub_, ProcessVideoPacket(_, _))
+        .Times(AnyNumber());
+    EXPECT_CALL(video_stub2_, ProcessVideoPacket(_, _))
+        .WillOnce(DoAll(
+            InvokeWithoutArgs(host_.get(), &ChromotingHost::Shutdown),
+            RunDoneTask()))
+        .RetiresOnSaturation();
+    EXPECT_CALL(video_stub2_, ProcessVideoPacket(_, _))
+        .Times(AnyNumber());
+  }
+
+  EXPECT_CALL(*connection_.get(), Disconnect())
+      .RetiresOnSaturation();
+  EXPECT_CALL(*connection2_.get(), Disconnect())
+      .RetiresOnSaturation();
+
+  SimulateClientConnection(0);
   message_loop_.Run();
 }
 
diff --git a/remoting/host/client_session.cc b/remoting/host/client_session.cc
new file mode 100755
index 0000000..f346e43
--- /dev/null
+++ b/remoting/host/client_session.cc
@@ -0,0 +1,66 @@
+// Copyright (c) 2011 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 "remoting/host/client_session.h"
+
+#include "base/scoped_ptr.h"
+#include "base/task.h"
+#include "remoting/host/user_authenticator.h"
+#include "remoting/proto/auth.pb.h"
+
+namespace remoting {
+
+ClientSession::ClientSession(
+    EventHandler* event_handler,
+    scoped_refptr<protocol::ConnectionToClient> connection)
+    : event_handler_(event_handler),
+      connection_(connection) {
+}
+
+ClientSession::~ClientSession() {
+}
+
+void ClientSession::SuggestResolution(
+    const protocol::SuggestResolutionRequest* msg, Task* done) {
+  done->Run();
+  delete done;
+}
+
+void ClientSession::BeginSessionRequest(
+    const protocol::LocalLoginCredentials* credentials, Task* done) {
+  DCHECK(event_handler_);
+
+  bool success = false;
+  scoped_ptr<UserAuthenticator> authenticator(UserAuthenticator::Create());
+  switch (credentials->type()) {
+    case protocol::PASSWORD:
+      success = authenticator->Authenticate(credentials->username(),
+                                            credentials->credential());
+      break;
+
+    default:
+      LOG(ERROR) << "Invalid credentials type " << credentials->type();
+      break;
+  }
+
+  if (success) {
+    event_handler_->LocalLoginSucceeded(connection_.get());
+  } else {
+    LOG(WARNING) << "Login failed for user " << credentials->username();
+    event_handler_->LocalLoginFailed(connection_.get());
+  }
+
+  done->Run();
+  delete done;
+}
+
+void ClientSession::Disconnect() {
+  connection_->Disconnect();
+}
+
+protocol::ConnectionToClient* ClientSession::connection() const {
+  return connection_.get();
+}
+
+}  // namespace remoting
diff --git a/remoting/host/client_session.h b/remoting/host/client_session.h
new file mode 100755
index 0000000..620923f
--- /dev/null
+++ b/remoting/host/client_session.h
@@ -0,0 +1,60 @@
+// Copyright (c) 2011 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.
+
+#ifndef REMOTING_HOST_CLIENT_SESSION_H_
+#define REMOTING_HOST_CLIENT_SESSION_H_
+
+#include "remoting/protocol/connection_to_client.h"
+#include "remoting/protocol/host_stub.h"
+
+namespace remoting {
+
+// A ClientSession keeps a reference to a connection to a client, and maintains
+// per-client state.
+class ClientSession : public protocol::HostStub,
+                      public base::RefCountedThreadSafe<ClientSession> {
+ public:
+  // Callback interface for passing events to the ChromotingHost.
+  class EventHandler {
+   public:
+    virtual ~EventHandler() {}
+
+    // Called to signal that local login has succeeded and ChromotingHost can
+    // proceed with the next step.
+    virtual void LocalLoginSucceeded(
+        scoped_refptr<protocol::ConnectionToClient> client) = 0;
+
+    // Called to signal that local login has failed.
+    virtual void LocalLoginFailed(
+        scoped_refptr<protocol::ConnectionToClient> client) = 0;
+  };
+
+  ClientSession(EventHandler* event_handler,
+                scoped_refptr<protocol::ConnectionToClient> connection);
+
+  // protocol::HostStub interface.
+  virtual void SuggestResolution(
+      const protocol::SuggestResolutionRequest* msg, Task* done);
+  virtual void BeginSessionRequest(
+      const protocol::LocalLoginCredentials* credentials, Task* done);
+
+  // Disconnect this client session.
+  void Disconnect();
+
+  protocol::ConnectionToClient* connection() const;
+
+ protected:
+  friend class base::RefCountedThreadSafe<ClientSession>;
+  ~ClientSession();
+
+ private:
+  EventHandler* event_handler_;
+  scoped_refptr<protocol::ConnectionToClient> connection_;
+
+  DISALLOW_COPY_AND_ASSIGN(ClientSession);
+};
+
+}  // namespace remoting
+
+#endif  // REMOTING_HOST_CLIENT_SESSION_H_
diff --git a/remoting/host/desktop_environment.cc b/remoting/host/desktop_environment.cc
index e6e1057..edc1792 100644
--- a/remoting/host/desktop_environment.cc
+++ b/remoting/host/desktop_environment.cc
@@ -5,10 +5,6 @@
 #include "remoting/host/desktop_environment.h"
 
 #include "remoting/host/capturer.h"
-#include "remoting/host/chromoting_host.h"
-#include "remoting/host/user_authenticator.h"
-#include "remoting/proto/auth.pb.h"
-#include "remoting/protocol/client_stub.h"
 #include "remoting/protocol/input_stub.h"
 
 using remoting::protocol::InputStub;
@@ -17,46 +13,11 @@ namespace remoting {
 
 DesktopEnvironment::DesktopEnvironment(Capturer* capturer,
                                        InputStub* input_stub)
-    : event_handler_(NULL),
-      capturer_(capturer),
+    : capturer_(capturer),
       input_stub_(input_stub) {
 }
 
 DesktopEnvironment::~DesktopEnvironment() {
 }
 
-void DesktopEnvironment::SuggestResolution(
-    const protocol::SuggestResolutionRequest* msg, Task* done) {
-  done->Run();
-  delete done;
-}
-
-void DesktopEnvironment::BeginSessionRequest(
-    const protocol::LocalLoginCredentials* credentials, Task* done) {
-  DCHECK(event_handler_);
-
-  bool success = false;
-  scoped_ptr<UserAuthenticator> authenticator(UserAuthenticator::Create());
-  switch (credentials->type()) {
-    case protocol::PASSWORD:
-      success = authenticator->Authenticate(credentials->username(),
-                                            credentials->credential());
-      break;
-
-    default:
-      LOG(ERROR) << "Invalid credentials type " << credentials->type();
-      break;
-  }
-
-  if (success) {
-    event_handler_->LocalLoginSucceeded();
-  } else {
-    LOG(WARNING) << "Login failed for user " << credentials->username();
-    event_handler_->LocalLoginFailed();
-  }
-
-  done->Run();
-  delete done;
-}
-
 }  // namespace remoting
diff --git a/remoting/host/desktop_environment.h b/remoting/host/desktop_environment.h
index 90bab71..7d5846b 100644
--- a/remoting/host/desktop_environment.h
+++ b/remoting/host/desktop_environment.h
@@ -7,7 +7,6 @@
 
 #include "base/basictypes.h"
 #include "base/scoped_ptr.h"
-#include "remoting/protocol/host_stub.h"
 
 namespace remoting {
 
@@ -17,44 +16,16 @@ class InputStub;
 
 class Capturer;
 
-class DesktopEnvironment : public protocol::HostStub {
+class DesktopEnvironment {
  public:
-  // Callback interface for passing events to the ChromotingHost.
-  class EventHandler {
-   public:
-    virtual ~EventHandler() {}
-
-    // Called to signal that local login has succeeded and ChromotingHost can
-    // proceed with the next step.
-    virtual void LocalLoginSucceeded() = 0;
-
-    // Called to signal that local login has failed.
-    virtual void LocalLoginFailed() = 0;
-  };
-
   DesktopEnvironment(Capturer* capturer, protocol::InputStub* input_stub);
   virtual ~DesktopEnvironment();
 
   Capturer* capturer() const { return capturer_.get(); }
   protocol::InputStub* input_stub() const { return input_stub_.get(); }
-  // Called by ChromotingHost constructor
-  void set_event_handler(EventHandler* event_handler) {
-    event_handler_ = event_handler;
-  }
-
-  // protocol::HostStub interface.
-  virtual void SuggestResolution(
-      const protocol::SuggestResolutionRequest* msg, Task* done);
-  virtual void BeginSessionRequest(
-      const protocol::LocalLoginCredentials* credentials, Task* done);
-
- protected:
-  // Allow access by DesktopEnvironmentFake for unittest.
-  EventHandler* event_handler_;
 
  private:
-  // Capturer to be used by ScreenRecorder. Once the ScreenRecorder is
-  // constructed this is set to NULL.
+  // Capturer to be used by ScreenRecorder.
   scoped_ptr<Capturer> capturer_;
 
   // InputStub in the host executes input events received from the client.
diff --git a/remoting/host/desktop_environment_fake.cc b/remoting/host/desktop_environment_fake.cc
deleted file mode 100644
index 36da27a7..0000000
--- a/remoting/host/desktop_environment_fake.cc
+++ /dev/null
@@ -1,31 +0,0 @@
-// Copyright (c) 2011 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 "remoting/host/desktop_environment_fake.h"
-
-#include "base/task.h"
-
-using remoting::protocol::InputStub;
-
-namespace remoting {
-
-DesktopEnvironmentFake::DesktopEnvironmentFake(Capturer* capturer,
-                                               InputStub* input_stub)
-    : DesktopEnvironment(capturer, input_stub) {
-}
-
-DesktopEnvironmentFake::~DesktopEnvironmentFake() {}
-
-void DesktopEnvironmentFake::BeginSessionRequest(
-    const protocol::LocalLoginCredentials* credentials, Task* done) {
-  DCHECK(event_handler_);
-
-  // For unit-test, don't require a valid Unix user/password for connection.
-  event_handler_->LocalLoginSucceeded();
-
-  done->Run();
-  delete done;
-}
-
-}  // namespace remoting
diff --git a/remoting/host/desktop_environment_fake.h b/remoting/host/desktop_environment_fake.h
deleted file mode 100644
index 66ec320..0000000
--- a/remoting/host/desktop_environment_fake.h
+++ /dev/null
@@ -1,29 +0,0 @@
-// Copyright (c) 2011 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.
-
-#ifndef REMOTING_HOST_DESKTOP_ENVIRONMENT_FAKE_H_
-#define REMOTING_HOST_DESKTOP_ENVIRONMENT_FAKE_H_
-
-#include "remoting/host/desktop_environment.h"
-
-namespace remoting {
-
-// A DesktopEnvironmentFake allows connection to proceed for unit-testing,
-// without needing a valid user/password for the system.
-class DesktopEnvironmentFake : public DesktopEnvironment {
- public:
-  DesktopEnvironmentFake(Capturer* capturer, protocol::InputStub* input_stub);
-  virtual ~DesktopEnvironmentFake();
-
-  // Overridden to do no authentication.
-  virtual void BeginSessionRequest(
-      const protocol::LocalLoginCredentials* credentials, Task* done);
-
- private:
-  DISALLOW_COPY_AND_ASSIGN(DesktopEnvironmentFake);
-};
-
-}  // namespace remoting
-
-#endif  // REMOTING_HOST_DESKTOP_ENVIRONMENT_FAKE_H_
diff --git a/remoting/protocol/connection_to_client.cc b/remoting/protocol/connection_to_client.cc
index 11101cd..4ec9a0b 100644
--- a/remoting/protocol/connection_to_client.cc
+++ b/remoting/protocol/connection_to_client.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2010 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 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.
 
@@ -23,12 +23,11 @@ static const size_t kAverageUpdateStream = 10;
 
 ConnectionToClient::ConnectionToClient(MessageLoop* message_loop,
                                        EventHandler* handler,
-                                       HostStub* host_stub,
                                        InputStub* input_stub)
     : client_authenticated_(false),
       loop_(message_loop),
       handler_(handler),
-      host_stub_(host_stub),
+      host_stub_(NULL),
       input_stub_(input_stub) {
   DCHECK(loop_);
   DCHECK(handler_);
@@ -76,6 +75,14 @@ ClientStub* ConnectionToClient::client_stub() {
   return client_stub_.get();
 }
 
+protocol::HostStub* ConnectionToClient::host_stub() {
+  return host_stub_;
+}
+
+void ConnectionToClient::set_host_stub(protocol::HostStub* host_stub) {
+  host_stub_ = host_stub;
+}
+
 void ConnectionToClient::OnSessionStateChange(protocol::Session::State state) {
   if (state == protocol::Session::CONNECTED) {
     client_stub_.reset(new ClientControlSender(session_->control_channel()));
@@ -145,5 +152,9 @@ void ConnectionToClient::OnClientAuthenticated() {
     client_stub_->OnAuthenticated();
 }
 
+bool ConnectionToClient::client_authenticated() {
+  return client_authenticated_;
+}
+
 }  // namespace protocol
 }  // namespace remoting
diff --git a/remoting/protocol/connection_to_client.h b/remoting/protocol/connection_to_client.h
index d2721f3..c9e4555 100644
--- a/remoting/protocol/connection_to_client.h
+++ b/remoting/protocol/connection_to_client.h
@@ -1,4 +1,4 @@
-// Copyright (c) 2010 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 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.
 
@@ -50,11 +50,8 @@ class ConnectionToClient :
   // It is guranteed that |handler| is called only on the |message_loop|.
   ConnectionToClient(MessageLoop* message_loop,
                      EventHandler* handler,
-                     HostStub* host_stub,
                      InputStub* input_stub);
 
-  virtual ~ConnectionToClient();
-
   virtual void Init(Session* session);
 
   // Returns the connection in use.
@@ -72,9 +69,19 @@ class ConnectionToClient :
   // Return pointer to ClientStub.
   virtual ClientStub* client_stub();
 
+  virtual HostStub* host_stub();
+  virtual void set_host_stub(HostStub* host_stub);
+
   // Called when the host accepts the client authentication.
   void OnClientAuthenticated();
 
+  // Whether the client has been authenticated.
+  bool client_authenticated();
+
+ protected:
+  friend class base::RefCountedThreadSafe<ConnectionToClient>;
+  virtual ~ConnectionToClient();
+
  private:
   // Callback for protocol Session.
   void OnSessionStateChange(Session::State state);
diff --git a/remoting/protocol/connection_to_client_unittest.cc b/remoting/protocol/connection_to_client_unittest.cc
index fcc732c1..7a21096 100644
--- a/remoting/protocol/connection_to_client_unittest.cc
+++ b/remoting/protocol/connection_to_client_unittest.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2010 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 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.
 
@@ -27,8 +27,8 @@ class ConnectionToClientTest : public testing::Test {
     session_->set_message_loop(&message_loop_);
 
     // Allocate a ClientConnection object with the mock objects.
-    viewer_ = new ConnectionToClient(&message_loop_, &handler_,
-                                     &host_stub_, &input_stub_);
+    viewer_ = new ConnectionToClient(&message_loop_, &handler_, &input_stub_);
+    viewer_->set_host_stub(&host_stub_);
     viewer_->Init(session_);
     EXPECT_CALL(handler_, OnConnectionOpened(viewer_.get()));
     session_->state_change_callback()->Run(
diff --git a/remoting/protocol/host_message_dispatcher.cc b/remoting/protocol/host_message_dispatcher.cc
index f4391a7..9b1ee28 100644
--- a/remoting/protocol/host_message_dispatcher.cc
+++ b/remoting/protocol/host_message_dispatcher.cc
@@ -1,4 +1,4 @@
-// Copyright (c) 2010 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 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.
 
@@ -7,6 +7,7 @@
 #include "remoting/proto/control.pb.h"
 #include "remoting/proto/event.pb.h"
 #include "remoting/proto/internal.pb.h"
+#include "remoting/protocol/connection_to_client.h"
 #include "remoting/protocol/host_message_dispatcher.h"
 #include "remoting/protocol/host_stub.h"
 #include "remoting/protocol/input_stub.h"
@@ -48,25 +49,22 @@ void HostMessageDispatcher::Initialize(
 
 void HostMessageDispatcher::OnControlMessageReceived(
     ControlMessage* message, Task* done_task) {
+  // BeginSessionRequest is always allowed.
+  if (message->has_begin_session_request()) {
+    host_stub_->BeginSessionRequest(
+        &message->begin_session_request().credentials(), done_task);
+    return;
+  }
   if (!host_stub_->authenticated()) {
-    // When the client has not authenticated with the host, we restrict the
-    // control messages that we support.
-    if (message->has_begin_session_request()) {
-      host_stub_->BeginSessionRequest(
-          &message->begin_session_request().credentials(), done_task);
-      return;
-    } else {
-      LOG(WARNING) << "Invalid control message received "
-                   << "(client not authenticated).";
-    }
+    // When the client has not authenticated with the host, no other messages
+    // are allowed.
+    LOG(WARNING) << "Invalid control message received "
+                 << "(client not authenticated).";
   } else {
     // TODO(sergeyu): Add message validation.
     if (message->has_suggest_resolution()) {
       host_stub_->SuggestResolution(&message->suggest_resolution(), done_task);
       return;
-    } else if (message->has_begin_session_request()) {
-      LOG(WARNING) << "BeginSessionRequest sent after client already "
-                   << "authorized.";
     } else {
       LOG(WARNING) << "Invalid control message received.";
     }
diff --git a/remoting/protocol/host_stub.h b/remoting/protocol/host_stub.h
index c4e9b46..737154d 100644
--- a/remoting/protocol/host_stub.h
+++ b/remoting/protocol/host_stub.h
@@ -1,4 +1,4 @@
-// Copyright (c) 2010 The Chromium Authors. All rights reserved.
+// Copyright (c) 2011 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.
 
@@ -16,8 +16,8 @@ class Task;
 namespace remoting {
 namespace protocol {
 
-class SuggestResolutionRequest;
 class LocalLoginCredentials;
+class SuggestResolutionRequest;
 
 class HostStub {
  public:
diff --git a/remoting/protocol/protocol_mock_objects.cc b/remoting/protocol/protocol_mock_objects.cc
index 5deb40f..85aad65a 100644
--- a/remoting/protocol/protocol_mock_objects.cc
+++ b/remoting/protocol/protocol_mock_objects.cc
@@ -11,7 +11,8 @@ MockConnectionToClient::MockConnectionToClient(MessageLoop* message_loop,
                                                EventHandler* handler,
                                                HostStub* host_stub,
                                                InputStub* input_stub)
-    : ConnectionToClient(message_loop, handler, host_stub, input_stub) {
+    : ConnectionToClient(message_loop, handler, input_stub) {
+  set_host_stub(host_stub);
 }
 
 MockConnectionToClient::~MockConnectionToClient() {}
diff --git a/remoting/remoting.gyp b/remoting/remoting.gyp
index 076763e..c402b27 100644
--- a/remoting/remoting.gyp
+++ b/remoting/remoting.gyp
@@ -207,6 +207,8 @@
         'host/chromoting_host.h',
         'host/chromoting_host_context.cc',
         'host/chromoting_host_context.h',
+        'host/client_session.cc',
+        'host/client_session.h',
         'host/desktop_environment.cc',
         'host/desktop_environment.h',
         'host/differ.h',
@@ -541,8 +543,6 @@
         'host/access_verifier_unittest.cc',
         'host/chromoting_host_context_unittest.cc',
         'host/chromoting_host_unittest.cc',
-        'host/desktop_environment_fake.cc',
-        'host/desktop_environment_fake.h',
         'host/differ_block_unittest.cc',
         'host/differ_unittest.cc',
         'host/heartbeat_sender_unittest.cc',