protocol: add PortForwardProtocol

doc/configuration.rst

@@ -2084,7 +2084,7 @@ Arguments:
 SSHDriver
 ~~~~~~~~~
 An :any:`SSHDriver` requires a `NetworkService`_ resource and allows the
-execution of commands and file upload via network.
+execution of commands, file upload and TCP port forwarding via network.
 It uses SSH's ``ServerAliveInterval`` option to detect failed connections.
 
 If a shared SSH connection to the target is already open, it will reuse it when
@@ -2099,6 +2099,7 @@ Binds to:
 Implements:
   - :any:`CommandProtocol`
   - :any:`FileTransferProtocol`
+  - :any:`PortForwardProtocol`
 
 .. code-block:: yaml
 
@@ -3922,10 +3923,13 @@ method removes a file by name.
 
 ADBDriver
 ~~~~~~~~~
-The :any:`ADBDriver` allows interaction with ADB devices. It allows the 
-execution of commands, transfer of files, and rebooting of the device.
+The :any:`ADBDriver` allows interaction with ADB devices. It allows the
+execution of commands, transfer of files, TCP port forwarding and rebooting of
+the device.
 
 It can interact with both USB and TCP adb devices.
+Port forwarding with a :any:`RemoteUSBADBDevice` requires SSH port forwarding
+access to the exporter.
 
 Binds to:
   iface:
@@ -3936,6 +3940,7 @@ Binds to:
 Implements:
   - :any:`CommandProtocol`
   - :any:`FileTransferProtocol`
+  - :any:`PortForwardProtocol`
   - :any:`ResetProtocol`
 
 .. _conf-strategies:

labgrid/driver/adb.py

@@ -1,10 +1,11 @@
+import contextlib
 import subprocess
 from enum import Enum
 
 import attr
 
 from ..factory import target_factory
-from ..protocol import CommandProtocol, FileTransferProtocol, ResetProtocol
+from ..protocol import CommandProtocol, FileTransferProtocol, PortForwardProtocol, ResetProtocol
 from ..resource.adb import NetworkADBDevice, RemoteUSBADBDevice, USBADBDevice
 from ..step import step
 from ..util.proxy import proxymanager
@@ -25,8 +26,8 @@ class ADBRebootMode(Enum):
 
 @target_factory.reg_driver
 @attr.s(eq=False)
-class ADBDriver(CommandMixin, Driver, CommandProtocol, FileTransferProtocol, ResetProtocol):
-    """ADB driver to execute commands, transfer files and reset devices via ADB."""
+class ADBDriver(CommandMixin, Driver, CommandProtocol, FileTransferProtocol, ResetProtocol, PortForwardProtocol):
+    """ADB driver for commands, file transfers, resets and port forwarding."""
 
     bindings = {"device": {"USBADBDevice", "RemoteUSBADBDevice", "NetworkADBDevice"}}
 
@@ -129,6 +130,65 @@ def get(self, filename: str, destination: str, timeout: float | None = None):
             check=True,
         )
 
+    # Port Forward Protocol
+
+    @contextlib.contextmanager
+    def _forward(self, command, listen_port, connect_port):
+        result = subprocess.run(
+            [
+                *self._base_command,
+                command,
+                "--no-rebind",
+                f"tcp:{listen_port}",
+                f"tcp:{connect_port}",
+            ],
+            stdout=subprocess.PIPE,
+            text=True,
+            timeout=ADB_TIMEOUT,
+            check=True,
+        )
+        try:
+            if listen_port == 0:
+                listen_port = int(result.stdout.strip())
+
+            yield listen_port
+        finally:
+            subprocess.run(
+                [*self._base_command, command, "--remove", f"tcp:{listen_port}"],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+                timeout=ADB_TIMEOUT,
+                check=False,
+            )
+
+    @Driver.check_active
+    @contextlib.contextmanager
+    def local_forward(self, remote_port: int, *, local_port: int = 0):
+        if isinstance(self.device, RemoteUSBADBDevice):
+            with (
+                self._forward("forward", 0, remote_port) as rp,
+                proxymanager.local_forward(self.device, "localhost", rp, local_port=local_port) as lp,
+            ):
+                yield lp
+            return
+
+        with self._forward("forward", local_port, remote_port) as lp:
+            yield lp
+
+    @Driver.check_active
+    @contextlib.contextmanager
+    def remote_forward(self, local_port: int, *, remote_port: int = 0):
+        if isinstance(self.device, RemoteUSBADBDevice):
+            with (
+                proxymanager.remote_forward(self.device, local_port) as lp,
+                self._forward("reverse", remote_port, lp) as rp,
+            ):
+                yield rp
+            return
+
+        with self._forward("reverse", remote_port, local_port) as rp:
+            yield rp
+
     # Reset Protocol
 
     @Driver.check_active

labgrid/driver/sshdriver.py

@@ -1,4 +1,4 @@
-"""The SSHDriver uses SSH as a transport to implement CommandProtocol and FileTransferProtocol"""
+"""The SSHDriver uses SSH as a transport to implement CommandProtocol, FileTransferProtocol and PortForwardProtocol"""
 import contextlib
 import os
 import re
@@ -8,12 +8,13 @@
 import subprocess
 import tempfile
 import time
+import warnings
 from functools import cached_property
 
 import attr
 
 from ..factory import target_factory
-from ..protocol import CommandProtocol, FileTransferProtocol
+from ..protocol import CommandProtocol, FileTransferProtocol, PortForwardProtocol
 from .commandmixin import CommandMixin
 from .common import Driver
 from ..step import step
@@ -26,10 +27,10 @@
 
 @target_factory.reg_driver
 @attr.s(eq=False)
-class SSHDriver(CommandMixin, Driver, CommandProtocol, FileTransferProtocol):
+class SSHDriver(CommandMixin, Driver, CommandProtocol, FileTransferProtocol, PortForwardProtocol):
     """SSHDriver - Driver to execute commands via SSH"""
     bindings = {"networkservice": "NetworkService", }
-    priorities = {CommandProtocol: 10, FileTransferProtocol: 10}
+    priorities = {CommandProtocol: 10, FileTransferProtocol: 10, PortForwardProtocol: 10}
     keyfile = attr.ib(default="", validator=attr.validators.instance_of(str))
     stderr_merge = attr.ib(default=False, validator=attr.validators.instance_of(bool))
     connection_timeout = attr.ib(default=float(get_ssh_connect_timeout()), validator=attr.validators.instance_of(float))
@@ -267,9 +268,9 @@ def _forward(self, forward):
                self.networkservice.address
                ]
         self.logger.debug("Running command: %s", cmd)
-        subprocess.run(cmd, check=True)
+        result = subprocess.run(cmd, check=True, stdout=subprocess.PIPE, text=True)
         try:
-            yield
+            yield result.stdout
         finally:
             cmd = [self._ssh, *self.ssh_prefix,
                    "-O", "cancel", forward,
@@ -282,7 +283,7 @@ def _forward(self, forward):
 
     @Driver.check_active
     @contextlib.contextmanager
-    def forward_local_port(self, remoteport, localport=None):
+    def local_forward(self, remote_port, *, local_port=0):
         """Forward a local port to a remote port on the target
 
         A context manager that keeps a local port forwarded to a remote port as
@@ -291,43 +292,71 @@ def forward_local_port(self, remoteport, localport=None):
         on the target device
 
         usage:
-            with ssh.forward_local_port(8080) as localport:
-                # Use localhost:localport here to connect to port 8080 on the
+            with ssh.local_forward(8080) as local_port:
+                # Use localhost:local_port here to connect to port 8080 on the
                 # target
 
         returns:
-        localport
+        local_port
         """
         if not self._check_keepalive():
             raise ExecutionError("Keepalive no longer running")
 
-        if localport is None:
-            localport = get_free_port()
+        # OpenSSH does not support dynamic port allocation for local forwards.
+        if local_port == 0:
+            local_port = get_free_port()
 
-        forward = f"-L{localport:d}:localhost:{remoteport:d}"
+        forward = f"-L{local_port:d}:localhost:{remote_port:d}"
         with self._forward(forward):
-            yield localport
+            yield local_port
 
     @Driver.check_active
     @contextlib.contextmanager
-    def forward_remote_port(self, remoteport, localport):
+    def forward_local_port(self, remoteport, localport=None):
+        warnings.warn(
+            "SSHDriver.forward_local_port() is deprecated, use local_forward() instead",
+            DeprecationWarning,
+            stacklevel=3,
+        )
+        local_port = 0 if localport is None else localport
+        with self.local_forward(remoteport, local_port=local_port) as local_port:
+            yield local_port
+
+    @Driver.check_active
+    @contextlib.contextmanager
+    def remote_forward(self, local_port, *, remote_port=0):
         """Forward a remote port on the target to a local port
 
         A context manager that keeps a remote port forwarded to a local port as
-        long as the context remains valid. A connection can be made to the
-        remote on the target device will be forwarded to the returned local
-        port on localhost
+        long as the context remains valid. Connections to the returned port on
+        the target device are forwarded to the local port on localhost.
 
         usage:
-            with ssh.forward_remote_port(8080, 8081) as localport:
+            with ssh.remote_forward(8081, remote_port=8080) as remote_port:
                 # Connections to port 8080 on the target will be redirected to
                 # localhost:8081
+
+        returns:
+        remote_port
         """
         if not self._check_keepalive():
             raise ExecutionError("Keepalive no longer running")
 
-        forward = f"-R{remoteport:d}:localhost:{localport:d}"
-        with self._forward(forward):
+        forward = f"-R{remote_port:d}:localhost:{local_port:d}"
+        with self._forward(forward) as stdout:
+            if remote_port == 0:
+                remote_port = int(stdout.strip())
+            yield remote_port
+
+    @Driver.check_active
+    @contextlib.contextmanager
+    def forward_remote_port(self, remoteport, localport):
+        warnings.warn(
+            "SSHDriver.forward_remote_port() is deprecated, use remote_forward() instead",
+            DeprecationWarning,
+            stacklevel=3,
+        )
+        with self.remote_forward(localport, remote_port=remoteport):
             yield
 
     @Driver.check_active

labgrid/protocol/__init__.py

@@ -3,6 +3,7 @@
 from .consoleprotocol import ConsoleProtocol
 from .linuxbootprotocol import LinuxBootProtocol
 from .powerprotocol import PowerProtocol
+from .portforwardprotocol import PortForwardProtocol
 from .filetransferprotocol import FileTransferProtocol
 from .infoprotocol import InfoProtocol
 from .digitaloutputprotocol import DigitalOutputProtocol

labgrid/protocol/portforwardprotocol.py

@@ -0,0 +1,20 @@
+import abc
+
+
+class PortForwardProtocol(abc.ABC):
+    """Interface for forwarding TCP ports between the local machine and target.
+
+    Methods return context managers which keep the forwarding active for the
+    duration of the context and yield the allocated listening port. A listening
+    port of zero requests automatic allocation.
+    """
+
+    @abc.abstractmethod
+    def local_forward(self, remote_port: int, *, local_port: int = 0):
+        """Forward a local port to a remote port on the target."""
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def remote_forward(self, local_port: int, *, remote_port: int = 0):
+        """Forward a remote port on the target to a local port."""
+        raise NotImplementedError

labgrid/remote/client.py

@@ -1310,6 +1310,16 @@ def _get_ssh(self):
             drv = self._get_driver_or_new(target, "SSHDriver", name=resource.name)
             return drv
 
+    def _get_port_forward(self):
+        place = self.get_acquired_place()
+        target = self._get_target(place)
+
+        try:
+            return target.get_driver("PortForwardProtocol", name=self.args.name)
+        except NoDriverFoundError:
+            self.logger.warning("no PortForwardProtocol driver found, falling back to implicit SSH port forwarding")
+            return self._get_ssh()
+
     def ssh(self):
         drv = self._get_ssh()
 
@@ -1341,16 +1351,16 @@ def forward(self):
             print("Nothing to forward", file=sys.stderr)
             return
 
-        drv = self._get_ssh()
+        drv = self._get_port_forward()
 
         with contextlib.ExitStack() as stack:
             for local, remote in self.args.local:
-                localport = stack.enter_context(drv.forward_local_port(remote, localport=local))
+                localport = stack.enter_context(drv.local_forward(remote, local_port=local or 0))
                 print(f"Forwarding local port {localport:d} to remote port {remote:d}")
 
             for local, remote in self.args.remote:
-                stack.enter_context(drv.forward_remote_port(remote, local))
-                print(f"Forwarding remote port {remote:d} to local port {local:d}")
+                allocated = stack.enter_context(drv.remote_forward(local, remote_port=remote))
+                print(f"Forwarding remote port {allocated:d} to local port {local:d}")
 
             try:
                 print("Waiting for CTRL+C...")

labgrid/util/proxy.py

@@ -1,3 +1,4 @@
+import contextlib
 import os
 
 from urllib.parse import urlsplit, urlunsplit, urlparse
@@ -25,6 +26,46 @@ def force_proxy(cls, force_proxy):
         assert isinstance(force_proxy, str)
         cls._force_proxy = force_proxy
 
+    @classmethod
+    @contextlib.contextmanager
+    def local_forward(cls, res, remote_host, remote_port, *, local_port=0):
+        assert isinstance(res, Resource)
+
+        extra = getattr(res, "extra", {})
+        host = extra.get("proxy") if extra.get("proxy_required") else res.host
+        connection = sshmanager.get(host)
+        local_port = connection.add_port_forward(
+            remote_host,
+            remote_port,
+            None if local_port == 0 else local_port,
+        )
+        try:
+            yield local_port
+        finally:
+            connection.remove_port_forward(remote_host, remote_port)
+
+    @classmethod
+    @contextlib.contextmanager
+    def remote_forward(cls, res, local_port, *, remote_port=0):
+        assert isinstance(res, Resource)
+
+        extra = getattr(res, "extra", {})
+        host = extra.get("proxy") if extra.get("proxy_required") else res.host
+        connection = sshmanager.get(host)
+        remote_port = connection.add_remote_port_forward(
+            remote_port,
+            local_port,
+            "localhost",
+        )
+        try:
+            yield remote_port
+        finally:
+            connection.remove_remote_port_forward(
+                remote_port,
+                local_port,
+                "localhost",
+            )
+
     @classmethod
     def get_host_and_port(cls, res, *, default_port=None, force_port=None):
         """get host and port for a proxy connection from a Resource