fix: handle transport exceptions from local connection

pjknkda · Feb 9, 2022 · 092713c · 092713c
1 parent 03e8b64
commit 092713c
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/rabbit_tunnel/__init__.py b/rabbit_tunnel/__init__.py
@@ -71,7 +71,11 @@ async def _reader_puller() -> None:
                 raise RuntimeError('Reader is not initialized')
 
             while True:
-                data = await reader.read(_READ_BUFFER)
+                try:
+                    data = await reader.read(_READ_BUFFER)
+                except (RuntimeError, ConnectionResetError, BrokenPipeError):
+                    break
+
                 if not data:
                     break
 
@@ -82,7 +86,7 @@ async def _reader_puller() -> None:
                         'data': data,
                     }))
                 except WsConnectionClosedError:
-                    return
+                    break
         finally:
             await _close('local-connection-closed')
 
@@ -126,7 +130,7 @@ async def _reader_puller() -> None:
                 try:
                     writer.write(msg['data'])
                     await writer.drain()
-                except ConnectionResetError:
+                except (RuntimeError, ConnectionResetError, BrokenPipeError):
                     await _close('local-connection-closed')
                     break