mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2026-02-11 15:40:54 +00:00
Migrate to new project structure.
This commit is contained in:
Alexander Koshkin
@ -1,6 +0,0 @@
|
||||
HEADER_TAGS = {
|
||||
"NEW_STREAM": 0,
|
||||
"MESSAGE": 2,
|
||||
"CLOSE": 4,
|
||||
"RESET": 6
|
||||
}
|
||||
@ -1,171 +0,0 @@
|
||||
import asyncio
|
||||
from .utils import encode_uvarint, decode_uvarint_from_stream
|
||||
from .mplex_stream import MplexStream
|
||||
from ..muxed_connection_interface import IMuxedConn
|
||||
|
||||
|
||||
class Mplex(IMuxedConn):
|
||||
# pylint: disable=too-many-instance-attributes
|
||||
"""
|
||||
reference: https://github.com/libp2p/go-mplex/blob/master/multiplex.go
|
||||
"""
|
||||
|
||||
def __init__(self, conn):
|
||||
"""
|
||||
create a new muxed connection
|
||||
:param conn: an instance of raw connection
|
||||
:param initiator: boolean to prevent multiplex with self
|
||||
"""
|
||||
self.raw_conn = conn
|
||||
self.initiator = conn.initiator
|
||||
|
||||
# Mapping from stream ID -> buffer of messages for that stream
|
||||
self.buffers = {}
|
||||
|
||||
self.stream_queue = asyncio.Queue()
|
||||
self.data_buffer = bytearray()
|
||||
|
||||
# The initiator of the raw connection need not read upon construction time.
|
||||
# It should read when the user decides that it wants to read from the constructed stream.
|
||||
if not self.initiator:
|
||||
asyncio.ensure_future(self.handle_incoming())
|
||||
|
||||
def close(self):
|
||||
"""
|
||||
close the stream muxer and underlying raw connection
|
||||
"""
|
||||
self.raw_conn.close()
|
||||
|
||||
def is_closed(self):
|
||||
"""
|
||||
check connection is fully closed
|
||||
:return: true if successful
|
||||
"""
|
||||
|
||||
async def read_buffer(self, stream_id):
|
||||
"""
|
||||
Read a message from stream_id's buffer, check raw connection for new messages
|
||||
:param stream_id: stream id of stream to read from
|
||||
:return: message read
|
||||
"""
|
||||
# Empty buffer or nonexistent stream
|
||||
# TODO: propagate up timeout exception and catch
|
||||
if stream_id not in self.buffers or self.buffers[stream_id].empty():
|
||||
await self.handle_incoming()
|
||||
if stream_id in self.buffers:
|
||||
return await self._read_buffer_exists(stream_id)
|
||||
|
||||
return None
|
||||
|
||||
async def _read_buffer_exists(self, stream_id):
|
||||
"""
|
||||
Reads from raw connection with the assumption that the message buffer for stream_id exsits
|
||||
:param stream_id: stream id of stream to read from
|
||||
:return: message read
|
||||
"""
|
||||
try:
|
||||
data = await asyncio.wait_for(self.buffers[stream_id].get(), timeout=5)
|
||||
return data
|
||||
except asyncio.TimeoutError:
|
||||
return None
|
||||
|
||||
async def open_stream(self, protocol_id, peer_id, multi_addr):
|
||||
"""
|
||||
creates a new muxed_stream
|
||||
:param protocol_id: protocol_id of stream
|
||||
:param stream_id: stream_id of stream
|
||||
:param peer_id: peer_id that stream connects to
|
||||
:param multi_addr: multi_addr that stream connects to
|
||||
:return: a new stream
|
||||
"""
|
||||
stream_id = self.raw_conn.next_stream_id()
|
||||
stream = MplexStream(stream_id, multi_addr, self)
|
||||
self.buffers[stream_id] = asyncio.Queue()
|
||||
return stream
|
||||
|
||||
async def accept_stream(self):
|
||||
"""
|
||||
accepts a muxed stream opened by the other end
|
||||
:return: the accepted stream
|
||||
"""
|
||||
# TODO update to pull out protocol_id from message
|
||||
protocol_id = "/echo/1.0.0"
|
||||
stream_id = await self.stream_queue.get()
|
||||
stream = MplexStream(stream_id, False, self)
|
||||
return stream, stream_id, protocol_id
|
||||
|
||||
async def send_message(self, flag, data, stream_id):
|
||||
"""
|
||||
sends a message over the connection
|
||||
:param header: header to use
|
||||
:param data: data to send in the message
|
||||
:param stream_id: stream the message is in
|
||||
:return: True if success
|
||||
"""
|
||||
# << by 3, then or with flag
|
||||
header = (stream_id << 3) | flag
|
||||
header = encode_uvarint(header)
|
||||
|
||||
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)
|
||||
|
||||
async def write_to_stream(self, _bytes):
|
||||
"""
|
||||
writes a byte array to a raw connection
|
||||
:param _bytes: byte array to write
|
||||
:return: length written
|
||||
"""
|
||||
self.raw_conn.writer.write(_bytes)
|
||||
await self.raw_conn.writer.drain()
|
||||
return len(_bytes)
|
||||
|
||||
async def handle_incoming(self):
|
||||
"""
|
||||
Read a message off of the raw connection and add it to the corresponding message buffer
|
||||
"""
|
||||
# TODO Deal with other types of messages using flag (currently _)
|
||||
# TODO call read_message in loop to handle case message for other stream was in conn
|
||||
|
||||
stream_id, _, message = await self.read_message()
|
||||
|
||||
if stream_id not in self.buffers:
|
||||
self.buffers[stream_id] = asyncio.Queue()
|
||||
await self.stream_queue.put(stream_id)
|
||||
|
||||
await self.buffers[stream_id].put(message)
|
||||
|
||||
async def read_chunk(self):
|
||||
"""
|
||||
Read a chunk of bytes off of the raw connection into data_buffer
|
||||
"""
|
||||
# unused now but possibly useful in the future
|
||||
try:
|
||||
chunk = await asyncio.wait_for(self.raw_conn.reader.read(-1), timeout=5)
|
||||
self.data_buffer += chunk
|
||||
except asyncio.TimeoutError:
|
||||
print('timeout!')
|
||||
return
|
||||
|
||||
async def read_message(self):
|
||||
"""
|
||||
Read a single message off of the raw connection
|
||||
:return: stream_id, flag, message contents
|
||||
"""
|
||||
try:
|
||||
header = await decode_uvarint_from_stream(self.raw_conn.reader)
|
||||
length = await decode_uvarint_from_stream(self.raw_conn.reader)
|
||||
message = await asyncio.wait_for(self.raw_conn.reader.read(length), timeout=5)
|
||||
except asyncio.TimeoutError:
|
||||
print("message malformed")
|
||||
return None, None, None
|
||||
|
||||
flag = header & 0x07
|
||||
stream_id = header >> 3
|
||||
|
||||
return stream_id, flag, message
|
||||
@ -1,123 +0,0 @@
|
||||
import asyncio
|
||||
from .constants import HEADER_TAGS
|
||||
from ..muxed_stream_interface import IMuxedStream
|
||||
|
||||
|
||||
class MplexStream(IMuxedStream):
|
||||
# pylint: disable=too-many-instance-attributes
|
||||
"""
|
||||
reference: https://github.com/libp2p/go-mplex/blob/master/stream.go
|
||||
"""
|
||||
|
||||
def __init__(self, stream_id, initiator, mplex_conn):
|
||||
"""
|
||||
create new MuxedStream in muxer
|
||||
:param stream_id: stream stream id
|
||||
:param initiator: boolean if this is an initiator
|
||||
:param mplex_conn: muxed connection of this muxed_stream
|
||||
"""
|
||||
self.stream_id = stream_id
|
||||
self.initiator = initiator
|
||||
self.mplex_conn = mplex_conn
|
||||
self.read_deadline = None
|
||||
self.write_deadline = None
|
||||
self.local_closed = False
|
||||
self.remote_closed = False
|
||||
self.stream_lock = asyncio.Lock()
|
||||
|
||||
def get_flag(self, action):
|
||||
"""
|
||||
get header flag based on action for mplex
|
||||
:param action: action type in str
|
||||
:return: int flag
|
||||
"""
|
||||
if self.initiator:
|
||||
return HEADER_TAGS[action]
|
||||
|
||||
return HEADER_TAGS[action] - 1
|
||||
|
||||
async def read(self):
|
||||
"""
|
||||
read messages associated with stream from buffer til end of file
|
||||
:return: bytes of input
|
||||
"""
|
||||
return await self.mplex_conn.read_buffer(self.stream_id)
|
||||
|
||||
async def write(self, data):
|
||||
"""
|
||||
write to stream
|
||||
:return: number of bytes written
|
||||
"""
|
||||
return await self.mplex_conn.send_message(self.get_flag("MESSAGE"), data, self.stream_id)
|
||||
|
||||
async def close(self):
|
||||
"""
|
||||
Closing a stream closes it for writing and closes the remote end for reading
|
||||
but allows writing in the other direction.
|
||||
:return: true if successful
|
||||
"""
|
||||
# TODO error handling with timeout
|
||||
# TODO understand better how mutexes are used from go repo
|
||||
await self.mplex_conn.send_message(self.get_flag("CLOSE"), None, self.stream_id)
|
||||
|
||||
remote_lock = ""
|
||||
async with self.stream_lock:
|
||||
if self.local_closed:
|
||||
return True
|
||||
self.local_closed = True
|
||||
remote_lock = self.remote_closed
|
||||
|
||||
if remote_lock:
|
||||
async with self.mplex_conn.conn_lock:
|
||||
self.mplex_conn.buffers.pop(self.stream_id)
|
||||
|
||||
return True
|
||||
|
||||
async def reset(self):
|
||||
"""
|
||||
closes both ends of the stream
|
||||
tells this remote side to hang up
|
||||
:return: true if successful
|
||||
"""
|
||||
# TODO understand better how mutexes are used here
|
||||
# TODO understand the difference between close and reset
|
||||
async with self.stream_lock:
|
||||
if self.remote_closed and self.local_closed:
|
||||
return True
|
||||
|
||||
if not self.remote_closed:
|
||||
await self.mplex_conn.send_message(self.get_flag("RESET"), None, self.stream_id)
|
||||
|
||||
self.local_closed = True
|
||||
self.remote_closed = True
|
||||
|
||||
async with self.mplex_conn.conn_lock:
|
||||
self.mplex_conn.buffers.pop(self.stream_id, None)
|
||||
|
||||
return True
|
||||
|
||||
# TODO deadline not in use
|
||||
def set_deadline(self, ttl):
|
||||
"""
|
||||
set deadline for muxed stream
|
||||
:return: True if successful
|
||||
"""
|
||||
self.read_deadline = ttl
|
||||
self.write_deadline = ttl
|
||||
return True
|
||||
|
||||
def set_read_deadline(self, ttl):
|
||||
"""
|
||||
set read deadline for muxed stream
|
||||
:return: True if successful
|
||||
"""
|
||||
self.read_deadline = ttl
|
||||
return True
|
||||
|
||||
def set_write_deadline(self, ttl):
|
||||
"""
|
||||
set write deadline for muxed stream
|
||||
:return: True if successful
|
||||
"""
|
||||
self.write_deadline = ttl
|
||||
return True
|
||||
@ -1,41 +0,0 @@
|
||||
import asyncio
|
||||
import struct
|
||||
|
||||
def encode_uvarint(number):
|
||||
"""Pack `number` into varint bytes"""
|
||||
buf = b''
|
||||
while True:
|
||||
towrite = number & 0x7f
|
||||
number >>= 7
|
||||
if number:
|
||||
buf += bytes((towrite | 0x80, ))
|
||||
else:
|
||||
buf += bytes((towrite, ))
|
||||
break
|
||||
return buf
|
||||
|
||||
def decode_uvarint(buff, index):
|
||||
shift = 0
|
||||
result = 0
|
||||
while True:
|
||||
i = buff[index]
|
||||
result |= (i & 0x7f) << shift
|
||||
shift += 7
|
||||
if not i & 0x80:
|
||||
break
|
||||
index += 1
|
||||
|
||||
return result, index + 1
|
||||
|
||||
async def decode_uvarint_from_stream(reader):
|
||||
shift = 0
|
||||
result = 0
|
||||
while True:
|
||||
byte = await asyncio.wait_for(reader.read(1), timeout=5)
|
||||
i = struct.unpack('>H', b'\x00' + byte)[0]
|
||||
result |= (i & 0x7f) << shift
|
||||
shift += 7
|
||||
if not i & 0x80:
|
||||
break
|
||||
|
||||
return result
|
||||
Reference in New Issue
Block a user