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 devil import base_error
13from devil import devil_env
14from devil.android.constants import file_system
15from devil.android.valgrind_tools import base_tool
16from devil.utils import cmd_helper
17
18
19def _GetProcessStartTime(pid):
20  return psutil.Process(pid).create_time
21
22
23class _FileLock(object):
24  """With statement-aware implementation of a file lock.
25
26  File locks are needed for cross-process synchronization when the
27  multiprocessing Python module is used.
28  """
29
30  def __init__(self, path):
31    self._fd = -1
32    self._path = path
33
34  def __enter__(self):
35    self._fd = os.open(self._path, os.O_RDONLY | os.O_CREAT)
36    if self._fd < 0:
37      raise Exception('Could not open file %s for reading' % self._path)
38    fcntl.flock(self._fd, fcntl.LOCK_EX)
39
40  def __exit__(self, _exception_type, _exception_value, traceback):
41    fcntl.flock(self._fd, fcntl.LOCK_UN)
42    os.close(self._fd)
43
44
45class HostForwarderError(base_error.BaseError):
46  """Exception for failures involving host_forwarder."""
47
48  def __init__(self, message):
49    super(HostForwarderError, self).__init__(message)
50
51
52class Forwarder(object):
53  """Thread-safe class to manage port forwards from the device to the host."""
54
55  _DEVICE_FORWARDER_FOLDER = (file_system.TEST_EXECUTABLE_DIR +
56                              '/forwarder/')
57  _DEVICE_FORWARDER_PATH = (file_system.TEST_EXECUTABLE_DIR +
58                            '/forwarder/device_forwarder')
59  _LOCK_PATH = '/tmp/chrome.forwarder.lock'
60  # Defined in host_forwarder_main.cc
61  _HOST_FORWARDER_LOG = '/tmp/host_forwarder_log'
62
63  _instance = None
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    if not tool:
83      tool = base_tool.BaseTool()
84    with _FileLock(Forwarder._LOCK_PATH):
85      instance = Forwarder._GetInstanceLocked(tool)
86      instance._InitDeviceLocked(device, tool)
87
88      device_serial = str(device)
89      redirection_commands = [
90          ['--adb=' + devil_env.config.FetchPath('adb'),
91           '--serial-id=' + device_serial,
92           '--map', str(device_port), str(host_port)]
93          for device_port, host_port in port_pairs]
94      logging.info('Forwarding using commands: %s', redirection_commands)
95
96      for redirection_command in redirection_commands:
97        try:
98          (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
99              [instance._host_forwarder_path] + redirection_command)
100        except OSError as e:
101          if e.errno == 2:
102            raise HostForwarderError(
103                'Unable to start host forwarder. '
104                'Make sure you have built host_forwarder.')
105          else: raise
106        if exit_code != 0:
107          Forwarder._KillDeviceLocked(device, tool)
108          # Log alive forwarders
109          ps_out = device.RunShellCommand(['ps'])
110          logging.info('Currently running device_forwarders:')
111          for line in ps_out:
112            if 'device_forwarder' in line:
113              logging.info('    %s', line)
114          raise HostForwarderError(
115              '%s exited with %d:\n%s' % (instance._host_forwarder_path,
116                                          exit_code, '\n'.join(output)))
117        tokens = output.split(':')
118        if len(tokens) != 2:
119          raise HostForwarderError(
120              'Unexpected host forwarder output "%s", '
121              'expected "device_port:host_port"' % output)
122        device_port = int(tokens[0])
123        host_port = int(tokens[1])
124        serial_with_port = (device_serial, device_port)
125        instance._device_to_host_port_map[serial_with_port] = host_port
126        instance._host_to_device_port_map[host_port] = serial_with_port
127        logging.info('Forwarding device port: %d to host port: %d.',
128                     device_port, host_port)
129
130  @staticmethod
131  def UnmapDevicePort(device_port, device):
132    """Unmaps a previously forwarded device port.
133
134    Args:
135      device: A DeviceUtils instance.
136      device_port: A previously forwarded port (through Map()).
137    """
138    with _FileLock(Forwarder._LOCK_PATH):
139      Forwarder._UnmapDevicePortLocked(device_port, device)
140
141  @staticmethod
142  def UnmapAllDevicePorts(device):
143    """Unmaps all the previously forwarded ports for the provided device.
144
145    Args:
146      device: A DeviceUtils instance.
147      port_pairs: A list of tuples (device_port, host_port) to unmap.
148    """
149    with _FileLock(Forwarder._LOCK_PATH):
150      if not Forwarder._instance:
151        return
152      adb_serial = str(device)
153      if adb_serial not in Forwarder._instance._initialized_devices:
154        return
155      port_map = Forwarder._GetInstanceLocked(
156          None)._device_to_host_port_map
157      for (device_serial, device_port) in port_map.keys():
158        if adb_serial == device_serial:
159          Forwarder._UnmapDevicePortLocked(device_port, device)
160      # There are no more ports mapped, kill the device_forwarder.
161      tool = base_tool.BaseTool()
162      Forwarder._KillDeviceLocked(device, tool)
163
164  @staticmethod
165  def DevicePortForHostPort(host_port):
166    """Returns the device port that corresponds to a given host port."""
167    with _FileLock(Forwarder._LOCK_PATH):
168      _, device_port = Forwarder._GetInstanceLocked(
169          None)._host_to_device_port_map.get(host_port)
170      return device_port
171
172  @staticmethod
173  def RemoveHostLog():
174    if os.path.exists(Forwarder._HOST_FORWARDER_LOG):
175      os.unlink(Forwarder._HOST_FORWARDER_LOG)
176
177  @staticmethod
178  def GetHostLog():
179    if not os.path.exists(Forwarder._HOST_FORWARDER_LOG):
180      return ''
181    with file(Forwarder._HOST_FORWARDER_LOG, 'r') as f:
182      return f.read()
183
184  @staticmethod
185  def _GetInstanceLocked(tool):
186    """Returns the singleton instance.
187
188    Note that the global lock must be acquired before calling this method.
189
190    Args:
191      tool: Tool class to use to get wrapper, if necessary, for executing the
192            forwarder (see valgrind_tools.py).
193    """
194    if not Forwarder._instance:
195      Forwarder._instance = Forwarder(tool)
196    return Forwarder._instance
197
198  def __init__(self, tool):
199    """Constructs a new instance of Forwarder.
200
201    Note that Forwarder is a singleton therefore this constructor should be
202    called only once.
203
204    Args:
205      tool: Tool class to use to get wrapper, if necessary, for executing the
206            forwarder (see valgrind_tools.py).
207    """
208    assert not Forwarder._instance
209    self._tool = tool
210    self._initialized_devices = set()
211    self._device_to_host_port_map = dict()
212    self._host_to_device_port_map = dict()
213    self._host_forwarder_path = devil_env.config.FetchPath('forwarder_host')
214    assert os.path.exists(self._host_forwarder_path), 'Please build forwarder2'
215    self._InitHostLocked()
216
217  @staticmethod
218  def _UnmapDevicePortLocked(device_port, device):
219    """Internal method used by UnmapDevicePort().
220
221    Note that the global lock must be acquired before calling this method.
222    """
223    instance = Forwarder._GetInstanceLocked(None)
224    serial = str(device)
225    serial_with_port = (serial, device_port)
226    if not serial_with_port in instance._device_to_host_port_map:
227      logging.error('Trying to unmap non-forwarded port %d', device_port)
228      return
229    redirection_command = ['--adb=' + devil_env.config.FetchPath('adb'),
230                           '--serial-id=' + serial,
231                           '--unmap', str(device_port)]
232    logging.info('Undo forwarding using command: %s', redirection_command)
233    (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
234        [instance._host_forwarder_path] + redirection_command)
235    if exit_code != 0:
236      logging.error(
237          '%s exited with %d:\n%s',
238          instance._host_forwarder_path, exit_code, '\n'.join(output))
239    host_port = instance._device_to_host_port_map[serial_with_port]
240    del instance._device_to_host_port_map[serial_with_port]
241    del instance._host_to_device_port_map[host_port]
242
243  @staticmethod
244  def _GetPidForLock():
245    """Returns the PID used for host_forwarder initialization.
246
247    The PID of the "sharder" is used to handle multiprocessing. The "sharder"
248    is the initial process that forks that is the parent process.
249    """
250    return os.getpgrp()
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      pid_file.truncate()
274
275  def _InitDeviceLocked(self, device, tool):
276    """Initializes the device_forwarder daemon for a specific device (once).
277
278    Note that the global lock must be acquired before calling this method. This
279    method kills any existing device_forwarder daemon on the device that could
280    be stale, pushes the latest version of the daemon (to the device) and starts
281    it.
282
283    Args:
284      device: A DeviceUtils instance.
285      tool: Tool class to use to get wrapper, if necessary, for executing the
286            forwarder (see valgrind_tools.py).
287    """
288    device_serial = str(device)
289    if device_serial in self._initialized_devices:
290      return
291    Forwarder._KillDeviceLocked(device, tool)
292    forwarder_device_path_on_host = devil_env.config.FetchPath(
293        'forwarder_device', device=device)
294    forwarder_device_path_on_device = (
295        Forwarder._DEVICE_FORWARDER_FOLDER
296        if os.path.isdir(forwarder_device_path_on_host)
297        else Forwarder._DEVICE_FORWARDER_PATH)
298    device.PushChangedFiles([(
299        forwarder_device_path_on_host,
300        forwarder_device_path_on_device)])
301
302    cmd = '%s %s' % (tool.GetUtilWrapper(), Forwarder._DEVICE_FORWARDER_PATH)
303    device.RunShellCommand(
304        cmd, env={'LD_LIBRARY_PATH': Forwarder._DEVICE_FORWARDER_FOLDER},
305        check_return=True)
306    self._initialized_devices.add(device_serial)
307
308  def _KillHostLocked(self):
309    """Kills the forwarder process running on the host.
310
311    Note that the global lock must be acquired before calling this method.
312    """
313    logging.info('Killing host_forwarder.')
314    (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
315        [self._host_forwarder_path, '--kill-server'])
316    if exit_code != 0:
317      (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
318          ['pkill', '-9', 'host_forwarder'])
319      if exit_code != 0:
320        raise HostForwarderError(
321            '%s exited with %d:\n%s' % (self._host_forwarder_path, exit_code,
322                                        '\n'.join(output)))
323
324  @staticmethod
325  def _KillDeviceLocked(device, tool):
326    """Kills the forwarder process running on the device.
327
328    Note that the global lock must be acquired before calling this method.
329
330    Args:
331      device: Instance of DeviceUtils for talking to the device.
332      tool: Wrapper tool (e.g. valgrind) that can be used to execute the device
333            forwarder (see valgrind_tools.py).
334    """
335    logging.info('Killing device_forwarder.')
336    Forwarder._instance._initialized_devices.discard(str(device))
337    if not device.FileExists(Forwarder._DEVICE_FORWARDER_PATH):
338      return
339
340    cmd = '%s %s --kill-server' % (tool.GetUtilWrapper(),
341                                   Forwarder._DEVICE_FORWARDER_PATH)
342    device.RunShellCommand(
343        cmd, env={'LD_LIBRARY_PATH': Forwarder._DEVICE_FORWARDER_FOLDER},
344        check_return=True)
345