From 9efc08cb2fe5e99af8dfa060fed86db6facd4fa9 Mon Sep 17 00:00:00 2001
From: Dmitriy Ryajov <dryajov@gmail.com>
Date: Fri, 21 Feb 2020 18:09:33 -0600
Subject: [PATCH] fix: #83

---
 libp2p/stream/bufferstream.nim |  2 +-
 tests/testbufferstream.nim     | 20 ++++++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/libp2p/stream/bufferstream.nim b/libp2p/stream/bufferstream.nim
index 8b9933a..681fb8e 100644
--- a/libp2p/stream/bufferstream.nim
+++ b/libp2p/stream/bufferstream.nim
@@ -106,7 +106,7 @@ proc pushTo*(s: BufferStream, data: seq[byte]) {.async.} =
   try:
     await s.lock.acquire()
     var index = 0
-    while true:
+    while not s.closed():
       while index < data.len and s.readBuf.len < s.maxSize:
         s.readBuf.addLast(data[index])
         inc(index)
diff --git a/tests/testbufferstream.nim b/tests/testbufferstream.nim
index 46ed711..2a02ca3 100644
--- a/tests/testbufferstream.nim
+++ b/tests/testbufferstream.nim
@@ -440,3 +440,23 @@ suite "BufferStream":
 
     check:
       waitFor(pipeTest()) == true
+
+  test "shouldn't get stuck on close":
+    proc test(): Future[bool] {.async.} =
+      proc createMessage(tmplate: string, size: int): seq[byte] =
+        result = newSeq[byte](size)
+        for i in 0 ..< len(result):
+          result[i] = byte(tmplate[i mod len(tmplate)])
+
+      var stream = newBufferStream()
+      var message = createMessage("MESSAGE", DefaultBufferSize * 2 + 1)
+      var fut = stream.pushTo(message)
+      await stream.close()
+      try:
+        await wait(fut, 100.milliseconds)
+        result = true
+      except AsyncTimeoutError:
+        result = false
+
+      check:
+        waitFor(test()) == true
\ No newline at end of file