1# Copyright 2022 The Chromium Authors 2# Use of this source code is governed by a BSD-style license that can be 3# found in the LICENSE file. 4"""Test server set up.""" 5 6import logging 7import os 8import sys 9import subprocess 10 11from typing import List, Optional, Tuple 12 13from common import DIR_SRC_ROOT, run_ffx_command 14from compatible_utils import get_ssh_prefix 15 16sys.path.append(os.path.join(DIR_SRC_ROOT, 'build', 'util', 'lib', 'common')) 17# pylint: disable=import-error,wrong-import-position 18import chrome_test_server_spawner 19# pylint: enable=import-error,wrong-import-position 20 21 22def port_forward(host_port_pair: str, host_port: int) -> int: 23 """Establishes a port forwarding SSH task to a localhost TCP endpoint 24 hosted at port |local_port|. Blocks until port forwarding is established. 25 26 Returns the remote port number.""" 27 28 ssh_prefix = get_ssh_prefix(host_port_pair) 29 30 # Allow a tunnel to be established. 31 subprocess.run(ssh_prefix + ['echo', 'true'], check=True) 32 33 forward_cmd = [ 34 '-O', 35 'forward', # Send SSH mux control signal. 36 '-R', 37 '0:localhost:%d' % host_port, 38 '-v', # Get forwarded port info from stderr. 39 '-NT' # Don't execute command; don't allocate terminal. 40 ] 41 forward_proc = subprocess.run(ssh_prefix + forward_cmd, 42 capture_output=True, 43 check=False, 44 text=True) 45 if forward_proc.returncode != 0: 46 raise Exception( 47 'Got an error code when requesting port forwarding: %d' % 48 forward_proc.returncode) 49 50 output = forward_proc.stdout 51 parsed_port = int(output.splitlines()[0].strip()) 52 logging.debug('Port forwarding established (local=%d, device=%d)', 53 host_port, parsed_port) 54 return parsed_port 55 56 57# Disable pylint errors since the subclass is not from this directory. 58# pylint: disable=invalid-name,missing-function-docstring 59class SSHPortForwarder(chrome_test_server_spawner.PortForwarder): 60 """Implementation of chrome_test_server_spawner.PortForwarder that uses 61 SSH's remote port forwarding feature to forward ports.""" 62 63 def __init__(self, host_port_pair: str) -> None: 64 self._host_port_pair = host_port_pair 65 66 # Maps the host (server) port to the device port number. 67 self._port_mapping = {} 68 69 def Map(self, port_pairs: List[Tuple[int, int]]) -> None: 70 for p in port_pairs: 71 _, host_port = p 72 self._port_mapping[host_port] = \ 73 port_forward(self._host_port_pair, host_port) 74 75 def GetDevicePortForHostPort(self, host_port: int) -> int: 76 return self._port_mapping[host_port] 77 78 def Unmap(self, device_port: int) -> None: 79 for host_port, entry in self._port_mapping.items(): 80 if entry == device_port: 81 ssh_prefix = get_ssh_prefix(self._host_port_pair) 82 unmap_cmd = [ 83 '-NT', '-O', 'cancel', '-R', 84 '0:localhost:%d' % host_port 85 ] 86 ssh_proc = subprocess.run(ssh_prefix + unmap_cmd, check=False) 87 if ssh_proc.returncode != 0: 88 raise Exception('Error %d when unmapping port %d' % 89 (ssh_proc.returncode, device_port)) 90 del self._port_mapping[host_port] 91 return 92 93 raise Exception('Unmap called for unknown port: %d' % device_port) 94 95 96# pylint: enable=invalid-name,missing-function-docstring 97 98 99def setup_test_server(target_id: Optional[str], test_concurrency: int)\ 100 -> Tuple[chrome_test_server_spawner.SpawningServer, str]: 101 """Provisions a test server and configures |target_id| to use it. 102 103 Args: 104 target_id: The target to which port forwarding to the test server will 105 be established. 106 test_concurrency: The number of parallel test jobs that will be run. 107 108 Returns a tuple of a SpawningServer object and the local url to use on 109 |target_id| to reach the test server.""" 110 111 logging.debug('Starting test server.') 112 113 host_port_pair = run_ffx_command(cmd=('target', 'get-ssh-address'), 114 target_id=target_id, 115 capture_output=True).stdout.strip() 116 117 # The TestLauncher can launch more jobs than the limit specified with 118 # --test-launcher-jobs so the max number of spawned test servers is set to 119 # twice that limit here. See https://crbug.com/913156#c19. 120 spawning_server = chrome_test_server_spawner.SpawningServer( 121 0, SSHPortForwarder(host_port_pair), test_concurrency * 2) 122 123 forwarded_port = port_forward(host_port_pair, spawning_server.server_port) 124 spawning_server.Start() 125 126 logging.debug('Test server listening for connections (port=%d)', 127 spawning_server.server_port) 128 logging.debug('Forwarded port is %d', forwarded_port) 129 130 return (spawning_server, 'http://localhost:%d' % forwarded_port) 131