From 921bfb65ccd60d770b3cc99c8353c2c1a0d21265 Mon Sep 17 00:00:00 2001
From: mhchia <kevin.mh.chia@gmail.com>
Date: Tue, 20 Aug 2019 23:54:33 +0800
Subject: [PATCH] Verify the remote pubkey and peer_id

- Add `from_bytes` in RSAPublicKey and Secp256k1PublicKey
- Add `pubkey_from_protobuf` to parse pubkey from protobuf
- Verify key and peer_id in `InsecureSession.run_handshake`
---
 libp2p/crypto/rsa.py                  | 10 +++++++++
 libp2p/crypto/secp256k1.py            |  5 +++++
 libp2p/crypto/utils.py                | 16 ++++++++++++++
 libp2p/security/base_session.py       |  6 ++++++
 libp2p/security/insecure/transport.py | 31 +++++++++++++++++++++------
 5 files changed, 62 insertions(+), 6 deletions(-)
 create mode 100644 libp2p/crypto/utils.py

diff --git a/libp2p/crypto/rsa.py b/libp2p/crypto/rsa.py
index ceea008..a0f8439 100644
--- a/libp2p/crypto/rsa.py
+++ b/libp2p/crypto/rsa.py
@@ -11,6 +11,11 @@ class RSAPublicKey(PublicKey):
     def to_bytes(self) -> bytes:
         return self.impl.export_key("DER")
 
+    @classmethod
+    def from_bytes(cls, key_bytes: bytes) -> "RSAPublicKey":
+        rsakey = RSA.import_key(key_bytes)
+        return cls(rsakey)
+
     def get_type(self) -> KeyType:
         return KeyType.RSA
 
@@ -30,6 +35,11 @@ class RSAPrivateKey(PrivateKey):
     def to_bytes(self) -> bytes:
         return self.impl.export_key("DER")
 
+    @classmethod
+    def from_bytes(cls, key_bytes: bytes) -> "RSAPrivateKey":
+        rsakey = RSA.import_key(key_bytes)
+        return cls(rsakey)
+
     def get_type(self) -> KeyType:
         return KeyType.RSA
 
diff --git a/libp2p/crypto/secp256k1.py b/libp2p/crypto/secp256k1.py
index 79ffc9d..e2d5fb2 100644
--- a/libp2p/crypto/secp256k1.py
+++ b/libp2p/crypto/secp256k1.py
@@ -10,6 +10,11 @@ class Secp256k1PublicKey(PublicKey):
     def to_bytes(self) -> bytes:
         return self.impl.format()
 
+    @classmethod
+    def from_bytes(cls, key_bytes: bytes) -> "Secp256k1PublicKey":
+        secp256k1_pubkey = coincurve.PublicKey(key_bytes)
+        return cls(secp256k1_pubkey)
+
     def get_type(self) -> KeyType:
         return KeyType.Secp256k1
 
diff --git a/libp2p/crypto/utils.py b/libp2p/crypto/utils.py
new file mode 100644
index 0000000..78b8b7b
--- /dev/null
+++ b/libp2p/crypto/utils.py
@@ -0,0 +1,16 @@
+from .keys import PublicKey
+from .pb import crypto_pb2 as protobuf
+from .rsa import RSAPublicKey
+from .secp256k1 import Secp256k1PublicKey
+
+
+def pubkey_from_protobuf(pubkey_pb: protobuf.PublicKey) -> PublicKey:
+    if pubkey_pb.key_type == protobuf.RSA:
+        return RSAPublicKey.from_bytes(pubkey_pb.data)
+    # TODO: Test against secp256k1 keys
+    elif pubkey_pb.key_type == protobuf.Secp256k1:
+        return Secp256k1PublicKey.from_bytes(pubkey_pb.data)
+    else:
+        raise ValueError(
+            f"unsupported key_type={pubkey_pb.key_type}, data={pubkey_pb.data!r}"
+        )
diff --git a/libp2p/security/base_session.py b/libp2p/security/base_session.py
index a41c52b..c4a7275 100644
--- a/libp2p/security/base_session.py
+++ b/libp2p/security/base_session.py
@@ -13,6 +13,12 @@ class BaseSession(ISecureConn):
     is only meant to be used in clases that derive from it.
     """
 
+    local_peer: ID
+    local_private_key: PrivateKey
+    conn: IRawConnection
+    remote_peer_id: ID
+    remote_permanent_pubkey: PublicKey
+
     def __init__(
         self, transport: BaseSecureTransport, conn: IRawConnection, peer_id: ID
     ) -> None:
diff --git a/libp2p/security/insecure/transport.py b/libp2p/security/insecure/transport.py
index c16713b..1048791 100644
--- a/libp2p/security/insecure/transport.py
+++ b/libp2p/security/insecure/transport.py
@@ -1,4 +1,5 @@
 from libp2p.crypto.keys import PublicKey
+from libp2p.crypto.utils import pubkey_from_protobuf
 from libp2p.network.connection.raw_connection_interface import IRawConnection
 from libp2p.peer.id import ID
 from libp2p.security.base_session import BaseSession
@@ -23,13 +24,31 @@ class InsecureSession(BaseSession):
         encoded_msg_bytes = encode_fixedint_prefixed(msg_bytes)
         await self.write(encoded_msg_bytes)
 
-        msg_bytes_other_side = await read_fixedint_prefixed(self.conn)
-        msg_other_side = plaintext_pb2.Exchange()
-        msg_other_side.ParseFromString(msg_bytes_other_side)
+        remote_msg_bytes = await read_fixedint_prefixed(self.conn)
+        remote_msg = plaintext_pb2.Exchange()
+        remote_msg.ParseFromString(remote_msg_bytes)
 
-        # TODO: Verify public key with peer id
-        # TODO: Store public key
-        self.remote_peer_id = ID(msg_other_side.id)
+        # Verify if the given `pubkey` matches the given `peer_id`
+        try:
+            remote_pubkey = pubkey_from_protobuf(remote_msg.pubkey)
+        except ValueError as error:
+            raise SecurityUpgradeFailure(
+                f"unknown protocol of remote_msg.pubkey={remote_msg.pubkey}"
+            ) from error
+        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:
+            raise SecurityUpgradeFailure(
+                "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}"
+            )
+
+        # Nothing is wrong. Store the `pubkey` and `peer_id` in the session.
+        self.remote_permanent_pubkey = remote_pubkey
+        self.remote_peer_id = ID(remote_msg.id)
+
+        # TODO: Store `pubkey` and `peer_id` to `PeerStore`
 
 
 class InsecureTransport(BaseSecureTransport):