Work in terms of the IRawConnection abstraction

This commit is contained in:
Alex Stokes 2019-08-24 22:52:09 +02:00
parent eb5ef39399
commit a363ba97d1
No known key found for this signature in database
GPG Key ID: 51CE1721B245C086

View File

@ -1,5 +1,5 @@
from dataclasses import dataclass from dataclasses import dataclass
from typing import Optional, Tuple, cast from typing import Optional, Tuple
import multihash import multihash
@ -16,7 +16,6 @@ from libp2p.crypto.keys import PrivateKey, PublicKey
from libp2p.crypto.serialization import deserialize_public_key from libp2p.crypto.serialization import deserialize_public_key
from libp2p.io.msgio import encode as encode_message from libp2p.io.msgio import encode as encode_message
from libp2p.io.msgio import read_next_message from libp2p.io.msgio import read_next_message
from libp2p.network.connection.raw_connection import RawConnection
from libp2p.network.connection.raw_connection_interface import IRawConnection from libp2p.network.connection.raw_connection_interface import IRawConnection
from libp2p.peer.id import ID as PeerID from libp2p.peer.id import ID as PeerID
from libp2p.security.base_session import BaseSession from libp2p.security.base_session import BaseSession
@ -72,8 +71,7 @@ class SecureSession(BaseSession):
async def _read_msg(self) -> bytes: async def _read_msg(self) -> bytes:
# TODO do we need to serialize reads? # TODO do we need to serialize reads?
# TODO do not expose reader msg = await read_next_message(self.conn)
msg = await read_next_message(cast(RawConnection, self.conn).reader)
return self.remote_encrypter.decrypt_if_valid(msg) return self.remote_encrypter.decrypt_if_valid(msg)
async def write(self, data: bytes) -> None: async def write(self, data: bytes) -> None:
@ -84,9 +82,7 @@ class SecureSession(BaseSession):
encrypted_data = self.local_encrypter.encrypt(data) encrypted_data = self.local_encrypter.encrypt(data)
tag = self.local_encrypter.authenticate(encrypted_data) tag = self.local_encrypter.authenticate(encrypted_data)
msg = encode_message(encrypted_data + tag) msg = encode_message(encrypted_data + tag)
# TODO clean up how we write messages await self.conn.write(msg)
await cast(RawConnection, self.conn).writer.write(msg)
await cast(RawConnection, self.conn).writer.drain()
@dataclass(frozen=True) @dataclass(frozen=True)
@ -160,12 +156,8 @@ class SessionParameters:
async def _response_to_msg(conn: IRawConnection, msg: bytes) -> bytes: async def _response_to_msg(conn: IRawConnection, msg: bytes) -> bytes:
# TODO clean up ``IRawConnection`` so that we don't have to break await conn.write(encode_message(msg))
# the abstraction return await read_next_message(conn)
await cast(RawConnection, conn).writer.write(encode_message(msg))
await cast(RawConnection, conn).writer.drain()
return await read_next_message(cast(RawConnection, conn).reader)
def _mk_multihash_sha256(data: bytes) -> bytes: def _mk_multihash_sha256(data: bytes) -> bytes: