[telemetry] Read MSRs from a separate privileged process.

tools/telemetry/telemetry/core/platform/win_platform_backend.py

Index: tools/telemetry/telemetry/core/platform/win_platform_backend.py
diff --git a/tools/telemetry/telemetry/core/platform/win_platform_backend.py b/tools/telemetry/telemetry/core/platform/win_platform_backend.py
index 6e91b3a91beae84d4966a64eca143a4c9afeb465..c81e3e756a51550f6416dcbd0a5bf36f58bcd1d6 100644
--- a/tools/telemetry/telemetry/core/platform/win_platform_backend.py
+++ b/tools/telemetry/telemetry/core/platform/win_platform_backend.py
@@ -2,19 +2,28 @@
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
+import atexit
 import collections
 import contextlib
 import ctypes
+import os
 import platform
 import re
+import socket
+import struct
 import subprocess
+import sys
 import time
+import zipfile
 
 from telemetry import decorators
 from telemetry.core import exceptions
+from telemetry.core import util
 from telemetry.core.platform import desktop_platform_backend
 from telemetry.core.platform import platform_backend
 from telemetry.core.platform.power_monitor import msr_power_monitor
+from telemetry.util import cloud_storage
+from telemetry.util import path
 
 try:
   import pywintypes  # pylint: disable=F0401
@@ -34,20 +43,64 @@ except ImportError:
   win32security = None
 
 
+def _InstallWinRing0():
+  """WinRing0 is used for reading MSRs."""
+  executable_dir = os.path.dirname(sys.executable)
 
-def IsCurrentProcessElevated():
-  handle = win32process.GetCurrentProcess()
-  with contextlib.closing(
-      win32security.OpenProcessToken(handle, win32con.TOKEN_QUERY)) as token:
-    return bool(
-        win32security.GetTokenInformation(token, win32security.TokenElevation))
+  python_is_64_bit = sys.maxsize > 2 ** 32
+  dll_file_name = 'WinRing0x64.dll' if python_is_64_bit else 'WinRing0.dll'
+  dll_path = os.path.join(executable_dir, dll_file_name)
+
+  os_is_64_bit = 'PROGRAMFILES(X86)' in os.environ
+  driver_file_name = 'WinRing0x64.sys' if os_is_64_bit else 'WinRing0.sys'
+  driver_path = os.path.join(executable_dir, driver_file_name)
+
+  # Check for WinRing0 and download if needed.
+  if not (os.path.exists(dll_path) and os.path.exists(driver_path)):
+    win_binary_dir = os.path.join(path.GetTelemetryDir(), 'bin', 'win')
+    zip_path = os.path.join(win_binary_dir, 'winring0.zip')
+    cloud_storage.GetIfChanged(zip_path, bucket=cloud_storage.PUBLIC_BUCKET)
+    try:
+      with zipfile.ZipFile(zip_path, 'r') as zip_file:
+        # Install DLL.
+        if not os.path.exists(dll_path):
+          zip_file.extract(dll_file_name, executable_dir)
+        # Install kernel driver.
+        if not os.path.exists(driver_path):
+          zip_file.extract(driver_file_name, executable_dir)
+    finally:
+      os.remove(zip_path)
+
+
+def TerminateProcess(process_handle):
+  if not process_handle:
+    return
+  if win32process.GetExitCodeProcess(process_handle) == win32con.STILL_ACTIVE:
+    win32process.TerminateProcess(process_handle, 0)
+  process_handle.close()
 
 
 class WinPlatformBackend(desktop_platform_backend.DesktopPlatformBackend):
   def __init__(self):
     super(WinPlatformBackend, self).__init__()
+    self._msr_server_handle = None
+    self._msr_server_port = None
     self._power_monitor = msr_power_monitor.MsrPowerMonitor(self)
 
+  def __del__(self):
+    self.close()
+
+  def close(self):
+    self.CloseMsrServer()
+
+  def CloseMsrServer(self):
+    if not self._msr_server_handle:
+      return
+
+    TerminateProcess(self._msr_server_handle)
+    self._msr_server_handle = None
+    self._msr_server_port = None
+
   # pylint: disable=W0613
   def StartRawDisplayFrameRateMeasurement(self):
     raise NotImplementedError()
@@ -224,12 +277,23 @@ class WinPlatformBackend(desktop_platform_backend.DesktopPlatformBackend):
         ctypes.byref(performance_info), performance_info.size)
     return performance_info
 
+  def IsCurrentProcessElevated(self):
+    if self.GetOSVersionName() < platform_backend.VISTA:
+      # TOKEN_QUERY is not defined before Vista. All processes are elevated.
+      return True
+
+    handle = win32process.GetCurrentProcess()
+    with contextlib.closing(
+        win32security.OpenProcessToken(handle, win32con.TOKEN_QUERY)) as token:
+      return bool(win32security.GetTokenInformation(
+          token, win32security.TokenElevation))
+
   def LaunchApplication(
       self, application, parameters=None, elevate_privilege=False):
     """Launch an application. Returns a PyHANDLE object."""
 
     parameters = ' '.join(parameters) if parameters else ''
-    if elevate_privilege and not IsCurrentProcessElevated():
+    if elevate_privilege and not self.IsCurrentProcessElevated():
       # Use ShellExecuteEx() instead of subprocess.Popen()/CreateProcess() to
       # elevate privileges. A new console will be created if the new process has
       # different permissions than this process.
@@ -259,3 +323,38 @@ class WinPlatformBackend(desktop_platform_backend.DesktopPlatformBackend):
 
   def StopMonitoringPower(self):
     return self._power_monitor.StopMonitoringPower()
+
+  def _StartMsrServerIfNeeded(self):
+    if self._msr_server_handle:
+      return
+
+    _InstallWinRing0()
+    self._msr_server_port = util.GetUnreservedAvailableLocalPort()
+    # It might be flaky to get a port number without reserving it atomically,
+    # but if the server process chooses a port, we have no way of getting it.
+    # The stdout of the elevated process isn't accessible.
+    parameters = (
+        os.path.join(os.path.dirname(__file__), 'msr_server_win.py'),
+        str(self._msr_server_port),
+    )
+    self._msr_server_handle = self.LaunchApplication(
+        sys.executable, parameters, elevate_privilege=True)
+    # Wait for server to start.
+    try:
+      socket.create_connection(('127.0.0.1', self._msr_server_port), 5).close()
+    except socket.error:
+      self.CloseMsrServer()
+    atexit.register(TerminateProcess, self._msr_server_handle)
+
+  def ReadMsr(self, msr_number):
+    self._StartMsrServerIfNeeded()
+    if not self._msr_server_handle:
+      raise OSError('Unable to start MSR server.')
+
+    sock = socket.create_connection(('127.0.0.1', self._msr_server_port), 0.1)
+    try:
+      sock.sendall(struct.pack('I', msr_number))
+      response = sock.recv(8)
+    finally:
+      sock.close()
+    return struct.unpack('Q', response)[0]