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# pylint: disable=W0212
6
7import fcntl
8import logging
9import os
10import psutil
11
12from pylib import cmd_helper
13from pylib import constants
14from pylib import valgrind_tools
15
16# TODO(jbudorick) Remove once telemetry gets switched over.
17import pylib.android_commands
18import pylib.device.device_utils
19
20
21def _GetProcessStartTime(pid):
22  return psutil.Process(pid).create_time
23
24
25class _FileLock(object):
26  """With statement-aware implementation of a file lock.
27
28  File locks are needed for cross-process synchronization when the
29  multiprocessing Python module is used.
30  """
31  def __init__(self, path):
32    self._fd = -1
33    self._path = path
34
35  def __enter__(self):
36    self._fd = os.open(self._path, os.O_RDONLY | os.O_CREAT)
37    if self._fd < 0:
38      raise Exception('Could not open file %s for reading' % self._path)
39    fcntl.flock(self._fd, fcntl.LOCK_EX)
40
41  def __exit__(self, _exception_type, _exception_value, traceback):
42    fcntl.flock(self._fd, fcntl.LOCK_UN)
43    os.close(self._fd)
44
45
46class Forwarder(object):
47  """Thread-safe class to manage port forwards from the device to the host."""
48
49  _DEVICE_FORWARDER_FOLDER = (constants.TEST_EXECUTABLE_DIR +
50                              '/forwarder/')
51  _DEVICE_FORWARDER_PATH = (constants.TEST_EXECUTABLE_DIR +
52                            '/forwarder/device_forwarder')
53  _LOCK_PATH = '/tmp/chrome.forwarder.lock'
54  _MULTIPROCESSING_ENV_VAR = 'CHROME_FORWARDER_USE_MULTIPROCESSING'
55  # Defined in host_forwarder_main.cc
56  _HOST_FORWARDER_LOG = '/tmp/host_forwarder_log'
57
58  _instance = None
59
60  @staticmethod
61  def UseMultiprocessing():
62    """Tells the forwarder that multiprocessing is used."""
63    os.environ[Forwarder._MULTIPROCESSING_ENV_VAR] = '1'
64
65  @staticmethod
66  def Map(port_pairs, device, tool=None):
67    """Runs the forwarder.
68
69    Args:
70      port_pairs: A list of tuples (device_port, host_port) to forward. Note
71                 that you can specify 0 as a device_port, in which case a
72                 port will by dynamically assigned on the device. You can
73                 get the number of the assigned port using the
74                 DevicePortForHostPort method.
75      device: A DeviceUtils instance.
76      tool: Tool class to use to get wrapper, if necessary, for executing the
77            forwarder (see valgrind_tools.py).
78
79    Raises:
80      Exception on failure to forward the port.
81    """
82    # TODO(jbudorick) Remove once telemetry gets switched over.
83    if isinstance(device, pylib.android_commands.AndroidCommands):
84      device = pylib.device.device_utils.DeviceUtils(device)
85    if not tool:
86      tool = valgrind_tools.CreateTool(None, device)
87    with _FileLock(Forwarder._LOCK_PATH):
88      instance = Forwarder._GetInstanceLocked(tool)
89      instance._InitDeviceLocked(device, tool)
90
91      device_serial = str(device)
92      redirection_commands = [
93          ['--serial-id=' + device_serial, '--map', str(device_port),
94           str(host_port)] for device_port, host_port in port_pairs]
95      logging.info('Forwarding using commands: %s', redirection_commands)
96
97      for redirection_command in redirection_commands:
98        try:
99          (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
100              [instance._host_forwarder_path] + redirection_command)
101        except OSError as e:
102          if e.errno == 2:
103            raise Exception('Unable to start host forwarder. Make sure you have'
104                            ' built host_forwarder.')
105          else: raise
106        if exit_code != 0:
107          Forwarder._KillDeviceLocked(device, tool)
108          raise Exception('%s exited with %d:\n%s' % (
109              instance._host_forwarder_path, exit_code, '\n'.join(output)))
110        tokens = output.split(':')
111        if len(tokens) != 2:
112          raise Exception('Unexpected host forwarder output "%s", '
113                          'expected "device_port:host_port"' % output)
114        device_port = int(tokens[0])
115        host_port = int(tokens[1])
116        serial_with_port = (device_serial, device_port)
117        instance._device_to_host_port_map[serial_with_port] = host_port
118        instance._host_to_device_port_map[host_port] = serial_with_port
119        logging.info('Forwarding device port: %d to host port: %d.',
120                     device_port, host_port)
121
122  @staticmethod
123  def UnmapDevicePort(device_port, device):
124    """Unmaps a previously forwarded device port.
125
126    Args:
127      device: A DeviceUtils instance.
128      device_port: A previously forwarded port (through Map()).
129    """
130    # TODO(jbudorick) Remove once telemetry gets switched over.
131    if isinstance(device, pylib.android_commands.AndroidCommands):
132      device = pylib.device.device_utils.DeviceUtils(device)
133    with _FileLock(Forwarder._LOCK_PATH):
134      Forwarder._UnmapDevicePortLocked(device_port, device)
135
136  @staticmethod
137  def UnmapAllDevicePorts(device):
138    """Unmaps all the previously forwarded ports for the provided device.
139
140    Args:
141      device: A DeviceUtils instance.
142      port_pairs: A list of tuples (device_port, host_port) to unmap.
143    """
144    # TODO(jbudorick) Remove once telemetry gets switched over.
145    if isinstance(device, pylib.android_commands.AndroidCommands):
146      device = pylib.device.device_utils.DeviceUtils(device)
147    with _FileLock(Forwarder._LOCK_PATH):
148      if not Forwarder._instance:
149        return
150      adb_serial = str(device)
151      if adb_serial not in Forwarder._instance._initialized_devices:
152        return
153      port_map = Forwarder._GetInstanceLocked(
154          None)._device_to_host_port_map
155      for (device_serial, device_port) in port_map.keys():
156        if adb_serial == device_serial:
157          Forwarder._UnmapDevicePortLocked(device_port, device)
158      # There are no more ports mapped, kill the device_forwarder.
159      tool = valgrind_tools.CreateTool(None, device)
160      Forwarder._KillDeviceLocked(device, tool)
161
162  @staticmethod
163  def DevicePortForHostPort(host_port):
164    """Returns the device port that corresponds to a given host port."""
165    with _FileLock(Forwarder._LOCK_PATH):
166      (_device_serial, device_port) = Forwarder._GetInstanceLocked(
167          None)._host_to_device_port_map.get(host_port)
168      return device_port
169
170  @staticmethod
171  def RemoveHostLog():
172    if os.path.exists(Forwarder._HOST_FORWARDER_LOG):
173      os.unlink(Forwarder._HOST_FORWARDER_LOG)
174
175  @staticmethod
176  def GetHostLog():
177    if not os.path.exists(Forwarder._HOST_FORWARDER_LOG):
178      return ''
179    with file(Forwarder._HOST_FORWARDER_LOG, 'r') as f:
180      return f.read()
181
182  @staticmethod
183  def _GetInstanceLocked(tool):
184    """Returns the singleton instance.
185
186    Note that the global lock must be acquired before calling this method.
187
188    Args:
189      tool: Tool class to use to get wrapper, if necessary, for executing the
190            forwarder (see valgrind_tools.py).
191    """
192    if not Forwarder._instance:
193      Forwarder._instance = Forwarder(tool)
194    return Forwarder._instance
195
196  def __init__(self, tool):
197    """Constructs a new instance of Forwarder.
198
199    Note that Forwarder is a singleton therefore this constructor should be
200    called only once.
201
202    Args:
203      tool: Tool class to use to get wrapper, if necessary, for executing the
204            forwarder (see valgrind_tools.py).
205    """
206    assert not Forwarder._instance
207    self._tool = tool
208    self._initialized_devices = set()
209    self._device_to_host_port_map = dict()
210    self._host_to_device_port_map = dict()
211    self._host_forwarder_path = os.path.join(
212        constants.GetOutDirectory(), 'host_forwarder')
213    assert os.path.exists(self._host_forwarder_path), 'Please build forwarder2'
214    self._device_forwarder_path_on_host = os.path.join(
215        constants.GetOutDirectory(), 'forwarder_dist')
216    self._InitHostLocked()
217
218  @staticmethod
219  def _UnmapDevicePortLocked(device_port, device):
220    """Internal method used by UnmapDevicePort().
221
222    Note that the global lock must be acquired before calling this method.
223    """
224    instance = Forwarder._GetInstanceLocked(None)
225    serial = str(device)
226    serial_with_port = (serial, device_port)
227    if not serial_with_port in instance._device_to_host_port_map:
228      logging.error('Trying to unmap non-forwarded port %d' % device_port)
229      return
230    redirection_command = ['--serial-id=' + serial, '--unmap', str(device_port)]
231    (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
232        [instance._host_forwarder_path] + redirection_command)
233    if exit_code != 0:
234      logging.error('%s exited with %d:\n%s' % (
235          instance._host_forwarder_path, exit_code, '\n'.join(output)))
236    host_port = instance._device_to_host_port_map[serial_with_port]
237    del instance._device_to_host_port_map[serial_with_port]
238    del instance._host_to_device_port_map[host_port]
239
240  @staticmethod
241  def _GetPidForLock():
242    """Returns the PID used for host_forwarder initialization.
243
244    In case multi-process sharding is used, the PID of the "sharder" is used.
245    The "sharder" is the initial process that forks that is the parent process.
246    By default, multi-processing is not used. In that case the PID of the
247    current process is returned.
248    """
249    use_multiprocessing = Forwarder._MULTIPROCESSING_ENV_VAR in os.environ
250    return os.getpgrp() if use_multiprocessing else os.getpid()
251
252  def _InitHostLocked(self):
253    """Initializes the host forwarder daemon.
254
255    Note that the global lock must be acquired before calling this method. This
256    method kills any existing host_forwarder process that could be stale.
257    """
258    # See if the host_forwarder daemon was already initialized by a concurrent
259    # process or thread (in case multi-process sharding is not used).
260    pid_for_lock = Forwarder._GetPidForLock()
261    fd = os.open(Forwarder._LOCK_PATH, os.O_RDWR | os.O_CREAT)
262    with os.fdopen(fd, 'r+') as pid_file:
263      pid_with_start_time = pid_file.readline()
264      if pid_with_start_time:
265        (pid, process_start_time) = pid_with_start_time.split(':')
266        if pid == str(pid_for_lock):
267          if process_start_time == str(_GetProcessStartTime(pid_for_lock)):
268            return
269      self._KillHostLocked()
270      pid_file.seek(0)
271      pid_file.write(
272          '%s:%s' % (pid_for_lock, str(_GetProcessStartTime(pid_for_lock))))
273
274  def _InitDeviceLocked(self, device, tool):
275    """Initializes the device_forwarder daemon for a specific device (once).
276
277    Note that the global lock must be acquired before calling this method. This
278    method kills any existing device_forwarder daemon on the device that could
279    be stale, pushes the latest version of the daemon (to the device) and starts
280    it.
281
282    Args:
283      device: A DeviceUtils instance.
284      tool: Tool class to use to get wrapper, if necessary, for executing the
285            forwarder (see valgrind_tools.py).
286    """
287    device_serial = str(device)
288    if device_serial in self._initialized_devices:
289      return
290    Forwarder._KillDeviceLocked(device, tool)
291    device.PushChangedFiles(
292        self._device_forwarder_path_on_host,
293        Forwarder._DEVICE_FORWARDER_FOLDER)
294    cmd = '%s %s' % (tool.GetUtilWrapper(), Forwarder._DEVICE_FORWARDER_PATH)
295    (exit_code, output) = device.old_interface.GetAndroidToolStatusAndOutput(
296        cmd, lib_path=Forwarder._DEVICE_FORWARDER_FOLDER)
297    if exit_code != 0:
298      raise Exception(
299          'Failed to start device forwarder:\n%s' % '\n'.join(output))
300    self._initialized_devices.add(device_serial)
301
302  def _KillHostLocked(self):
303    """Kills the forwarder process running on the host.
304
305    Note that the global lock must be acquired before calling this method.
306    """
307    logging.info('Killing host_forwarder.')
308    (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
309        [self._host_forwarder_path, '--kill-server'])
310    if exit_code != 0:
311      (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
312          ['pkill', '-9', 'host_forwarder'])
313      if exit_code != 0:
314        raise Exception('%s exited with %d:\n%s' % (
315              self._host_forwarder_path, exit_code, '\n'.join(output)))
316
317  @staticmethod
318  def _KillDeviceLocked(device, tool):
319    """Kills the forwarder process running on the device.
320
321    Note that the global lock must be acquired before calling this method.
322
323    Args:
324      device: Instance of DeviceUtils for talking to the device.
325      tool: Wrapper tool (e.g. valgrind) that can be used to execute the device
326            forwarder (see valgrind_tools.py).
327    """
328    logging.info('Killing device_forwarder.')
329    Forwarder._instance._initialized_devices.discard(str(device))
330    if not device.FileExists(Forwarder._DEVICE_FORWARDER_PATH):
331      return
332
333    cmd = '%s %s --kill-server' % (tool.GetUtilWrapper(),
334                                   Forwarder._DEVICE_FORWARDER_PATH)
335    device.old_interface.GetAndroidToolStatusAndOutput(
336        cmd, lib_path=Forwarder._DEVICE_FORWARDER_FOLDER)
337