2018-11-21 10:46:18 +08:00
|
|
|
import asyncio
|
2018-11-12 07:03:04 +08:00
|
|
|
|
2019-08-05 11:22:44 +08:00
|
|
|
from libp2p.stream_muxer.abc import IMuxedConn, IMuxedStream
|
2018-11-01 05:31:00 +08:00
|
|
|
|
2019-08-05 10:20:30 +08:00
|
|
|
from .constants import HeaderTags
|
2019-08-02 18:28:04 +08:00
|
|
|
|
|
|
|
|
2018-11-21 09:28:41 +08:00
|
|
|
class MplexStream(IMuxedStream):
|
2018-11-01 06:02:00 +08:00
|
|
|
"""
|
|
|
|
reference: https://github.com/libp2p/go-mplex/blob/master/stream.go
|
|
|
|
"""
|
|
|
|
|
2019-08-02 18:28:04 +08:00
|
|
|
stream_id: int
|
|
|
|
initiator: bool
|
2019-08-05 10:20:30 +08:00
|
|
|
mplex_conn: IMuxedConn
|
2019-08-05 17:02:18 +08:00
|
|
|
read_deadline: int
|
|
|
|
write_deadline: int
|
2019-08-02 18:28:04 +08:00
|
|
|
local_closed: bool
|
|
|
|
remote_closed: bool
|
|
|
|
stream_lock: asyncio.Lock
|
|
|
|
|
2019-08-15 18:31:18 +08:00
|
|
|
_buf: bytearray
|
2019-08-07 15:23:20 +08:00
|
|
|
|
2019-08-05 10:20:30 +08:00
|
|
|
def __init__(self, stream_id: int, initiator: bool, mplex_conn: IMuxedConn) -> None:
|
2018-11-12 07:03:04 +08:00
|
|
|
"""
|
|
|
|
create new MuxedStream in muxer
|
|
|
|
:param stream_id: stream stream id
|
|
|
|
:param initiator: boolean if this is an initiator
|
2018-11-21 10:46:18 +08:00
|
|
|
:param mplex_conn: muxed connection of this muxed_stream
|
2018-11-12 07:03:04 +08:00
|
|
|
"""
|
2018-11-13 02:02:49 +08:00
|
|
|
self.stream_id = stream_id
|
2018-11-12 07:03:04 +08:00
|
|
|
self.initiator = initiator
|
2018-11-21 10:46:18 +08:00
|
|
|
self.mplex_conn = mplex_conn
|
2018-11-13 02:02:49 +08:00
|
|
|
self.read_deadline = None
|
|
|
|
self.write_deadline = None
|
2018-11-12 07:03:04 +08:00
|
|
|
self.local_closed = False
|
|
|
|
self.remote_closed = False
|
2018-11-21 10:46:18 +08:00
|
|
|
self.stream_lock = asyncio.Lock()
|
2019-08-15 18:31:18 +08:00
|
|
|
self._buf = bytearray()
|
2018-11-12 07:03:04 +08:00
|
|
|
|
2019-08-12 14:25:17 +08:00
|
|
|
async def read(self, n: int = -1) -> bytes:
|
2018-11-12 07:03:04 +08:00
|
|
|
"""
|
2019-08-11 23:49:58 +08:00
|
|
|
Read up to n bytes. Read possibly returns fewer than `n` bytes,
|
|
|
|
if there are not enough bytes in the Mplex buffer.
|
2019-08-12 14:25:17 +08:00
|
|
|
If `n == -1`, read until EOF.
|
2019-08-07 15:23:20 +08:00
|
|
|
:param n: number of bytes to read
|
2019-08-11 23:49:58 +08:00
|
|
|
:return: bytes actually read
|
2018-11-12 07:03:04 +08:00
|
|
|
"""
|
2019-08-12 14:25:17 +08:00
|
|
|
if n < 0 and n != -1:
|
|
|
|
raise ValueError("`n` can only be -1 if it is negative")
|
2019-08-11 23:49:58 +08:00
|
|
|
# If the buffer is empty at first, blocking wait for data.
|
|
|
|
if len(self._buf) == 0:
|
2019-08-15 18:31:18 +08:00
|
|
|
self._buf.extend(await self.mplex_conn.read_buffer(self.stream_id))
|
2019-08-12 14:25:17 +08:00
|
|
|
# Sanity check: `self._buf` should never be empty here.
|
2019-08-11 23:49:58 +08:00
|
|
|
if self._buf is None or len(self._buf) == 0:
|
2019-08-12 14:25:17 +08:00
|
|
|
raise Exception("`self._buf` should never be empty here")
|
2019-08-11 23:49:58 +08:00
|
|
|
|
2019-08-12 14:25:17 +08:00
|
|
|
# FIXME: If `n == -1`, we should blocking read until EOF, instead of returning when
|
|
|
|
# no message is available.
|
|
|
|
# If `n >= 0`, read up to `n` bytes.
|
|
|
|
# Else, read until no message is available.
|
|
|
|
while len(self._buf) < n or n == -1:
|
2019-08-11 23:49:58 +08:00
|
|
|
new_bytes = await self.mplex_conn.read_buffer_nonblocking(self.stream_id)
|
|
|
|
if new_bytes is None:
|
|
|
|
# Nothing to read in the `MplexConn` buffer
|
|
|
|
break
|
2019-08-15 18:31:18 +08:00
|
|
|
self._buf.extend(new_bytes)
|
|
|
|
payload: bytearray
|
2019-08-12 14:25:17 +08:00
|
|
|
if n == -1:
|
|
|
|
payload = self._buf
|
|
|
|
else:
|
|
|
|
payload = self._buf[:n]
|
|
|
|
self._buf = self._buf[len(payload) :]
|
2019-08-15 18:31:18 +08:00
|
|
|
return bytes(payload)
|
2018-11-01 05:31:00 +08:00
|
|
|
|
2019-08-02 18:28:04 +08:00
|
|
|
async def write(self, data: bytes) -> int:
|
2018-11-12 07:03:04 +08:00
|
|
|
"""
|
|
|
|
write to stream
|
|
|
|
:return: number of bytes written
|
|
|
|
"""
|
2019-08-14 05:36:42 +08:00
|
|
|
flag = (
|
|
|
|
HeaderTags.MessageInitiator
|
|
|
|
if self.initiator
|
|
|
|
else HeaderTags.MessageReceiver
|
|
|
|
)
|
2019-08-02 17:14:43 +08:00
|
|
|
return await self.mplex_conn.send_message(flag, data, self.stream_id)
|
2018-11-01 05:31:00 +08:00
|
|
|
|
2019-08-02 18:28:04 +08:00
|
|
|
async def close(self) -> bool:
|
2018-11-12 07:03:04 +08:00
|
|
|
"""
|
2018-11-21 10:46:18 +08:00
|
|
|
Closing a stream closes it for writing and closes the remote end for reading
|
|
|
|
but allows writing in the other direction.
|
2018-11-12 07:03:04 +08:00
|
|
|
:return: true if successful
|
|
|
|
"""
|
2018-11-21 13:41:13 +08:00
|
|
|
# TODO error handling with timeout
|
|
|
|
# TODO understand better how mutexes are used from go repo
|
2019-08-02 17:14:43 +08:00
|
|
|
flag = HeaderTags.CloseInitiator if self.initiator else HeaderTags.CloseReceiver
|
|
|
|
await self.mplex_conn.send_message(flag, None, self.stream_id)
|
2018-11-12 07:03:04 +08:00
|
|
|
|
2019-08-02 18:28:04 +08:00
|
|
|
remote_lock = False
|
2018-11-21 10:46:18 +08:00
|
|
|
async with self.stream_lock:
|
|
|
|
if self.local_closed:
|
|
|
|
return True
|
|
|
|
self.local_closed = True
|
|
|
|
remote_lock = self.remote_closed
|
2018-11-12 07:03:04 +08:00
|
|
|
|
2018-11-21 10:46:18 +08:00
|
|
|
if remote_lock:
|
2019-08-02 18:28:04 +08:00
|
|
|
# FIXME: mplex_conn has no conn_lock!
|
|
|
|
async with self.mplex_conn.conn_lock: # type: ignore
|
|
|
|
# FIXME: Don't access to buffers directly
|
|
|
|
self.mplex_conn.buffers.pop(self.stream_id) # type: ignore
|
2018-11-12 07:03:04 +08:00
|
|
|
|
|
|
|
return True
|
2018-11-01 05:31:00 +08:00
|
|
|
|
2019-08-02 18:28:04 +08:00
|
|
|
async def reset(self) -> bool:
|
2018-11-01 05:31:00 +08:00
|
|
|
"""
|
|
|
|
closes both ends of the stream
|
|
|
|
tells this remote side to hang up
|
2018-11-12 07:03:04 +08:00
|
|
|
:return: true if successful
|
2018-11-01 05:31:00 +08:00
|
|
|
"""
|
2018-11-21 13:41:13 +08:00
|
|
|
# TODO understand better how mutexes are used here
|
|
|
|
# TODO understand the difference between close and reset
|
2018-11-21 10:46:18 +08:00
|
|
|
async with self.stream_lock:
|
|
|
|
if self.remote_closed and self.local_closed:
|
|
|
|
return True
|
|
|
|
|
|
|
|
if not self.remote_closed:
|
2019-08-14 05:36:42 +08:00
|
|
|
flag = (
|
|
|
|
HeaderTags.ResetInitiator
|
|
|
|
if self.initiator
|
|
|
|
else HeaderTags.ResetInitiator
|
|
|
|
)
|
2019-08-02 17:14:43 +08:00
|
|
|
await self.mplex_conn.send_message(flag, None, self.stream_id)
|
2018-11-21 10:46:18 +08:00
|
|
|
|
|
|
|
self.local_closed = True
|
|
|
|
self.remote_closed = True
|
|
|
|
|
2019-08-02 18:28:04 +08:00
|
|
|
# FIXME: mplex_conn has no conn_lock!
|
|
|
|
async with self.mplex_conn.conn_lock: # type: ignore
|
|
|
|
# FIXME: Don't access to buffers directly
|
|
|
|
self.mplex_conn.buffers.pop(self.stream_id, None) # type: ignore
|
2018-11-21 10:46:18 +08:00
|
|
|
|
|
|
|
return True
|
2018-11-01 05:31:00 +08:00
|
|
|
|
2018-11-12 07:03:04 +08:00
|
|
|
# TODO deadline not in use
|
2019-08-05 17:02:18 +08:00
|
|
|
def set_deadline(self, ttl: int) -> bool:
|
2018-11-01 05:31:00 +08:00
|
|
|
"""
|
|
|
|
set deadline for muxed stream
|
2018-11-12 07:03:04 +08:00
|
|
|
:return: True if successful
|
2018-11-01 05:31:00 +08:00
|
|
|
"""
|
2018-11-12 07:03:04 +08:00
|
|
|
self.read_deadline = ttl
|
|
|
|
self.write_deadline = ttl
|
|
|
|
return True
|
|
|
|
|
2019-08-05 17:02:18 +08:00
|
|
|
def set_read_deadline(self, ttl: int) -> bool:
|
2018-11-12 07:03:04 +08:00
|
|
|
"""
|
|
|
|
set read deadline for muxed stream
|
|
|
|
:return: True if successful
|
|
|
|
"""
|
|
|
|
self.read_deadline = ttl
|
|
|
|
return True
|
|
|
|
|
2019-08-05 17:02:18 +08:00
|
|
|
def set_write_deadline(self, ttl: int) -> bool:
|
2018-11-12 07:03:04 +08:00
|
|
|
"""
|
|
|
|
set write deadline for muxed stream
|
|
|
|
:return: True if successful
|
|
|
|
"""
|
|
|
|
self.write_deadline = ttl
|
|
|
|
return True
|