From 092713c361244cc15cc4c4ef2b861a852fdcc4a0 Mon Sep 17 00:00:00 2001
From: Jungkook Park <pjknkda@gmail.com>
Date: Wed, 9 Feb 2022 09:56:24 +0900
Subject: [PATCH] fix: handle transport exceptions from local connection

---
 rabbit_tunnel/__init__.py | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/rabbit_tunnel/__init__.py b/rabbit_tunnel/__init__.py
index 52b0d79..5dbc3ea 100644
--- 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