Use orjson encoder for websocket messages (#4854)

I missed that these need to have dumps passed
home-assistant · Jan 30, 2024 · 375789b · 375789b
1 parent 140b769
commit 375789b
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/supervisor/api/proxy.py b/supervisor/api/proxy.py
@@ -14,6 +14,7 @@
 
 from ..coresys import CoreSysAttributes
 from ..exceptions import APIError, HomeAssistantAPIError, HomeAssistantAuthError
+from ..utils.json import json_dumps
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -145,7 +146,8 @@ async def _websocket_client(self) -> ClientWebSocketResponse:
                 {
                     "type": "auth",
                     "access_token": self.sys_homeassistant.api.access_token,
-                }
+                },
+                dumps=json_dumps,
             )
 
             data = await client.receive_json()
@@ -202,7 +204,8 @@ async def websocket(self, request: web.Request):
         # handle authentication
         try:
             await server.send_json(
-                {"type": "auth_required", "ha_version": self.sys_homeassistant.version}
+                {"type": "auth_required", "ha_version": self.sys_homeassistant.version},
+                dumps=json_dumps,
             )
 
             # Check API access
@@ -215,14 +218,16 @@ async def websocket(self, request: web.Request):
             if not addon or not addon.access_homeassistant_api:
                 _LOGGER.warning("Unauthorized WebSocket access!")
                 await server.send_json(
-                    {"type": "auth_invalid", "message": "Invalid access"}
+                    {"type": "auth_invalid", "message": "Invalid access"},
+                    dumps=json_dumps,
                 )
                 return server
 
             _LOGGER.info("WebSocket access from %s", addon.slug)
 
             await server.send_json(
-                {"type": "auth_ok", "ha_version": self.sys_homeassistant.version}
+                {"type": "auth_ok", "ha_version": self.sys_homeassistant.version},
+                dumps=json_dumps,
             )
         except (RuntimeError, ValueError) as err:
             _LOGGER.error("Can't initialize handshake: %s", err)

diff --git a/supervisor/homeassistant/websocket.py b/supervisor/homeassistant/websocket.py
@@ -26,6 +26,7 @@
     HomeAssistantWSError,
     HomeAssistantWSNotSupported,
 )
+from ..utils.json import json_dumps
 from .const import CLOSING_STATES, WSEvent, WSType
 
 MIN_VERSION = {
@@ -74,7 +75,7 @@ async def async_send_message(self, message: dict[str, Any]) -> None:
         self._message_id += 1
         _LOGGER.debug("Sending: %s", message)
         try:
-            await self._client.send_json(message)
+            await self._client.send_json(message, dumps=json_dumps)
         except ConnectionError as err:
             raise HomeAssistantWSConnectionError(err) from err
 
@@ -85,7 +86,7 @@ async def async_send_command(self, message: dict[str, Any]) -> dict | None:
         self._futures[message["id"]] = self._loop.create_future()
         _LOGGER.debug("Sending: %s", message)
         try:
-            await self._client.send_json(message)
+            await self._client.send_json(message, dumps=json_dumps)
         except ConnectionError as err:
             raise HomeAssistantWSConnectionError(err) from err
 
@@ -163,7 +164,9 @@ async def connect_with_auth(
 
         hello_message = await client.receive_json()
 
-        await client.send_json({ATTR_TYPE: WSType.AUTH, ATTR_ACCESS_TOKEN: token})
+        await client.send_json(
+            {ATTR_TYPE: WSType.AUTH, ATTR_ACCESS_TOKEN: token}, dumps=json_dumps
+        )
 
         auth_ok_message = await client.receive_json()