Continue central error handling for OAuth (home-assistant#180588)

Author: zweckj

homeassistant/components/cloud/account_link.py

@@ -113,6 +113,12 @@ def domain(self) -> str:
         """Domain that is providing the implementation."""
         return DOMAIN
 
+    @property
+    @override
+    def service_domain(self) -> str:
+        """Domain of the service the tokens are for."""
+        return self.service
+
     @override
     async def async_generate_authorize_url(self, flow_id: str) -> str:
         """Generate a url for the user to authorize."""

homeassistant/exceptions.py

@@ -3,7 +3,7 @@
 from collections.abc import Callable, Generator, Sequence
 from typing import TYPE_CHECKING, Any, override
 
-from aiohttp import ClientResponse, ClientResponseError, RequestInfo
+from aiohttp import ClientError, ClientResponse, ClientResponseError, RequestInfo
 from multidict import MultiMapping
 
 from .util.event_type import EventType
@@ -253,8 +253,20 @@ class ConfigEntryAuthFailed(IntegrationError):
     """Error to indicate that config entry could not authenticate."""
 
 
-class OAuth2TokenRequestError(ClientResponseError, HomeAssistantError):
-    """Error to indicate that the OAuth 2.0 flow could not refresh token."""
+class OAuth2TokenRequestBaseError(ConfigEntryNotReady):
+    """Base class for the errors a failed OAuth 2.0 token request raises.
+
+    Catch this to handle every token request failure; the subclasses differ in
+    whether a status was received and what should happen to the config entry.
+    """
+
+
+class OAuth2TokenRequestError(ClientResponseError, OAuth2TokenRequestBaseError):
+    """Error to indicate that the OAuth 2.0 flow could not refresh token.
+
+    Inherits ConfigEntryNotReady so setup retries without the integration having to
+    map it. Catch it explicitly to handle it differently.
+    """
 
     def __init__(
         self,
@@ -275,15 +287,32 @@ def __init__(
             message=message,
             headers=headers,
         )
-        HomeAssistantError.__init__(self)
+        OAuth2TokenRequestBaseError.__init__(self)
         self.domain = domain
         self.translation_domain = "homeassistant"
         self.translation_key = "oauth2_helper_refresh_failed"
         self.translation_placeholders = {"domain": domain}
         self.generate_message = True
 
 
-class OAuth2TokenRequestTransientError(OAuth2TokenRequestError, ConfigEntryNotReady):
+class OAuth2TokenRequestConnectionError(ClientError, OAuth2TokenRequestBaseError):
+    """Recoverable error to indicate the token request yielded no usable token.
+
+    Covers a request that never got a response and one whose response could not
+    be used, neither of which has a status to tell the causes apart.
+    """
+
+    def __init__(self, *, domain: str) -> None:
+        """Initialize OAuth2TokenRequestConnectionError."""
+        OAuth2TokenRequestBaseError.__init__(self)
+        self.domain = domain
+        self.translation_domain = "homeassistant"
+        self.translation_key = "oauth2_helper_refresh_transient"
+        self.translation_placeholders = {"domain": domain}
+        self.generate_message = True
+
+
+class OAuth2TokenRequestTransientError(OAuth2TokenRequestError):
     """Recoverable error to indicate flow could not refresh token.
 
     Inherits ConfigEntryNotReady so setup retries without the integration having to

homeassistant/helpers/config_entry_oauth2_flow.py

@@ -17,7 +17,7 @@
 import logging
 import secrets
 import time
-from typing import Any, cast, override
+from typing import Any, NoReturn, cast, override
 
 from aiohttp import ClientError, ClientResponseError, client, hdrs, web
 from habluetooth import BluetoothServiceInfoBleak
@@ -30,6 +30,7 @@
 from homeassistant.core import DOMAIN as HOMEASSISTANT_DOMAIN, HomeAssistant, callback
 from homeassistant.exceptions import (
     ImplementationUnavailableError,
+    OAuth2TokenRequestConnectionError,
     OAuth2TokenRequestError,
     OAuth2TokenRequestReauthError,
     OAuth2TokenRequestTransientError,
@@ -105,6 +106,28 @@
 )
 
 
+def _raise_mapped_token_error(err: ClientError, domain: str) -> NoReturn:
+    """Re-raise a failed token request as the matching OAuth2 token error."""
+    if not isinstance(err, ClientResponseError):
+        # Nothing was received, so there is no status to tell the causes apart.
+        _LOGGER.debug("Token request for %s got no response: %s", domain, err)
+        raise OAuth2TokenRequestConnectionError(domain=domain) from err
+
+    kwargs: dict[str, Any] = {
+        "request_info": err.request_info,
+        "history": err.history,
+        "status": err.status,
+        "message": err.message,
+        "headers": err.headers,
+        "domain": domain,
+    }
+    if err.status == HTTPStatus.TOO_MANY_REQUESTS or 500 <= err.status <= 599:
+        raise OAuth2TokenRequestTransientError(**kwargs) from err
+    if 400 <= err.status <= 499:
+        raise OAuth2TokenRequestReauthError(**kwargs) from err
+    raise OAuth2TokenRequestError(**kwargs) from err
+
+
 @callback
 def async_get_redirect_uri(hass: HomeAssistant) -> str:
     """Return the redirect uri."""
@@ -163,11 +186,30 @@ async def async_resolve_external_data(self, external_data: Any) -> dict:
         config entry data.
         """
 
+    @property
+    def service_domain(self) -> str:
+        """Domain of the service the tokens are for.
+
+        Defaults to the implementation itself, but an implementation that obtains
+        tokens on behalf of other integrations has to name the one it serves.
+        """
+        return self.domain
+
     async def async_refresh_token(self, token: dict) -> dict:
         """Refresh a token and update expires info."""
-        new_token = await self._async_refresh_token(token)
+        try:
+            new_token = await self._async_refresh_token(token)
+        except OAuth2TokenRequestError, OAuth2TokenRequestConnectionError:
+            raise
+        except ClientError as err:
+            # Implementations that issue their own token request may not map their
+            # failures, so callers would see a raw aiohttp error instead.
+            _raise_mapped_token_error(err, self.service_domain)
         # Force int for non-compliant oauth2 providers
-        new_token["expires_in"] = int(new_token["expires_in"])
+        try:
+            new_token["expires_in"] = int(new_token["expires_in"])
+        except (KeyError, TypeError, ValueError) as err:
+            raise OAuth2TokenRequestConnectionError(domain=self.service_domain) from err
         new_token["expires_at"] = time.time() + new_token["expires_in"]
         return new_token
 
@@ -268,6 +310,11 @@ async def _async_refresh_token(self, token: dict) -> dict:
             }
         )
 
+        # Merging a response without one would keep the stale access token while
+        # extending its expiry, so the session would never recover.
+        if not new_token.get("access_token"):
+            raise OAuth2TokenRequestConnectionError(domain=self.service_domain)
+
         return {**token, **new_token}
 
     async def _token_request(self, data: dict) -> dict:
@@ -306,38 +353,13 @@ async def _token_request(self, data: dict) -> dict:
                     detail,
                 )
             resp.raise_for_status()
+            return cast(dict, await resp.json())
         except ClientResponseError as err:
-            if err.status == HTTPStatus.TOO_MANY_REQUESTS or 500 <= err.status <= 599:
-                # Recoverable error
-                raise OAuth2TokenRequestTransientError(
-                    request_info=err.request_info,
-                    history=err.history,
-                    status=err.status,
-                    message=err.message,
-                    headers=err.headers,
-                    domain=self._domain,
-                ) from err
-            if 400 <= err.status <= 499:
-                # Non-recoverable error
-                raise OAuth2TokenRequestReauthError(
-                    request_info=err.request_info,
-                    history=err.history,
-                    status=err.status,
-                    message=err.message,
-                    headers=err.headers,
-                    domain=self._domain,
-                ) from err
-
-            raise OAuth2TokenRequestError(
-                request_info=err.request_info,
-                history=err.history,
-                status=err.status,
-                message=err.message,
-                headers=err.headers,
-                domain=self._domain,
-            ) from err
-
-        return cast(dict, await resp.json())
+            _raise_mapped_token_error(err, self.service_domain)
+        except ClientError as err:
+            # Bare TimeoutError is left alone so an enclosing asyncio.timeout still
+            # aborts with oauth_timeout; aiohttp's own timeouts are ClientErrors.
+            _raise_mapped_token_error(err, self.service_domain)
 
 
 class LocalOAuth2ImplementationWithPkce(LocalOAuth2Implementation):
@@ -844,6 +866,15 @@ async def async_ensure_token_valid(self) -> None:
                 self.config_entry.async_start_reauth_if_available(self.hass)
                 raise
 
+            # Checked before storing, so reads can trust what is on the entry.
+            if any(
+                new_token.get(field) in (None, "")
+                for field in ("access_token", "expires_at")
+            ):
+                raise OAuth2TokenRequestConnectionError(
+                    domain=self.implementation.service_domain
+                )
+
             self.hass.config_entries.async_update_entry(
                 self.config_entry, data={**self.config_entry.data, "token": new_token}
             )

script/hassfest/quality_scale_validation/test_before_setup.py

@@ -14,6 +14,13 @@
     "ConfigEntryError",
 }
 
+# Helpers that raise one of the above on the caller's behalf, so an integration
+# awaiting them satisfies the rule without repeating the mapping itself.
+_VALID_AWAITED_CALLS = {
+    "async_config_entry_first_refresh",
+    "async_ensure_token_valid",
+}
+
 
 def _get_exception_name(expression: ast.expr) -> str:
     """Get the name of the exception being raised."""
@@ -58,17 +65,20 @@ def _raises_exception(integration: Integration) -> bool:
     return False
 
 
-def _calls_first_refresh(async_setup_entry_function: ast.AsyncFunctionDef) -> bool:
-    """Check that a async_config_entry_first_refresh within `async_setup_entry`."""
-    for node in ast.walk(async_setup_entry_function):
-        if (
-            isinstance(node, ast.Call)
-            and isinstance(node.func, ast.Attribute)
-            and node.func.attr == "async_config_entry_first_refresh"
-        ):
-            return True
+def _awaits_raising_helper(async_setup_entry_function: ast.AsyncFunctionDef) -> bool:
+    """Check that `async_setup_entry` awaits a helper that raises on its behalf.
 
-    return False
+    The call only has to sit somewhere inside an await, so gathering several of
+    them still counts, while an unawaited call does not.
+    """
+    return any(
+        isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Attribute)
+        and node.func.attr in _VALID_AWAITED_CALLS
+        for await_node in ast.walk(async_setup_entry_function)
+        if isinstance(await_node, ast.Await)
+        for node in ast.walk(await_node)
+    )
 
 
 def _get_setup_entry_function(module: ast.Module) -> ast.AsyncFunctionDef | None:
@@ -90,6 +100,8 @@ def validate(
     if not (async_setup_entry := _get_setup_entry_function(init)):
         return [f"Could not find `async_setup_entry` in {init_file}"]
 
-    if not (_calls_first_refresh(async_setup_entry) or _raises_exception(integration)):
+    if not (
+        _awaits_raising_helper(async_setup_entry) or _raises_exception(integration)
+    ):
         return [f"Integration does not raise one of {_VALID_EXCEPTIONS}"]
     return None

tests/components/cloud/test_account_link.py

@@ -6,7 +6,7 @@
 from time import time
 from unittest.mock import AsyncMock, Mock, patch
 
-from aiohttp import ClientResponseError, RequestInfo
+from aiohttp import ClientError, ClientResponseError, RequestInfo
 import pytest
 from yarl import URL
 
@@ -16,6 +16,7 @@
 from homeassistant.core import HomeAssistant
 from homeassistant.data_entry_flow import FlowResultType
 from homeassistant.exceptions import (
+    OAuth2TokenRequestConnectionError,
     OAuth2TokenRequestError,
     OAuth2TokenRequestReauthError,
     OAuth2TokenRequestTransientError,
@@ -304,3 +305,24 @@ async def test_refresh_token_error(
 
     assert exc_info.value.status == status
     assert exc_info.value.domain == "test"
+
+
+async def test_refresh_token_connection_error(hass: HomeAssistant) -> None:
+    """Test a failure without a response reports the service, not the cloud domain."""
+    hass.data[DATA_CLOUD] = None
+    impl = account_link.CloudOAuth2Implementation(hass, "test")
+
+    with (
+        patch(
+            "hass_nabucasa.account_link.async_fetch_access_token",
+            side_effect=ClientError("Cannot connect"),
+        ),
+        pytest.raises(OAuth2TokenRequestConnectionError) as exc_info,
+    ):
+        await impl.async_refresh_token(
+            {"refresh_token": "mock-refresh", "access_token": "mock-access"}
+        )
+
+    assert impl.domain == "cloud"
+    assert exc_info.value.domain == "test"
+    assert exc_info.value.translation_placeholders == {"domain": "test"}