From 094b0f6247fea1a001a0c7079892a8d3c803725f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ond=C5=99ej=20Gajdu=C5=A1ek?= <ogajduse@redhat.com>
Date: Thu, 27 Jun 2024 15:51:30 +0200
Subject: [PATCH] Enhance port mapping logic, add exposed ports to Host object

---
 broker/hosts.py                   |  1 +
 broker/providers/container.py     | 20 ++++++++++++++------
 tests/providers/test_container.py |  1 +
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/broker/hosts.py b/broker/hosts.py
index 036b8178..1149469f 100644
--- a/broker/hosts.py
+++ b/broker/hosts.py
@@ -180,6 +180,7 @@ def to_dict(self):
             "os_distribution",
             "os_distribution_version",
             "reported_devices",
+            "exposed_ports",
         )
         ret_dict = {
             "name": getattr(self, "name", None),
diff --git a/broker/providers/container.py b/broker/providers/container.py
index b8b27aad..37b460a8 100644
--- a/broker/providers/container.py
+++ b/broker/providers/container.py
@@ -1,4 +1,5 @@
 """Container provider implementation."""
+
 from functools import cache
 import getpass
 import inspect
@@ -154,6 +155,13 @@ def _port_mapping(self, image, **kwargs):
             22:1337 23:1335.
         """
         mapping = {}
+        # create mapping for all exposed ports in the image
+        if settings.container.auto_map_ports or kwargs.get("auto_map_ports"):
+            mapping = {
+                k: v or None
+                for k, v in self.runtime.image_info(image)["config"]["ExposedPorts"].items()
+            }
+        # add any ports that were passed as arguments
         if ports := kwargs.pop("ports", None):
             if isinstance(ports, str):
                 for _map in ports.split():
@@ -166,11 +174,6 @@ def _port_mapping(self, image, **kwargs):
                     else:
                         s = "tcp"
                     mapping[f"{p}/{s}"] = int(h) if h else None
-        elif settings.container.auto_map_ports:
-            mapping = {
-                k: v or None
-                for k, v in self.runtime.image_info(image)["config"]["ExposedPorts"].items()
-            }
         return mapping
 
     def _cont_inst_by_name(self, cont_name):
@@ -200,13 +203,18 @@ def construct_host(self, provider_params, host_classes, **kwargs):
         logger.debug(cont_attrs)
         hostname = cont_inst.id[:12]
         if port := self._find_ssh_port(cont_attrs["ports"]):
-            hostname = f"{hostname}:{port}"
+            hostname = f"{self.runtime.host}:{port}"
         if not hostname:
             raise Exception(f"Could not determine container hostname:\n{cont_attrs}")
         name = cont_attrs["name"]
         logger.debug(f"hostname: {hostname}, name: {name}, host type: host")
         host_inst = host_classes["host"](**{**kwargs, "hostname": hostname, "name": name})
         self._set_attributes(host_inst, broker_args=kwargs, cont_inst=cont_inst)
+        # add the container's port mapping to the host instance only if there are any ports open
+        if cont_attrs.get("ports"):
+            host_inst.exposed_ports = {
+                f"{k.split('/')[0]}": v[0]["HostPort"] for k, v in cont_attrs["ports"].items()
+            }
         return host_inst
 
     def provider_help(
diff --git a/tests/providers/test_container.py b/tests/providers/test_container.py
index 9ab3aae6..1d6957e2 100644
--- a/tests/providers/test_container.py
+++ b/tests/providers/test_container.py
@@ -24,6 +24,7 @@ def __init__(self, **kwargs):
             "images": [MockStub({"tags": "ch-d:ubi8"})],  # self.runtime.images
             "containers": [MockStub({"tags": "f37d3058317f"})],  # self.runtime.containers
             "name": "f37d3058317f",  # self.runtime.get_attrs(cont_inst)["name"]
+            "ports": MockStub({'22/tcp': [{'HostIp': '', 'HostPort': '1337'}]}),  # self.runtime.get_attrs(cont_inst)["ports"]
         }
         if "job_id" in kwargs:
             # we need to load in an image object