xref: /aosp_15_r20/external/angle/build/fuchsia/test/test_server.py (revision 8975f5c5ed3d1c378011245431ada316dfb6f244)
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, get_free_local_port, get_ssh_address
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 _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
45def _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
56def _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
66def 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
80def 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
95def 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
108def 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
126def 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
140class 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
174def 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)
204