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
5import logging
6import os
7import re
8import sys
9import time
10
11import android_commands
12import cmd_helper
13import constants
14import ports
15
16from pylib import pexpect
17
18class Forwarder(object):
19  """Class to manage port forwards from the device to the host."""
20
21  _DEVICE_FORWARDER_PATH = constants.TEST_EXECUTABLE_DIR + '/device_forwarder'
22
23  # Unix Abstract socket path:
24  _DEVICE_ADB_CONTROL_PORT = 'chrome_device_forwarder'
25  _TIMEOUT_SECS = 30
26
27  def __init__(self, adb, port_pairs, tool, host_name, build_type):
28    """Forwards TCP ports on the device back to the host.
29
30    Works like adb forward, but in reverse.
31
32    Args:
33      adb: Instance of AndroidCommands for talking to the device.
34      port_pairs: A list of tuples (device_port, host_port) to forward. Note
35                 that you can specify 0 as a device_port, in which case a
36                 port will by dynamically assigned on the device. You can
37                 get the number of the assigned port using the
38                 DevicePortForHostPort method.
39      tool: Tool class to use to get wrapper, if necessary, for executing the
40            forwarder (see valgrind_tools.py).
41      host_name: Address to forward to, must be addressable from the
42                 host machine. Usually use loopback '127.0.0.1'.
43      build_type: 'Release' or 'Debug'.
44
45    Raises:
46      Exception on failure to forward the port.
47    """
48    self._adb = adb
49    self._host_to_device_port_map = dict()
50    self._host_process = None
51    self._device_process = None
52    self._adb_forward_process = None
53
54    self._host_adb_control_port = ports.AllocateTestServerPort()
55    if not self._host_adb_control_port:
56      raise Exception('Failed to allocate a TCP port in the host machine.')
57    adb.PushIfNeeded(
58        os.path.join(constants.CHROME_DIR, 'out', build_type,
59                     'device_forwarder'),
60        Forwarder._DEVICE_FORWARDER_PATH)
61    self._host_forwarder_path = os.path.join(constants.CHROME_DIR,
62                                             'out',
63                                             build_type,
64                                             'host_forwarder')
65    forward_string = ['%d:%d:%s' %
66                      (device, host, host_name) for device, host in port_pairs]
67    logging.info('Forwarding ports: %s', forward_string)
68    timeout_sec = 5
69    host_pattern = 'host_forwarder.*' + ' '.join(forward_string)
70    # TODO(felipeg): Rather than using a blocking kill() here, the device
71    # forwarder could try to bind the Unix Domain Socket until it succeeds or
72    # while it fails because the socket is already bound (with appropriate
73    # timeout handling obviously).
74    self._KillHostForwarderBlocking(host_pattern, timeout_sec)
75    self._KillDeviceForwarderBlocking(timeout_sec)
76    self._adb_forward_process = pexpect.spawn(
77        'adb', ['-s',
78                adb._adb.GetSerialNumber(),
79                'forward',
80                'tcp:%s' % self._host_adb_control_port,
81                'localabstract:%s' % Forwarder._DEVICE_ADB_CONTROL_PORT])
82    self._device_process = pexpect.spawn(
83        'adb', ['-s',
84                adb._adb.GetSerialNumber(),
85                'shell',
86                '%s %s -D --adb_sock=%s' % (
87                    tool.GetUtilWrapper(),
88                    Forwarder._DEVICE_FORWARDER_PATH,
89                    Forwarder._DEVICE_ADB_CONTROL_PORT)])
90
91    device_success_re = re.compile('Starting Device Forwarder.')
92    device_failure_re = re.compile('.*:ERROR:(.*)')
93    index = self._device_process.expect([device_success_re,
94                                         device_failure_re,
95                                         pexpect.EOF,
96                                         pexpect.TIMEOUT],
97                                        Forwarder._TIMEOUT_SECS)
98    if index == 1:
99      # Failure
100      error_msg = str(self._device_process.match.group(1))
101      logging.error(self._device_process.before)
102      self._CloseProcess()
103      raise Exception('Failed to start Device Forwarder with Error: %s' %
104                      error_msg)
105    elif index == 2:
106      logging.error(self._device_process.before)
107      self._CloseProcess()
108      raise Exception('Unexpected EOF while trying to start Device Forwarder.')
109    elif index == 3:
110      logging.error(self._device_process.before)
111      self._CloseProcess()
112      raise Exception('Timeout while trying start Device Forwarder')
113
114    self._host_process = pexpect.spawn(self._host_forwarder_path,
115                                       ['--adb_port=%s' % (
116                                           self._host_adb_control_port)] +
117                                       forward_string)
118
119    # Read the output of the command to determine which device ports where
120    # forwarded to which host ports (necessary if
121    host_success_re = re.compile('Forwarding device port (\d+) to host (\d+):')
122    host_failure_re = re.compile('Couldn\'t start forwarder server for port '
123                                 'spec: (\d+):(\d+)')
124    for pair in port_pairs:
125      index = self._host_process.expect([host_success_re,
126                                         host_failure_re,
127                                         pexpect.EOF,
128                                         pexpect.TIMEOUT],
129                                        Forwarder._TIMEOUT_SECS)
130      if index == 0:
131        # Success
132        device_port = int(self._host_process.match.group(1))
133        host_port = int(self._host_process.match.group(2))
134        self._host_to_device_port_map[host_port] = device_port
135        logging.info("Forwarding device port: %d to host port: %d." %
136                     (device_port, host_port))
137      elif index == 1:
138        # Failure
139        device_port = int(self._host_process.match.group(1))
140        host_port = int(self._host_process.match.group(2))
141        self._CloseProcess()
142        raise Exception('Failed to forward port %d to %d' % (device_port,
143                                                             host_port))
144      elif index == 2:
145        logging.error(self._host_process.before)
146        self._CloseProcess()
147        raise Exception('Unexpected EOF while trying to forward ports %s' %
148                        port_pairs)
149      elif index == 3:
150        logging.error(self._host_process.before)
151        self._CloseProcess()
152        raise Exception('Timeout while trying to forward ports %s' % port_pairs)
153
154  def _KillHostForwarderBlocking(self, host_pattern, timeout_sec):
155    """Kills any existing host forwarders using the provided pattern.
156
157    Note that this waits until the process terminates.
158    """
159    cmd_helper.RunCmd(['pkill', '-f', host_pattern])
160    elapsed = 0
161    wait_period = 0.1
162    while not cmd_helper.RunCmd(['pgrep', '-f', host_pattern]) and (
163        elapsed < timeout_sec):
164      time.sleep(wait_period)
165      elapsed += wait_period
166    if elapsed >= timeout_sec:
167        raise Exception('Timed out while killing ' + host_pattern)
168
169  def _KillDeviceForwarderBlocking(self, timeout_sec):
170    """Kills any existing device forwarders.
171
172    Note that this waits until the process terminates.
173    """
174    processes_killed = self._adb.KillAllBlocking(
175        'device_forwarder', timeout_sec)
176    if not processes_killed:
177      pids = self._adb.ExtractPid('device_forwarder')
178      if pids:
179        raise Exception('Timed out while killing device_forwarder')
180
181  def _CloseProcess(self):
182    if self._host_process:
183      self._host_process.close()
184    if self._device_process:
185      self._device_process.close()
186    if self._adb_forward_process:
187      self._adb_forward_process.close()
188    self._host_process = None
189    self._device_process = None
190    self._adb_forward_process = None
191
192  def DevicePortForHostPort(self, host_port):
193    """Get the device port that corresponds to a given host port."""
194    return self._host_to_device_port_map.get(host_port)
195
196  def Close(self):
197    """Terminate the forwarder process."""
198    self._CloseProcess()
199