remoting/host/win/wts_terminal_monitor.cc

   1 // Copyright (c) 2012 The Chromium Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style license that can be
   3 // found in the LICENSE file.
   4
   5 #include "remoting/host/win/wts_terminal_monitor.h"
   6
   7 #include <windows.h>
   8 #include <winternl.h>
   9 #include <wtsapi32.h>
  10
  11 #include "base/basictypes.h"
  12 #include "base/files/file_path.h"
  13 #include "base/lazy_instance.h"
  14 #include "base/native_library.h"
  15 #include "base/scoped_native_library.h"
  16 #include "base/utf_string_conversions.h"
  17 #include "net/base/ip_endpoint.h"
  18
  19 namespace {
  20
  21 // Used to query the endpoint of an attached RDP client.
  22 const WINSTATIONINFOCLASS kWinStationRemoteAddress =
  23     static_cast<WINSTATIONINFOCLASS>(29);
  24
  25 // WinStationRemoteAddress information class returns the following structure.
  26 // Note that its layout is different from sockaddr_in/sockaddr_in6. For
  27 // instance both |ipv4| and |ipv6| structures are 4 byte aligned so there is
  28 // additional 2 byte padding after |sin_family|.
  29 struct RemoteAddress {
  30   unsigned short sin_family;
  31   union {
  32     struct {
  33       USHORT sin_port;
  34       ULONG in_addr;
  35       UCHAR sin_zero[8];
  36     } ipv4;
  37     struct {
  38       USHORT sin6_port;
  39       ULONG sin6_flowinfo;
  40       USHORT sin6_addr[8];
  41       ULONG sin6_scope_id;
  42     } ipv6;
  43   };
  44 };
  45
  46 // Loads winsta.dll dynamically and resolves the address of
  47 // the winsta!WinStationQueryInformationW() function.
  48 class WinstaLoader {
  49  public:
  50   WinstaLoader();
  51   ~WinstaLoader();
  52
  53   // Returns the address and port of the RDP client attached to |session_id|.
  54   bool GetRemoteAddress(uint32 session_id, RemoteAddress* address);
  55
  56  private:
  57   // Handle of dynamically loaded winsta.dll.
  58   base::ScopedNativeLibrary winsta_;
  59
  60   // Points to winsta!WinStationQueryInformationW().
  61   PWINSTATIONQUERYINFORMATIONW win_station_query_information_;
  62
  63   DISALLOW_COPY_AND_ASSIGN(WinstaLoader);
  64 };
  65
  66 static base::LazyInstance<WinstaLoader> g_winsta = LAZY_INSTANCE_INITIALIZER;
  67
  68 WinstaLoader::WinstaLoader() :
  69   winsta_(base::FilePath(base::GetNativeLibraryName(UTF8ToUTF16("winsta")))) {
  70
  71   // Resolve the function pointer.
  72   win_station_query_information_ =
  73       static_cast<PWINSTATIONQUERYINFORMATIONW>(
  74           winsta_.GetFunctionPointer("WinStationQueryInformationW"));
  75 }
  76
  77 WinstaLoader::~WinstaLoader() {
  78 }
  79
  80 bool WinstaLoader::GetRemoteAddress(uint32 session_id, RemoteAddress* address) {
  81   ULONG length;
  82   return win_station_query_information_(WTS_CURRENT_SERVER_HANDLE,
  83                                         session_id,
  84                                         kWinStationRemoteAddress,
  85                                         address,
  86                                         sizeof(*address),
  87                                         &length) != FALSE;
  88 }
  89
  90 }  // namespace
  91
  92 namespace remoting {
  93
  94 // Session id that does not represent any session.
  95 const uint32 kInvalidSessionId = 0xffffffffu;
  96
  97 WtsTerminalMonitor::~WtsTerminalMonitor() {
  98 }
  99
 100 // static
 101 bool WtsTerminalMonitor::GetEndpointForSessionId(uint32 session_id,
 102                                                  net::IPEndPoint* endpoint) {
 103   // Fast path for the case when |session_id| is currently attached to
 104   // the physical console.
 105   if (session_id == WTSGetActiveConsoleSessionId()) {
 106     *endpoint = net::IPEndPoint();
 107     return true;
 108   }
 109
 110   RemoteAddress address;
 111   // WinStationQueryInformationW() fails if no RDP client is attached to
 112   // |session_id|.
 113   if (!g_winsta.Get().GetRemoteAddress(session_id, &address))
 114     return false;
 115
 116   // Convert the RemoteAddress structure into sockaddr_in/sockaddr_in6.
 117   switch (address.sin_family) {
 118     case AF_INET: {
 119       sockaddr_in ipv4 = { 0 };
 120       ipv4.sin_family = AF_INET;
 121       ipv4.sin_port = address.ipv4.sin_port;
 122       ipv4.sin_addr.S_un.S_addr = address.ipv4.in_addr;
 123       return endpoint->FromSockAddr(
 124           reinterpret_cast<struct sockaddr*>(&ipv4), sizeof(ipv4));
 125     }
 126
 127     case AF_INET6: {
 128       sockaddr_in6 ipv6 = { 0 };
 129       ipv6.sin6_family = AF_INET6;
 130       ipv6.sin6_port = address.ipv6.sin6_port;
 131       ipv6.sin6_flowinfo = address.ipv6.sin6_flowinfo;
 132       memcpy(&ipv6.sin6_addr, address.ipv6.sin6_addr, sizeof(ipv6.sin6_addr));
 133       ipv6.sin6_scope_id = address.ipv6.sin6_scope_id;
 134       return endpoint->FromSockAddr(
 135           reinterpret_cast<struct sockaddr*>(&ipv6), sizeof(ipv6));
 136     }
 137
 138     default:
 139       return false;
 140   }
 141 }
 142
 143 // static
 144 uint32 WtsTerminalMonitor::GetSessionIdForEndpoint(
 145     const net::IPEndPoint& client_endpoint) {
 146   // Use the fast path if the caller wants to get id of the session attached to
 147   // the physical console.
 148   if (client_endpoint == net::IPEndPoint())
 149     return WTSGetActiveConsoleSessionId();
 150
 151   // Enumerate all sessions and try to match the client endpoint.
 152   WTS_SESSION_INFO* session_info;
 153   DWORD session_info_count;
 154   if (!WTSEnumerateSessions(WTS_CURRENT_SERVER_HANDLE, 0, 1, &session_info,
 155                             &session_info_count)) {
 156     LOG_GETLASTERROR(ERROR) << "Failed to enumerate all sessions";
 157     return kInvalidSessionId;
 158   }
 159   for (DWORD i = 0; i < session_info_count; ++i) {
 160     net::IPEndPoint endpoint;
 161     if (GetEndpointForSessionId(session_info[i].SessionId, &endpoint) &&
 162         endpoint == client_endpoint) {
 163       uint32 session_id = session_info[i].SessionId;
 164       WTSFreeMemory(session_info);
 165       return session_id;
 166     }
 167   }
 168
 169   // |client_endpoint| is not associated with any session.
 170   WTSFreeMemory(session_info);
 171   return kInvalidSessionId;
 172 }
 173
 174 WtsTerminalMonitor::WtsTerminalMonitor() {
 175 }
 176
 177 }  // namespace remoting