Repository: canonical/operator
Pull request #905 (2023-02-16 03:43:35)

Support for open-port, close-port, and opened-ports
Adds support for open-port, close-port and opened-ports as per the [OP029 spec](https://docs.google.com/document/d/1xHLzBgcXtLsBpMsjuqUW2pI1tP4D8lgSUD_sg3HClng/edit#).

Also adds testing backend support for the same for use in charm tests.

Fixes https://github.com/canonical/operator/issues/179

jnsgruk: Looking tidy so far! 😎
benhoyt: Going to merge without a second approval: spec's been approved by Gustavo, John's taken a close look at the PR and approved, and others have cast their eyes over it too -- I think we're good.

Patch:
diff --git a/ops/model.py b/ops/model.py
index 0db760f..bc6c63a 100644
--- a/ops/model.py
+++ b/ops/model.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 
 """Representations of Juju's model, application, unit, and other entities."""
+
+import dataclasses
 import datetime
 import enum
 import ipaddress
@@ -600,6 +602,61 @@ class Unit:
             owner='unit')
         return Secret(self._backend, id=id, label=label, content=content)
 
+    def open_port(self, protocol: typing.Literal['tcp', 'udp', 'icmp'],
+                  port: Optional[int] = None):
+        """Open a port with the given protocol for this unit.
+
+        Calling this registers intent with Juju that the application should be
+        accessed on the given port, but the port isn't actually opened
+        externally until the operator runs "juju expose".
+
+        On Kubernetes sidecar charms, the ports opened are not strictly
+        per-unit: Juju will open the union of ports from all units.
+        However, normally charms should make the same open_port() call from
+        every unit.
+
+        Args:
+            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.
+        """
+        self._backend.open_port(protocol.lower(), port)
+
+    def close_port(self, protocol: typing.Literal['tcp', 'udp', 'icmp'],
+                   port: Optional[int] = None):
+        """Close a port with the given protocol for this unit.
+
+        On Kubernetes sidecar charms, Juju will only close the port once the
+        last unit that opened that port has closed it. However, this is
+        usually not an issue; normally charms should make the same
+        close_port() call from every unit.
+
+        Args:
+            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.
+        """
+        self._backend.close_port(protocol.lower(), port)
+
+    def opened_ports(self) -> Set['OpenedPort']:
+        """Return a list of opened ports for this unit."""
+        return self._backend.opened_ports()
+
+
+@dataclasses.dataclass(frozen=True)
+class OpenedPort:
+    """Represents a port opened by :meth:`Unit.open_port`."""
+
+    """The IP protocol: 'tcp', 'udp', or 'icmp'."""
+    protocol: typing.Literal['tcp', 'udp', 'icmp']
+
+    """The port number. Will be None if protocol is 'icmp'."""
+    port: Optional[int]
+
 
 class LazyMapping(Mapping[str, str], ABC):
     """Represents a dict that isn't populated until it is accessed.
@@ -2998,6 +3055,44 @@ class _ModelBackend:
             args.extend(['--revision', str(revision)])
         self._run_for_secret('secret-remove', *args)
 
+    def open_port(self, protocol: str, port: Optional[int] = None):
+        arg = f'{port}/{protocol}' if port is not None else protocol
+        self._run('open-port', arg)
+
+    def close_port(self, protocol: str, port: Optional[int] = None):
+        arg = f'{port}/{protocol}' if port is not None else protocol
+        self._run('close-port', arg)
+
+    def opened_ports(self) -> Set[OpenedPort]:
+        # We could use "opened-ports --format=json", but it's not really
+        # structured; it's just an array of strings which are the lines of the
+        # text output, like ["icmp","8081/udp"]. So it's probably just as
+        # likely to change as the text output, and doesn't seem any better.
+        output = typing.cast(str, self._run('opened-ports', return_output=True))
+        ports: Set[OpenedPort] = set()
+        for line in output.splitlines():
+            line = line.strip()
+            if not line:
+                continue
+            port = self._parse_opened_port(line)
+            if port is not None:
+                ports.add(port)
+        return ports
+
+    @classmethod
+    def _parse_opened_port(cls, port_str: str) -> Optional[OpenedPort]:
+        if port_str == 'icmp':
+            return OpenedPort('icmp', None)
+        port_range, slash, protocol = port_str.partition('/')
+        if not slash or protocol not in ['tcp', 'udp']:
+            logger.warning('Unexpected opened-ports protocol: %s', port_str)
+            return None
+        port, hyphen, _ = port_range.partition('-')
+        if hyphen:
+            logger.warning('Ignoring opened-ports port range: %s', port_str)
+        protocol_lit = typing.cast(typing.Literal['tcp', 'udp'], protocol)
+        return OpenedPort(protocol_lit, int(port))
+
 
 class _ModelBackendValidator:
     """Provides facilities for validating inputs and formatting them for model backends."""


Test patch:
diff --git a/ops/testing.py b/ops/testing.py
index 7f05bae..9a215f9 100755
--- a/ops/testing.py
+++ b/ops/testing.py
@@ -39,6 +39,7 @@ from typing import (
     Iterable,
     Iterator,
     List,
+    Literal,
     Mapping,
     Optional,
     Set,
@@ -56,7 +57,7 @@ from ops.charm import CharmBase, CharmMeta, RelationRole
 from ops.model import RelationNotFoundError
 
 if TYPE_CHECKING:
-    from typing_extensions import Literal, TypedDict
+    from typing_extensions import TypedDict
 
     from ops.model import UnitOrApplication
 
@@ -1547,6 +1548,7 @@ class _TestingModelBackend:
         self._planned_units: Optional[int] = None
         self._hook_is_running = ''
         self._secrets: List[_Secret] = []
+        self._opened_ports: Set[model.OpenedPort] = set()
 
     def _validate_relation_access(self, relation_name: str, relations: List[model.Relation]):
         """Ensures that the named relation exists/has been added.
@@ -2059,6 +2061,33 @@ class _TestingModelBackend:
         else:
             self._secrets = [s for s in self._secrets if s.id != id]
 
+    def open_port(self, protocol: str, port: Optional[int] = None):
+        self._check_protocol_and_port(protocol, port)
+        protocol_lit = cast(Literal['tcp', 'udp', 'icmp'], protocol)
+        self._opened_ports.add(model.OpenedPort(protocol_lit, port))
+
+    def close_port(self, protocol: str, port: Optional[int] = None):
+        self._check_protocol_and_port(protocol, port)
+        protocol_lit = cast(Literal['tcp', 'udp', 'icmp'], protocol)
+        self._opened_ports.discard(model.OpenedPort(protocol_lit, port))
+
+    def opened_ports(self) -> Set[model.OpenedPort]:
+        return set(self._opened_ports)
+
+    def _check_protocol_and_port(self, protocol: str, port: Optional[int]):
+        # Simulate the error messages we get from Juju (not that charm tests
+        # should be testing details of error messages).
+        if protocol == 'icmp':
+            if port is not None:
+                raise model.ModelError(f'ERROR protocol "{protocol}" doesn\'t support any ports; got "{port}"\n')  # NOQA: test_quote_backslashes
+        elif protocol in ['tcp', 'udp']:
+            if port is None:
+                raise model.ModelError(f'ERROR invalid port "{protocol}": strconv.Atoi: parsing "{protocol}": invalid syntax\n')  # NOQA: test_quote_backslashes
+            if not (1 <= port <= 65535):
+                raise model.ModelError(f'ERROR port range bounds must be between 1 and 65535, got {port}-{port}\n')  # NOQA: test_quote_backslashes
+        else:
+            raise model.ModelError(f'ERROR invalid protocol "{protocol}", expected "tcp", "udp", or "icmp"\n')  # NOQA: test_quote_backslashes
+
 
 @_copy_docstrings(pebble.Client)
 class _TestingPebbleClient:
diff --git a/test/test_model.py b/test/test_model.py
index 35df7ea..703cd7a 100755
--- a/test/test_model.py
+++ b/test/test_model.py
@@ -3165,5 +3165,100 @@ class TestSecretClass(unittest.TestCase):
         ])
 
 
+class TestPorts(unittest.TestCase):
+    def setUp(self):
+        self.model = ops.model.Model(ops.charm.CharmMeta(), ops.model._ModelBackend('myapp/0'))
+        self.unit = self.model.unit
+
+    def test_open_port(self):
+        fake_script(self, 'open-port', 'exit 0')
+
+        self.unit.open_port('tcp', 8080)
+        self.unit.open_port('UDP', 4000)
+        self.unit.open_port('icmp')
+
+        self.assertEqual(fake_script_calls(self, clear=True), [
+            ['open-port', '8080/tcp'],
+            ['open-port', '4000/udp'],
+            ['open-port', 'icmp'],
+        ])
+
+    def test_open_port_error(self):
+        fake_script(self, 'open-port', "echo 'ERROR bad protocol' >&2; exit 1")
+
+        with self.assertRaises(model.ModelError) as cm:
+            self.unit.open_port('ftp', 8080)
+        self.assertEqual(str(cm.exception), 'ERROR bad protocol\n')
+
+        self.assertEqual(fake_script_calls(self, clear=True), [
+            ['open-port', '8080/ftp'],
+        ])
+
+    def test_close_port(self):
+        fake_script(self, 'close-port', 'exit 0')
+
+        self.unit.close_port('tcp', 8080)
+        self.unit.close_port('UDP', 4000)
+        self.unit.close_port('icmp')
+
+        self.assertEqual(fake_script_calls(self, clear=True), [
+            ['close-port', '8080/tcp'],
+            ['close-port', '4000/udp'],
+            ['close-port', 'icmp'],
+        ])
+
+    def test_close_port_error(self):
+        fake_script(self, 'close-port', "echo 'ERROR bad protocol' >&2; exit 1")
+
+        with self.assertRaises(model.ModelError) as cm:
+            self.unit.close_port('ftp', 8080)
+        self.assertEqual(str(cm.exception), 'ERROR bad protocol\n')
+
+        self.assertEqual(fake_script_calls(self, clear=True), [
+            ['close-port', '8080/ftp'],
+        ])
+
+    def test_opened_ports(self):
+        fake_script(self, 'opened-ports', """echo 8080/tcp; echo icmp""")
+
+        ports_set = self.unit.opened_ports()
+        self.assertIsInstance(ports_set, set)
+        ports = sorted(ports_set, key=lambda p: (p.protocol, p.port))
+        self.assertEqual(len(ports), 2)
+        self.assertIsInstance(ports[0], model.OpenedPort)
+        self.assertEqual(ports[0].protocol, 'icmp')
+        self.assertIsNone(ports[0].port)
+        self.assertIsInstance(ports[1], model.OpenedPort)
+        self.assertEqual(ports[1].protocol, 'tcp')
+        self.assertEqual(ports[1].port, 8080)
+
+        self.assertEqual(fake_script_calls(self, clear=True), [
+            ['opened-ports', ''],
+        ])
+
+    def test_opened_ports_warnings(self):
+        fake_script(self, 'opened-ports', """echo 8080/tcp; echo 1234/ftp; echo 1000-2000/udp""")
+
+        with self.assertLogs('ops.model', level='WARNING') as cm:
+            ports_set = self.unit.opened_ports()
+        self.assertEqual(len(cm.output), 2)
+        self.assertRegex(cm.output[0], r'WARNING:ops.model:.*protocol.*')
+        self.assertRegex(cm.output[1], r'WARNING:ops.model:.*range.*')
+
+        self.assertIsInstance(ports_set, set)
+        ports = sorted(ports_set, key=lambda p: (p.protocol, p.port))
+        self.assertEqual(len(ports), 2)
+        self.assertIsInstance(ports[0], model.OpenedPort)
+        self.assertEqual(ports[0].protocol, 'tcp')
+        self.assertEqual(ports[0].port, 8080)
+        self.assertIsInstance(ports[1], model.OpenedPort)
+        self.assertEqual(ports[1].protocol, 'udp')
+        self.assertEqual(ports[1].port, 1000)
+
+        self.assertEqual(fake_script_calls(self, clear=True), [
+            ['opened-ports', ''],
+        ])
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/test/test_testing.py b/test/test_testing.py
index 1a3c4c1..ed92f99 100644
--- a/test/test_testing.py
+++ b/test/test_testing.py
@@ -4572,3 +4572,63 @@ class EventRecorder(CharmBase):
 
     def record_event(self, event):
         self.events.append(event)
+
+
+class TestPorts(unittest.TestCase):
+    def test_ports(self):
+        harness = Harness(CharmBase, meta='name: webapp')
+        self.addCleanup(harness.cleanup)
+        unit = harness.model.unit
+
+        unit.open_port('tcp', 8080)
+        unit.open_port('udp', 4000)
+        unit.open_port('icmp')
+
+        ports_set = unit.opened_ports()
+        self.assertIsInstance(ports_set, set)
+        ports = sorted(ports_set, key=lambda p: (p.protocol, p.port))
+        self.assertEqual(len(ports), 3)
+        self.assertIsInstance(ports[0], model.OpenedPort)
+        self.assertEqual(ports[0].protocol, 'icmp')
+        self.assertIsNone(ports[0].port)
+        self.assertIsInstance(ports[1], model.OpenedPort)
+        self.assertEqual(ports[1].protocol, 'tcp')
+        self.assertEqual(ports[1].port, 8080)
+        self.assertIsInstance(ports[2], model.OpenedPort)
+        self.assertEqual(ports[2].protocol, 'udp')
+        self.assertEqual(ports[2].port, 4000)
+
+        unit.close_port('tcp', 8080)
+        unit.close_port('tcp', 8080)  # closing same port again has no effect
+        unit.close_port('udp', 4000)
+
+        ports_set = unit.opened_ports()
+        self.assertIsInstance(ports_set, set)
+        ports = sorted(ports_set, key=lambda p: (p.protocol, p.port))
+        self.assertEqual(len(ports), 1)
+        self.assertIsInstance(ports[0], model.OpenedPort)
+        self.assertEqual(ports[0].protocol, 'icmp')
+        self.assertIsNone(ports[0].port)
+
+        unit.close_port('icmp')
+
+        ports_set = unit.opened_ports()
+        self.assertEqual(ports_set, set())
+
+    def test_errors(self):
+        harness = Harness(CharmBase, meta='name: webapp')
+        self.addCleanup(harness.cleanup)
+        unit = harness.model.unit
+
+        with self.assertRaises(model.ModelError):
+            unit.open_port('icmp', 8080)  # icmp cannot have port
+        with self.assertRaises(model.ModelError):
+            unit.open_port('ftp', 8080)  # invalid protocol
+        with self.assertRaises(model.ModelError):
+            unit.open_port('tcp')  # tcp must have port
+        with self.assertRaises(model.ModelError):
+            unit.open_port('udp')  # udp must have port
+        with self.assertRaises(model.ModelError):
+            unit.open_port('tcp', 0)  # port out of range
+        with self.assertRaises(model.ModelError):
+            unit.open_port('tcp', 65536)  # port out of range