Add websocket proxy support (#286)

* Add websocket proxy support

* forward

* update proxy code

* fix import

* fix import

* fix

* reorder

* fix setup

* fix code

* stage al

* fix lint

* convert it into object

* fix lint

* fix url

* fix routing

* update log output

* fix future

* add loop

* Update log messages & error handling

* fix error message

* Update logging

* improve handling

* better error handling

* Fix server read

* fix cancel reader

hassio/api/__init__.py

@@ -8,6 +8,7 @@ from .addons import APIAddons
 from .homeassistant import APIHomeAssistant
 from .host import APIHost
 from .network import APINetwork
+from .proxy import APIProxy
 from .supervisor import APISupervisor
 from .security import APISecurity
 from .snapshots import APISnapshots
@@ -63,9 +64,9 @@ class RestAPI(object):
             '/supervisor/options', api_supervisor.options)
         self.webapp.router.add_get('/supervisor/logs', api_supervisor.logs)
 
-    def register_homeassistant(self, dock_homeassistant):
+    def register_homeassistant(self, homeassistant):
         """Register homeassistant function."""
-        api_hass = APIHomeAssistant(self.config, self.loop, dock_homeassistant)
+        api_hass = APIHomeAssistant(self.config, self.loop, homeassistant)
 
         self.webapp.router.add_get('/homeassistant/info', api_hass.info)
         self.webapp.router.add_get('/homeassistant/logs', api_hass.logs)
@@ -75,12 +76,21 @@ class RestAPI(object):
         self.webapp.router.add_post('/homeassistant/stop', api_hass.stop)
         self.webapp.router.add_post('/homeassistant/start', api_hass.start)
         self.webapp.router.add_post('/homeassistant/check', api_hass.check)
+
+    def register_proxy(self, homeassistant, websession):
+        """Register HomeAssistant API Proxy."""
+        api_proxy = APIProxy(self.loop, homeassistant, websession)
+
+        self.webapp.router.add_get(
+            '/homeassistant/api/websocket', api_proxy.websocket)
+        self.webapp.router.add_get(
+            '/homeassistant/websocket', api_proxy.websocket)
         self.webapp.router.add_post(
-            '/homeassistant/api/{path:.+}', api_hass.api)
+            '/homeassistant/api/{path:.+}', api_proxy.api)
         self.webapp.router.add_get(
-            '/homeassistant/api/{path:.+}', api_hass.api)
+            '/homeassistant/api/{path:.+}', api_proxy.api)
         self.webapp.router.add_get(
-            '/homeassistant/api', api_hass.api)
+            '/homeassistant/api', api_proxy.api)
 
     def register_addons(self, addons):
         """Register homeassistant function."""

hassio/api/homeassistant.py

@@ -2,18 +2,13 @@
 import asyncio
 import logging
 
-import aiohttp
-from aiohttp import web
-from aiohttp.web_exceptions import HTTPBadGateway
-from aiohttp.hdrs import CONTENT_TYPE
-import async_timeout
 import voluptuous as vol
 
 from .util import api_process, api_process_raw, api_validate
 from ..const import (
     ATTR_VERSION, ATTR_LAST_VERSION, ATTR_DEVICES, ATTR_IMAGE, ATTR_CUSTOM,
     ATTR_BOOT, ATTR_PORT, ATTR_PASSWORD, ATTR_SSL, ATTR_WATCHDOG,
-    CONTENT_TYPE_BINARY, HEADER_HA_ACCESS)
+    CONTENT_TYPE_BINARY)
 from ..validate import HASS_DEVICES, NETWORK_PORT
 
 _LOGGER = logging.getLogger(__name__)
@@ -46,45 +41,6 @@ class APIHomeAssistant(object):
         self.loop = loop
         self.homeassistant = homeassistant
 
-    async def homeassistant_proxy(self, path, request, timeout=300):
-        """Return a client request with proxy origin for Home-Assistant."""
-        url = f"{self.homeassistant.api_url}/api/{path}"
-
-        try:
-            data = None
-            headers = {}
-            method = getattr(
-                self.homeassistant.websession, request.method.lower())
-
-            # read data
-            with async_timeout.timeout(30, loop=self.loop):
-                data = await request.read()
-
-            if data:
-                headers.update({CONTENT_TYPE: request.content_type})
-
-            # need api password?
-            if self.homeassistant.api_password:
-                headers = {HEADER_HA_ACCESS: self.homeassistant.api_password}
-
-            # reset headers
-            if not headers:
-                headers = None
-
-            client = await method(
-                url, data=data, headers=headers, timeout=timeout
-            )
-
-            return client
-
-        except aiohttp.ClientError as err:
-            _LOGGER.error("Client error on api %s request %s.", path, err)
-
-        except asyncio.TimeoutError:
-            _LOGGER.error("Client timeout error on api request %s.", path)
-
-        raise HTTPBadGateway()
-
     @api_process
     async def info(self, request):
         """Return host information."""
@@ -169,44 +125,3 @@ class APIHomeAssistant(object):
             raise RuntimeError(message)
 
         return True
-
-    async def api(self, request):
-        """Proxy API request to Home-Assistant."""
-        path = request.match_info.get('path', '')
-        _LOGGER.info("Proxy /api/%s request", path)
-
-        # API stream
-        if path.startswith("stream"):
-            client = await self.homeassistant_proxy(
-                path, request, timeout=None)
-
-            response = web.StreamResponse()
-            response.content_type = request.headers.get(CONTENT_TYPE)
-            try:
-                await response.prepare(request)
-                while True:
-                    data = await client.content.read(10)
-                    if not data:
-                        await response.write_eof()
-                        break
-                    response.write(data)
-
-            except aiohttp.ClientError:
-                await response.write_eof()
-
-            except asyncio.TimeoutError:
-                pass
-
-            finally:
-                client.close()
-
-        # Normal request
-        else:
-            client = await self.homeassistant_proxy(path, request)
-
-            data = await client.read()
-            return web.Response(
-                body=data,
-                status=client.status,
-                content_type=client.content_type
-            )

hassio/api/proxy.py

@@ -0,0 +1,195 @@
+"""Utils for HomeAssistant Proxy."""
+import asyncio
+import logging
+
+import aiohttp
+from aiohttp import web
+from aiohttp.web_exceptions import HTTPBadGateway
+from aiohttp.hdrs import CONTENT_TYPE
+import async_timeout
+
+from ..const import HEADER_HA_ACCESS
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class APIProxy(object):
+    """API Proxy for Home-Assistant."""
+
+    def __init__(self, loop, homeassistant, websession):
+        """Initialize api proxy."""
+        self.loop = loop
+        self.homeassistant = homeassistant
+        self.websession = websession
+
+    async def _api_client(self, request, path, timeout=300):
+        """Return a client request with proxy origin for Home-Assistant."""
+        url = f"{self.homeassistant.api_url}/api/{path}"
+
+        try:
+            data = None
+            headers = {}
+            method = getattr(self.websession, request.method.lower())
+
+            # read data
+            with async_timeout.timeout(30, loop=self.loop):
+                data = await request.read()
+
+            if data:
+                headers.update({CONTENT_TYPE: request.content_type})
+
+            # need api password?
+            if self.homeassistant.api_password:
+                headers = {HEADER_HA_ACCESS: self.homeassistant.api_password}
+
+            # reset headers
+            if not headers:
+                headers = None
+
+            client = await method(
+                url, data=data, headers=headers, timeout=timeout
+            )
+
+            return client
+
+        except aiohttp.ClientError as err:
+            _LOGGER.error("Client error on API %s request %s.", path, err)
+
+        except asyncio.TimeoutError:
+            _LOGGER.error("Client timeout error on API request %s.", path)
+
+        raise HTTPBadGateway()
+
+    async def api(self, request):
+        """Proxy HomeAssistant API Requests."""
+        path = request.match_info.get('path', '')
+
+        # API stream
+        if path.startswith("stream"):
+            _LOGGER.info("Home-Assistant Event-Stream start")
+            client = await self._api_client(request, path, timeout=None)
+
+            response = web.StreamResponse()
+            response.content_type = request.headers.get(CONTENT_TYPE)
+            try:
+                await response.prepare(request)
+                while True:
+                    data = await client.content.read(10)
+                    if not data:
+                        await response.write_eof()
+                        break
+                    response.write(data)
+
+            except aiohttp.ClientError:
+                await response.write_eof()
+
+            except asyncio.CancelledError:
+                pass
+
+            finally:
+                client.close()
+
+            _LOGGER.info("Home-Assistant Event-Stream close")
+
+        # Normal request
+        else:
+            _LOGGER.info("Home-Assistant '/api/%s' request", path)
+            client = await self._api_client(request, path)
+
+            data = await client.read()
+            return web.Response(
+                body=data,
+                status=client.status,
+                content_type=client.content_type
+            )
+
+    async def _websocket_client(self):
+        """Initialize a websocket api connection."""
+        url = f"{self.homeassistant.api_url}/api/websocket"
+
+        try:
+            client = await self.websession.ws_connect(
+                url, heartbeat=60)
+
+            # handle authentication
+            for _ in range(2):
+                data = await client.receive_json()
+                if data.get('type') == 'auth_ok':
+                    return client
+                elif data.get('type') == 'auth_required':
+                    await client.send_json({
+                        'type': 'auth',
+                        'api_password': self.homeassistant.api_password,
+                    })
+
+            _LOGGER.error("Authentication to Home-Assistant websocket")
+
+        except (aiohttp.ClientError, RuntimeError) as err:
+            _LOGGER.error("Client error on websocket API %s.", err)
+
+        raise HTTPBadGateway()
+
+    async def websocket(self, request):
+        """Initialize a websocket api connection."""
+        _LOGGER.info("Home-Assistant Websocket API request initialze")
+
+        # init server
+        server = web.WebSocketResponse(heartbeat=60)
+        await server.prepare(request)
+
+        # handle authentication
+        await server.send_json({'type': 'auth_required'})
+        await server.receive_json()  # get internal token
+        await server.send_json({'type': 'auth_ok'})
+
+        # init connection to hass
+        client = await self._websocket_client()
+
+        _LOGGER.info("Home-Assistant Websocket API request running")
+        try:
+            client_read = None
+            server_read = None
+            while not server.closed and not client.closed:
+                if not client_read:
+                    client_read = asyncio.ensure_future(
+                        client.receive_str(), loop=self.loop)
+                if not server_read:
+                    server_read = asyncio.ensure_future(
+                        server.receive_str(), loop=self.loop)
+
+                # wait until data need to be processed
+                await asyncio.wait(
+                    [client_read, server_read],
+                    loop=self.loop, return_when=asyncio.FIRST_COMPLETED
+                )
+
+                # server
+                if server_read.done() and not client.closed:
+                    server_read.exception()
+                    await client.send_str(server_read.result())
+                    server_read = None
+
+                # client
+                if client_read.done() and not server.closed:
+                    client_read.exception()
+                    await server.send_str(client_read.result())
+                    client_read = None
+
+        except asyncio.CancelledError:
+            pass
+
+        except RuntimeError as err:
+            _LOGGER.info("Home-Assistant Websocket API error: %s", err)
+
+        finally:
+            if client_read:
+                client_read.cancel()
+            if server_read:
+                server_read.cancel()
+
+            # close connections
+            await client.close()
+            await server.close()
+
+        _LOGGER.info("Home-Assistant Websocket API connection is closed")
+        return server

hassio/core.py

@@ -91,6 +91,7 @@ class HassIO(object):
             self.supervisor, self.snapshots, self.addons, self.host_control,
             self.updater)
         self.api.register_homeassistant(self.homeassistant)
+        self.api.register_proxy(self.homeassistant, self.websession)
         self.api.register_addons(self.addons)
         self.api.register_security()
         self.api.register_snapshots(self.snapshots)