feat: support port ranges and endpoints

ops/_private/harness.py

@@ -3000,15 +3000,39 @@ def secret_remove(self, id: str, *, revision: int | None = None) -> None:
         else:
             self._secrets = [s for s in self._secrets if not self._secret_ids_are_equal(s.id, id)]
 
-    def open_port(self, protocol: str, port: int | None = None):
+    def open_port(
+        self,
+        protocol: str,
+        port: int | None = None,
+        *,
+        to_port: int | None = None,
+        endpoints: Sequence[str] = '*',
+    ):
         self._check_protocol_and_port(protocol, port)
         protocol_lit = cast('Literal["tcp", "udp", "icmp"]', protocol)
-        self._opened_ports.add(model.Port(protocol_lit, port))
-
-    def close_port(self, protocol: str, port: int | None = None):
+        endpoints = tuple(endpoints) if endpoints != '*' else '*'
+        # This only really works for the happy path.
+        # We should really be checking for overlapping port ranges (and erroring),
+        # and merging with existing entries for endpoints, but since harness is deprecated,
+        # and this is only needed for charms adopting this new feature, we can get away with it.
+        self._opened_ports.add(model.Port(protocol_lit, port, to_port, endpoints=endpoints))
+
+    def close_port(
+        self,
+        protocol: str,
+        port: int | None = None,
+        *,
+        to_port: int | None = None,
+        endpoints: Sequence[str] = '*',
+    ):
         self._check_protocol_and_port(protocol, port)
         protocol_lit = cast('Literal["tcp", "udp", "icmp"]', protocol)
-        self._opened_ports.discard(model.Port(protocol_lit, port))
+        endpoints = tuple(endpoints) if endpoints != '*' else '*'
+        # This only really works for the happy path.
+        # We should really be checking for overlapping port ranges (and erroring),
+        # and merging with existing entries for endpoints, but since harness is deprecated,
+        # and this is only needed for charms adopting this new feature, we can get away with it.
+        self._opened_ports.discard(model.Port(protocol_lit, port, to_port, endpoints=endpoints))
 
     def opened_ports(self) -> set[model.Port]:
         return set(self._opened_ports)

ops/hookcmds/_port.py

@@ -29,22 +29,22 @@ def close_port(
     *,
     to_port: int | None = None,
     endpoints: str | Iterable[str] | None = None,
-) -> None: ...
+) -> str | None: ...
 @overload
 def close_port(
     protocol: str | None,
     port: int,
     *,
     to_port: int | None = None,
     endpoints: str | Iterable[str] | None = None,
-) -> None: ...
+) -> str | None: ...
 def close_port(
     protocol: str | None = None,
     port: int | None = None,
     *,
     to_port: int | None = None,
     endpoints: str | Iterable[str] | None = None,
-):
+) -> str | None:
     """Register a request to close a port or port range.
 
     For more details, see:
@@ -58,13 +58,16 @@ def close_port(
     if port is None:
         if protocol is None:
             raise TypeError('You must provide a port or protocol.')
+        if to_port is not None:
+            raise TypeError('to_port cannot be specified if port is not specified')
         args.append(protocol)
     else:
         port_arg = f'{port}-{to_port}' if to_port is not None else str(port)
         if protocol is not None:
             port_arg = f'{port_arg}/{protocol}'
         args.append(port_arg)
-    run('close-port', *args)
+    result = run('close-port', *args).strip()
+    return result or None
 
 
 @overload
@@ -74,22 +77,22 @@ def open_port(
     *,
     to_port: int | None = None,
     endpoints: str | Iterable[str] | None = None,
-) -> None: ...
+) -> str | None: ...
 @overload
 def open_port(
     protocol: str | None,
     port: int,
     *,
     to_port: int | None = None,
     endpoints: str | Iterable[str] | None = None,
-) -> None: ...
+) -> str | None: ...
 def open_port(
     protocol: str | None = None,
     port: int | None = None,
     *,
     to_port: int | None = None,
     endpoints: str | Iterable[str] | None = None,
-):
+) -> str | None:
     """Register a request to open a port or port range.
 
     For more details, see:
@@ -113,13 +116,21 @@ def open_port(
     if port is None:
         if protocol is None:
             raise TypeError('You must provide a port or protocol.')
+        if to_port is not None:
+            raise TypeError('to_port can only be specified if port is also specified')
         args.append(protocol)
     else:
         port_arg = f'{port}-{to_port}' if to_port is not None else str(port)
         if protocol is not None:
             port_arg = f'{port_arg}/{protocol}'
         args.append(port_arg)
-    run('open-port', *args)
+    # In the happy case (already open or opened successfully) open-ports exits silently with 0.
+    # If open-port exits with an error code, then run will raise an Error.
+    # Specifying a non-existent endpoint exits with 0, but prints an error message,
+    # **and does not open the port**. In this case, we return the error message.
+    # Ops will use this to raise an error.
+    result = run('open-port', *args).strip()
+    return result or None
 
 
 def opened_ports(*, endpoints: bool = False) -> list[Port]:
@@ -146,7 +157,8 @@ def opened_ports(*, endpoints: bool = False) -> list[Port]:
     # '8000-8999/tcp' or '8000-8999/udp' (where the two numbers can be any ports)
     # '8000-8999' (where these could be any port number)
     # If ``--endpoints`` is used, then each port will be followed by a
-    # (possibly empty) tuple of endpoints.
+    # (non-empty) tuple of endpoints, e.g. '8000-8999/tcp (ep1,ep2)' or '80/tcp (*)'.
+    # (*) indicates that the port applies to all endpoints.
     for port in result:
         if endpoints:
             port, port_endpoints = port.rsplit(' ', 1)

ops/hookcmds/_types.py

@@ -242,9 +242,6 @@ def _from_dict(cls, d: dict[str, Any]) -> Network:
         return cls(bind_addresses=bind, egress_subnets=egress, ingress_addresses=ingress)
 
 
-# Note that we intend to merge this with model.py's `Port` in the future, and
-# that does not have `kw_only=True`. That means that we should not use it here,
-# either, so that merging can be backwards compatible.
 @dataclasses.dataclass(frozen=True)
 class Port:
     """A port that Juju has opened for the charm."""

ops/model.py

@@ -37,7 +37,7 @@
 import warnings
 import weakref
 from abc import ABC, abstractmethod
-from collections.abc import Callable, Generator, Iterable, Mapping, MutableMapping
+from collections.abc import Callable, Generator, Iterable, Mapping, MutableMapping, Sequence
 from pathlib import Path, PurePath
 from typing import (
     Any,
@@ -719,7 +719,11 @@ def add_secret(
         )
 
     def open_port(
-        self, protocol: typing.Literal['tcp', 'udp', 'icmp'], port: int | None = None
+        self,
+        protocol: typing.Literal['tcp', 'udp', 'icmp'],
+        port: int | tuple[int, int | None] | None = None,
+        *,
+        endpoints: Sequence[str] = '*',
     ) -> None:
         """Open a port with the given protocol for this unit.
 
@@ -736,18 +740,30 @@ def open_port(
             protocol: String representing the protocol; must be one of
                 'tcp', 'udp', or 'icmp' (lowercase is recommended, but
                 uppercase is also supported).
-            port: The port to open. Required for TCP and UDP; not allowed
-                for ICMP.
+            port: The port to open. Required for TCP and UDP; not allowed for ICMP.
+                May be a tuple of two integers to specify a port range.
+            endpoints: The endpoints for which to open the port.
+                '*' means to open the port for all endpoints.
 
         Raises:
             ModelError: If ``port`` is provided when ``protocol`` is 'icmp'
                 or ``port`` is not provided when ``protocol`` is 'tcp' or
                 'udp'.
         """
-        self._backend.open_port(protocol.lower(), port)
+        if isinstance(port, tuple):
+            port, to_port = port
+        else:
+            port, to_port = port, None
+        if not endpoints:
+            raise TypeError('endpoints must be a non-empty string or sequence of strings')
+        self._backend.open_port(protocol.lower(), port, to_port=to_port, endpoints=endpoints)
 
     def close_port(
-        self, protocol: typing.Literal['tcp', 'udp', 'icmp'], port: int | None = None
+        self,
+        protocol: typing.Literal['tcp', 'udp', 'icmp'],
+        port: int | tuple[int, int | None] | None = None,
+        *,
+        endpoints: Sequence[str] = '*',
     ) -> None:
         """Close a port with the given protocol for this unit.
 
@@ -765,21 +781,29 @@ def close_port(
             protocol: String representing the protocol; must be one of
                 'tcp', 'udp', or 'icmp' (lowercase is recommended, but
                 uppercase is also supported).
-            port: The port to open. Required for TCP and UDP; not allowed
-                for ICMP.
+            port: The port to open. Required for TCP and UDP; not allowed for ICMP.
+                May be a tuple of two integers to specify a port range.
+            endpoints: The endpoints for which to close the port.
+                '*' means to close the port for all endpoints.
 
         Raises:
             ModelError: If ``port`` is provided when ``protocol`` is 'icmp'
                 or ``port`` is not provided when ``protocol`` is 'tcp' or
                 'udp'.
         """
-        self._backend.close_port(protocol.lower(), port)
+        if isinstance(port, tuple):
+            port, to_port = port
+        else:
+            port, to_port = port, None
+        if not endpoints:
+            raise TypeError('endpoints must be a non-empty string or sequence of strings')
+        self._backend.close_port(protocol.lower(), port, to_port=to_port, endpoints=endpoints)
 
     def opened_ports(self) -> set[Port]:
         """Return a list of opened ports for this unit."""
         return self._backend.opened_ports()
 
-    def set_ports(self, *ports: int | Port) -> None:
+    def set_ports(self, *ports: int | tuple[int, int | None] | Port) -> None:
         """Set the open ports for this unit, closing any others that are open.
 
         Some behaviour, such as whether the port is opened or closed externally without
@@ -800,16 +824,19 @@ def set_ports(self, *ports: int | Port) -> None:
                 ``port`` is not ``None``, or where ``protocol`` is 'tcp' or 'udp' and ``port``
                 is ``None``.
         """
-        # Normalise to get easier comparisons.
-        existing = {(port.protocol, port.port) for port in self._backend.opened_ports()}
+        existing = self._backend.opened_ports()
         desired = {
-            ('tcp', port) if isinstance(port, int) else (port.protocol, port.port)
+            Port('tcp', port)
+            if isinstance(port, int)
+            else Port('tcp', port[0], to_port=port[1])
+            if isinstance(port, tuple)
+            else port
             for port in ports
         }
-        for protocol, port in existing - desired:
-            self._backend.close_port(protocol, port)
-        for protocol, port in desired - existing:
-            self._backend.open_port(protocol, port)
+        for p in existing - desired:
+            self._backend.close_port(p.protocol, p.port, to_port=p.to_port, endpoints=p.endpoints)
+        for p in desired - existing:
+            self._backend.open_port(p.protocol, p.port, to_port=p.to_port, endpoints=p.endpoints)
 
     def reboot(self, now: bool = False) -> None:
         """Reboot the host machine.
@@ -846,6 +873,21 @@ class Port:
     port: int | None
     """The port number. Will be ``None`` if protocol is ``'icmp'``."""
 
+    to_port: int | None = None
+    """The end of the port range, if a range was specified.
+
+    Will be ``None`` if a single port was specified (or if protocol is ``'icmp'``).
+    """
+
+    _: dataclasses.KW_ONLY
+
+    endpoints: tuple[str, ...] | Literal['*'] = '*'
+    """The endpoints for which the port is open.
+
+    Will be ``"*"`` if open for all endpoints,
+    or a tuple of endpoint names if specified for particular endpoints.
+    """
+
 
 OpenedPort = Port
 """Alias to Port for backwards compatibility.
@@ -4037,26 +4079,59 @@ def secret_remove(self, id: str, *, revision: int | None = None):
         with self._wrap_hookcmd('secret-remove', id=id, revision=revision):
             hookcmds.secret_remove(id, revision=revision)
 
-    def open_port(self, protocol: str, port: int | None = None):
-        with self._wrap_hookcmd('open-port', protocol=protocol, port=port):
-            hookcmds.open_port(protocol, port)
+    def open_port(
+        self,
+        protocol: str,
+        port: int | None = None,
+        *,
+        to_port: int | None = None,
+        endpoints: Sequence[str] = '*',
+    ):
+        with self._wrap_hookcmd(
+            'open-port', protocol=protocol, port=port, to_port=to_port, endpoints=endpoints
+        ):
+            result = hookcmds.open_port(protocol, port, to_port=to_port, endpoints=endpoints)
+        if result is not None:
+            raise ModelError(result)
 
-    def close_port(self, protocol: str, port: int | None = None):
-        with self._wrap_hookcmd('close-port', protocol=protocol, port=port):
-            hookcmds.close_port(protocol, port)
+    def close_port(
+        self,
+        protocol: str,
+        port: int | None = None,
+        *,
+        to_port: int | None = None,
+        endpoints: Sequence[str] = '*',
+    ):
+        with self._wrap_hookcmd(
+            'close-port', protocol=protocol, port=port, to_port=to_port, endpoints=endpoints
+        ):
+            result = hookcmds.close_port(protocol, port, to_port=to_port, endpoints=endpoints)
+        if result is not None:
+            raise ModelError(result)
 
     def opened_ports(self) -> set[Port]:
-        with self._wrap_hookcmd('opened-ports'):
-            results = hookcmds.opened_ports()
+        with self._wrap_hookcmd('opened-ports', endpoints=True):
+            result = hookcmds.opened_ports(endpoints=True)
         ports: set[Port] = set()
-        for raw_port in results:
-            if raw_port.protocol not in ('tcp', 'udp', 'icmp'):
-                logger.warning('Unexpected opened-ports protocol: %s', raw_port.protocol)
+        for port in result:
+            if port.protocol not in ('tcp', 'udp', 'icmp'):
+                logger.warning('Unexpected opened-ports protocol: %s', port.protocol)
+                continue
+            if not port.endpoints:
+                logger.warning('opened-ports result with no endpoints: %s', port)
                 continue
-            if raw_port.to_port is not None:
-                logger.warning('Ignoring opened-ports port range: %s', raw_port)
-            port = Port(raw_port.protocol or 'tcp', raw_port.port)
-            ports.add(port)
+            match port.endpoints:
+                case ['*']:
+                    model_endpoints = '*'
+                case _:
+                    model_endpoints = tuple(port.endpoints)
+            model_port = Port(
+                port.protocol,
+                port.port,
+                to_port=port.to_port,
+                endpoints=model_endpoints,
+            )
+            ports.add(model_port)
         return ports
 
     def reboot(self, now: bool = False):