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 import device_errors
15from devil.android.constants import file_system
16from devil.android.valgrind_tools import base_tool
17from devil.utils import cmd_helper
18
19
20def _GetProcessStartTime(pid):
21  return psutil.Process(pid).create_time
22
23
24class _FileLock(object):
25  """With statement-aware implementation of a file lock.
26
27  File locks are needed for cross-process synchronization when the
28  multiprocessing Python module is used.
29  """
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 HostForwarderError(base_error.BaseError):
47  """Exception for failures involving host_forwarder."""
48
49  def __init__(self, message):
50    super(HostForwarderError, self).__init__(message)
51
52
53class Forwarder(object):
54  """Thread-safe class to manage port forwards from the device to the host."""
55
56  _DEVICE_FORWARDER_FOLDER = (file_system.TEST_EXECUTABLE_DIR +
57                              '/forwarder/')
58  _DEVICE_FORWARDER_PATH = (file_system.TEST_EXECUTABLE_DIR +
59                            '/forwarder/device_forwarder')
60  _LOCK_PATH = '/tmp/chrome.forwarder.lock'
61  # Defined in host_forwarder_main.cc
62  _HOST_FORWARDER_LOG = '/tmp/host_forwarder_log'
63
64  _instance = None
65
66  @staticmethod
67  def Map(port_pairs, device, tool=None):
68    """Runs the forwarder.
69
70    Args:
71      port_pairs: A list of tuples (device_port, host_port) to forward. Note
72                 that you can specify 0 as a device_port, in which case a
73                 port will by dynamically assigned on the device. You can
74                 get the number of the assigned port using the
75                 DevicePortForHostPort method.
76      device: A DeviceUtils instance.
77      tool: Tool class to use to get wrapper, if necessary, for executing the
78            forwarder (see valgrind_tools.py).
79
80    Raises:
81      Exception on failure to forward the port.
82    """
83    if not tool:
84      tool = base_tool.BaseTool()
85    with _FileLock(Forwarder._LOCK_PATH):
86      instance = Forwarder._GetInstanceLocked(tool)
87      instance._InitDeviceLocked(device, tool)
88
89      device_serial = str(device)
90      redirection_commands = [
91          ['--adb=' + devil_env.config.FetchPath('adb'),
92           '--serial-id=' + device_serial,
93           '--map', str(device_port), str(host_port)]
94          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 HostForwarderError(
104                'Unable to start host forwarder. '
105                'Make sure you have built host_forwarder.')
106          else: raise
107        if exit_code != 0:
108          Forwarder._KillDeviceLocked(device, tool)
109          # Log alive forwarders
110          ps_out = device.RunShellCommand(['ps'])
111          logging.info('Currently running device_forwarders:')
112          for line in ps_out:
113            if 'device_forwarder' in line:
114              logging.info('    %s', line)
115          raise HostForwarderError(
116              '%s exited with %d:\n%s' % (instance._host_forwarder_path,
117                                          exit_code, '\n'.join(output)))
118        tokens = output.split(':')
119        if len(tokens) != 2:
120          raise HostForwarderError(
121              'Unexpected host forwarder output "%s", '
122              'expected "device_port:host_port"' % output)
123        device_port = int(tokens[0])
124        host_port = int(tokens[1])
125        serial_with_port = (device_serial, device_port)
126        instance._device_to_host_port_map[serial_with_port] = host_port
127        instance._host_to_device_port_map[host_port] = serial_with_port
128        logging.info('Forwarding device port: %d to host port: %d.',
129                     device_port, host_port)
130
131  @staticmethod
132  def UnmapDevicePort(device_port, device):
133    """Unmaps a previously forwarded device port.
134
135    Args:
136      device: A DeviceUtils instance.
137      device_port: A previously forwarded port (through Map()).
138    """
139    with _FileLock(Forwarder._LOCK_PATH):
140      Forwarder._UnmapDevicePortLocked(device_port, device)
141
142  @staticmethod
143  def UnmapAllDevicePorts(device):
144    """Unmaps all the previously forwarded ports for the provided device.
145
146    Args:
147      device: A DeviceUtils instance.
148      port_pairs: A list of tuples (device_port, host_port) to unmap.
149    """
150    with _FileLock(Forwarder._LOCK_PATH):
151      if not Forwarder._instance:
152        return
153      adb_serial = str(device)
154      if adb_serial not in Forwarder._instance._initialized_devices:
155        return
156      port_map = Forwarder._GetInstanceLocked(
157          None)._device_to_host_port_map
158      for (device_serial, device_port) in port_map.keys():
159        if adb_serial == device_serial:
160          Forwarder._UnmapDevicePortLocked(device_port, device)
161      # There are no more ports mapped, kill the device_forwarder.
162      tool = base_tool.BaseTool()
163      Forwarder._KillDeviceLocked(device, tool)
164
165  @staticmethod
166  def DevicePortForHostPort(host_port):
167    """Returns the device port that corresponds to a given host port."""
168    with _FileLock(Forwarder._LOCK_PATH):
169      _, device_port = Forwarder._GetInstanceLocked(
170          None)._host_to_device_port_map.get(host_port)
171      return device_port
172
173  @staticmethod
174  def RemoveHostLog():
175    if os.path.exists(Forwarder._HOST_FORWARDER_LOG):
176      os.unlink(Forwarder._HOST_FORWARDER_LOG)
177
178  @staticmethod
179  def GetHostLog():
180    if not os.path.exists(Forwarder._HOST_FORWARDER_LOG):
181      return ''
182    with file(Forwarder._HOST_FORWARDER_LOG, 'r') as f:
183      return f.read()
184
185  @staticmethod
186  def _GetInstanceLocked(tool):
187    """Returns the singleton instance.
188
189    Note that the global lock must be acquired before calling this method.
190
191    Args:
192      tool: Tool class to use to get wrapper, if necessary, for executing the
193            forwarder (see valgrind_tools.py).
194    """
195    if not Forwarder._instance:
196      Forwarder._instance = Forwarder(tool)
197    return Forwarder._instance
198
199  def __init__(self, tool):
200    """Constructs a new instance of Forwarder.
201
202    Note that Forwarder is a singleton therefore this constructor should be
203    called only once.
204
205    Args:
206      tool: Tool class to use to get wrapper, if necessary, for executing the
207            forwarder (see valgrind_tools.py).
208    """
209    assert not Forwarder._instance
210    self._tool = tool
211    self._initialized_devices = set()
212    self._device_to_host_port_map = dict()
213    self._host_to_device_port_map = dict()
214    self._host_forwarder_path = devil_env.config.FetchPath('forwarder_host')
215    assert os.path.exists(self._host_forwarder_path), 'Please build forwarder2'
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 = ['--adb=' + devil_env.config.FetchPath('adb'),
231                           '--serial-id=' + serial,
232                           '--unmap', str(device_port)]
233    logging.info('Undo forwarding using command: %s', redirection_command)
234    (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
235        [instance._host_forwarder_path] + redirection_command)
236    if exit_code != 0:
237      logging.error(
238          '%s exited with %d:\n%s',
239          instance._host_forwarder_path, exit_code, '\n'.join(output))
240    host_port = instance._device_to_host_port_map[serial_with_port]
241    del instance._device_to_host_port_map[serial_with_port]
242    del instance._host_to_device_port_map[host_port]
243
244  @staticmethod
245  def _GetPidForLock():
246    """Returns the PID used for host_forwarder initialization.
247
248    The PID of the "sharder" is used to handle multiprocessing. The "sharder"
249    is the initial process that forks that is the parent process.
250    """
251    return os.getpgrp()
252
253  def _InitHostLocked(self):
254    """Initializes the host forwarder daemon.
255
256    Note that the global lock must be acquired before calling this method. This
257    method kills any existing host_forwarder process that could be stale.
258    """
259    # See if the host_forwarder daemon was already initialized by a concurrent
260    # process or thread (in case multi-process sharding is not used).
261    pid_for_lock = Forwarder._GetPidForLock()
262    fd = os.open(Forwarder._LOCK_PATH, os.O_RDWR | os.O_CREAT)
263    with os.fdopen(fd, 'r+') as pid_file:
264      pid_with_start_time = pid_file.readline()
265      if pid_with_start_time:
266        (pid, process_start_time) = pid_with_start_time.split(':')
267        if pid == str(pid_for_lock):
268          if process_start_time == str(_GetProcessStartTime(pid_for_lock)):
269            return
270      self._KillHostLocked()
271      pid_file.seek(0)
272      pid_file.write(
273          '%s:%s' % (pid_for_lock, str(_GetProcessStartTime(pid_for_lock))))
274      pid_file.truncate()
275
276  def _InitDeviceLocked(self, device, tool):
277    """Initializes the device_forwarder daemon for a specific device (once).
278
279    Note that the global lock must be acquired before calling this method. This
280    method kills any existing device_forwarder daemon on the device that could
281    be stale, pushes the latest version of the daemon (to the device) and starts
282    it.
283
284    Args:
285      device: A DeviceUtils instance.
286      tool: Tool class to use to get wrapper, if necessary, for executing the
287            forwarder (see valgrind_tools.py).
288    """
289    device_serial = str(device)
290    if device_serial in self._initialized_devices:
291      return
292    try:
293      Forwarder._KillDeviceLocked(device, tool)
294    except device_errors.CommandFailedError:
295      logging.warning('Failed to kill device forwarder. Rebooting.')
296      device.Reboot()
297    forwarder_device_path_on_host = devil_env.config.FetchPath(
298        'forwarder_device', device=device)
299    forwarder_device_path_on_device = (
300        Forwarder._DEVICE_FORWARDER_FOLDER
301        if os.path.isdir(forwarder_device_path_on_host)
302        else Forwarder._DEVICE_FORWARDER_PATH)
303    device.PushChangedFiles([(
304        forwarder_device_path_on_host,
305        forwarder_device_path_on_device)])
306
307    cmd = '%s %s' % (tool.GetUtilWrapper(), Forwarder._DEVICE_FORWARDER_PATH)
308    device.RunShellCommand(
309        cmd, env={'LD_LIBRARY_PATH': Forwarder._DEVICE_FORWARDER_FOLDER},
310        check_return=True)
311    self._initialized_devices.add(device_serial)
312
313  def _KillHostLocked(self):
314    """Kills the forwarder process running on the host.
315
316    Note that the global lock must be acquired before calling this method.
317    """
318    logging.info('Killing host_forwarder.')
319    (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
320        [self._host_forwarder_path, '--kill-server'])
321    if exit_code != 0:
322      (exit_code, output) = cmd_helper.GetCmdStatusAndOutput(
323          ['pkill', '-9', 'host_forwarder'])
324      if exit_code != 0:
325        raise HostForwarderError(
326            '%s exited with %d:\n%s' % (self._host_forwarder_path, exit_code,
327                                        '\n'.join(output)))
328
329  @staticmethod
330  def _KillDeviceLocked(device, tool):
331    """Kills the forwarder process running on the device.
332
333    Note that the global lock must be acquired before calling this method.
334
335    Args:
336      device: Instance of DeviceUtils for talking to the device.
337      tool: Wrapper tool (e.g. valgrind) that can be used to execute the device
338            forwarder (see valgrind_tools.py).
339    """
340    logging.info('Killing device_forwarder.')
341    Forwarder._instance._initialized_devices.discard(str(device))
342    if not device.FileExists(Forwarder._DEVICE_FORWARDER_PATH):
343      return
344
345    cmd = '%s %s --kill-server' % (tool.GetUtilWrapper(),
346                                   Forwarder._DEVICE_FORWARDER_PATH)
347    device.RunShellCommand(
348        cmd, env={'LD_LIBRARY_PATH': Forwarder._DEVICE_FORWARDER_FOLDER},
349        check_return=True)
350