From d3b79b002e267cb02ce11e453d6dcd5d6021b712 Mon Sep 17 00:00:00 2001
From: Dmitriy Ryajov <dryajov@gmail.com>
Date: Sat, 23 May 2020 10:52:33 -0600
Subject: [PATCH] better exceptions and don't fail writes

---
 libp2p/stream/chronosstream.nim | 35 +++++++++++++++++----------------
 1 file changed, 18 insertions(+), 17 deletions(-)

diff --git a/libp2p/stream/chronosstream.nim b/libp2p/stream/chronosstream.nim
index c1bca76..2ebe529 100644
--- a/libp2p/stream/chronosstream.nim
+++ b/libp2p/stream/chronosstream.nim
@@ -21,12 +21,12 @@ proc newChronosStream*(client: StreamTransport): ChronosStream =
   result.client = client
   result.closeEvent = newAsyncEvent()
 
-
 template withExceptions(body: untyped) =
   try:
     body
   except TransportIncompleteError:
-    raise newLPStreamIncompleteError()
+    # for all intents and purposes this is an EOF
+    raise newLPStreamEOFError()
   except TransportLimitError:
     raise newLPStreamLimitError()
   except TransportUseClosedError:
@@ -60,9 +60,12 @@ method write*(s: ChronosStream, msg: seq[byte]) {.async.} =
     return
 
   withExceptions:
-    # Returns 0 sometimes when write fails - but there's not much we can do here?
-    if (await s.client.write(msg)) != msg.len:
-      raise (ref LPStreamError)(msg: "Write couldn't finish writing")
+    var writen = 0
+    while not s.client.closed and writen < msg.len:
+      writen += await s.client.write(msg[writen..<msg.len])
+
+    if writen < msg.len:
+      raise (ref LPStreamClosedError)(msg: "Write couldn't finish writing")
 
 method closed*(s: ChronosStream): bool {.inline.} =
   result = s.client.closed
@@ -71,16 +74,14 @@ method atEof*(s: ChronosStream): bool {.inline.} =
   s.client.atEof()
 
 method close*(s: ChronosStream) {.async.} =
-  if not s.isClosed:
-    s.isClosed = true
-    trace "shutting chronos stream", address = $s.client.remoteAddress()
-    if not s.client.closed():
-      try:
-        await s.client.closeWait()
-      except CancelledError as exc:
-        raise exc
-      except CatchableError as exc:
-        # Shouldn't happen, but we can't be sure
-        warn "error while closing connection", msg = exc.msg
+  try:
+    if not s.isClosed:
+      s.isClosed = true
 
-    s.closeEvent.fire()
+      trace "shutting down chronos stream", address = $s.client.remoteAddress()
+      if not s.client.closed():
+        await s.client.closeWait()
+
+      s.closeEvent.fire()
+  except CatchableError as exc:
+    trace "error closing chronosstream", exc = exc.msg