feat: allow tmate ip addresses to firewall.

src-docs/firewall.py.md

@@ -50,12 +50,15 @@ Get the host IP address for the corresponding LXD network.
 
 ---
 
-<a href="../src/firewall.py#L69"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
+<a href="../src/firewall.py#L96"><img align="right" style="float:right;" src="https://img.shields.io/badge/-source-cccccc?style=flat-square"></a>
 
 ### <kbd>function</kbd> `refresh_firewall`
 
 ```python
-refresh_firewall(denylist: List[FirewallEntry])
+refresh_firewall(
+    denylist: List[FirewallEntry],
+    allowlist: Optional[List[FirewallEntry]] = None
+)
 ```
 
 Refresh the firewall configuration. 

src/charm.py

@@ -910,8 +910,11 @@ def _refresh_firewall(self):
                 FirewallEntry.decode(entry.strip())
                 for entry in firewall_denylist_config.split(",")
             ]
+        allowlist = [
+            FirewallEntry.decode(str(entry.host)) for entry in self._state.ssh_debug_infos
+        ]
         firewall = Firewall("lxdbr0")
-        firewall.refresh_firewall(denylist)
+        firewall.refresh_firewall(denylist=denylist, allowlist=allowlist)
         logger.debug(
             "firewall update, current firewall: %s",
             execute_command(["/usr/sbin/nft", "list", "ruleset"]),
@@ -930,6 +933,7 @@ def _apt_install(self, packages: Sequence[str]) -> None:
 
     def _on_debug_ssh_relation_changed(self, _: ops.RelationChangedEvent) -> None:
         """Handle debug ssh relation changed event."""
+        self._refresh_firewall()
         runner_manager = self._get_runner_manager()
         runner_manager.flush(flush_busy=False)
         self._reconcile_runners(runner_manager)

src/firewall.py

@@ -66,7 +66,38 @@ def get_host_ip(self) -> str:
         )
         return str(ipaddress.IPv4Interface(address.strip()).ip)
 
-    def refresh_firewall(self, denylist: typing.List[FirewallEntry]):
+    def _exclude_network(
+        self,
+        networks: typing.List[ipaddress.IPv4Network],
+        excludes: typing.List[ipaddress.IPv4Network],
+    ) -> typing.Set[ipaddress.IPv4Network]:
+        """Excludes the network segment from a pool of networks.
+
+        Args:
+            networks: The network pool to apply.
+            excludes: The networks to exclude from the pool.
+
+        Returns:
+            The network pool without the network segments in excludes.
+        """
+        total_excluded_networks = set(networks)
+
+        for exclude in excludes:
+            scoped_excluded_networks: set[ipaddress.IPv4Network] = set()
+            for net in total_excluded_networks:
+                if net.overlaps(exclude):
+                    scoped_excluded_networks.update(net.address_exclude(exclude))
+                else:
+                    scoped_excluded_networks.add(net)
+            total_excluded_networks = scoped_excluded_networks
+
+        return total_excluded_networks
+
+    def refresh_firewall(
+        self,
+        denylist: typing.List[FirewallEntry],
+        allowlist: typing.List[FirewallEntry] | None = None,
+    ):
         """Refresh the firewall configuration.
 
         Args:
@@ -128,23 +159,23 @@ def refresh_firewall(self, denylist: typing.List[FirewallEntry]):
                 "state": "enabled",
             },
         ]
-        host_network = ipaddress.IPv4Network(host_ip)
-        for entry in denylist:
-            entry_network = ipaddress.IPv4Network(entry.ip_range)
-            try:
-                excluded = list(entry_network.address_exclude(host_network))
-            except ValueError:
-                excluded = [entry_network]
-            egress_rules.extend(
-                [
-                    {
-                        "action": "reject",
-                        "destination": str(ip),
-                        "state": "enabled",
-                    }
-                    for ip in excluded
-                ]
-            )
+
+        allowed_ips = [
+            ipaddress.IPv4Network(host_ip),
+            *[ipaddress.IPv4Network(entry.ip_range) for entry in (allowlist or [])],
+        ]
+        ips_to_deny = [ipaddress.IPv4Network(entry.ip_range) for entry in denylist]
+        denied_ips = self._exclude_network(networks=ips_to_deny, excludes=allowed_ips)
+        egress_rules.extend(
+            [
+                {
+                    "action": "reject",
+                    "destination": str(ip),
+                    "state": "enabled",
+                }
+                for ip in denied_ips
+            ]
+        )
         acl_config["egress"] = egress_rules
         execute_command(
             ["lxc", "network", "acl", "edit", self._ACL_RULESET_NAME],

tests/integration/conftest.py

@@ -294,7 +294,7 @@ async def tmate_ssh_server_app_fixture(
     """tmate-ssh-server charm application related to GitHub-Runner app charm."""
     tmate_app: Application = await model.deploy("tmate-ssh-server", channel="edge")
     await app_no_wait.relate("debug-ssh", f"{tmate_app.name}:debug-ssh")
-    await model.wait_for_idle(status=ACTIVE, timeout=60 * 30)
+    await model.wait_for_idle(apps=[tmate_app.name], status=ACTIVE, timeout=60 * 30)
 
     return tmate_app
 

tests/integration/test_debug_ssh.py

@@ -9,25 +9,41 @@
 from github.Repository import Repository
 from github.WorkflowRun import WorkflowRun
 from juju.application import Application
+from juju.model import Model
 
 from tests.integration.helpers import dispatch_workflow, get_job_logs, get_workflow_runs
+from tests.status_name import ACTIVE
 
 logger = logging.getLogger(__name__)
 
 SSH_DEBUG_WORKFLOW_FILE_NAME = "workflow_dispatch_ssh_debug.yaml"
 
 
 async def test_ssh_debug(
+    model: Model,
     app_no_wait: Application,
     github_repository: Repository,
     test_github_branch: Branch,
     tmate_ssh_server_unit_ip: str,
 ):
     """
-    arrange: given an integrated GitHub-Runner charm and tmate-ssh-server charm.
+    arrange: given an integrated GitHub-Runner charm and tmate-ssh-server charm with a denylist
+        covering ip ranges of tmate-ssh-server.
     act: when canonical/action-tmate is triggered.
     assert: the ssh connection info from action-log and tmate-ssh-server matches.
     """
+    await app_no_wait.set_config(
+        {
+            "denylist": (
+                "0.0.0.0/8,10.0.0.0/8,100.64.0.0/10,169.254.0.0/16,"
+                "172.16.0.0/12,192.0.0.0/24,192.0.2.0/24,192.88.99.0/24,192.168.0.0/16,"
+                "198.18.0.0/15,198.51.100.0/24,203.0.113.0/24,224.0.0.0/4,233.252.0.0/24,"
+                "240.0.0.0/4"
+            ),
+        }
+    )
+    await model.wait_for_idle(status=ACTIVE, timeout=60 * 30)
+
     # trigger tmate action
     logger.info("Dispatching workflow_dispatch_ssh_debug.yaml workflow.")
     start_time = datetime.now(timezone.utc)

tests/unit/test_charm.py

@@ -19,6 +19,7 @@
     RunnerError,
     SubprocessError,
 )
+from firewall import FirewallEntry
 from github_type import GitHubRunnerStatus
 from runner_manager import RunnerInfo, RunnerManagerConfig
 from runner_type import GithubOrg, GithubRepo, VirtualMachineResources
@@ -372,8 +373,8 @@ def test_check_runners_action_with_errors(self, run, wt, mkdir, rm):
         )
 
     @patch("charm.RunnerManager")
-    @patch("pathlib.Path.mkdir")
     @patch("pathlib.Path.write_text")
+    @patch("pathlib.Path.mkdir")
     @patch("subprocess.run")
     def test_on_flush_runners_action(self, run, wt, mkdir, rm):
         mock_event = MagicMock()
@@ -391,3 +392,61 @@ def test_on_flush_runners_action(self, run, wt, mkdir, rm):
         harness.charm._on_flush_runners_action(mock_event)
         mock_event.set_results.assert_called()
         mock_event.reset_mock()
+
+    @patch("charm.RunnerManager")
+    @patch("pathlib.Path.write_text")
+    @patch("pathlib.Path.mkdir")
+    @patch("subprocess.run")
+    @patch("charm.Firewall")
+    def test__refresh_firewall(self, mock_firewall, *args):
+        """
+        arrange: given multiple tmate-ssh-server units in relation.
+        act: when refresh_firewall is called.
+        assert: the unit ip addresses are included in allowlist.
+        """
+        harness = Harness(GithubRunnerCharm)
+        relation_id = harness.add_relation("debug-ssh", "tmate-ssh-server")
+        harness.add_relation_unit(relation_id, "tmate-ssh-server/0")
+        harness.add_relation_unit(relation_id, "tmate-ssh-server/1")
+        harness.add_relation_unit(relation_id, "tmate-ssh-server/2")
+        test_unit_ip_addresses = ["127.0.0.1", "127.0.0.2", "127.0.0.3"]
+
+        harness.update_relation_data(
+            relation_id,
+            "tmate-ssh-server/0",
+            {
+                "host": test_unit_ip_addresses[0],
+                "port": "10022",
+                "rsa_fingerprint": "SHA256:abcd",
+                "ed25519_fingerprint": "abcd",
+            },
+        )
+        harness.update_relation_data(
+            relation_id,
+            "tmate-ssh-server/1",
+            {
+                "host": test_unit_ip_addresses[1],
+                "port": "10022",
+                "rsa_fingerprint": "SHA256:abcd",
+                "ed25519_fingerprint": "abcd",
+            },
+        )
+        harness.update_relation_data(
+            relation_id,
+            "tmate-ssh-server/2",
+            {
+                "host": test_unit_ip_addresses[2],
+                "port": "10022",
+                "rsa_fingerprint": "SHA256:abcd",
+                "ed25519_fingerprint": "abcd",
+            },
+        )
+
+        harness.begin()
+
+        harness.charm._refresh_firewall()
+        mocked_firewall_instance = mock_firewall.return_value
+        allowlist = mocked_firewall_instance.refresh_firewall.call_args_list[0][1]["allowlist"]
+        assert all(
+            FirewallEntry(ip) in allowlist for ip in test_unit_ip_addresses
+        ), "Expected IP firewall entry not found in allowlist arg."

tests/unit/test_firewall.py

@@ -0,0 +1,97 @@
+# Copyright 2024 Canonical Ltd.
+# See LICENSE file for licensing details.
+
+"""Test cases for firewall module."""
+
+from ipaddress import IPv4Network
+
+import pytest
+
+from firewall import Firewall
+
+
+@pytest.mark.parametrize(
+    "domain_ranges, excludes_ranges, expected",
+    [
+        pytest.param(
+            [IPv4Network("10.0.0.0/8")],
+            [IPv4Network("10.136.12.36/32")],
+            {
+                IPv4Network("10.0.0.0/9"),
+                IPv4Network("10.192.0.0/10"),
+                IPv4Network("10.160.0.0/11"),
+                IPv4Network("10.144.0.0/12"),
+                IPv4Network("10.128.0.0/13"),
+                IPv4Network("10.140.0.0/14"),
+                IPv4Network("10.138.0.0/15"),
+                IPv4Network("10.137.0.0/16"),
+                IPv4Network("10.136.128.0/17"),
+                IPv4Network("10.136.64.0/18"),
+                IPv4Network("10.136.32.0/19"),
+                IPv4Network("10.136.16.0/20"),
+                IPv4Network("10.136.0.0/21"),
+                IPv4Network("10.136.8.0/22"),
+                IPv4Network("10.136.14.0/23"),
+                IPv4Network("10.136.13.0/24"),
+                IPv4Network("10.136.12.128/25"),
+                IPv4Network("10.136.12.64/26"),
+                IPv4Network("10.136.12.0/27"),
+                IPv4Network("10.136.12.48/28"),
+                IPv4Network("10.136.12.40/29"),
+                IPv4Network("10.136.12.32/30"),
+                IPv4Network("10.136.12.38/31"),
+                IPv4Network("10.136.12.37/32"),
+            },
+            id="exclude a single IP",
+        ),
+        pytest.param(
+            [IPv4Network("10.0.0.0/8"), IPv4Network("192.0.2.0/28")],
+            [IPv4Network("10.136.12.36/32")],
+            {
+                IPv4Network("10.0.0.0/9"),
+                IPv4Network("10.192.0.0/10"),
+                IPv4Network("10.160.0.0/11"),
+                IPv4Network("10.144.0.0/12"),
+                IPv4Network("10.128.0.0/13"),
+                IPv4Network("10.140.0.0/14"),
+                IPv4Network("10.138.0.0/15"),
+                IPv4Network("10.137.0.0/16"),
+                IPv4Network("10.136.128.0/17"),
+                IPv4Network("10.136.64.0/18"),
+                IPv4Network("10.136.32.0/19"),
+                IPv4Network("10.136.16.0/20"),
+                IPv4Network("10.136.0.0/21"),
+                IPv4Network("10.136.8.0/22"),
+                IPv4Network("10.136.14.0/23"),
+                IPv4Network("10.136.13.0/24"),
+                IPv4Network("10.136.12.128/25"),
+                IPv4Network("10.136.12.64/26"),
+                IPv4Network("10.136.12.0/27"),
+                IPv4Network("10.136.12.48/28"),
+                IPv4Network("10.136.12.40/29"),
+                IPv4Network("10.136.12.32/30"),
+                IPv4Network("10.136.12.38/31"),
+                IPv4Network("10.136.12.37/32"),
+                IPv4Network("192.0.2.0/28"),
+            },
+            id="exclude a single IP and one different subnet",
+        ),
+        pytest.param(
+            [IPv4Network("198.18.0.0/15"), IPv4Network("172.16.0.0/12")],
+            [IPv4Network("10.136.12.36/32")],
+            {IPv4Network("198.18.0.0/15"), IPv4Network("172.16.0.0/12")},
+            id="no matching networks",
+        ),
+    ],
+)
+def test__exclude_network(
+    domain_ranges: list[IPv4Network],
+    excludes_ranges: list[IPv4Network],
+    expected: set[IPv4Network],
+):
+    """
+    arrange: given domain networks and some IPs to exclude from the domains.
+    act: when _exclude_network is called.
+    assert: new ip networks are returned with excluded target IP ranges.
+    """
+    assert Firewall("test")._exclude_network(domain_ranges, excludes_ranges) == expected