Recreate aiohttp ClientSession after DNS plug-in load (#5862)

* Recreate aiohttp ClientSession after DNS plug-in load

Create a temporary ClientSession early in case we need to load version
information from the internet. This doesn't use the final DNS setup
and hence might fail to load in certain situations since we don't have
the fallback mechanims in place yet. But if the DNS container image
is present, we'll continue the setup and load the DNS plug-in. We then
can recreate the ClientSession such that it uses the DNS plug-in.

This works around an issue with aiodns, which today doesn't reload
`resolv.conf` automatically when it changes. This lead to Supervisor
using the initial `resolv.conf` as created by Docker. It meant that
we did not use the DNS plug-in (and its fallback capabilities) in
Supervisor. Also it meant that changes to the DNS setup at runtime
did not propagate to the aiohttp ClientSession (as observed in #5332).

* Mock aiohttp.ClientSession for all tests

Currently in several places pytest actually uses the aiohttp
ClientSession and reaches out to the internet. This is not ideal
for unit tests and should be avoided.

This creates several new fixtures to aid this effort: The `websession`
fixture simply returns a mocked aiohttp.ClientSession, which can be
used whenever a function is tested which needs the global websession.

A separate new fixture to mock the connectivity check named
`supervisor_internet` since this is often used through the Job
decorator which require INTERNET_SYSTEM.

And the `mock_update_data` uses the already existing update json
test data from the fixture directory instead of loading the data
from the internet.

* Log ClientSession nameserver information

When recreating the aiohttp ClientSession, log information what
nameservers exactly are going to be used.

* Refuse ClientSession initialization when API is available

Previous attempts to reinitialize the ClientSession have shown
use of the ClientSession after it was closed due to API requets
being handled in parallel to the reinitialization (see #5851).
Make sure this is not possible by refusing to reinitialize the
ClientSession when the API is available.

* Fix pytests

Also sure we don't create aiohttp ClientSession objects unnecessarily.

* Apply suggestions from code review

Co-authored-by: Jan Čermák <sairon@users.noreply.github.com>

---------

Co-authored-by: Jan Čermák <sairon@users.noreply.github.com>

Author: agners

pyproject.toml

@@ -230,6 +230,9 @@ filterwarnings = [
     "ignore:pkg_resources is deprecated as an API:DeprecationWarning:dirhash",
     "ignore::pytest.PytestUnraisableExceptionWarning",
 ]
+markers = [
+    "no_mock_init_websession: disable the autouse mock of init_websession for this test",
+]
 
 [tool.ruff]
 lint.select = [

supervisor/api/middleware/security.py

@@ -20,7 +20,7 @@
     ROLE_DEFAULT,
     ROLE_HOMEASSISTANT,
     ROLE_MANAGER,
-    CoreState,
+    VALID_API_STATES,
 )
 from ...coresys import CoreSys, CoreSysAttributes
 from ...utils import version_is_new_enough
@@ -200,11 +200,7 @@ async def block_bad_requests(self, request: Request, handler: Callable) -> Respo
     @middleware
     async def system_validation(self, request: Request, handler: Callable) -> Response:
         """Check if core is ready to response."""
-        if self.sys_core.state not in (
-            CoreState.STARTUP,
-            CoreState.RUNNING,
-            CoreState.FREEZE,
-        ):
+        if self.sys_core.state not in VALID_API_STATES:
             return api_return_error(
                 message=f"System is not ready with state: {self.sys_core.state}"
             )

supervisor/const.py

@@ -552,3 +552,12 @@ def from_dict(cls, data: dict[str, dict[str, str | None]]) -> Self:
     CoreState.STARTUP,
     CoreState.SETUP,
 ]
+
+# States in which the API can be used (enforced by system_validation())
+VALID_API_STATES = frozenset(
+    {
+        CoreState.STARTUP,
+        CoreState.RUNNING,
+        CoreState.FREEZE,
+    }
+)

supervisor/core.py

@@ -124,6 +124,19 @@ async def setup(self):
         """Start setting up supervisor orchestration."""
         await self.set_state(CoreState.SETUP)
 
+        # Initialize websession early. At this point we'll use the Docker DNS proxy
+        # at 127.0.0.11, which does not have the fallback feature and hence might
+        # fail in certain environments. But a websession is required to get the
+        # initial version information after a device wipe or otherwise empty state
+        # (e.g. CI environment, Supervised).
+        #
+        # An OS installation has the plug-in container images pre-installed, so we
+        # setup can continue even if this early websession fails to connect to the
+        # internet. We'll reinitialize the websession when the DNS plug-in is up to
+        # make sure the DNS plug-in along with its fallback capabilities is used
+        # (see #5857).
+        await self.coresys.init_websession()
+
         # Check internet on startup
         await self.sys_supervisor.check_connectivity()
 

supervisor/coresys.py

@@ -21,6 +21,7 @@
     ENV_SUPERVISOR_MACHINE,
     MACHINE_ID,
     SERVER_SOFTWARE,
+    VALID_API_STATES,
 )
 
 if TYPE_CHECKING:
@@ -68,7 +69,6 @@ def __init__(self):
 
         # External objects
         self._loop: asyncio.BaseEventLoop = asyncio.get_running_loop()
-        self._websession: aiohttp.ClientSession = aiohttp.ClientSession()
 
         # Global objects
         self._config: CoreConfig = CoreConfig()
@@ -100,11 +100,7 @@ def __init__(self):
         self._security: Security | None = None
         self._bus: Bus | None = None
         self._mounts: MountManager | None = None
-
-        # Set default header for aiohttp
-        self._websession._default_headers = MappingProxyType(
-            {aiohttp.hdrs.USER_AGENT: SERVER_SOFTWARE}
-        )
+        self._websession: aiohttp.ClientSession | None = None
 
         # Task factory attributes
         self._set_task_context: list[Callable[[Context], Context]] = []
@@ -114,6 +110,33 @@ async def load_config(self) -> Self:
         await self.config.read_data()
         return self
 
+    async def init_websession(self) -> None:
+        """Initialize global aiohttp ClientSession."""
+        if self.core.state in VALID_API_STATES:
+            # Make sure we don't reinitialize the session if the API is running (see #5851)
+            raise RuntimeError(
+                "Initializing ClientSession is not safe when API is running"
+            )
+
+        if self._websession:
+            await self._websession.close()
+
+        resolver = aiohttp.AsyncResolver()
+
+        # pylint: disable=protected-access
+        _LOGGER.debug(
+            "Initializing ClientSession with AsyncResolver. Using nameservers %s",
+            resolver._resolver.nameservers,
+        )
+        connector = aiohttp.TCPConnector(loop=self.loop, resolver=resolver)
+
+        session = aiohttp.ClientSession(
+            headers=MappingProxyType({aiohttp.hdrs.USER_AGENT: SERVER_SOFTWARE}),
+            connector=connector,
+        )
+
+        self._websession = session
+
     async def init_machine(self):
         """Initialize machine information."""
 
@@ -165,6 +188,8 @@ def loop(self) -> asyncio.BaseEventLoop:
     @property
     def websession(self) -> aiohttp.ClientSession:
         """Return websession object."""
+        if self._websession is None:
+            raise RuntimeError("WebSession not setup yet")
         return self._websession
 
     @property

supervisor/plugins/dns.py

@@ -177,7 +177,13 @@ async def load(self) -> None:
 
         # Update supervisor
         await self._write_resolv(HOST_RESOLV)
-        await self.sys_supervisor.check_connectivity()
+
+        # Reinitializing aiohttp.ClientSession after DNS setup makes sure that
+        # aiodns is using the right DNS servers (see #5857).
+        # At this point it should be fairly safe to replace the session since
+        # we only use the session synchronously during setup and not thorugh the
+        # API which previously caused issues (see #5851).
+        await self.coresys.init_websession()
 
     async def install(self) -> None:
         """Install CoreDNS."""

tests/addons/test_manager.py

@@ -190,7 +190,7 @@ async def test_addon_shutdown_error(
 
 
 async def test_addon_uninstall_removes_discovery(
-    coresys: CoreSys, install_addon_ssh: Addon
+    coresys: CoreSys, install_addon_ssh: Addon, websession: MagicMock
 ):
     """Test discovery messages removed when addon uninstalled."""
     assert coresys.discovery.list_messages == []
@@ -203,7 +203,6 @@ async def test_addon_uninstall_removes_discovery(
     assert coresys.discovery.list_messages == [message]
 
     coresys.homeassistant.api.ensure_access_token = AsyncMock()
-    coresys.websession.delete = MagicMock()
 
     await coresys.addons.uninstall(TEST_ADDON_SLUG)
     await asyncio.sleep(0)

tests/api/test_auth.py

@@ -1,14 +1,15 @@
 """Test auth API."""
 
 from datetime import UTC, datetime, timedelta
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from aiohttp.test_utils import TestClient
 import pytest
 
 from supervisor.addons.addon import Addon
 from supervisor.coresys import CoreSys
 
+from tests.common import MockResponse
 from tests.const import TEST_ADDON_SLUG
 
 LIST_USERS_RESPONSE = [
@@ -78,7 +79,10 @@ def fixture_mock_check_login(coresys: CoreSys):
 
 
 async def test_password_reset(
-    api_client: TestClient, coresys: CoreSys, caplog: pytest.LogCaptureFixture
+    api_client: TestClient,
+    coresys: CoreSys,
+    caplog: pytest.LogCaptureFixture,
+    websession: MagicMock,
 ):
     """Test password reset api."""
     coresys.homeassistant.api.access_token = "abc123"
@@ -87,15 +91,12 @@ async def test_password_reset(
         days=1
     )
 
-    mock_websession = AsyncMock()
-    mock_websession.post.return_value.__aenter__.return_value.status = 200
-    with patch("supervisor.coresys.aiohttp.ClientSession.post") as post:
-        post.return_value.__aenter__.return_value.status = 200
-        resp = await api_client.post(
-            "/auth/reset", json={"username": "john", "password": "doe"}
-        )
-        assert resp.status == 200
-        assert "Successful password reset for 'john'" in caplog.text
+    websession.post = MagicMock(return_value=MockResponse(status=200))
+    resp = await api_client.post(
+        "/auth/reset", json={"username": "john", "password": "doe"}
+    )
+    assert resp.status == 200
+    assert "Successful password reset for 'john'" in caplog.text
 
 
 async def test_list_users(

tests/api/test_backups.py

@@ -276,6 +276,7 @@ async def test_api_backup_restore_background(
     backup_type: str,
     options: dict[str, Any],
     tmp_supervisor_data: Path,
+    supervisor_internet: AsyncMock,
 ):
     """Test background option on backup/restore APIs."""
     await coresys.core.set_state(CoreState.RUNNING)
@@ -472,6 +473,7 @@ async def test_restore_immediate_errors(
     api_client: TestClient,
     coresys: CoreSys,
     mock_partial_backup: Backup,
+    supervisor_internet: AsyncMock,
 ):
     """Test restore errors that return immediately even in background mode."""
     await coresys.core.set_state(CoreState.RUNNING)
@@ -1010,6 +1012,7 @@ async def test_restore_backup_from_location(
     coresys: CoreSys,
     tmp_supervisor_data: Path,
     local_location: str | None,
+    supervisor_internet: AsyncMock,
 ):
     """Test restoring a backup from a specific location."""
     await coresys.core.set_state(CoreState.RUNNING)
@@ -1059,6 +1062,7 @@ async def test_restore_backup_from_location(
 async def test_restore_backup_unencrypted_after_encrypted(
     api_client: TestClient,
     coresys: CoreSys,
+    supervisor_internet: AsyncMock,
 ):
     """Test restoring an unencrypted backup after an encrypted backup and vis-versa."""
     enc_tar = copy(get_fixture_path("test_consolidate.tar"), coresys.config.path_backup)
@@ -1131,6 +1135,7 @@ async def test_restore_homeassistant_adds_env(
     docker: DockerAPI,
     backup_type: str,
     postbody: dict[str, Any],
+    supervisor_internet: AsyncMock,
 ):
     """Test restoring home assistant from backup adds env to container."""
     event = asyncio.Event()
@@ -1328,6 +1333,7 @@ async def test_missing_file_removes_location_from_cache(
     url_path: str,
     body: dict[str, Any] | None,
     backup_file: str,
+    supervisor_internet: AsyncMock,
 ):
     """Test finding a missing file removes the location from cache."""
     await coresys.core.set_state(CoreState.RUNNING)
@@ -1387,6 +1393,7 @@ async def test_missing_file_removes_backup_from_cache(
     url_path: str,
     body: dict[str, Any] | None,
     backup_file: str,
+    supervisor_internet: AsyncMock,
 ):
     """Test finding a missing file removes the backup from cache if its the only one."""
     await coresys.core.set_state(CoreState.RUNNING)
@@ -1412,7 +1419,9 @@ async def test_missing_file_removes_backup_from_cache(
 
 @pytest.mark.usefixtures("tmp_supervisor_data")
 async def test_immediate_list_after_missing_file_restore(
-    api_client: TestClient, coresys: CoreSys
+    api_client: TestClient,
+    coresys: CoreSys,
+    supervisor_internet: AsyncMock,
 ):
     """Test race with reload for missing file on restore does not error."""
     await coresys.core.set_state(CoreState.RUNNING)

tests/api/test_discovery.py

@@ -84,12 +84,14 @@ async def test_api_list_discovery(
 
 @pytest.mark.parametrize("api_client", [TEST_ADDON_SLUG], indirect=True)
 async def test_api_send_del_discovery(
-    api_client: TestClient, coresys: CoreSys, install_addon_ssh: Addon
+    api_client: TestClient,
+    coresys: CoreSys,
+    install_addon_ssh: Addon,
+    websession: MagicMock,
 ):
     """Test adding and removing discovery."""
     install_addon_ssh.data["discovery"] = ["test"]
     coresys.homeassistant.api.ensure_access_token = AsyncMock()
-    coresys.websession.post = MagicMock()
 
     resp = await api_client.post("/discovery", json={"service": "test", "config": {}})
     assert resp.status == 200

tests/api/test_security.py

@@ -1,5 +1,7 @@
 """Test Supervisor API."""
 
+from unittest.mock import AsyncMock
+
 import pytest
 
 from supervisor.coresys import CoreSys
@@ -36,7 +38,9 @@ async def test_api_security_options_pwned(api_client, coresys: CoreSys):
 
 
 @pytest.mark.asyncio
-async def test_api_integrity_check(api_client, coresys: CoreSys):
+async def test_api_integrity_check(
+    api_client, coresys: CoreSys, supervisor_internet: AsyncMock
+):
     """Test security integrity check."""
     coresys.security.content_trust = False
 

tests/api/test_store.py

@@ -2,7 +2,7 @@
 
 import asyncio
 from pathlib import Path
-from unittest.mock import MagicMock, PropertyMock, patch
+from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch
 
 from aiohttp import ClientResponse
 from aiohttp.test_utils import TestClient
@@ -92,7 +92,9 @@ async def test_api_store_repositories_repository(
     assert result["data"]["slug"] == repository.slug
 
 
-async def test_api_store_add_repository(api_client: TestClient, coresys: CoreSys):
+async def test_api_store_add_repository(
+    api_client: TestClient, coresys: CoreSys, supervisor_internet: AsyncMock
+) -> None:
     """Test POST /store/repositories REST API."""
     with (
         patch("supervisor.store.repository.Repository.load", return_value=None),

tests/api/test_supervisor.py

@@ -2,7 +2,7 @@
 
 # pylint: disable=protected-access
 import time
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from aiohttp.test_utils import TestClient
 from blockbuster import BlockingError
@@ -34,7 +34,7 @@ async def test_api_supervisor_options_debug(api_client: TestClient, coresys: Cor
 
 
 async def test_api_supervisor_options_add_repository(
-    api_client: TestClient, coresys: CoreSys
+    api_client: TestClient, coresys: CoreSys, supervisor_internet: AsyncMock
 ):
     """Test add a repository via POST /supervisor/options REST API."""
     assert REPO_URL not in coresys.store.repository_urls
@@ -231,7 +231,9 @@ async def test_api_supervisor_fallback_log_capture(
         capture_exception.assert_called_once()
 
 
-async def test_api_supervisor_reload(api_client: TestClient):
+async def test_api_supervisor_reload(
+    api_client: TestClient, supervisor_internet: AsyncMock, websession: MagicMock
+):
     """Test supervisor reload."""
     resp = await api_client.post("/supervisor/reload")
     assert resp.status == 200