From 96edf96e76343deabd1e4c15bd73ca71d4776304 Mon Sep 17 00:00:00 2001
From: Christophe de Carvalho Pereira Martins <christophe.dcpm@gmail.com>
Date: Sun, 18 Nov 2018 17:22:17 +0100
Subject: [PATCH] make stream close async

---
 muxer/mplex/muxed_connection.py | 8 ++++++--
 muxer/mplex/muxed_stream.py     | 4 ++--
 network/stream/net_stream.py    | 4 ++--
 3 files changed, 10 insertions(+), 6 deletions(-)

diff --git a/muxer/mplex/muxed_connection.py b/muxer/mplex/muxed_connection.py
index c2fd343..0a43640 100644
--- a/muxer/mplex/muxed_connection.py
+++ b/muxer/mplex/muxed_connection.py
@@ -83,8 +83,12 @@ class MuxedConn(IMuxedConn):
         # << by 3, then or with flag
         header = (stream_id << 3) | flag
         header = encode_uvarint(header)
-        data_length = encode_uvarint(len(data))
-        _bytes = header + data_length + data
+        if data is None:
+            data_length = encode_uvarint(0)
+            _bytes = header + data_length
+        else:
+            data_length = encode_uvarint(len(data))
+            _bytes = header + data_length + data
 
         return await self.write_to_stream(_bytes)
 
diff --git a/muxer/mplex/muxed_stream.py b/muxer/mplex/muxed_stream.py
index 1d6bd5f..11e04f4 100644
--- a/muxer/mplex/muxed_stream.py
+++ b/muxer/mplex/muxed_stream.py
@@ -49,7 +49,7 @@ class MuxedStream(IMuxedStream):
         """
         return await self.muxed_conn.send_message(self.get_flag("MESSAGE"), data, self.stream_id)
 
-    def close(self):
+    async def close(self):
         """
         close stream
         :return: true if successful
@@ -58,7 +58,7 @@ class MuxedStream(IMuxedStream):
         if self.local_closed and self.remote_closed:
             return True
 
-        self.muxed_conn.send_message(self.get_flag("CLOSE"), None, self.stream_id)
+        await self.muxed_conn.send_message(self.get_flag("CLOSE"), None, self.stream_id)
         self.muxed_conn.streams.pop(self.stream_id)
 
         self.local_closed = True
diff --git a/network/stream/net_stream.py b/network/stream/net_stream.py
index 0377c32..8cfb635 100644
--- a/network/stream/net_stream.py
+++ b/network/stream/net_stream.py
@@ -34,10 +34,10 @@ class NetStream(INetStream):
         """
         return await self.muxed_stream.write(data)
 
-    def close(self):
+    async def close(self):
         """
         close stream
         :return: true if successful
         """
-        self.muxed_stream.close()
+        await self.muxed_stream.close()
         return True