• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python
2#
3# Copyright 2015 The Chromium Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6#
7# This is a script developers can use to set-up their workstation to let
8# Telemetry read the CPU's Model Specific Registers in order to get power
9# measurements. It can check if reading from MSRs is possible as any user, but
10# must run as root to make changes. Not all changes are sticky, so one has to
11# re-run this script after each reboot.
12#
13# This script is currently Debian/Ubuntu specific.
14
15import os
16import subprocess
17import sys
18
19MSR_DEV_FILE_PATH = '/dev/cpu/0/msr'
20RDMSR_PATH = '/usr/sbin/rdmsr'
21
22def _Usage(prog_name):
23  """Print a help message."""
24  print 'Run "%s" as a regular user to check if reading from the MSR ' \
25      'is possible.' % prog_name
26  print 'Run "%s enable" as root to automatically set up reading from ' \
27      'the MSR.' % prog_name
28
29
30def _CheckMsrKernelModule():
31  """Return whether the 'msr' kernel module is loaded."""
32  proc = subprocess.Popen('/sbin/lsmod', stdout=subprocess.PIPE)
33  stdout = proc.communicate()[0]
34  ret = proc.wait()
35  if ret != 0:
36    raise OSError('lsmod failed')
37
38  if not any([line.startswith('msr ') for line in stdout.splitlines()]):
39    print 'Error: MSR module not loaded.'
40    return False
41
42  return True
43
44
45def _CheckMsrDevNodes():
46  """Check whether the MSR /dev files have the right permissions."""
47  if not os.path.exists(MSR_DEV_FILE_PATH):
48    print 'Error: %s does not exist.' % MSR_DEV_FILE_PATH
49    return False
50
51  if not os.access(MSR_DEV_FILE_PATH, os.R_OK):
52    print 'Error: Cannot read from %s' % MSR_DEV_FILE_PATH
53    return False
54
55  return True
56
57
58def _CheckRdmsr():
59  """Check and make sure /usr/sbin/rdmsr is set up correctly."""
60  if not os.access(RDMSR_PATH, os.X_OK):
61    print 'Error: %s missing or not executable.' % RDMSR_PATH
62    return False
63
64  proc = subprocess.Popen(['/sbin/getcap', RDMSR_PATH], stdout=subprocess.PIPE)
65  stdout = proc.communicate()[0]
66  ret = proc.wait()
67  if ret != 0:
68    raise OSError('getcap failed')
69
70  if not 'cap_sys_rawio+ep' in stdout:
71    print 'Error: /usr/sbin/rdmsr needs RAWIO capability.'
72    return False
73
74  return True
75
76
77def _RunAllChecks():
78  """Check to make sure it is possible to read from the MSRs."""
79  if os.geteuid() == 0:
80    print 'WARNING: Running as root, msr permission check likely inaccurate.'
81
82  has_dev_node = _CheckMsrDevNodes() if _CheckMsrKernelModule() else False
83  has_rdmsr = _CheckRdmsr()
84  return has_dev_node and has_rdmsr
85
86
87def _EnableMsr(prog_name):
88  """Do all the setup needed to pass _RunAllChecks().
89
90  Needs to run as root."""
91  if os.geteuid() != 0:
92    print 'Error: Must run "%s enable" as root.' % prog_name
93    return False
94
95  print 'Loading msr kernel module.'
96  ret = subprocess.call(['/sbin/modprobe', 'msr'])
97  if ret != 0:
98    print 'Error: Cannot load msr module.'
99    return False
100
101  print 'Running chmod on %s.' % MSR_DEV_FILE_PATH
102  ret = subprocess.call(['/bin/chmod', 'a+r', MSR_DEV_FILE_PATH])
103  if ret != 0:
104    print 'Error: Cannot chmod %s.' % MSR_DEV_FILE_PATH
105    return False
106
107  if not os.access(RDMSR_PATH, os.F_OK):
108    print 'Need to install the msr-tools package.'
109    ret = subprocess.call(['/usr/bin/apt-get', 'install', '-y', 'msr-tools'])
110    if ret != 0:
111      print 'Error: Did not successfully install msr-tools.'
112      return False
113
114  print 'Running setcap on %s.' % RDMSR_PATH
115  ret = subprocess.call(['/sbin/setcap', 'cap_sys_rawio+ep', RDMSR_PATH])
116  if ret != 0:
117    print 'Error: Cannot give /usr/sbin/rdmsr RAWIO capability.'
118    return False
119
120  return True
121
122
123def main(prog_name, argv):
124  if len(argv) == 0:
125    if _RunAllChecks():
126      print 'Check succeeded'
127      return 0
128
129    print 'Check failed, try running "%s enable" as root to fix.' % prog_name
130    return 1
131
132  if len(argv) == 1:
133    if argv[0] == 'enable':
134      return 0 if _EnableMsr(prog_name) else 1
135
136    print 'Error: Unknown sub-command %s' % argv[0]
137    _Usage(prog_name)
138    return 1
139
140  print 'Error: Bad number of arguments'
141  _Usage(prog_name)
142  return 1
143
144
145if '__main__' == __name__:
146  sys.exit(main(os.path.basename(sys.argv[0]), sys.argv[1:]))
147