Skip to content

Commit

Permalink
Make check_port an async function (#4677)
Browse files Browse the repository at this point in the history
* Make check_port asyncio

This requires to change the ingress_port property to a async method.

* Avoid using wait_for

* Add missing async

* Really await

* Set dynamic ingress port on add-on installation/update

* Fix pytest issue

* Rename async_check_port back to check_port

* Raise RuntimeError in case port is not set

* Make sure port gets set on add-on restore

* Drop unnecessary async

* Simplify check_port by using asyncio.get_running_loop()
  • Loading branch information
agners authored Dec 5, 2023
1 parent c2d4be3 commit 883e54f
Show file tree
Hide file tree
Showing 6 changed files with 49 additions and 35 deletions.
19 changes: 17 additions & 2 deletions supervisor/addons/addon.py
Original file line number Diff line number Diff line change
Expand Up @@ -395,7 +395,7 @@ def ingress_port(self) -> int | None:

port = self.data[ATTR_INGRESS_PORT]
if port == 0:
return self.sys_ingress.get_dynamic_port(self.slug)
raise RuntimeError(f"No port set for add-on {self.slug}")
return port

@property
Expand Down Expand Up @@ -539,7 +539,7 @@ async def watchdog_application(self) -> bool:

# TCP monitoring
if s_prefix == "tcp":
return await self.sys_run_in_executor(check_port, self.ip_address, port)
return await check_port(self.ip_address, port)

# lookup the correct protocol from config
if t_proto:
Expand Down Expand Up @@ -602,6 +602,16 @@ async def unload(self) -> None:
_LOGGER.info("Removing add-on data folder %s", self.path_data)
await remove_data(self.path_data)

async def _check_ingress_port(self):
"""Assign a ingress port if dynamic port selection is used."""
if not self.with_ingress:
return

if self.data[ATTR_INGRESS_PORT] == 0:
self.data[ATTR_INGRESS_PORT] = await self.sys_ingress.get_dynamic_port(
self.slug
)

@Job(
name="addon_install",
limit=JobExecutionLimit.GROUP_ONCE,
Expand Down Expand Up @@ -630,6 +640,8 @@ async def install(self) -> None:
self.sys_addons.data.uninstall(self)
raise AddonsError() from err

await self._check_ingress_port()

# Add to addon manager
self.sys_addons.local[self.slug] = self

Expand Down Expand Up @@ -716,6 +728,7 @@ async def update(self) -> asyncio.Task | None:
try:
_LOGGER.info("Add-on '%s' successfully updated", self.slug)
self.sys_addons.data.update(store)
await self._check_ingress_port()

# Cleanup
with suppress(DockerError):
Expand Down Expand Up @@ -756,6 +769,7 @@ async def rebuild(self) -> asyncio.Task | None:
raise AddonsError() from err

self.sys_addons.data.update(self.addon_store)
await self._check_ingress_port()
_LOGGER.info("Add-on '%s' successfully rebuilt", self.slug)

finally:
Expand Down Expand Up @@ -1199,6 +1213,7 @@ def _extract_tarfile():
_LOGGER.info("Restore/Update of image for addon %s", self.slug)
with suppress(DockerError):
await self.instance.update(version, restore_image)
self._check_ingress_port()

# Restore data and config
def _restore_data():
Expand Down
3 changes: 1 addition & 2 deletions supervisor/homeassistant/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,8 +140,7 @@ async def get_api_state(self) -> str | None:
return None

# Check if port is up
if not await self.sys_run_in_executor(
check_port,
if not await check_port(
self.sys_homeassistant.ip_address,
self.sys_homeassistant.api_port,
):
Expand Down
4 changes: 2 additions & 2 deletions supervisor/ingress.py
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,7 @@ def validate_session(self, session: str) -> bool:

return True

def get_dynamic_port(self, addon_slug: str) -> int:
async def get_dynamic_port(self, addon_slug: str) -> int:
"""Get/Create a dynamic port from range."""
if addon_slug in self.ports:
return self.ports[addon_slug]
Expand All @@ -163,7 +163,7 @@ def get_dynamic_port(self, addon_slug: str) -> int:
while (
port is None
or port in self.ports.values()
or check_port(self.sys_docker.network.gateway, port)
or await check_port(self.sys_docker.network.gateway, port)
):
port = random.randint(62000, 65500)

Expand Down
21 changes: 10 additions & 11 deletions supervisor/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,20 +39,19 @@ async def wrap_api(api, *args, **kwargs):
return wrap_api


def check_port(address: IPv4Address, port: int) -> bool:
async def check_port(address: IPv4Address, port: int) -> bool:
"""Check if port is mapped."""
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.settimeout(0.5)
sock.setblocking(False)
try:
result = sock.connect_ex((str(address), port))
sock.close()

# Check if the port is available
if result == 0:
return True
except OSError:
pass
return False
async with asyncio.timeout(0.5):
await asyncio.get_running_loop().sock_connect(sock, (str(address), port))
except (OSError, TimeoutError):
return False
finally:
if sock is not None:
sock.close()
return True


def check_exception_chain(err: Exception, object_type: Any) -> bool:
Expand Down
8 changes: 4 additions & 4 deletions tests/test_ingress.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,15 +50,15 @@ async def test_save_on_unload(coresys: CoreSys):
assert coresys.ingress.save_data.called


def test_dynamic_ports(coresys: CoreSys):
async def test_dynamic_ports(coresys: CoreSys):
"""Test dyanmic port handling."""
port_test1 = coresys.ingress.get_dynamic_port("test1")
port_test1 = await coresys.ingress.get_dynamic_port("test1")

assert port_test1
assert coresys.ingress.save_data.called
assert port_test1 == coresys.ingress.get_dynamic_port("test1")
assert port_test1 == await coresys.ingress.get_dynamic_port("test1")

port_test2 = coresys.ingress.get_dynamic_port("test2")
port_test2 = await coresys.ingress.get_dynamic_port("test2")

assert port_test2
assert port_test2 != port_test1
Expand Down
29 changes: 15 additions & 14 deletions tests/utils/test_check_port.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,15 @@
"""Check ports."""
from ipaddress import ip_address

from supervisor.utils import check_port


def test_exists_open_port():
"""Test a exists network port."""
assert check_port(ip_address("8.8.8.8"), 53)


def test_not_exists_port():
"""Test a not exists network service."""
assert not check_port(ip_address("192.0.2.1"), 53)
"""Check ports."""
from ipaddress import ip_address

from supervisor.coresys import CoreSys
from supervisor.utils import check_port


async def test_exists_open_port(coresys: CoreSys):
"""Test a exists network port."""
assert await check_port(ip_address("8.8.8.8"), 53)


async def test_not_exists_port(coresys: CoreSys):
"""Test a not exists network service."""
assert not await check_port(ip_address("192.0.2.1"), 53)

0 comments on commit 883e54f

Please sign in to comment.