michael@0: # Copyright (c) 2012 The Chromium Authors. All rights reserved.
michael@0: # Use of this source code is governed by a BSD-style license that can be
michael@0: # found in the LICENSE file.
michael@0: 
michael@0: import logging
michael@0: import os
michael@0: import re
michael@0: import sys
michael@0: import time
michael@0: 
michael@0: import android_commands
michael@0: import cmd_helper
michael@0: import constants
michael@0: import ports
michael@0: 
michael@0: from pylib import pexpect
michael@0: 
michael@0: class Forwarder(object):
michael@0:   """Class to manage port forwards from the device to the host."""
michael@0: 
michael@0:   _DEVICE_FORWARDER_PATH = constants.TEST_EXECUTABLE_DIR + '/device_forwarder'
michael@0: 
michael@0:   # Unix Abstract socket path:
michael@0:   _DEVICE_ADB_CONTROL_PORT = 'chrome_device_forwarder'
michael@0:   _TIMEOUT_SECS = 30
michael@0: 
michael@0:   def __init__(self, adb, port_pairs, tool, host_name, build_type):
michael@0:     """Forwards TCP ports on the device back to the host.
michael@0: 
michael@0:     Works like adb forward, but in reverse.
michael@0: 
michael@0:     Args:
michael@0:       adb: Instance of AndroidCommands for talking to the device.
michael@0:       port_pairs: A list of tuples (device_port, host_port) to forward. Note
michael@0:                  that you can specify 0 as a device_port, in which case a
michael@0:                  port will by dynamically assigned on the device. You can
michael@0:                  get the number of the assigned port using the
michael@0:                  DevicePortForHostPort method.
michael@0:       tool: Tool class to use to get wrapper, if necessary, for executing the
michael@0:             forwarder (see valgrind_tools.py).
michael@0:       host_name: Address to forward to, must be addressable from the
michael@0:                  host machine. Usually use loopback '127.0.0.1'.
michael@0:       build_type: 'Release' or 'Debug'.
michael@0: 
michael@0:     Raises:
michael@0:       Exception on failure to forward the port.
michael@0:     """
michael@0:     self._adb = adb
michael@0:     self._host_to_device_port_map = dict()
michael@0:     self._host_process = None
michael@0:     self._device_process = None
michael@0:     self._adb_forward_process = None
michael@0: 
michael@0:     self._host_adb_control_port = ports.AllocateTestServerPort()
michael@0:     if not self._host_adb_control_port:
michael@0:       raise Exception('Failed to allocate a TCP port in the host machine.')
michael@0:     adb.PushIfNeeded(
michael@0:         os.path.join(constants.CHROME_DIR, 'out', build_type,
michael@0:                      'device_forwarder'),
michael@0:         Forwarder._DEVICE_FORWARDER_PATH)
michael@0:     self._host_forwarder_path = os.path.join(constants.CHROME_DIR,
michael@0:                                              'out',
michael@0:                                              build_type,
michael@0:                                              'host_forwarder')
michael@0:     forward_string = ['%d:%d:%s' %
michael@0:                       (device, host, host_name) for device, host in port_pairs]
michael@0:     logging.info('Forwarding ports: %s', forward_string)
michael@0:     timeout_sec = 5
michael@0:     host_pattern = 'host_forwarder.*' + ' '.join(forward_string)
michael@0:     # TODO(felipeg): Rather than using a blocking kill() here, the device
michael@0:     # forwarder could try to bind the Unix Domain Socket until it succeeds or
michael@0:     # while it fails because the socket is already bound (with appropriate
michael@0:     # timeout handling obviously).
michael@0:     self._KillHostForwarderBlocking(host_pattern, timeout_sec)
michael@0:     self._KillDeviceForwarderBlocking(timeout_sec)
michael@0:     self._adb_forward_process = pexpect.spawn(
michael@0:         'adb', ['-s',
michael@0:                 adb._adb.GetSerialNumber(),
michael@0:                 'forward',
michael@0:                 'tcp:%s' % self._host_adb_control_port,
michael@0:                 'localabstract:%s' % Forwarder._DEVICE_ADB_CONTROL_PORT])
michael@0:     self._device_process = pexpect.spawn(
michael@0:         'adb', ['-s',
michael@0:                 adb._adb.GetSerialNumber(),
michael@0:                 'shell',
michael@0:                 '%s %s -D --adb_sock=%s' % (
michael@0:                     tool.GetUtilWrapper(),
michael@0:                     Forwarder._DEVICE_FORWARDER_PATH,
michael@0:                     Forwarder._DEVICE_ADB_CONTROL_PORT)])
michael@0: 
michael@0:     device_success_re = re.compile('Starting Device Forwarder.')
michael@0:     device_failure_re = re.compile('.*:ERROR:(.*)')
michael@0:     index = self._device_process.expect([device_success_re,
michael@0:                                          device_failure_re,
michael@0:                                          pexpect.EOF,
michael@0:                                          pexpect.TIMEOUT],
michael@0:                                         Forwarder._TIMEOUT_SECS)
michael@0:     if index == 1:
michael@0:       # Failure
michael@0:       error_msg = str(self._device_process.match.group(1))
michael@0:       logging.error(self._device_process.before)
michael@0:       self._CloseProcess()
michael@0:       raise Exception('Failed to start Device Forwarder with Error: %s' %
michael@0:                       error_msg)
michael@0:     elif index == 2:
michael@0:       logging.error(self._device_process.before)
michael@0:       self._CloseProcess()
michael@0:       raise Exception('Unexpected EOF while trying to start Device Forwarder.')
michael@0:     elif index == 3:
michael@0:       logging.error(self._device_process.before)
michael@0:       self._CloseProcess()
michael@0:       raise Exception('Timeout while trying start Device Forwarder')
michael@0: 
michael@0:     self._host_process = pexpect.spawn(self._host_forwarder_path,
michael@0:                                        ['--adb_port=%s' % (
michael@0:                                            self._host_adb_control_port)] +
michael@0:                                        forward_string)
michael@0: 
michael@0:     # Read the output of the command to determine which device ports where
michael@0:     # forwarded to which host ports (necessary if
michael@0:     host_success_re = re.compile('Forwarding device port (\d+) to host (\d+):')
michael@0:     host_failure_re = re.compile('Couldn\'t start forwarder server for port '
michael@0:                                  'spec: (\d+):(\d+)')
michael@0:     for pair in port_pairs:
michael@0:       index = self._host_process.expect([host_success_re,
michael@0:                                          host_failure_re,
michael@0:                                          pexpect.EOF,
michael@0:                                          pexpect.TIMEOUT],
michael@0:                                         Forwarder._TIMEOUT_SECS)
michael@0:       if index == 0:
michael@0:         # Success
michael@0:         device_port = int(self._host_process.match.group(1))
michael@0:         host_port = int(self._host_process.match.group(2))
michael@0:         self._host_to_device_port_map[host_port] = device_port
michael@0:         logging.info("Forwarding device port: %d to host port: %d." %
michael@0:                      (device_port, host_port))
michael@0:       elif index == 1:
michael@0:         # Failure
michael@0:         device_port = int(self._host_process.match.group(1))
michael@0:         host_port = int(self._host_process.match.group(2))
michael@0:         self._CloseProcess()
michael@0:         raise Exception('Failed to forward port %d to %d' % (device_port,
michael@0:                                                              host_port))
michael@0:       elif index == 2:
michael@0:         logging.error(self._host_process.before)
michael@0:         self._CloseProcess()
michael@0:         raise Exception('Unexpected EOF while trying to forward ports %s' %
michael@0:                         port_pairs)
michael@0:       elif index == 3:
michael@0:         logging.error(self._host_process.before)
michael@0:         self._CloseProcess()
michael@0:         raise Exception('Timeout while trying to forward ports %s' % port_pairs)
michael@0: 
michael@0:   def _KillHostForwarderBlocking(self, host_pattern, timeout_sec):
michael@0:     """Kills any existing host forwarders using the provided pattern.
michael@0: 
michael@0:     Note that this waits until the process terminates.
michael@0:     """
michael@0:     cmd_helper.RunCmd(['pkill', '-f', host_pattern])
michael@0:     elapsed = 0
michael@0:     wait_period = 0.1
michael@0:     while not cmd_helper.RunCmd(['pgrep', '-f', host_pattern]) and (
michael@0:         elapsed < timeout_sec):
michael@0:       time.sleep(wait_period)
michael@0:       elapsed += wait_period
michael@0:     if elapsed >= timeout_sec:
michael@0:         raise Exception('Timed out while killing ' + host_pattern)
michael@0: 
michael@0:   def _KillDeviceForwarderBlocking(self, timeout_sec):
michael@0:     """Kills any existing device forwarders.
michael@0: 
michael@0:     Note that this waits until the process terminates.
michael@0:     """
michael@0:     processes_killed = self._adb.KillAllBlocking(
michael@0:         'device_forwarder', timeout_sec)
michael@0:     if not processes_killed:
michael@0:       pids = self._adb.ExtractPid('device_forwarder')
michael@0:       if pids:
michael@0:         raise Exception('Timed out while killing device_forwarder')
michael@0: 
michael@0:   def _CloseProcess(self):
michael@0:     if self._host_process:
michael@0:       self._host_process.close()
michael@0:     if self._device_process:
michael@0:       self._device_process.close()
michael@0:     if self._adb_forward_process:
michael@0:       self._adb_forward_process.close()
michael@0:     self._host_process = None
michael@0:     self._device_process = None
michael@0:     self._adb_forward_process = None
michael@0: 
michael@0:   def DevicePortForHostPort(self, host_port):
michael@0:     """Get the device port that corresponds to a given host port."""
michael@0:     return self._host_to_device_port_map.get(host_port)
michael@0: 
michael@0:   def Close(self):
michael@0:     """Terminate the forwarder process."""
michael@0:     self._CloseProcess()