Update GnubbyAuthHandler to use the current session ID

remoting/host/security_key/gnubby_auth_handler_win_unittest.cc

 // Copyright 2016 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/security_key/gnubby_auth_handler.h"
 
+#include <cstdint>
 #include <memory>
 #include <string>
 
 #include "base/bind.h"
 #include "base/macros.h"
 #include "base/memory/weak_ptr.h"
 #include "base/message_loop/message_loop.h"
 #include "base/run_loop.h"
 #include "ipc/ipc_channel.h"
 #include "ipc/ipc_listener.h"
@@ skipping 21 matching lines @@
   void OperationComplete();
 
  protected:
   // Waits until the current |run_loop_| instance is signaled, then resets it.
   void WaitForOperationComplete();
 
   // Used as a callback given to the object under test, expected to be called
   // back when a security key request is received by it.
   void SendMessageToClient(int connection_id, const std::string& data);
 
+  // Used as a callback given to the object under test, returns the id of the
+  // Windows session which is being remoted.
+  uint32_t GetDesktopSessionId() const;
+
   // Creates a new gnubby connection on the object under test.
   void CreateGnubbyConnection(const std::string& channel_name);
 
+  // Sets |desktop_session_id_| to the id for the current Windows session.
+  bool SetDesktopSessionId();
+
   // Uses |fake_ipc_client| to connect to the initial IPC server channel, it
   // then validates internal state of the object under test and closes the
   // connection based on |close_connection|.
   void EstablishInitialIpcConnection(
       FakeRemoteSecurityKeyIpcClient* fake_ipc_client,
       int expected_connection_id,
       const std::string& channel_name,
       bool close_connection);
 
   // Sends a security key response message using |fake_ipc_server| and
@@ skipping 31 matching lines @@
 
   // Set as the default factory to create RemoteSecurityKeyIpcServerFactory
   // instances, this class will track each objects creation and allow the tests
   // to access it and use it for driving tests and validate state.
   FakeRemoteSecurityKeyIpcServerFactory ipc_server_factory_;
 
   // Used to validate the object under test uses the correct ID when
   // communicating over the IPC channel.
   int last_connection_id_received_ = -1;
 
+  // Used to validate that IPC connections are only allowed from a specific
+  // Windows session.
+  DWORD desktop_session_id_ = UINT32_MAX;
+
   // Stores the contents of the last IPC message received for validation.
   std::string last_message_received_;
 
  private:
   DISALLOW_COPY_AND_ASSIGN(GnubbyAuthHandlerWinTest);
 };
 
 GnubbyAuthHandlerWinTest::GnubbyAuthHandlerWinTest()
     : run_loop_(new base::RunLoop()) {
-  auth_handler_ = remoting::GnubbyAuthHandler::Create(base::Bind(
-      &GnubbyAuthHandlerWinTest::SendMessageToClient, base::Unretained(this)));
+  auth_handler_ = remoting::GnubbyAuthHandler::Create(
+      base::Bind(&GnubbyAuthHandlerWinTest::SendMessageToClient,
+                 base::Unretained(this)),
+      base::Bind(&GnubbyAuthHandlerWinTest::GetDesktopSessionId,
+                 base::Unretained(this)));
 }
 
 GnubbyAuthHandlerWinTest::~GnubbyAuthHandlerWinTest() {}
 
 void GnubbyAuthHandlerWinTest::OperationComplete() {
   run_loop_->Quit();
 }
 
 void GnubbyAuthHandlerWinTest::WaitForOperationComplete() {
   run_loop_->Run();
   run_loop_.reset(new base::RunLoop());
 }
 
 void GnubbyAuthHandlerWinTest::SendMessageToClient(int connection_id,
                                                    const std::string& data) {
   last_connection_id_received_ = connection_id;
   last_message_received_ = data;
   OperationComplete();
 }
 
+uint32_t GnubbyAuthHandlerWinTest::GetDesktopSessionId() const {
+  return desktop_session_id_;
+}
+
 void GnubbyAuthHandlerWinTest::CreateGnubbyConnection(
     const std::string& channel_name) {
   ASSERT_EQ(0u, auth_handler_->GetActiveConnectionCountForTest());
 
   remoting::SetRemoteSecurityKeyIpcChannelNameForTest(channel_name);
 
   // Create a new Gnubby IPC Server connection.
   auth_handler_->CreateGnubbyConnection();
   ASSERT_TRUE(IPC::Channel::IsNamedServerInitialized(channel_name));
+
+  ASSERT_TRUE(SetDesktopSessionId());
+}
+
+bool GnubbyAuthHandlerWinTest::SetDesktopSessionId() {
+  return ProcessIdToSessionId(GetCurrentProcessId(), &desktop_session_id_);
 }
 
 void GnubbyAuthHandlerWinTest::EstablishInitialIpcConnection(
     FakeRemoteSecurityKeyIpcClient* fake_ipc_client,
     int expected_connection_id,
     const std::string& channel_name,
     bool close_connection) {
   size_t expected_connection_count =
       auth_handler_->GetActiveConnectionCountForTest() + 1;
 
@@ skipping 330 matching lines @@
   // Verify the connection was cleaned up.
   ASSERT_FALSE(fake_ipc_server.get());
   ASSERT_FALSE(auth_handler_->IsValidConnectionId(kConnectionId1));
   ASSERT_EQ(0u, auth_handler_->GetActiveConnectionCountForTest());
 
   // Attempt to connect again after the error.
   EstablishInitialIpcConnection(&fake_ipc_client, kConnectionId2, channel_name,
                                 /*close_connection=*/true);
 }
 
+TEST_F(GnubbyAuthHandlerWinTest, IpcConnectionFailsFromInvalidSession) {
+  std::string channel_name(GetUniqueTestChannelName());
+  CreateGnubbyConnection(channel_name);
+
+  // Set the current session id to a 'different' session.
+  desktop_session_id_ = desktop_session_id_ + 1;
+
+  // Create a fake client and connect to the IPC server channel.
+  FakeRemoteSecurityKeyIpcClient fake_ipc_client(base::Bind(
+      &GnubbyAuthHandlerWinTest::OperationComplete, base::Unretained(this)));
+  ASSERT_TRUE(fake_ipc_client.ConnectViaIpc(channel_name));
+  // Wait for the error callback to be signaled.
+  WaitForOperationComplete();
+
+  // Verify the connection was not set up.
+  ASSERT_FALSE(auth_handler_->IsValidConnectionId(kConnectionId1));
+  ASSERT_EQ(0u, auth_handler_->GetActiveConnectionCountForTest());
+}
+
 }  // namespace remoting