From 07da14a7a7de467b9d5cfe46eeaba18c8de7a7e0 Mon Sep 17 00:00:00 2001
From: Tanguy <tanguy@status.im>
Date: Fri, 28 Jan 2022 19:05:07 +0100
Subject: [PATCH] Fix websocket EOF reading exception (#689)

---
 libp2p/transports/wstransport.nim | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/libp2p/transports/wstransport.nim b/libp2p/transports/wstransport.nim
index c0b14278e..4f6444294 100644
--- a/libp2p/transports/wstransport.nim
+++ b/libp2p/transports/wstransport.nim
@@ -49,11 +49,24 @@ proc new*(T: type WsStream,
   stream.initStream()
   return stream
 
+template mapExceptions(body: untyped) =
+  try:
+    body
+  except AsyncStreamIncompleteError:
+    raise newLPStreamEOFError()
+  except AsyncStreamUseClosedError:
+    raise newLPStreamEOFError()
+  except WSClosedError:
+    raise newLPStreamEOFError()
+  except AsyncStreamLimitError:
+    raise newLPStreamLimitError()
+
 method readOnce*(
   s: WsStream,
   pbytes: pointer,
   nbytes: int): Future[int] {.async.} =
-  let res = await s.session.recv(pbytes, nbytes)
+  let res = mapExceptions(await s.session.recv(pbytes, nbytes))
+
   if res == 0 and s.session.readyState == ReadyState.Closed:
     raise newLPStreamEOFError()
   return res
@@ -61,10 +74,7 @@ method readOnce*(
 method write*(
   s: WsStream,
   msg: seq[byte]): Future[void] {.async.} =
-  try:
-    await s.session.send(msg, Opcode.Binary)
-  except WSClosedError:
-    raise newLPStreamEOFError()
+  mapExceptions(await s.session.send(msg, Opcode.Binary))
 
 method closeImpl*(s: WsStream): Future[void] {.async.} =
   await s.session.close()