test_server.py (7941B)
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 6 import logging 7 import os 8 import sys 9 import subprocess 10 11 from typing import List, Optional, Tuple 12 13 from common import DIR_SRC_ROOT, get_free_local_port, get_ssh_address 14 from compatible_utils import get_ssh_prefix 15 16 sys.path.append(os.path.join(DIR_SRC_ROOT, 'build', 'util', 'lib', 'common')) 17 # pylint: disable=import-error,wrong-import-position 18 import chrome_test_server_spawner 19 # pylint: enable=import-error,wrong-import-position 20 21 22 def _run_ssh_tunnel(target_addr: str, command: str, 23 port_maps: List[str]) -> subprocess.CompletedProcess: 24 assert port_maps 25 26 ssh_prefix = get_ssh_prefix(target_addr) 27 28 # Allow a tunnel / control path to be established for the first time. 29 # The sshconfig https://crsrc.org/c/build/fuchsia/test/sshconfig used here 30 # persists the connection. 31 subprocess.run(ssh_prefix + ['echo', 'true'], check=True) 32 33 forward_proc = subprocess.run( 34 ssh_prefix + [ 35 '-O', 36 command, # Send SSH mux control signal. 37 '-NT' # Don't execute command; don't allocate terminal. 38 ] + port_maps, 39 capture_output=True, 40 check=True, 41 text=True) 42 return forward_proc 43 44 45 def _forward_command(fuchsia_port: int, host_port: int, 46 port_forwarding: bool) -> List[str]: 47 max_port = 65535 48 assert fuchsia_port is not None and 0 <= fuchsia_port <= max_port 49 assert host_port is not None and 0 < host_port <= max_port 50 if port_forwarding: 51 return ['-R', f'{fuchsia_port}:localhost:{host_port}'] 52 assert fuchsia_port != 0 53 return ['-L', f'{host_port}:localhost:{fuchsia_port}'] 54 55 56 def _forward_commands(ports: List[Tuple[int, int]], 57 port_forwarding: bool) -> List[str]: 58 assert ports 59 forward_cmd = [] 60 for port in ports: 61 assert port is not None 62 forward_cmd.extend(_forward_command(port[0], port[1], port_forwarding)) 63 return forward_cmd 64 65 66 def ports_forward(target_addr: str, 67 ports: List[Tuple[int, int]]) -> subprocess.CompletedProcess: 68 """Establishes a port forwarding SSH task to forward ports from the host to 69 the fuchsia endpoints specified by tuples of port numbers in format of 70 [fuchsia-port, host-port]. Setting fuchsia-port to 0 would allow the fuchsia 71 selecting a free port; host-port shouldn't be 0. 72 73 Blocks until port forwarding is established. 74 75 Returns the CompletedProcess of the SSH task.""" 76 return _run_ssh_tunnel(target_addr, 'forward', 77 _forward_commands(ports, True)) 78 79 80 def ports_backward( 81 target_addr: str, 82 ports: List[Tuple[int, int]]) -> subprocess.CompletedProcess: 83 """Establishes a reverse port forwarding SSH task to forward ports from the 84 fuchsia to the host endpoints specified by tuples of port numbers in format 85 of [fuchsia-port, host-port]. Both host-port and fuchsia-port shouldn't be 86 0. 87 88 Blocks until port forwarding is established. 89 90 Returns the CompletedProcess of the SSH task.""" 91 return _run_ssh_tunnel(target_addr, 'forward', 92 _forward_commands(ports, False)) 93 94 95 def port_forward(target_addr: str, host_port: int) -> int: 96 """Establishes a port forwarding SSH task to a host TCP endpoint at port 97 |host_port|. Blocks until port forwarding is established. 98 99 Returns the fuchsia port number.""" 100 101 forward_proc = ports_forward(target_addr, [(0, host_port)]) 102 parsed_port = int(forward_proc.stdout.splitlines()[0].strip()) 103 logging.debug('Port forwarding established (local=%d, device=%d)', 104 host_port, parsed_port) 105 return parsed_port 106 107 108 def port_backward(target_addr: str, 109 fuchsia_port: int, 110 host_port: int = 0) -> int: 111 """Establishes a reverse port forwarding SSH task to a fuchsia TCP endpoint 112 at port |fuchsia_port| from the host at port |host_port|. If |host_port| is 113 None or 0, a local free port will be selected. 114 Blocks until reverse port forwarding is established. 115 116 Returns the local port number.""" 117 118 if not host_port: 119 host_port = get_free_local_port() 120 ports_backward(target_addr, [(fuchsia_port, host_port)]) 121 logging.debug('Reverse port forwarding established (local=%d, device=%d)', 122 host_port, fuchsia_port) 123 return host_port 124 125 126 def cancel_port_forwarding(target_addr: str, fuchsia_port: int, host_port: int, 127 port_forwarding: bool) -> None: 128 """Cancels an existing port forwarding, if port_forwarding is false, it will 129 be treated as reverse port forwarding. 130 Note, the ports passing in here need to exactly match the ports used to 131 setup the port forwarding, i.e. if ports_forward([0, 8080]) was issued, even 132 it returned an allocated port, cancel_port_forwarding(..., 0, 8080, ...) 133 should still be used to cancel the port forwarding.""" 134 _run_ssh_tunnel(target_addr, 'cancel', 135 _forward_command(fuchsia_port, host_port, port_forwarding)) 136 137 138 # Disable pylint errors since the subclass is not from this directory. 139 # pylint: disable=invalid-name,missing-function-docstring 140 class SSHPortForwarder(chrome_test_server_spawner.PortForwarder): 141 """Implementation of chrome_test_server_spawner.PortForwarder that uses 142 SSH's remote port forwarding feature to forward ports.""" 143 144 def __init__(self, target_addr: str) -> None: 145 self._target_addr = target_addr 146 147 # Maps the host (server) port to the device port number. 148 self._port_mapping = {} 149 150 def Map(self, port_pairs: List[Tuple[int, int]]) -> None: 151 for p in port_pairs: 152 fuchsia_port, host_port = p 153 assert fuchsia_port == 0, \ 154 'Port forwarding with a fixed fuchsia-port is unsupported yet.' 155 self._port_mapping[host_port] = \ 156 port_forward(self._target_addr, host_port) 157 158 def GetDevicePortForHostPort(self, host_port: int) -> int: 159 return self._port_mapping[host_port] 160 161 def Unmap(self, device_port: int) -> None: 162 for host_port, fuchsia_port in self._port_mapping.items(): 163 if fuchsia_port == device_port: 164 cancel_port_forwarding(self._target_addr, 0, host_port, True) 165 del self._port_mapping[host_port] 166 return 167 168 raise Exception('Unmap called for unknown port: %d' % device_port) 169 170 171 # pylint: enable=invalid-name,missing-function-docstring 172 173 174 def setup_test_server(target_id: Optional[str], test_concurrency: int)\ 175 -> Tuple[chrome_test_server_spawner.SpawningServer, str]: 176 """Provisions a test server and configures |target_id| to use it. 177 178 Args: 179 target_id: The target to which port forwarding to the test server will 180 be established. 181 test_concurrency: The number of parallel test jobs that will be run. 182 183 Returns a tuple of a SpawningServer object and the local url to use on 184 |target_id| to reach the test server.""" 185 186 logging.debug('Starting test server.') 187 188 target_addr = get_ssh_address(target_id) 189 190 # The TestLauncher can launch more jobs than the limit specified with 191 # --test-launcher-jobs so the max number of spawned test servers is set to 192 # twice that limit here. See https://crbug.com/913156#c19. 193 spawning_server = chrome_test_server_spawner.SpawningServer( 194 0, SSHPortForwarder(target_addr), test_concurrency * 2) 195 196 forwarded_port = port_forward(target_addr, spawning_server.server_port) 197 spawning_server.Start() 198 199 logging.debug('Test server listening for connections (port=%d)', 200 spawning_server.server_port) 201 logging.debug('Forwarded port is %d', forwarded_port) 202 203 return (spawning_server, 'http://localhost:%d' % forwarded_port)