Dispatch serialization of keys based on key type

- Add some tests to check high-level roundtrip

libp2p/crypto/keys.py

@@ -33,6 +33,9 @@ class Key(ABC):
         """
         ...
 
+    def __eq__(self, other: "Key") -> bool:
+        return self.impl == other.impl
+
 
 class PublicKey(Key):
     """
@@ -61,14 +64,18 @@ class PublicKey(Key):
         """
         return self._serialize_to_protobuf().SerializeToString()
 
+    @classmethod
+    def deserialize_from_protobuf(cls, protobuf_data: bytes) -> protobuf.PublicKey:
+        protobuf_key = protobuf.PublicKey()
+        protobuf_key.ParseFromString(protobuf_data)
+        return protobuf_key
+
 
 class PrivateKey(Key):
     """
     A ``PrivateKey`` represents a cryptographic private key.
     """
 
-    protobuf_constructor = protobuf.PrivateKey
-
     @abstractmethod
     def sign(self, data: bytes) -> bytes:
         ...
@@ -92,6 +99,21 @@ class PrivateKey(Key):
         """
         return self._serialize_to_protobuf().SerializeToString()
 
+    def _protobuf_from_serialization(self, data: bytes) -> protobuf.PrivateKey:
+        """
+        Return the protobuf representation of this ``Key``.
+        """
+        key_type = self.get_type().value
+        data = self.to_bytes()
+        protobuf_key = protobuf.PrivateKey(key_type=key_type, data=data)
+        return protobuf_key
+
+    @classmethod
+    def deserialize_from_protobuf(cls, protobuf_data: bytes) -> protobuf.PrivateKey:
+        protobuf_key = protobuf.PrivateKey()
+        protobuf_key.ParseFromString(protobuf_data)
+        return protobuf_key
+
 
 @dataclass(frozen=True)
 class KeyPair:

libp2p/crypto/secp256k1.py

@@ -11,9 +11,14 @@ class Secp256k1PublicKey(PublicKey):
         return self.impl.format()
 
     @classmethod
-    def from_bytes(cls, key_bytes: bytes) -> "Secp256k1PublicKey":
-        secp256k1_pubkey = coincurve.PublicKey(key_bytes)
-        return cls(secp256k1_pubkey)
+    def from_bytes(cls, data: bytes) -> "Secp256k1PublicKey":
+        impl = coincurve.PublicKey(data)
+        return cls(impl)
+
+    @classmethod
+    def deserialize(cls, data: bytes) -> "Secp256k1PublicKey":
+        protobuf_key = cls.deserialize_from_protobuf(data)
+        return cls.from_bytes(protobuf_key.data)
 
     def get_type(self) -> KeyType:
         return KeyType.Secp256k1
@@ -34,6 +39,16 @@ class Secp256k1PrivateKey(PrivateKey):
     def to_bytes(self) -> bytes:
         return self.impl.secret
 
+    @classmethod
+    def from_bytes(cls, data: bytes) -> "Secp256k1PrivateKey":
+        impl = coincurve.PrivateKey(data)
+        return cls(impl)
+
+    @classmethod
+    def deserialize(cls, data: bytes) -> "Secp256k1PrivateKey":
+        protobuf_key = cls.deserialize_from_protobuf(data)
+        return cls.from_bytes(protobuf_key.data)
+
     def get_type(self) -> KeyType:
         return KeyType.Secp256k1
 

libp2p/crypto/serialization.py

@@ -0,0 +1,22 @@
+from libp2p.crypto.keys import KeyType, PrivateKey, PublicKey
+from libp2p.crypto.secp256k1 import Secp256k1PrivateKey, Secp256k1PublicKey
+
+key_type_to_public_key_deserializer = {
+    KeyType.Secp256k1.value: Secp256k1PublicKey.from_bytes
+}
+
+key_type_to_private_key_deserializer = {
+    KeyType.Secp256k1.value: Secp256k1PrivateKey.from_bytes
+}
+
+
+def deserialize_public_key(data: bytes) -> PublicKey:
+    f = PublicKey.deserialize_from_protobuf(data)
+    deserializer = key_type_to_public_key_deserializer[f.key_type]
+    return deserializer(f.data)
+
+
+def deserialize_private_key(data: bytes) -> PrivateKey:
+    f = PrivateKey.deserialize_from_protobuf(data)
+    deserializer = key_type_to_private_key_deserializer[f.key_type]
+    return deserializer(f.data)

libp2p/security/secio/transport.py

@@ -14,6 +14,7 @@ from libp2p.crypto.authenticated_encryption import MacAndCipher as Encrypter
 from libp2p.crypto.ecc import ECCPublicKey
 from libp2p.crypto.key_exchange import create_ephemeral_key_pair
 from libp2p.crypto.keys import PrivateKey, PublicKey
+from libp2p.crypto.serialization import deserialize_public_key
 from libp2p.io.msgio import encode as encode_message
 from libp2p.io.msgio import read_next_message
 from libp2p.network.connection.raw_connection_interface import IRawConnection
@@ -112,8 +113,7 @@ class Proposal:
 
         nonce = protobuf.rand
         public_key_protobuf_bytes = protobuf.public_key
-        # TODO (ralexstokes) handle genericity in the deserialization
-        public_key = PublicKey.deserialize(public_key_protobuf_bytes)
+        public_key = deserialize_public_key(public_key_protobuf_bytes)
         exchanges = protobuf.exchanges
         ciphers = protobuf.ciphers
         hashes = protobuf.hashes

tests/crypto/secp256k1.py

@@ -0,0 +1,22 @@
+from libp2p.crypto.secp256k1 import create_new_key_pair
+from libp2p.crypto.serialization import deserialize_private_key, deserialize_public_key
+
+
+def test_public_key_serialize_deserialize_round_trip():
+    key_pair = create_new_key_pair()
+    public_key = key_pair.public_key
+
+    public_key_bytes = public_key.serialize()
+    another_public_key = deserialize_public_key(public_key_bytes)
+
+    assert public_key == another_public_key
+
+
+def test_private_key_serialize_deserialize_round_trip():
+    key_pair = create_new_key_pair()
+    private_key = key_pair.private_key
+
+    private_key_bytes = private_key.serialize()
+    another_private_key = deserialize_private_key(private_key_bytes)
+
+    assert private_key == another_private_key