varun-r-mallya/py-libp2p
1
0
Fork 0
mirror of https://github.com/varun-r-mallya/py-libp2p.git synced 2025-12-31 20:36:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

Stream rearchitecture (#126)

Browse Source 
* Add generic protocol handler

* Add generic protocol handler to stream muxing pipeline

* Modify conn_handler to only deal with connections

* mplex accept stream architecture changes

* Add create generic protocol handler

* Fix minor bugs

* who would win 4 devs or one not

* Debugging

* rearch with handle_incoming infinite loop, seems to work, needs cleanup"

* passing linting, still needs cleanup

* fixing linting again; code still needs cleanup

* fixing tests; code still needs cleanup

* adding test cleanup and task cleanup, removing prints

* linting, and cleanup complete

* storing connections based on peer id

* remove dead code

* remove unnecessary peer_id
This commit is contained in:
Robert Zajac
2019-02-24 20:58:23 -05:00
committed by GitHub
parent 17c778de15
commit 82840b5e6c
14 changed files with 367 additions and 120 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
libp2p/__init__.py
View File

@ -1,7 +1,7 @@
from Crypto.PublicKey import RSA
import asyncio
import multiaddr
from Crypto.PublicKey import RSA
from .peer.peerstore import PeerStore
from .peer.id import id_from_public_key
from .network.swarm import Swarm
@ -10,6 +10,16 @@ from .transport.upgrader import TransportUpgrader
from .transport.tcp.tcp import TCP
async def cleanup_done_tasks():
while True:
for task in asyncio.all_tasks():
if task.done():
await task
# Need not run often
# Some sleep necessary to context switch
await asyncio.sleep(3)
async def new_node(
id_opt=None, transport_opt=None,
muxer_opt=None, sec_opt=None, peerstore=None):
@ -35,4 +45,7 @@ async def new_node(
# TODO routing unimplemented
host = BasicHost(swarm)
# Kick off cleanup job
asyncio.ensure_future(cleanup_done_tasks())
return host

56
libp2p/network/swarm.py
View File

@ -1,12 +1,13 @@
import asyncio
from libp2p.protocol_muxer.multiselect_client import MultiselectClient
from libp2p.protocol_muxer.multiselect import Multiselect
from libp2p.peer.id import id_b58_decode
from .network_interface import INetwork
from .stream.net_stream import NetStream
from .connection.raw_connection import RawConnection
class Swarm(INetwork):
# pylint: disable=too-many-instance-attributes, cell-var-from-loop
@ -23,6 +24,9 @@ class Swarm(INetwork):
self.multiselect = Multiselect()
self.multiselect_client = MultiselectClient()
# Create generic protocol handler
self.generic_protocol_handler = create_generic_protocol_handler(self)
def get_peer_id(self):
return self.self_id
@ -58,10 +62,10 @@ class Swarm(INetwork):
muxed_conn = self.connections[peer_id]
else:
# Transport dials peer (gets back a raw conn)
raw_conn = await self.transport.dial(multiaddr)
raw_conn = await self.transport.dial(multiaddr, self.self_id)
# Use upgrader to upgrade raw conn to muxed conn
muxed_conn = self.upgrader.upgrade_connection(raw_conn)
muxed_conn = self.upgrader.upgrade_connection(raw_conn, self.generic_protocol_handler)
# Store muxed connection in connections
self.connections[peer_id] = muxed_conn
@ -87,7 +91,7 @@ class Swarm(INetwork):
# Use muxed conn to open stream, which returns
# a muxed stream
# TODO: Remove protocol id from being passed into muxed_conn
muxed_stream = await muxed_conn.open_stream(protocol_ids[0], peer_id, multiaddr)
muxed_stream = await muxed_conn.open_stream(protocol_ids[0], multiaddr)
# Perform protocol muxing to determine protocol to use
selected_protocol = await self.multiselect_client.select_one_of(protocol_ids, muxed_stream)
@ -117,26 +121,21 @@ class Swarm(INetwork):
return True
async def conn_handler(reader, writer):
# Read in first message (should be peer_id of initiator) and ack
peer_id = id_b58_decode((await reader.read(1024)).decode())
writer.write("received peer id".encode())
await writer.drain()
# Upgrade reader/write to a net_stream and pass \
# to appropriate stream handler (using multiaddr)
raw_conn = RawConnection(multiaddr.value_for_protocol('ip4'),
multiaddr.value_for_protocol('tcp'), reader, writer, False)
muxed_conn = self.upgrader.upgrade_connection(raw_conn)
muxed_conn = self.upgrader.upgrade_connection(raw_conn, \
self.generic_protocol_handler)
# TODO: Remove protocol id from muxed_conn accept stream or
# move protocol muxing into accept_stream
muxed_stream, _, _ = await muxed_conn.accept_stream()
# Perform protocol muxing to determine protocol to use
selected_protocol, handler = await self.multiselect.negotiate(muxed_stream)
net_stream = NetStream(muxed_stream)
net_stream.set_protocol(selected_protocol)
# Give to stream handler
# TODO: handle case when stream handler is set
# TODO: handle case of multiple protocols over same raw connection
await handler(net_stream)
# Store muxed_conn with peer id
self.connections[peer_id] = muxed_conn
try:
# Success
@ -155,6 +154,23 @@ class Swarm(INetwork):
# TODO: Support more than one transport
self.transport = transport
def create_generic_protocol_handler(swarm):
"""
Create a generic protocol handler from the given swarm. We use swarm
to extract the multiselect module so that generic_protocol_handler
can use multiselect when generic_protocol_handler is called
from a different class
"""
multiselect = swarm.multiselect
async def generic_protocol_handler(muxed_stream):
# Perform protocol muxing to determine protocol to use
_, handler = await multiselect.negotiate(muxed_stream)
# Give to stream handler
asyncio.ensure_future(handler(muxed_stream))
return generic_protocol_handler
class SwarmException(Exception):
pass

1
libp2p/protocol_muxer/multiselect.py
View File

@ -31,7 +31,6 @@ class Multiselect(IMultiselectMuxer):
:return: selected protocol name, handler function
:raise Exception: negotiation failed exception
"""
# Create a communicator to handle all communication across the stream
communicator = MultiselectCommunicator(stream)

92
libp2p/stream_muxer/mplex/mplex.py
View File

@ -1,6 +1,6 @@
import asyncio
from .utils import encode_uvarint, decode_uvarint_from_stream
from .utils import encode_uvarint, decode_uvarint_from_stream, get_flag
from .mplex_stream import MplexStream
from ..muxed_connection_interface import IMuxedConn
@ -11,25 +11,28 @@ class Mplex(IMuxedConn):
reference: https://github.com/libp2p/go-mplex/blob/master/multiplex.go
"""
def __init__(self, conn):
def __init__(self, conn, generic_protocol_handler):
"""
create a new muxed connection
:param conn: an instance of raw connection
:param initiator: boolean to prevent multiplex with self
:param generic_protocol_handler: generic protocol handler
for new muxed streams
"""
super(Mplex, self).__init__(conn, generic_protocol_handler)
self.raw_conn = conn
self.initiator = conn.initiator
# Store generic protocol handler
self.generic_protocol_handler = generic_protocol_handler
# 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(None))
# Kick off reading
asyncio.ensure_future(self.handle_incoming())
def close(self):
"""
@ -49,39 +52,29 @@ class Mplex(IMuxedConn):
: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(stream_id)
# TODO: pass down timeout from user and use that
if stream_id in self.buffers:
return await self._read_buffer_exists(stream_id)
try:
data = await asyncio.wait_for(self.buffers[stream_id].get(), timeout=3)
return data
except asyncio.TimeoutError:
return None
# Stream not created yet
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):
async def open_stream(self, protocol_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()
await self.send_message(get_flag(self.initiator, "NEW_STREAM"), None, stream_id)
return stream
async def accept_stream(self):
@ -89,11 +82,9 @@ class Mplex(IMuxedConn):
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
asyncio.ensure_future(self.generic_protocol_handler(stream))
async def send_message(self, flag, data, stream_id):
"""
@ -126,38 +117,29 @@ class Mplex(IMuxedConn):
await self.raw_conn.writer.drain()
return len(_bytes)
async def handle_incoming(self, my_stream_id):
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 _)
continue_reading = True
i = 0
while continue_reading:
i += 1
stream_id, _, message = await self.read_message()
continue_reading = (stream_id is not None and
stream_id != my_stream_id and
my_stream_id is not None)
while True:
stream_id, flag, 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)
if stream_id is not None and flag is not None and message is not None:
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)
if flag is get_flag(True, "NEW_STREAM"):
# new stream detected on connection
await self.accept_stream()
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
if message:
await self.buffers[stream_id].put(message)
# Force context switch
await asyncio.sleep(0)
async def read_message(self):
"""
@ -167,7 +149,7 @@ class Mplex(IMuxedConn):
# Timeout is set to a relatively small value to alleviate wait time to exit
# loop in handle_incoming
timeout = .1
timeout = 0.1
try:
header = await decode_uvarint_from_stream(self.raw_conn.reader, timeout)
length = await decode_uvarint_from_stream(self.raw_conn.reader, timeout)

21
libp2p/stream_muxer/mplex/mplex_stream.py
View File

@ -1,6 +1,6 @@
import asyncio
from .constants import HEADER_TAGS
from .utils import get_flag
from ..muxed_stream_interface import IMuxedStream
@ -26,17 +26,6 @@ class MplexStream(IMuxedStream):
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
@ -49,7 +38,8 @@ class MplexStream(IMuxedStream):
write to stream
:return: number of bytes written
"""
return await self.mplex_conn.send_message(self.get_flag("MESSAGE"), data, self.stream_id)
return await self.mplex_conn.send_message(
get_flag(self.initiator, "MESSAGE"), data, self.stream_id)
async def close(self):
"""
@ -59,7 +49,7 @@ class MplexStream(IMuxedStream):
"""
# 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)
await self.mplex_conn.send_message(get_flag(self.initiator, "CLOSE"), None, self.stream_id)
remote_lock = ""
async with self.stream_lock:
@ -87,7 +77,8 @@ class MplexStream(IMuxedStream):
return True
if not self.remote_closed:
await self.mplex_conn.send_message(self.get_flag("RESET"), None, self.stream_id)
await self.mplex_conn.send_message(
get_flag(self.initiator, "RESET"), None, self.stream_id)
self.local_closed = True
self.remote_closed = True

12
libp2p/stream_muxer/mplex/utils.py
View File

@ -1,5 +1,6 @@
import asyncio
import struct
from .constants import HEADER_TAGS
def encode_uvarint(number):
@ -41,3 +42,14 @@ async def decode_uvarint_from_stream(reader, timeout):
break
return result
def get_flag(initiator, action):
"""
get header flag based on action for mplex
:param action: action type in str
:return: int flag
"""
if initiator or HEADER_TAGS[action] == 0:
return HEADER_TAGS[action]
return HEADER_TAGS[action] - 1

13
libp2p/stream_muxer/muxed_connection_interface.py
View File

@ -6,6 +6,15 @@ class IMuxedConn(ABC):
reference: https://github.com/libp2p/go-stream-muxer/blob/master/muxer.go
"""
@abstractmethod
def __init__(self, conn, generic_protocol_handler):
"""
create a new muxed connection
:param conn: an instance of raw connection
:param generic_protocol_handler: generic protocol handler
for new muxed streams
"""
@abstractmethod
def close(self):
"""
@ -21,12 +30,10 @@ class IMuxedConn(ABC):
"""
@abstractmethod
def open_stream(self, protocol_id, peer_id, multi_addr):
def open_stream(self, protocol_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
"""

14
libp2p/transport/tcp/tcp.py
View File

@ -3,6 +3,7 @@ import asyncio
import multiaddr
from libp2p.network.connection.raw_connection import RawConnection
from libp2p.peer.id import id_b58_encode
from ..listener_interface import IListener
from ..transport_interface import ITransport
@ -63,10 +64,11 @@ class TCP(ITransport):
self.server = None
return True
async def dial(self, multiaddr, options=None):
async def dial(self, multiaddr, self_id, options=None):
"""
dial a transport to peer listening on multiaddr
:param multiaddr: multiaddr of peer
:param self_id: peer_id of the dialer (to send to receier)
:param options: optional object
:return: True if successful
"""
@ -75,6 +77,16 @@ class TCP(ITransport):
reader, writer = await asyncio.open_connection(host, port)
# First: send our peer ID so receiver knows it
writer.write(id_b58_encode(self_id).encode())
await writer.drain()
# Await ack for peer id
ack = (await reader.read(1024)).decode()
if ack != "received peer id":
raise Exception("Receiver did not receive peer id")
return RawConnection(host, port, reader, writer, True)
def create_listener(self, handler_function, options=None):

3
libp2p/transport/transport_interface.py
View File

@ -4,10 +4,11 @@ from abc import ABC, abstractmethod
class ITransport(ABC):
@abstractmethod
def dial(self, multiaddr, options=None):
def dial(self, multiaddr, self_id, options=None):
"""
dial a transport to peer listening on multiaddr
:param multiaddr: multiaddr of peer
:param self_id: peer_id of the dialer (to send to receier)
:param options: optional object
:return: list of multiaddrs
"""

4
libp2p/transport/upgrader.py
View File

@ -17,11 +17,11 @@ class TransportUpgrader:
def upgrade_security(self):
pass
def upgrade_connection(self, conn):
def upgrade_connection(self, conn, generic_protocol_handler):
"""
upgrade raw connection to muxed connection
"""
# For PoC, no security, default to mplex
# TODO do exchange to determine multiplexer
return Mplex(conn)
return Mplex(conn, generic_protocol_handler)