From 5ab5036504cb8c4e16e80361b829347091f06b24 Mon Sep 17 00:00:00 2001
From: Pascal Vizeli <pascal.vizeli@syshack.ch>
Date: Mon, 15 Oct 2018 13:01:52 +0200
Subject: [PATCH] Fix proxy handling with failing connection (#760)

* Fix proxy handling with failing connection

* fix lint

* Fix exception handling

* clenaup error handling

* Fix type error

* Fix event stream

* Fix stream handling

* Fix

* Fix lint

* Handle

* Update proxy.py

* fix lint
---
 hassio/api/proxy.py | 44 ++++++++++++++++++++++----------------------
 requirements.txt    |  8 ++++----
 2 files changed, 26 insertions(+), 26 deletions(-)

diff --git a/hassio/api/proxy.py b/hassio/api/proxy.py
index af6d38af6..f67a1adc5 100644
--- a/hassio/api/proxy.py
+++ b/hassio/api/proxy.py
@@ -5,14 +5,15 @@ import logging
 
 import aiohttp
 from aiohttp import web
-from aiohttp.web_exceptions import (
-    HTTPBadGateway, HTTPInternalServerError, HTTPUnauthorized)
+from aiohttp.web_exceptions import HTTPBadGateway, HTTPUnauthorized
+from aiohttp.client_exceptions import ClientConnectorError
 from aiohttp.hdrs import CONTENT_TYPE, AUTHORIZATION
 import async_timeout
 
 from ..const import HEADER_HA_ACCESS
 from ..coresys import CoreSysAttributes
-from ..exceptions import HomeAssistantAuthError, HomeAssistantAPIError
+from ..exceptions import (
+    HomeAssistantAuthError, HomeAssistantAPIError, APIError)
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -82,19 +83,13 @@ class APIProxy(CoreSysAttributes):
             response.content_type = request.headers.get(CONTENT_TYPE)
             try:
                 await response.prepare(request)
-                while True:
-                    data = await client.content.read(10)
-                    if not data:
-                        break
+                async for data in client.content:
                     await response.write(data)
 
-            except aiohttp.ClientError:
+            except (aiohttp.ClientError, aiohttp.ClientPayloadError):
                 pass
 
-            finally:
-                client.close()
-                _LOGGER.info("Home Assistant EventStream close")
-
+            _LOGGER.info("Home Assistant EventStream close")
             return response
 
     async def api(self, request):
@@ -117,7 +112,7 @@ class APIProxy(CoreSysAttributes):
 
         try:
             client = await self.sys_websession_ssl.ws_connect(
-                url, heartbeat=60, verify_ssl=False)
+                url, heartbeat=30, verify_ssl=False)
 
             # Handle authentication
             data = await client.receive_json()
@@ -129,7 +124,7 @@ class APIProxy(CoreSysAttributes):
                 # Invalid protocol
                 _LOGGER.error(
                     "Got unexpected response from HA WebSocket: %s", data)
-                raise HTTPBadGateway()
+                raise APIError()
 
             if self.sys_homeassistant.refresh_token:
                 await self.sys_homeassistant.ensure_access_token()
@@ -156,19 +151,19 @@ class APIProxy(CoreSysAttributes):
 
             raise HomeAssistantAuthError()
 
-        except (RuntimeError, ValueError) as err:
+        except (RuntimeError, ValueError, ClientConnectorError) as err:
             _LOGGER.error("Client error on WebSocket API %s.", err)
         except HomeAssistantAuthError as err:
             _LOGGER.error("Failed authentication to Home Assistant WebSocket")
 
-        raise HTTPBadGateway()
+        raise APIError()
 
     async def websocket(self, request):
         """Initialize a WebSocket API connection."""
         _LOGGER.info("Home Assistant WebSocket API request initialize")
 
         # init server
-        server = web.WebSocketResponse(heartbeat=60)
+        server = web.WebSocketResponse(heartbeat=30)
         await server.prepare(request)
 
         # handle authentication
@@ -200,10 +195,13 @@ class APIProxy(CoreSysAttributes):
             })
         except (RuntimeError, ValueError) as err:
             _LOGGER.error("Can't initialize handshake: %s", err)
-            raise HTTPInternalServerError() from None
+            return server
 
         # init connection to hass
-        client = await self._websocket_client()
+        try:
+            client = await self._websocket_client()
+        except APIError:
+            return server
 
         _LOGGER.info("Home Assistant WebSocket API request running")
         try:
@@ -238,7 +236,7 @@ class APIProxy(CoreSysAttributes):
         except asyncio.CancelledError:
             pass
 
-        except RuntimeError as err:
+        except (RuntimeError, ConnectionError, TypeError) as err:
             _LOGGER.info("Home Assistant WebSocket API error: %s", err)
 
         finally:
@@ -248,8 +246,10 @@ class APIProxy(CoreSysAttributes):
                 server_read.cancel()
 
             # close connections
-            await client.close()
-            await server.close()
+            if not client.closed:
+                await client.close()
+            if not server.closed:
+                await server.close()
 
         _LOGGER.info("Home Assistant WebSocket API connection is closed")
         return server
diff --git a/requirements.txt b/requirements.txt
index eae7f825f..4cd773e76 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,11 +1,11 @@
 attr==0.3.1
-async_timeout==3.0.0
-aiohttp==3.4.0
+async_timeout==3.0.1
+aiohttp==3.4.4
 docker==3.5.0
-colorlog==3.1.2
+colorlog==3.1.4
 voluptuous==0.11.5
 gitpython==2.1.10
-pytz==2018.4
+pytz==2018.5
 pyudev==0.21.0
 pycryptodome==3.6.6
 cpe==1.2.1