• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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_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 ports_forward(host_port_pair: str,
23                  ports: List[Tuple[int, int]]) -> subprocess.CompletedProcess:
24    """Establishes a port forwarding SSH task to forward ports from fuchsia to
25    the local endpoints specified by tuples of port numbers. Blocks until port
26    forwarding is established.
27
28    Returns the CompletedProcess of the SSH task."""
29    assert len(ports) > 0
30
31    ssh_prefix = get_ssh_prefix(host_port_pair)
32
33    # Allow a tunnel to be established.
34    subprocess.run(ssh_prefix + ['echo', 'true'], check=True)
35
36    forward_cmd = [
37        '-O',
38        'forward',  # Send SSH mux control signal.
39        '-v',  # Get forwarded port info from stderr.
40        '-NT'  # Don't execute command; don't allocate terminal.
41    ]
42    for port in ports:
43        forward_cmd.extend(['-R', f'{port[0]}:localhost:{port[1]}'])
44    forward_proc = subprocess.run(ssh_prefix + forward_cmd,
45                                  capture_output=True,
46                                  check=False,
47                                  text=True)
48    if forward_proc.returncode != 0:
49        raise Exception(
50            'Got an error code when requesting port forwarding: %d' %
51            forward_proc.returncode)
52    return forward_proc
53
54
55def port_forward(host_port_pair: str, host_port: int) -> int:
56    """Establishes a port forwarding SSH task to a localhost TCP endpoint
57    hosted at port |local_port|. Blocks until port forwarding is established.
58
59    Returns the remote port number."""
60
61    forward_proc = ports_forward(host_port_pair, [(0, host_port)])
62    parsed_port = int(forward_proc.stdout.splitlines()[0].strip())
63    logging.debug('Port forwarding established (local=%d, device=%d)',
64                  host_port, parsed_port)
65    return parsed_port
66
67
68# Disable pylint errors since the subclass is not from this directory.
69# pylint: disable=invalid-name,missing-function-docstring
70class SSHPortForwarder(chrome_test_server_spawner.PortForwarder):
71    """Implementation of chrome_test_server_spawner.PortForwarder that uses
72    SSH's remote port forwarding feature to forward ports."""
73
74    def __init__(self, host_port_pair: str) -> None:
75        self._host_port_pair = host_port_pair
76
77        # Maps the host (server) port to the device port number.
78        self._port_mapping = {}
79
80    def Map(self, port_pairs: List[Tuple[int, int]]) -> None:
81        for p in port_pairs:
82            _, host_port = p
83            self._port_mapping[host_port] = \
84                port_forward(self._host_port_pair, host_port)
85
86    def GetDevicePortForHostPort(self, host_port: int) -> int:
87        return self._port_mapping[host_port]
88
89    def Unmap(self, device_port: int) -> None:
90        for host_port, entry in self._port_mapping.items():
91            if entry == device_port:
92                ssh_prefix = get_ssh_prefix(self._host_port_pair)
93                unmap_cmd = [
94                    '-NT', '-O', 'cancel', '-R',
95                    '0:localhost:%d' % host_port
96                ]
97                ssh_proc = subprocess.run(ssh_prefix + unmap_cmd, check=False)
98                if ssh_proc.returncode != 0:
99                    raise Exception('Error %d when unmapping port %d' %
100                                    (ssh_proc.returncode, device_port))
101                del self._port_mapping[host_port]
102                return
103
104        raise Exception('Unmap called for unknown port: %d' % device_port)
105
106
107# pylint: enable=invalid-name,missing-function-docstring
108
109
110def setup_test_server(target_id: Optional[str], test_concurrency: int)\
111         -> Tuple[chrome_test_server_spawner.SpawningServer, str]:
112    """Provisions a test server and configures |target_id| to use it.
113
114    Args:
115        target_id: The target to which port forwarding to the test server will
116            be established.
117        test_concurrency: The number of parallel test jobs that will be run.
118
119    Returns a tuple of a SpawningServer object and the local url to use on
120    |target_id| to reach the test server."""
121
122    logging.debug('Starting test server.')
123
124    host_port_pair = get_ssh_address(target_id)
125
126    # The TestLauncher can launch more jobs than the limit specified with
127    # --test-launcher-jobs so the max number of spawned test servers is set to
128    # twice that limit here. See https://crbug.com/913156#c19.
129    spawning_server = chrome_test_server_spawner.SpawningServer(
130        0, SSHPortForwarder(host_port_pair), test_concurrency * 2)
131
132    forwarded_port = port_forward(host_port_pair, spawning_server.server_port)
133    spawning_server.Start()
134
135    logging.debug('Test server listening for connections (port=%d)',
136                  spawning_server.server_port)
137    logging.debug('Forwarded port is %d', forwarded_port)
138
139    return (spawning_server, 'http://localhost:%d' % forwarded_port)
140