diff --git a/doc/man/exporter.rst b/doc/man/exporter.rst index 598a3a0f2..b0d3b4cdf 100644 --- a/doc/man/exporter.rst +++ b/doc/man/exporter.rst @@ -41,6 +41,8 @@ OPTIONS use fully qualified domain name as default for hostname -d, --debug enable debug mode +--port-range START-END + allocate TCP ports from this inclusive range of at least 1000 ports -i / --isolated ~~~~~~~~~~~~~~~ @@ -72,6 +74,15 @@ In some networks the fully qualified domain name may be needed to reach resource on an exporter. This option changes the default to fqdn when no --hostname is explicitly set. +--port-range +~~~~~~~~~~~~ +Limit automatic TCP port allocation to an inclusive range, for example ``--port-range=4000-5000``. +This applies to all calls to ``get_free_port()`` in the exporter process, including ser2net and ADB servers. +The range must contain at least 1000 ports between 1 and 65535. +Ports are tried in random order, and occupied ports are skipped. +Allocation fails if no port in the range is available. +Without this option, the operating system chooses an available port. + CONFIGURATION ------------- The exporter uses a YAML configuration file which defines groups of related diff --git a/doc/usage.rst b/doc/usage.rst index 8e5419ecf..8caed10c8 100644 --- a/doc/usage.rst +++ b/doc/usage.rst @@ -832,6 +832,15 @@ resources. See :doc:`man/client`, :doc:`man/device-config` and :doc:`man/exporter` for more information. +Exporter port range +~~~~~~~~~~~~~~~~~~~ + +On networks with a fixed set of allowed TCP ports, the exporter option ``--port-range=4000-5000`` +restricts automatic port allocation to that inclusive range. +The range applies to serial exports and ADB servers, rather than an individual resource. +It must contain at least 1000 ports, which are tried in random order on each allocation. +See :doc:`man/exporter` for details. + Advanced CLI features ~~~~~~~~~~~~~~~~~~~~~ diff --git a/labgrid/remote/exporter.py b/labgrid/remote/exporter.py index 50fa700b4..aed8c72f7 100755 --- a/labgrid/remote/exporter.py +++ b/labgrid/remote/exporter.py @@ -29,7 +29,7 @@ from .config import ResourceConfig from .common import ResourceEntry, get_client_credentials, queue_as_aiter from .generated import labgrid_coordinator_pb2, labgrid_coordinator_pb2_grpc -from ..util import get_free_port, labgrid_version +from ..util import get_free_port, helper, labgrid_version exports: Dict[str, Type[ResourceEntry]] = {} @@ -1190,6 +1190,18 @@ def _stop(): await exporter.run() +def parse_port_range(value): + try: + start, end = (int(port) for port in value.split("-")) + except ValueError as exc: + raise argparse.ArgumentTypeError("expected START-END") from exc + if not 1 <= start <= end <= 65535: + raise argparse.ArgumentTypeError("ports must satisfy 1 <= START <= END <= 65535") + if end - start + 1 < 1000: + raise argparse.ArgumentTypeError("port range must contain at least 1000 ports") + return range(start, end + 1) + + def main(): parser = argparse.ArgumentParser() parser.add_argument( @@ -1233,12 +1245,19 @@ def main(): help="enable isolated mode (always request SSH forwards)", ) parser.add_argument("--pystuck", action="store_true", help="enable pystuck") + parser.add_argument( + "--port-range", + metavar="START-END", + type=parse_port_range, + help="allocate TCP ports from this inclusive range (at least 1000 ports)", + ) parser.add_argument( "--pystuck-port", metavar="PORT", type=int, default=6667, help="use a different pystuck port than 6667" ) parser.add_argument("resources", metavar="RESOURCES", type=str, help="resource config file name") args = parser.parse_args() + helper.set_port_range(args.port_range) logging.basicConfig(level=logging.DEBUG if args.debug else logging.INFO) diff --git a/labgrid/util/helper.py b/labgrid/util/helper.py index 3c5ea4559..a1bac0591 100644 --- a/labgrid/util/helper.py +++ b/labgrid/util/helper.py @@ -3,6 +3,7 @@ import logging import os import pty +import random import re import select import subprocess @@ -15,11 +16,33 @@ from .timeout import Timeout re_vt100 = re.compile(r"(\x1b\[|\x9b)[^@-_a-z]*[@-_a-z]|\x1b[@-_a-z]") +_port_range = None + + +def set_port_range(port_range): + """Set the default range used when allocating free TCP ports.""" + global _port_range + _port_range = port_range + def get_free_port(): """Helper function to always return an unused port.""" with closing(socket(AF_INET, SOCK_STREAM)) as s: - s.bind(('', 0)) + if _port_range is not None: + ports = random.sample(_port_range, len(_port_range)) + while ports: + port = ports.pop() + try: + s.bind(('', port)) + except OSError as exc: + if exc.errno != errno.EADDRINUSE: + raise + else: + break + else: + raise OSError(errno.EADDRINUSE, "No free ports in the configured port range") + else: + s.bind(('', 0)) return s.getsockname()[1] diff --git a/man/labgrid-exporter.1 b/man/labgrid-exporter.1 index dcfbcd7a5..5589aa816 100644 --- a/man/labgrid-exporter.1 +++ b/man/labgrid-exporter.1 @@ -71,6 +71,9 @@ use fully qualified domain name as default for hostname .TP .B \-d\fP,\fB \-\-debug enable debug mode +.TP +.BI \-\-port\-range \ START\-END +allocate TCP ports from this inclusive range of at least 1000 ports .UNINDENT .SS \-i / \-\-isolated .sp @@ -98,6 +101,14 @@ override this default with another name (or an IP address). In some networks the fully qualified domain name may be needed to reach resources on an exporter. This option changes the default to fqdn when no \-\-hostname is explicitly set. +.SS \-\-port\-range +.sp +Limit automatic TCP port allocation to an inclusive range, for example \fB\-\-port\-range=4000\-5000\fP\&. +This applies to all calls to \fBget_free_port()\fP in the exporter process, including ser2net and ADB servers. +The range must contain at least 1000 ports between 1 and 65535. +Ports are tried in random order, and occupied ports are skipped. +Allocation fails if no port in the range is available. +Without this option, the operating system chooses an available port. .SS CONFIGURATION .sp The exporter uses a YAML configuration file which defines groups of related diff --git a/tests/test_port_range.py b/tests/test_port_range.py new file mode 100644 index 000000000..935a896d1 --- /dev/null +++ b/tests/test_port_range.py @@ -0,0 +1,126 @@ +import asyncio +import errno +import random +import sys +from unittest.mock import Mock, call + +import pytest + +from labgrid.remote import exporter +from labgrid.util import helper, ssh + + +@pytest.fixture +def port_socket(monkeypatch): + monkeypatch.setattr(helper, "_port_range", None, raising=False) + monkeypatch.setattr(helper, "random", random, raising=False) + sock = Mock() + monkeypatch.setattr(helper, "socket", Mock(return_value=sock)) + return sock + + +def test_get_free_port_default(port_socket): + port_socket.getsockname.return_value = ("", 32768) + + assert helper.get_free_port() == 32768 + + port_socket.bind.assert_called_once_with(("", 0)) + port_socket.close.assert_called_once_with() + + +def test_get_free_port_range(monkeypatch, port_socket): + monkeypatch.setattr(helper, "_port_range", range(4000, 4003)) + sample = Mock(return_value=[4001, 4002, 4000]) + monkeypatch.setattr(helper.random, "sample", sample) + port_socket.bind.side_effect = [OSError(errno.EADDRINUSE, "occupied"), None] + port_socket.getsockname.return_value = ("", 4002) + + assert helper.get_free_port() == 4002 + + sample.assert_called_once_with(range(4000, 4003), 3) + assert port_socket.bind.call_args_list == [call(("", 4000)), call(("", 4002))] + port_socket.close.assert_called_once_with() + + +def test_get_free_port_range_exhausted(monkeypatch, port_socket): + monkeypatch.setattr(helper, "_port_range", range(4000, 4002)) + monkeypatch.setattr(helper.random, "sample", lambda ports, count: list(ports)) + port_socket.bind.side_effect = OSError(errno.EADDRINUSE, "occupied") + + with pytest.raises(OSError) as excinfo: + helper.get_free_port() + + assert excinfo.value.errno == errno.EADDRINUSE + assert port_socket.bind.call_args_list == [call(("", 4001)), call(("", 4000))] + port_socket.close.assert_called_once_with() + + +def test_get_free_port_range_resampled(monkeypatch, port_socket): + monkeypatch.setattr(helper, "_port_range", range(4000, 4002)) + sample = Mock(side_effect=[[4000, 4001], [4001, 4000]]) + monkeypatch.setattr(helper.random, "sample", sample) + port_socket.getsockname.side_effect = [("", 4001), ("", 4000)] + + assert helper.get_free_port() == 4001 + assert helper.get_free_port() == 4000 + + assert sample.call_count == 2 + assert port_socket.bind.call_args_list == [call(("", 4001)), call(("", 4000))] + + +@pytest.mark.parametrize("error", [errno.EACCES, errno.EMFILE]) +def test_get_free_port_range_bind_error(monkeypatch, port_socket, error): + monkeypatch.setattr(helper, "_port_range", range(4000, 4002)) + monkeypatch.setattr(helper.random, "sample", lambda ports, count: list(ports)) + port_socket.bind.side_effect = OSError(error, "bind failed") + + with pytest.raises(OSError) as excinfo: + helper.get_free_port() + + assert excinfo.value.errno == error + port_socket.bind.assert_called_once_with(("", 4001)) + port_socket.close.assert_called_once_with() + + +@pytest.mark.parametrize("value", ["1-1000", "4000-4999", "64536-65535", "1-65535"]) +def test_parse_port_range(value): + start, end = (int(part) for part in value.split("-")) + assert exporter.parse_port_range(value) == range(start, end + 1) + + +@pytest.mark.parametrize( + "value", + ["4000", "4000:5000", "a-b", "4000-5000-6000", "0-1000", "64536-65536", "5000-4000", "4000-4998"], +) +def test_exporter_invalid_port_range(monkeypatch, capsys, value): + monkeypatch.setattr(sys, "argv", ["labgrid-exporter", f"--port-range={value}", "exports.yaml"]) + + with pytest.raises(SystemExit) as excinfo: + exporter.main() + + assert excinfo.value.code == 2 + assert "argument --port-range:" in capsys.readouterr().err + + +@pytest.mark.parametrize("args, expected", [(["--port-range=4000-5000"], range(4000, 5001)), ([], None)]) +def test_exporter_port_range(monkeypatch, port_socket, args, expected): + monkeypatch.setattr(helper, "_port_range", range(6000, 7001)) + monkeypatch.setattr(sys, "argv", ["labgrid-exporter", *args, "exports.yaml"]) + monkeypatch.setattr(exporter, "reexec", False) + monkeypatch.setattr(exporter.asyncio, "run", lambda coroutine, **kwargs: coroutine.close()) + monkeypatch.setattr(exporter.asyncio, "set_event_loop", lambda loop: None) + loop = asyncio.new_event_loop() + monkeypatch.setattr(exporter.asyncio, "new_event_loop", lambda: loop) + try: + exporter.main() + finally: + loop.close() + + assert helper._port_range == expected + monkeypatch.setattr(helper.random, "sample", lambda ports, count: list(ports)) + port_socket.getsockname.return_value = ("", 5000) + + assert exporter.get_free_port() == 5000 + assert ssh.get_free_port() == 5000 + port = 5000 if expected is not None else 0 + assert port_socket.bind.call_args_list == [call(("", port)), call(("", port))]