Security: SecureSession

Make security sessions(secio, noise) share the same implementation
`BaseSession` to avoid duplicate implementation of buffered read.

libp2p/security/insecure/io.py

@@ -0,0 +1,19 @@
+from libp2p.io.abc import MsgReadWriteCloser, ReadWriteCloser
+from libp2p.utils import encode_fixedint_prefixed, read_fixedint_prefixed
+
+
+class PlaintextHandshakeReadWriter(MsgReadWriteCloser):
+    conn: ReadWriteCloser
+
+    def __init__(self, conn: ReadWriteCloser) -> None:
+        self.conn = conn
+
+    async def read_msg(self) -> bytes:
+        return await read_fixedint_prefixed(self.conn)
+
+    async def write_msg(self, msg: bytes) -> None:
+        encoded_msg_bytes = encode_fixedint_prefixed(msg)
+        await self.conn.write(encoded_msg_bytes)
+
+    async def close(self) -> None:
+        await self.conn.close()

libp2p/security/insecure/transport.py

@@ -13,8 +13,8 @@ from libp2p.security.base_transport import BaseSecureTransport
 from libp2p.security.exceptions import HandshakeFailure
 from libp2p.security.secure_conn_interface import ISecureConn
 from libp2p.typing import TProtocol
-from libp2p.utils import encode_fixedint_prefixed, read_fixedint_prefixed
 
+from .io import PlaintextHandshakeReadWriter
 from .pb import plaintext_pb2
 
 # Reference: https://github.com/libp2p/go-libp2p-core/blob/master/sec/insecure/insecure.go
@@ -44,60 +44,66 @@ class InsecureSession(BaseSession):
     async def close(self) -> None:
         await self.conn.close()
 
-    async def run_handshake(self) -> None:
-        """Raise `HandshakeFailure` when handshake failed."""
-        msg = make_exchange_message(self.local_private_key.get_public_key())
-        msg_bytes = msg.SerializeToString()
-        encoded_msg_bytes = encode_fixedint_prefixed(msg_bytes)
-        try:
-            await self.write(encoded_msg_bytes)
-        except RawConnError as e:
-            raise HandshakeFailure("connection closed") from e
 
-        try:
-            remote_msg_bytes = await read_fixedint_prefixed(self.conn)
-        except RawConnError as e:
-            raise HandshakeFailure("connection closed") from e
-        remote_msg = plaintext_pb2.Exchange()
-        remote_msg.ParseFromString(remote_msg_bytes)
-        received_peer_id = ID(remote_msg.id)
+async def run_handshake(
+    local_peer: ID,
+    local_private_key: PrivateKey,
+    conn: IRawConnection,
+    is_initiator: bool,
+    remote_peer_id: ID,
+) -> ISecureConn:
+    """Raise `HandshakeFailure` when handshake failed."""
+    msg = make_exchange_message(local_private_key.get_public_key())
+    msg_bytes = msg.SerializeToString()
+    read_writer = PlaintextHandshakeReadWriter(conn)
+    try:
+        await read_writer.write_msg(msg_bytes)
+    except RawConnError as e:
+        raise HandshakeFailure("connection closed") from e
 
-        # Verify if the receive `ID` matches the one we originally initialize the session.
-        # We only need to check it when we are the initiator, because only in that condition
-        # we possibly knows the `ID` of the remote.
-        if self.is_initiator and self.remote_peer_id != received_peer_id:
-            raise HandshakeFailure(
-                "remote peer sent unexpected peer ID. "
-                f"expected={self.remote_peer_id} received={received_peer_id}"
-            )
+    try:
+        remote_msg_bytes = await read_writer.read_msg()
+    except RawConnError as e:
+        raise HandshakeFailure("connection closed") from e
+    remote_msg = plaintext_pb2.Exchange()
+    remote_msg.ParseFromString(remote_msg_bytes)
+    received_peer_id = ID(remote_msg.id)
 
-        # Verify if the given `pubkey` matches the given `peer_id`
-        try:
-            received_pubkey = deserialize_public_key(
-                remote_msg.pubkey.SerializeToString()
-            )
-        except ValueError as e:
-            raise HandshakeFailure(
-                f"unknown `key_type` of remote_msg.pubkey={remote_msg.pubkey}"
-            ) from e
-        except MissingDeserializerError as error:
-            raise HandshakeFailure() from error
-        peer_id_from_received_pubkey = ID.from_pubkey(received_pubkey)
-        if peer_id_from_received_pubkey != received_peer_id:
-            raise HandshakeFailure(
-                "peer id and pubkey from the remote mismatch: "
-                f"received_peer_id={received_peer_id}, remote_pubkey={received_pubkey}, "
-                f"peer_id_from_received_pubkey={peer_id_from_received_pubkey}"
-            )
+    # Verify if the receive `ID` matches the one we originally initialize the session.
+    # We only need to check it when we are the initiator, because only in that condition
+    # we possibly knows the `ID` of the remote.
+    if is_initiator and remote_peer_id != received_peer_id:
+        raise HandshakeFailure(
+            "remote peer sent unexpected peer ID. "
+            f"expected={remote_peer_id} received={received_peer_id}"
+        )
 
-        # Nothing is wrong. Store the `pubkey` and `peer_id` in the session.
-        self.remote_permanent_pubkey = received_pubkey
-        # Only need to set peer's id when we don't know it before,
-        # i.e. we are not the connection initiator.
-        if not self.is_initiator:
-            self.remote_peer_id = received_peer_id
+    # Verify if the given `pubkey` matches the given `peer_id`
+    try:
+        received_pubkey = deserialize_public_key(remote_msg.pubkey.SerializeToString())
+    except ValueError as e:
+        raise HandshakeFailure(
+            f"unknown `key_type` of remote_msg.pubkey={remote_msg.pubkey}"
+        ) from e
+    except MissingDeserializerError as error:
+        raise HandshakeFailure() from error
+    peer_id_from_received_pubkey = ID.from_pubkey(received_pubkey)
+    if peer_id_from_received_pubkey != received_peer_id:
+        raise HandshakeFailure(
+            "peer id and pubkey from the remote mismatch: "
+            f"received_peer_id={received_peer_id}, remote_pubkey={received_pubkey}, "
+            f"peer_id_from_received_pubkey={peer_id_from_received_pubkey}"
+        )
 
-        # TODO: Store `pubkey` and `peer_id` to `PeerStore`
+    secure_conn = InsecureSession(
+        local_peer, local_private_key, conn, is_initiator, received_peer_id
+    )
+
+    # Nothing is wrong. Store the `pubkey` and `peer_id` in the session.
+    secure_conn.remote_permanent_pubkey = received_pubkey
+    # TODO: Store `pubkey` and `peer_id` to `PeerStore`
+
+    return secure_conn
 
 
 class InsecureTransport(BaseSecureTransport):
@@ -113,9 +119,9 @@ class InsecureTransport(BaseSecureTransport):
 
         :return: secure connection object (that implements secure_conn_interface)
         """
-        session = InsecureSession(self.local_peer, self.local_private_key, conn, False)
-        await session.run_handshake()
-        return session
+        return await run_handshake(
+            self.local_peer, self.local_private_key, conn, False, None
+        )
 
     async def secure_outbound(self, conn: IRawConnection, peer_id: ID) -> ISecureConn:
         """
@@ -124,11 +130,9 @@ class InsecureTransport(BaseSecureTransport):
 
         :return: secure connection object (that implements secure_conn_interface)
         """
-        session = InsecureSession(
+        return await run_handshake(
             self.local_peer, self.local_private_key, conn, True, peer_id
         )
-        await session.run_handshake()
-        return session
 
 
 def make_exchange_message(pubkey: PublicKey) -> plaintext_pb2.Exchange: