diff --git a/doc/configuration.rst b/doc/configuration.rst index 72ac124fd..a3d56c7ab 100644 --- a/doc/configuration.rst +++ b/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: diff --git a/labgrid/driver/adb.py b/labgrid/driver/adb.py index 8db3ce2b5..151a2bf56 100644 --- a/labgrid/driver/adb.py +++ b/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 diff --git a/labgrid/driver/sshdriver.py b/labgrid/driver/sshdriver.py index 110a4f707..1d60aae57 100644 --- a/labgrid/driver/sshdriver.py +++ b/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 diff --git a/labgrid/protocol/__init__.py b/labgrid/protocol/__init__.py index 0ac225622..a7ba040c0 100644 --- a/labgrid/protocol/__init__.py +++ b/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 diff --git a/labgrid/protocol/portforwardprotocol.py b/labgrid/protocol/portforwardprotocol.py new file mode 100644 index 000000000..43ba72215 --- /dev/null +++ b/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 diff --git a/labgrid/remote/client.py b/labgrid/remote/client.py index 1576d39f1..a02224e19 100755 --- a/labgrid/remote/client.py +++ b/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...") diff --git a/labgrid/util/proxy.py b/labgrid/util/proxy.py index 9df168dcf..dfce3d5fe 100644 --- a/labgrid/util/proxy.py +++ b/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 diff --git a/labgrid/util/ssh.py b/labgrid/util/ssh.py index 62bbf4cbb..b688f7acd 100644 --- a/labgrid/util/ssh.py +++ b/labgrid/util/ssh.py @@ -139,7 +139,7 @@ class SSHConnection: validator=attr.validators.instance_of(str) ) _l_forwards = attr.ib(init=False, default=attr.Factory(dict)) - _r_forwards = attr.ib(init=False, default=attr.Factory(set)) + _r_forwards = attr.ib(init=False, default=attr.Factory(dict)) def __attrs_post_init__(self): self._logger = logging.getLogger(f"{self}") @@ -185,11 +185,11 @@ def _run_socket_command(self, command, forward=None): complete_cmd.append(item) complete_cmd.append(self.host) self._logger.debug("Running control command: %s", " ".join(complete_cmd)) - subprocess.check_call( + return subprocess.check_output( complete_cmd, stdin=subprocess.DEVNULL, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, + stderr=subprocess.PIPE, + text=True, timeout=2, ) @@ -381,8 +381,16 @@ def add_remote_port_forward(self, remote_port, local_port, remote_bind=None): forward = f"-R{remote_bind}:{remote_port:d}:localhost:{local_port:d}" - self._run_socket_command("forward", [forward]) - self._r_forwards.add(forward) + stdout = self._run_socket_command("forward", [forward]) + if remote_port == 0: + try: + remote_port = int(stdout.strip()) + except ValueError: + self._run_socket_command("cancel", [forward]) + raise + + self._r_forwards[remote_port, local_port, remote_bind] = forward + return remote_port @_check_connected def remove_remote_port_forward(self, remote_port, local_port, remote_bind=None): @@ -390,9 +398,7 @@ def remove_remote_port_forward(self, remote_port, local_port, remote_bind=None): if remote_bind is None: remote_bind = "*" - forward = f"-R{remote_bind}:{remote_port:d}:localhost:{local_port:d}" - - self._r_forwards.remove(forward) + forward = self._r_forwards.pop((remote_port, local_port, remote_bind)) self._run_socket_command("cancel", [forward]) def connect(self): @@ -536,7 +542,7 @@ def cleanup(self): self._run_socket_command("cancel", [f"-L{local_port}:{destination}"]) self._l_forwards.clear() # cancel remote forwards - for forward in self._r_forwards: + for forward in self._r_forwards.values(): self._run_socket_command("cancel", [forward]) self._r_forwards.clear() self.disconnect() diff --git a/tests/test_client.py b/tests/test_client.py index 8aa7835c9..c3aefbacc 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -642,3 +642,45 @@ def test_same_name_resources(place, exporter, tmpdir): spawn.expect(pexpect.EOF) spawn.close() assert spawn.exitstatus == 0, spawn.before.strip() + +def test_forward(place, tmpdir): + # Port forwarding requires an acquired place. + with pexpect.spawn('python -m labgrid.remote.client -p test acquire') as spawn: + spawn.expect(pexpect.EOF) + spawn.close() + assert spawn.exitstatus == 0, spawn.before.strip() + + adb = tmpdir.join('adb') + adb.write( + """#!/bin/sh +echo 42000 +""" + ) + adb.chmod(0o755) + + config = tmpdir.join('config.yaml') + config.write( + f""" +targets: + main: + resources: + RemotePlace: + name: test + NetworkADBDevice: + host: localhost + port: 5555 + drivers: + ADBDriver: {{}} +tools: + adb: "{adb}" +""" + ) + + with pexpect.spawn(f'python -m labgrid.remote.client -c {config} -p test forward -L 8080') as spawn: + spawn.expect('Forwarding local port 42000 to remote port 8080') + spawn.expect(r'Waiting for CTRL\+C...') + spawn.sendcontrol('c') + spawn.expect('Exiting...') + spawn.expect(pexpect.EOF) + spawn.close() + assert spawn.exitstatus == 0, spawn.before.strip() diff --git a/tests/test_sshdriver.py b/tests/test_sshdriver.py index 4c233a834..73044e6e1 100644 --- a/tests/test_sshdriver.py +++ b/tests/test_sshdriver.py @@ -165,7 +165,7 @@ def test_local_port_forward(ssh_localhost, tmpdir): remoteport = get_free_port() test_string = "Hello World" - with ssh_localhost.forward_local_port(remoteport) as localport: + with ssh_localhost.local_forward(remoteport) as localport: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as send_socket: server_socket.bind(("127.0.0.1", remoteport)) @@ -179,13 +179,23 @@ def test_local_port_forward(ssh_localhost, tmpdir): assert client_socket.recv(16).decode("utf-8") == test_string +def test_deprecated_local_port_forward(ssh_driver_mocked_and_activated, mocker): + s = ssh_driver_mocked_and_activated + s.local_forward = mocker.MagicMock() + s.local_forward.return_value.__enter__.return_value = 42000 + + with pytest.warns(DeprecationWarning, match="use local_forward"): + with s.forward_local_port(remoteport=1234, localport=None) as local_port: + assert local_port == 42000 + s.local_forward.assert_called_once_with(1234, local_port=0) + @pytest.mark.sshusername -def test_local_remote_forward(ssh_localhost, tmpdir): +def test_remote_port_forward(ssh_localhost, tmpdir): remoteport = get_free_port() localport = get_free_port() test_string = "Hello World" - with ssh_localhost.forward_remote_port(remoteport, localport): + with ssh_localhost.remote_forward(localport, remote_port=remoteport): with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as send_socket: server_socket.bind(("127.0.0.1", localport)) @@ -199,6 +209,34 @@ def test_local_remote_forward(ssh_localhost, tmpdir): assert client_socket.recv(16).decode("utf-8") == test_string +@pytest.mark.sshusername +def test_remote_port_forward_auto_allocation(ssh_localhost, tmpdir): + localport = get_free_port() + test_string = "Hello World" + + with ssh_localhost.remote_forward(localport) as remoteport: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server_socket: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as send_socket: + server_socket.bind(("127.0.0.1", localport)) + server_socket.listen(1) + + send_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + send_socket.connect(("127.0.0.1", remoteport)) + + client_socket, address = server_socket.accept() + send_socket.send(test_string.encode('utf-8')) + + assert client_socket.recv(16).decode("utf-8") == test_string + +def test_deprecated_remote_port_forward(ssh_driver_mocked_and_activated, mocker): + s = ssh_driver_mocked_and_activated + s.remote_forward = mocker.MagicMock() + + with pytest.warns(DeprecationWarning, match="use remote_forward"): + with s.forward_remote_port(remoteport=42000, localport=1234) as result: + assert result is None + s.remote_forward.assert_called_once_with(1234, remote_port=42000) + @pytest.mark.sshusername def test_unix_socket_forward(ssh_localhost, tmpdir): diff --git a/tests/test_util.py b/tests/test_util.py index 29ec2555f..74fb2dc41 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -182,7 +182,25 @@ def test_sshconnection_port_remote_forward_add_remove(connection_localhost): assert client_socket.recv(16).decode("utf-8") == test_string connection_localhost.remove_remote_port_forward(rport, lport) - assert connection_localhost._r_forwards == set() + assert connection_localhost._r_forwards == {} + +@pytest.mark.localsshmanager +def test_sshconnection_port_remote_forward_auto_allocation(connection_localhost): + lport = get_free_port() + test_string = "Hello World" + + rport = connection_localhost.add_remote_port_forward(0, lport, "localhost") + server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + server_socket.bind(("127.0.0.1", lport)) + server_socket.listen(1) + send_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + send_socket.connect(("127.0.0.1", rport)) + client_socket, address = server_socket.accept() + send_socket.send(test_string.encode("utf-8")) + + assert client_socket.recv(16).decode("utf-8") == test_string + connection_localhost.remove_remote_port_forward(rport, lport, "localhost") + assert connection_localhost._r_forwards == {} @pytest.mark.localsshmanager