From c1eacf221fa5331375d6ef7acc769a3c5a054714 Mon Sep 17 00:00:00 2001
From: mhchia <kevin.mh.chia@gmail.com>
Date: Wed, 21 Aug 2019 23:56:04 +0800
Subject: [PATCH] PR feedback

- Check if the received peer id matches the one we initialize the
session with.
- Move the check inside `run_handshake`
---
 libp2p/security/insecure/transport.py | 34 ++++++++++++++++-----------
 1 file changed, 20 insertions(+), 14 deletions(-)

diff --git a/libp2p/security/insecure/transport.py b/libp2p/security/insecure/transport.py
index 6894cbd..2aad45c 100644
--- a/libp2p/security/insecure/transport.py
+++ b/libp2p/security/insecure/transport.py
@@ -28,26 +28,38 @@ class InsecureSession(BaseSession):
         remote_msg_bytes = await read_fixedint_prefixed(self.conn)
         remote_msg = plaintext_pb2.Exchange()
         remote_msg.ParseFromString(remote_msg_bytes)
+        received_peer_id = ID(remote_msg.id)
+
+        # 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.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}"
+            )
 
         # Verify if the given `pubkey` matches the given `peer_id`
         try:
-            remote_pubkey = pubkey_from_protobuf(remote_msg.pubkey)
+            received_pubkey = pubkey_from_protobuf(remote_msg.pubkey)
         except ValueError:
             raise HandshakeFailure(
                 f"unknown `key_type` of remote_msg.pubkey={remote_msg.pubkey}"
             )
-        remote_peer_id = ID(remote_msg.id)
-        remote_peer_id_from_pubkey = ID.from_pubkey(remote_pubkey)
-        if remote_peer_id_from_pubkey != remote_peer_id:
+        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"remote_peer_id={remote_peer_id}, remote_pubkey={remote_pubkey}, "
-                f"remote_peer_id_from_pubkey={remote_peer_id_from_pubkey}"
+                f"received_peer_id={received_peer_id}, remote_pubkey={received_pubkey}, "
+                f"peer_id_from_received_pubkey={peer_id_from_received_pubkey}"
             )
 
         # Nothing is wrong. Store the `pubkey` and `peer_id` in the session.
-        self.remote_permanent_pubkey = remote_pubkey
-        self.remote_peer_id = remote_peer_id
+        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.initiator:
+            self.remote_peer_id = received_peer_id
 
         # TODO: Store `pubkey` and `peer_id` to `PeerStore`
 
@@ -76,12 +88,6 @@ class InsecureTransport(BaseSecureTransport):
         """
         session = InsecureSession(self, conn, peer_id)
         await session.run_handshake()
-        received_peer_id = session.get_remote_peer()
-        if session.remote_permanent_pubkey is not None and received_peer_id != peer_id:
-            raise HandshakeFailure(
-                "remote peer sent unexpected peer ID. "
-                f"expected={peer_id} received={received_peer_id}"
-            )
         return session