Support read/write of noise msg and payload

This commit is contained in:
mhchia
2020-02-15 12:12:12 +08:00
parent 1f881e0464
commit 8a4ebd4cbb
5 changed files with 158 additions and 33 deletions

View File

@ -9,24 +9,8 @@ from libp2p.peer.id import ID
from libp2p.security.secure_conn_interface import ISecureConn
from .connection import NoiseConnection
# FIXME: Choose a serious bound number.
NUM_BYTES_TO_READ = 2048
# TODO: Merged into `BasePattern`?
class PreHandshakeConnection:
conn: IRawConnection
def __init__(self, conn: IRawConnection) -> None:
self.conn = conn
async def write_msg(self, data: bytes) -> None:
# TODO:
await self.conn.write(data)
async def read_msg(self) -> bytes:
return await self.conn.read(NUM_BYTES_TO_READ)
from .exceptions import HandshakeHasNotFinished
from .io import NoiseHandshakeReadWriter
class IPattern(ABC):
@ -66,25 +50,24 @@ class PatternXX(BasePattern):
async def handshake_inbound(self, conn: IRawConnection) -> ISecureConn:
noise_state = self.create_noise_state()
handshake_conn = PreHandshakeConnection(conn)
noise_state.set_as_responder()
noise_state.start_handshake()
msg_0_encrypted = await handshake_conn.read_msg()
read_writer = NoiseHandshakeReadWriter(conn, noise_state)
# TODO: Parse and save the payload from the other side.
_ = noise_state.read_message(msg_0_encrypted)
_ = await read_writer.read_msg()
# TODO: Send our payload.
our_payload = b"server"
msg_1_encrypted = noise_state.write_message(our_payload)
await handshake_conn.write_msg(msg_1_encrypted)
await read_writer.write_msg(our_payload)
msg_2_encrypted = await handshake_conn.read_msg()
# TODO: Parse and save another payload from the other side.
_ = noise_state.read_message(msg_2_encrypted)
_ = await read_writer.read_msg()
# TODO: Add a specific exception
if not noise_state.handshake_finished:
raise Exception
raise HandshakeHasNotFinished(
"handshake done but it is not marked as finished in `noise_state`"
)
# FIXME: `remote_peer` should be derived from the messages.
return NoiseConnection(self.local_peer, self.libp2p_privkey, None, conn, False)
@ -93,19 +76,17 @@ class PatternXX(BasePattern):
self, conn: IRawConnection, remote_peer: ID
) -> ISecureConn:
noise_state = self.create_noise_state()
handshake_conn = PreHandshakeConnection(conn)
read_writer = NoiseHandshakeReadWriter(conn, noise_state)
noise_state.set_as_initiator()
noise_state.start_handshake()
msg_0 = noise_state.write_message()
await handshake_conn.write_msg(msg_0)
msg_1_encrypted = await handshake_conn.read_msg()
await read_writer.write_msg(b"")
# TODO: Parse and save the payload from the other side.
_ = noise_state.read_message(msg_1_encrypted)
_ = await read_writer.read_msg()
# TODO: Send our payload.
our_payload = b"client"
msg_2_encrypted = noise_state.write_message(our_payload)
await handshake_conn.write_msg(msg_2_encrypted)
await read_writer.write_msg(our_payload)
# TODO: Add a specific exception
if not noise_state.handshake_finished: