Avoid using message-based IO in the plaintext protocol

Can reuse the machinery in `secio` but need to generalize the
"buffering" there
This commit is contained in:
Alex Stokes 2019-09-04 09:28:53 -07:00
parent f86ba7283d
commit b214f88f75
No known key found for this signature in database
GPG Key ID: 51CE1721B245C086
3 changed files with 28 additions and 22 deletions

View File

@ -1,7 +1,6 @@
from typing import Optional from typing import Optional
from libp2p.crypto.keys import PrivateKey, PublicKey from libp2p.crypto.keys import PrivateKey, PublicKey
from libp2p.io.msgio import MsgIOReadWriter
from libp2p.peer.id import ID from libp2p.peer.id import ID
from libp2p.security.secure_conn_interface import ISecureConn from libp2p.security.secure_conn_interface import ISecureConn
@ -14,7 +13,6 @@ class BaseSession(ISecureConn):
local_peer: ID local_peer: ID
local_private_key: PrivateKey local_private_key: PrivateKey
conn: MsgIOReadWriter
remote_peer_id: ID remote_peer_id: ID
remote_permanent_pubkey: PublicKey remote_permanent_pubkey: PublicKey
@ -22,27 +20,14 @@ class BaseSession(ISecureConn):
self, self,
local_peer: ID, local_peer: ID,
local_private_key: PrivateKey, local_private_key: PrivateKey,
conn: MsgIOReadWriter,
peer_id: Optional[ID] = None, peer_id: Optional[ID] = None,
) -> None: ) -> None:
self.local_peer = local_peer self.local_peer = local_peer
self.local_private_key = local_private_key self.local_private_key = local_private_key
self.remote_peer_id = peer_id self.remote_peer_id = peer_id
self.remote_permanent_pubkey = None self.remote_permanent_pubkey = None
self.conn = conn
self.initiator = peer_id is not None self.initiator = peer_id is not None
async def write(self, data: bytes) -> int:
await self.conn.write(data)
return len(data)
async def read(self, n: int = -1) -> bytes:
return await self.conn.read(n)
async def close(self) -> None:
await self.conn.close()
def get_local_peer(self) -> ID: def get_local_peer(self) -> ID:
return self.local_peer return self.local_peer

View File

@ -1,7 +1,9 @@
from libp2p.crypto.keys import PublicKey from typing import Optional
from libp2p.crypto.keys import PrivateKey, PublicKey
from libp2p.crypto.pb import crypto_pb2 from libp2p.crypto.pb import crypto_pb2
from libp2p.crypto.utils import pubkey_from_protobuf from libp2p.crypto.utils import pubkey_from_protobuf
from libp2p.io.msgio import MsgIOReadWriter from libp2p.io.abc import ReadWriteCloser
from libp2p.network.connection.raw_connection_interface import IRawConnection from libp2p.network.connection.raw_connection_interface import IRawConnection
from libp2p.peer.id import ID from libp2p.peer.id import ID
from libp2p.security.base_session import BaseSession from libp2p.security.base_session import BaseSession
@ -20,6 +22,26 @@ PLAINTEXT_PROTOCOL_ID = TProtocol("/plaintext/2.0.0")
class InsecureSession(BaseSession): class InsecureSession(BaseSession):
def __init__(
self,
local_peer: ID,
local_private_key: PrivateKey,
conn: ReadWriteCloser,
peer_id: Optional[ID] = None,
) -> None:
super().__init__(local_peer, local_private_key, peer_id)
self.conn = conn
async def write(self, data: bytes) -> int:
await self.conn.write(data)
return len(data)
async def read(self, n: int = -1) -> bytes:
return await self.conn.read(n)
async def close(self) -> None:
await self.conn.close()
async def run_handshake(self) -> None: async def run_handshake(self) -> None:
msg = make_exchange_message(self.local_private_key.get_public_key()) msg = make_exchange_message(self.local_private_key.get_public_key())
msg_bytes = msg.SerializeToString() msg_bytes = msg.SerializeToString()
@ -77,8 +99,7 @@ class InsecureTransport(BaseSecureTransport):
for an inbound connection (i.e. we are not the initiator) for an inbound connection (i.e. we are not the initiator)
:return: secure connection object (that implements secure_conn_interface) :return: secure connection object (that implements secure_conn_interface)
""" """
msg_io = MsgIOReadWriter(conn) session = InsecureSession(self.local_peer, self.local_private_key, conn)
session = InsecureSession(self.local_peer, self.local_private_key, msg_io)
await session.run_handshake() await session.run_handshake()
return session return session
@ -88,9 +109,8 @@ class InsecureTransport(BaseSecureTransport):
for an inbound connection (i.e. we are the initiator) for an inbound connection (i.e. we are the initiator)
:return: secure connection object (that implements secure_conn_interface) :return: secure connection object (that implements secure_conn_interface)
""" """
msg_io = MsgIOReadWriter(conn)
session = InsecureSession( session = InsecureSession(
self.local_peer, self.local_private_key, msg_io, peer_id self.local_peer, self.local_private_key, conn, peer_id
) )
await session.run_handshake() await session.run_handshake()
return session return session

View File

@ -58,7 +58,8 @@ class SecureSession(BaseSession):
remote_encryption_parameters: AuthenticatedEncryptionParameters, remote_encryption_parameters: AuthenticatedEncryptionParameters,
conn: MsgIOReadWriter, conn: MsgIOReadWriter,
) -> None: ) -> None:
super().__init__(local_peer, local_private_key, conn, remote_peer) super().__init__(local_peer, local_private_key, remote_peer)
self.conn = conn
self.local_encryption_parameters = local_encryption_parameters self.local_encryption_parameters = local_encryption_parameters
self.remote_encryption_parameters = remote_encryption_parameters self.remote_encryption_parameters = remote_encryption_parameters