Move port-manager code to custom PortManager classes

Move all port-managing methods from Platform to the dedicated
PortManager classes.

- Add LocalPortManager default implementation
- Complete AndroidAdbPortManager implementation
- Complete SshPortManager implementation
- Add new NumberParser.port_number_zero helper
- Make build_ssh_cmd method public so it can be used from the
  SshPortManager
- Add more test PortManagers

Change-Id: I5aaccdf0325172586ad42ae345556f68be3d1e14
Reviewed-on: https://chromium-review.googlesource.com/c/crossbench/+/6650994
Commit-Queue: Camillo Bruni <[email protected]>
Reviewed-by: Patrick Thier <[email protected]>

Author: camillobruni

crossbench/browsers/chromium/devtools.py

@@ -32,7 +32,7 @@ def connect(self) -> None:
     if self._ws and self._ws.connected:
       return
     try:
-      self._devtools_port = self._platform.forward_devtools_port(
+      self._devtools_port = self._platform.ports.forward_devtools(
           local_port=self._requested_local_port,
           remote_identifier=self._remote_devtools_identifier)
       self._ws = websocket.WebSocket()
@@ -59,7 +59,7 @@ def _disconnect_internal(self) -> None:
     self._ws = None
     if self._devtools_port:
       try:
-        self._platform.stop_port_forward(self._devtools_port)
+        self._platform.ports.stop_forward(self._devtools_port)
       except Exception as e:  # pylint: disable=broad-except
         # Best effort to remove forwarding, log if it fails but don't crash
         logging.warning(

crossbench/parse.py

@@ -645,6 +645,13 @@ def port_number(cls,
                   parse_str: bool = True) -> int:
     return cls.int_range(1, 65535, name, parse_str)(value)
 
+  @classmethod
+  def port_number_zero(cls,
+                       value: Any,
+                       name: str = "port",
+                       parse_str: bool = True) -> int:
+    return cls.int_range(0, 65535, name, parse_str)(value)
+
 
 class LateArgumentError(argparse.ArgumentTypeError):
   """Signals argument parse errors after parser.parse_args().

crossbench/plt/android_adb.py

@@ -16,8 +16,8 @@
 from snippet_uiautomator import uiautomator
 from typing_extensions import override
 
-from crossbench.flags.base import Flags, FlagsData
 from crossbench import path as pth
+from crossbench.flags.base import Flags, FlagsData
 from crossbench.parse import NumberParser
 from crossbench.plt.arch import MachineArch
 from crossbench.plt.base import SubprocessError
@@ -476,7 +476,50 @@ class AndroidAdbPortManager(PortManager):
 
   def __init__(self, platform: AndroidAdbPlatform, adb: Adb) -> None:
     super().__init__(platform)
-    self._adb = adb
+    self._adb: Adb = adb
+
+  @property
+  def host_platform(self) -> Platform:
+    return self._platform.host_platform
+
+  @override
+  def forward(self, local_port: int, remote_port: int) -> int:
+    local_port = NumberParser.positive_zero_int(local_port, "local_port")
+    remote_port = NumberParser.port_number(remote_port, "remote_port")
+    local_port = self._adb.forward(
+        local_port, remote_port, local_protocol="tcp", remote_protocol="tcp")
+    logging.debug("Forwarded Remote Port: %s:%s <= %s:%s",
+                  self.host_platform.name, local_port, self, remote_port)
+    return local_port
+
+  @override
+  def forward_devtools(self, local_port: int, remote_identifier: str) -> int:
+    local_port = NumberParser.positive_zero_int(local_port, "local_port")
+    local_port = self._adb.forward(
+        local=local_port,
+        remote=remote_identifier,
+        local_protocol="tcp",
+        remote_protocol="localabstract")
+    logging.debug("Forwarded DevTools Port: %s:%s <= %s:%s",
+                  self.host_platform.name, local_port, self, remote_identifier)
+    return local_port
+
+  @override
+  def stop_forward(self, local_port: int) -> None:
+    self._adb.forward_remove(local_port, protocol="tcp")
+
+  @override
+  def reverse_forward(self, remote_port: int, local_port: int) -> int:
+    remote_port = NumberParser.positive_zero_int(remote_port, "remote_port")
+    local_port = NumberParser.port_number(local_port, "local_port")
+    remote_port = self._adb.reverse(remote_port, local_port, protocol="tcp")
+    logging.debug("Forwarded Local Port: %s:%s => %s:%s", self.host_platform,
+                  local_port, self, remote_port)
+    return remote_port
+
+  @override
+  def stop_reverse_forward(self, remote_port: int) -> None:
+    self._adb.reverse_remove(remote_port, protocol="tcp")
 
 
 class AndroidAdbPlatform(RemotePosixPlatform):
@@ -680,45 +723,6 @@ def sh_stdout_bytes(self,
     return self.adb.shell_stdout_bytes(
         *args, shell=shell, stdin=stdin, env=env, quiet=quiet, check=check)
 
-  @override
-  def port_forward(self, local_port: int, remote_port: int) -> int:
-    local_port = NumberParser.positive_zero_int(local_port, "local_port")
-    remote_port = NumberParser.port_number(remote_port, "remote_port")
-    local_port = self.adb.forward(local_port, remote_port)
-    logging.debug("Forwarded Remote Port: %s:%s <= %s:%s",
-                  self._host_platform.name, local_port, self, remote_port)
-    return local_port
-
-  @override
-  def forward_devtools_port(self, local_port: int,
-                            remote_identifier: str) -> int:
-    local_port = NumberParser.positive_zero_int(local_port, "local_port")
-    local_port = self.adb.forward(
-        local=local_port,
-        remote=remote_identifier,
-        local_protocol="tcp",
-        remote_protocol="localabstract")
-    logging.debug("Forwarded DevTools Port: %s:%s <= %s:%s",
-                  self._host_platform.name, local_port, self, remote_identifier)
-    return local_port
-
-  @override
-  def stop_port_forward(self, local_port: int) -> None:
-    self.adb.forward_remove(local_port, protocol="tcp")
-
-  @override
-  def reverse_port_forward(self, remote_port: int, local_port: int) -> int:
-    remote_port = NumberParser.positive_zero_int(remote_port, "remote_port")
-    local_port = NumberParser.port_number(local_port, "local_port")
-    remote_port = self.adb.reverse(remote_port, local_port, protocol="tcp")
-    logging.debug("Forwarded Local Port: %s:%s => %s:%s", self._host_platform,
-                  local_port, self, remote_port)
-    return remote_port
-
-  @override
-  def stop_reverse_port_forward(self, remote_port: int) -> None:
-    self.adb.reverse_remove(remote_port, protocol="tcp")
-
   @override
   def pull(self, from_path: pth.AnyPath,
            to_path: pth.LocalPath) -> pth.LocalPath:

crossbench/plt/base.py

@@ -35,7 +35,8 @@
 from crossbench.plt import proc_helper
 from crossbench.plt.arch import MachineArch
 from crossbench.plt.bin import Binary
-from crossbench.plt.port_manager import PortManager, PortScope
+from crossbench.plt.port_manager import (LocalPortManager, PortManager,
+                                         PortScope)
 from crossbench.plt.remote import RemotePopen
 
 if TYPE_CHECKING:
@@ -112,7 +113,7 @@ def __init__(self) -> None:
     self._default_port_manager: PortManager = self._create_port_manager()
 
   def _create_port_manager(self) -> PortManager:
-    return PortManager(self)
+    return LocalPortManager(self)
 
   def assert_is_local(self) -> None:
     if self.is_local:
@@ -577,54 +578,6 @@ def default_tmp_dir(self) -> pth.AnyPath:
   def ports(self) -> PortScope:
     return self._default_port_manager.scope
 
-  def port_forward(self, local_port: int, remote_port: int) -> int:
-    """ Forwards a device remote_port to a local port."""
-    # TODO: Migrate forwarding methods to custom PortManager
-    if remote_port != local_port:
-      raise ValueError("Cannot forward a remote port on a local platform.")
-    parse.NumberParser.port_number(local_port, "local_port")
-    self.assert_is_local()
-    return local_port
-
-  def forward_devtools_port(self, local_port: int,
-                            remote_identifier: str) -> int:
-    """Forwards a DevTools debugging port from a remote target to a local port.
-
-    Args:
-      local_port: The local port number to forward to. If 0, a free
-                  port will be chosen by the system.
-      remote_identifier: A string identifying the remote DevTools socket or
-                         service. For Android, this is typically a
-                         localabstract socket name like
-                         "chrome_devtools_remote".
-                         For other platforms, it might be a remote port number
-                         or other service identifier.
-
-    Returns:
-      The local port number that was actually used for forwarding.
-    """
-    raise NotImplementedError(
-        f"forward_devtools_port not implemented for {self}")
-
-  def stop_port_forward(self, local_port: int) -> None:
-    # TODO: Migrate forwarding methods to custom PortManager
-    del local_port
-    self.assert_is_local()
-
-  def reverse_port_forward(self, remote_port: int, local_port: int) -> int:
-    """ Forwards a local port to a device port."""
-    # TODO: Migrate forwarding methods to custom PortManager
-    if remote_port != local_port:
-      raise ValueError("Cannot forward a remote port on a local platform.")
-    parse.NumberParser.port_number(remote_port, "remote_port")
-    self.assert_is_local()
-    return remote_port
-
-  def stop_reverse_port_forward(self, remote_port: int) -> None:
-    # TODO: Migrate forwarding methods to custom PortManager
-    del remote_port
-    self.assert_is_local()
-
   def is_port_used(self, port: int) -> bool:
     self.assert_is_local()
     for conn in psutil.net_connections(kind="inet"):

crossbench/plt/linux_ssh.py

@@ -4,19 +4,15 @@
 
 from __future__ import annotations
 
-import atexit
-import datetime as dt
-import logging
 import shlex
-import subprocess
 from typing import TYPE_CHECKING, Any, Optional
 
 from typing_extensions import override
 
-from crossbench import parse
 from crossbench.plt.arch import MachineArch
 from crossbench.plt.linux import RemoteLinuxPlatform
-from crossbench.plt.ssh import SshPlatformMixin, SshPortManager
+from crossbench.plt.ssh import SshPlatformMixin
+from crossbench.plt.ssh_port_manager import SshPortManager
 
 if TYPE_CHECKING:
   from crossbench.path import AnyPath, LocalPath
@@ -26,18 +22,13 @@
 
 class LinuxSshPlatform(SshPlatformMixin, RemoteLinuxPlatform):
 
-  PORT_FORWARDING_TIMEOUT = dt.timedelta(seconds=10)
 
   def __init__(self, host_platform: Platform, host: str, port: int,
                ssh_port: int, ssh_user: str) -> None:
     super().__init__(host_platform, host, port, ssh_port, ssh_user)
     self._machine: MachineArch | None = None
     self._system_details: dict[str, Any] | None = None
     self._cpu_details: dict[str, Any] | None = None
-    # TOOO: create custom PortManager for linux-ssh
-    self._port_forward_popens: dict[int, subprocess.Popen] = {}
-    self._reverse_port_forward_popens: dict[int, subprocess.Popen] = {}
-    atexit.register(self._stop_all_port_forward)
 
   def _create_port_manager(self) -> PortManager:
     return SshPortManager(self)
@@ -47,7 +38,7 @@ def _create_port_manager(self) -> PortManager:
   def name(self) -> str:
     return "linux_ssh"
 
-  def _build_ssh_cmd(self, *args: CmdArg, shell: bool = False) -> ListCmdArgs:
+  def build_ssh_cmd(self, *args: CmdArg, shell: bool = False) -> ListCmdArgs:
     self.validate_shell_args(args, shell)
     ssh_cmd: ListCmdArgs = [
         "ssh", "-p", f"{self._ssh_port}", f"{self._ssh_user}@{self._host}"
@@ -64,7 +55,7 @@ def _build_ssh_cmd(self, *args: CmdArg, shell: bool = False) -> ListCmdArgs:
 
   @override
   def build_shell_cmd(self, *args: CmdArg, shell: bool = False) -> ListCmdArgs:
-    return self._build_ssh_cmd(*args, shell=shell)
+    return self.build_ssh_cmd(*args, shell=shell)
 
   def processes(self,
                 attrs: Optional[list[str]] = None) -> list[dict[str, Any]]:
@@ -99,62 +90,3 @@ def pull(self, from_path: AnyPath, to_path: LocalPath) -> LocalPath:
     ]
     self._host_platform.sh_stdout(*scp_cmd)
     return to_path
-
-  def port_forward(self, local_port: int, remote_port: int) -> int:
-    local_port, remote_port = self._validate_forwarding_ports(
-        local_port, remote_port)
-    self._port_forward_popens[local_port] = self.host_platform.popen(
-        *self._build_ssh_cmd("-NL", f"{local_port}:localhost:{remote_port}"))
-    self.host_platform.wait_for_port(local_port, self.PORT_FORWARDING_TIMEOUT)
-    logging.debug("Forwarded Remote Port: %s:%s <= %s:%s", self._host_platform,
-                  local_port, self, remote_port)
-    return local_port
-
-  def _validate_forwarding_ports(self, local_port: int,
-                                 remote_port: int) -> tuple[int, int]:
-    local_port = parse.NumberParser.positive_zero_int(local_port, "local_port")
-    remote_port = parse.NumberParser.port_number(remote_port, "remote_port")
-    if not local_port:
-      local_port = self.host_platform.get_free_port()
-    if local_port in self._port_forward_popens:
-      raise RuntimeError(f"Cannot forward local port {local_port} twice.")
-    return local_port, remote_port
-
-  def stop_port_forward(self, local_port: int) -> None:
-    self._port_forward_popens.pop(local_port).terminate()
-
-  def reverse_port_forward(self, remote_port: int, local_port: int) -> int:
-    # TODO: this should likely match with adb, where we support 0
-    # for auto-allocating a remote_port
-    remote_port, local_port = self._validate_reverse_forwarding_ports(
-        remote_port, local_port)
-    self._reverse_port_forward_popens[remote_port] = self.host_platform.popen(
-        *self._build_ssh_cmd("-NR", f"{remote_port}:localhost:{local_port}"))
-    self.wait_for_port(remote_port, self.PORT_FORWARDING_TIMEOUT)
-    logging.debug("Forwarded Local Port: %s:%s => %s:%s", self._host_platform,
-                  local_port, self, remote_port)
-    return remote_port
-
-  def _validate_reverse_forwarding_ports(self, remote_port: int,
-                                         local_port: int) -> tuple[int, int]:
-    remote_port = parse.NumberParser.port_number(remote_port, "remote_port")
-    local_port = parse.NumberParser.positive_zero_int(local_port, "local_port")
-    if not local_port:
-      local_port = self.host_platform.get_free_port()
-    if remote_port in self._reverse_port_forward_popens:
-      raise RuntimeError(f"Cannot forward remote port {remote_port} twice.")
-    return remote_port, local_port
-
-  def stop_reverse_port_forward(self, remote_port: int) -> None:
-    self._reverse_port_forward_popens.pop(remote_port).terminate()
-
-  def _stop_all_port_forward(self) -> None:
-    for port in list(self._port_forward_popens.keys()):
-      self.stop_port_forward(port)
-    for port in list(self._reverse_port_forward_popens.keys()):
-      self.stop_reverse_port_forward(port)
-
-    assert not self._port_forward_popens, (
-        "Did not stop all port forwarding processes.")
-    assert not self._reverse_port_forward_popens, (
-        "Did not stop all reverse port forwarding processes.")