Fix restoring login from stored access token. (#10)

* Rename authentication token to proper access token

* Bump matrix-nio to 0.21.1 to take advantage of whoami

* Fix restoring login from stored access token.

* Extra logging.

* Bump matrix-nio  to 0.21.2 for bugfixes.

Author: PaarthShah

custom_components/matrix/__init__.py

@@ -19,7 +19,7 @@
     LoginError,
     Response,
     UploadError,
-    UploadResponse,
+    UploadResponse, WhoamiError, WhoamiResponse,
 )
 import voluptuous as vol
 
@@ -166,7 +166,7 @@ def __init__(
         self.hass = hass
 
         self._session_filepath = config_file
-        self._auth_tokens: JsonObjectType = {}
+        self._access_tokens: JsonObjectType = {}
 
         self._homeserver = homeserver
         self._verify_tls = verify_ssl
@@ -192,7 +192,7 @@ async def stop_client(event: HassEvent) -> None:
 
         async def handle_startup(event: HassEvent) -> None:
             """Run once when Home Assistant finished startup."""
-            self._auth_tokens = await self._get_auth_tokens()
+            self._access_tokens = await self._get_auth_tokens()
             await self._login()
             await self._join_rooms()
             # Sync once so that we don't respond to past events.
@@ -283,44 +283,55 @@ async def _join_rooms(self) -> None:
         await asyncio.wait(rooms)
 
     async def _get_auth_tokens(self) -> JsonObjectType:
-        """Read sorted authentication tokens from disk."""
+        """Read sorted access tokens from disk."""
         try:
             return load_json_object(self._session_filepath)
         except HomeAssistantError as ex:
             _LOGGER.warning(
-                "Loading authentication tokens from file '%s' failed: %s",
+                "Loading access tokens from file '%s' failed: %s",
                 self._session_filepath,
                 str(ex),
             )
             return {}
 
     async def _store_auth_token(self, token: str) -> None:
-        """Store authentication token to session and persistent storage."""
-        self._auth_tokens[self._mx_id] = token
+        """Store access token to session and persistent storage."""
+        self._access_tokens[self._mx_id] = token
 
         await self.hass.async_add_executor_job(
-            save_json, self._session_filepath, self._auth_tokens, True  # private=True
+            save_json, self._session_filepath, self._access_tokens, True  # private=True
         )
 
     async def _login(self) -> None:
         """Log in to the Matrix homeserver.
 
-        Attempts to use the stored authentication token.
+        Attempts to use the stored access token.
         If that fails, then tries using the password.
         If that also fails, raises LocalProtocolError.
         """
 
-        # If we have an authentication token
-        if (token := self._auth_tokens.get(self._mx_id)) is not None:
-            response = await self._client.login(token=token)
-            _LOGGER.debug("Logging in using stored token")
-
-            if isinstance(response, LoginError):
+        # If we have an access token
+        if (token := self._access_tokens.get(self._mx_id)) is not None:
+            _LOGGER.debug("Restoring login from stored access token")
+            self._client.restore_login(
+                user_id=self._client.user_id,
+                device_id=self._client.device_id,
+                access_token=token,
+            )
+            response = await self._client.whoami()
+            if isinstance(response, WhoamiError):
                 _LOGGER.warning(
-                    "Login by token failed: %s, %s",
+                    "Restoring login from access token failed: %s, %s",
                     response.status_code,
                     response.message,
                 )
+                self._client.access_token = ""  # Force a soft-logout if the homeserver didn't.
+            elif isinstance(response, WhoamiResponse):
+                _LOGGER.debug(
+                    "Successfully restored login from access token: user_id '%s', device_id '%s'",
+                    response.user_id,
+                    response.device_id,
+                )
 
         # If the token login did not succeed
         if not self._client.logged_in:

custom_components/matrix/manifest.json

@@ -6,6 +6,6 @@
   "iot_class": "cloud_push",
   "issue_tracker": "https://github.com/PaarthShah/matrix-nio-hacs/issues",
   "loggers": ["matrix_client"],
-  "requirements": ["matrix-nio==0.20.2", "Pillow==9.5.0"],
-  "version": "v1.0.0"
+  "requirements": ["matrix-nio==0.21.2", "Pillow==10.0.0"],
+  "version": "v1.0.2"
 }

requirements.txt

@@ -2,5 +2,5 @@ colorlog==6.7.0
 homeassistant>=2023.7.2
 pip>=21.0,<23.2
 ruff==0.0.278
-matrix-nio==0.21.0
-Pillow==10.0.0
+matrix-nio==0.21.2
+Pillow==10.0.0

tests/conftest.py

@@ -26,7 +26,7 @@
     LoginError,
     LoginResponse,
     Response,
-    UploadResponse,
+    UploadResponse, WhoamiError,
 )
 from pytest_homeassistant_custom_component.common import async_capture_events
 
@@ -49,6 +49,7 @@
 TEST_JOINABLE_ROOMS = ["!RoomIdString:example.com", "#RoomAliasString:example.com"]
 TEST_BAD_ROOM = "!UninvitedRoom:example.com"
 TEST_MXID = "@user:example.com"
+TEST_DEVICE_ID = "FAKEID"
 TEST_PASSWORD = "password"
 TEST_TOKEN = "access_token"
 
@@ -63,8 +64,6 @@ def my_enable_custom_integrations(enable_custom_integrations):
 class _MockAsyncClient(AsyncClient):
     """Mock class to simulate MatrixBot._client's I/O methods."""
 
-    logged_in: bool = False
-
     async def close(self):
         return None
 
@@ -76,17 +75,24 @@ async def join(self, room_id: RoomID):
 
     async def login(self, *args, **kwargs):
         if kwargs.get("password") == TEST_PASSWORD or kwargs.get("token") == TEST_TOKEN:
-            self.logged_in = True
             self.access_token = TEST_TOKEN
             return LoginResponse(
                 access_token=TEST_TOKEN,
                 device_id="test_device",
                 user_id=TEST_MXID,
             )
         else:
-            self.logged_in = False
+            self.access_token = ""
             return LoginError(message="LoginError", status_code="status_code")
 
+    async def whoami(self):
+        if self.access_token == TEST_TOKEN:
+            self.user_id = TEST_MXID
+            self.device_id = TEST_DEVICE_ID
+        else:
+            self.access_token = ""
+            return WhoamiError(message="Invalid access token passed.", status_code="M_UNKNOWN_TOKEN")
+
     async def room_send(self, *args, **kwargs):
         if not self.logged_in:
             raise LocalProtocolError

tests/test_login.py

@@ -19,24 +19,24 @@ async def test_login(matrix_bot: MatrixBot, mock_client):
     assert matrix_bot._client.logged_in
 
     # Test good login using a good password and bad stored access_token.
-    matrix_bot._auth_tokens = {TEST_MXID: "WrongToken"}
+    matrix_bot._access_tokens = {TEST_MXID: "WrongToken"}
     await matrix_bot._login()
     assert matrix_bot._client.logged_in
 
     # Test good login using a good access_token and bad password.
-    matrix_bot._auth_tokens = {TEST_MXID: TEST_TOKEN}
+    matrix_bot._access_tokens = {TEST_MXID: TEST_TOKEN}
     await matrix_bot._login()
     assert matrix_bot._client.logged_in
 
     # Test bad login using a bad password and bad stored access_token.
     matrix_bot._password = "WrongPassword"
-    matrix_bot._auth_tokens = {TEST_MXID: "WrongToken"}
+    matrix_bot._access_tokens = {TEST_MXID: "WrongToken"}
     with pytest.raises(ConfigEntryAuthFailed):
         await matrix_bot._login()
     assert not matrix_bot._client.logged_in
 
     # Test bad login using bad password and missing access token.
-    matrix_bot._auth_tokens = {}
+    matrix_bot._access_tokens = {}
     with pytest.raises(ConfigEntryAuthFailed):
         await matrix_bot._login()
     assert not matrix_bot._client.logged_in