Merge pull request #1107 from sumanjeet0012/universal_connectivity

Improve RSA Key Compatibility, Public Key Extraction, and Pubsub Connection Handling

Author: seetadev

libp2p/crypto/rsa.py

@@ -8,6 +8,9 @@
 from Crypto.Signature import (
     pkcs1_15,
 )
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.serialization import load_der_public_key
 
 from libp2p.crypto.exceptions import (
     CryptographyError,
@@ -59,7 +62,20 @@ def to_bytes(self) -> bytes:
 
     @classmethod
     def from_bytes(cls, key_bytes: bytes) -> "RSAPublicKey":
-        rsakey = RSA.import_key(key_bytes)
+        try:
+            rsakey = RSA.import_key(key_bytes)
+        except ValueError:
+            # PyCryptodome might fail on some PKIX formats, try using cryptography lib
+            try:
+                crypto_key = load_der_public_key(key_bytes, backend=default_backend())
+                # Re-export in PKCS1 format that PyCryptodome understands
+                pkcs1_bytes = crypto_key.public_bytes(
+                    encoding=serialization.Encoding.DER,
+                    format=serialization.PublicFormat.PKCS1,
+                )
+                rsakey = RSA.import_key(pkcs1_bytes)
+            except Exception:
+                raise
         validate_rsa_key_size(rsakey)
         return cls(rsakey)
 

libp2p/peer/id.py

@@ -7,6 +7,7 @@
 from libp2p.crypto.keys import (
     PublicKey,
 )
+from libp2p.crypto.serialization import deserialize_public_key
 
 # NOTE: On inlining...
 # See: https://github.com/libp2p/specs/issues/138
@@ -90,6 +91,32 @@ def from_pubkey(cls, key: PublicKey) -> "ID":
         mh_digest = multihash.digest(serialized_key, algo)
         return cls(mh_digest.encode())
 
+    def extract_public_key(self) -> PublicKey | None:
+        """
+        Extract the public key from this peer ID if it uses an identity multihash.
+
+        For Ed25519 and other small keys, the public key is embedded directly
+        in the peer ID using an identity multihash. For larger keys like RSA,
+        the peer ID uses a SHA-256 hash and the key cannot be extracted.
+
+        Returns:
+            The public key if it can be extracted, None otherwise.
+
+        """
+        try:
+            # Decode the multihash to check if it's an identity hash
+            mh_decoded = multihash.decode(self._bytes)
+
+            # Identity multihash func code is 0x00
+            if mh_decoded.func == IDENTITY_MULTIHASH_CODE:
+                # The digest is the serialized public key protobuf
+                return deserialize_public_key(mh_decoded.digest)
+            else:
+                # Not an identity hash, key cannot be extracted
+                return None
+        except Exception:
+            return None
+
 
 def sha256_digest(data: str | bytes) -> bytes:
     if isinstance(data, str):

libp2p/pubsub/pubsub.py

@@ -454,6 +454,11 @@ async def wait_until_ready(self) -> None:
         await self.event_handle_dead_peer_queue_started.wait()
 
     async def _handle_new_peer(self, peer_id: ID) -> None:
+        # Check if we already have a pubsub stream with this peer to avoid duplicates
+        if peer_id in self.peers:
+            logger.debug("Peer %s already has pubsub stream, skipping", peer_id)
+            return
+
         if self.is_peer_blacklisted(peer_id):
             logger.debug("Rejecting blacklisted peer %s", peer_id)
             return
@@ -521,11 +526,28 @@ async def handle_dead_peer_queue(self) -> None:
         """
         Continuously read from dead peer channel and close the stream
         between that peer and remove peer info from pubsub and pubsub router.
+        Only removes the peer if there are no remaining active connections.
         """
         async with self.dead_peer_receive_channel:
             self.event_handle_dead_peer_queue_started.set()
             async for peer_id in self.dead_peer_receive_channel:
-                # Remove Peer
+                # Check if peer still has active connections before removing
+                # This prevents premature removal when multiple connections exist
+                network = self.host.get_network()
+                remaining_connections = network.get_connections(peer_id)
+                if remaining_connections:
+                    # Filter out closed connections
+                    active_connections = [
+                        c for c in remaining_connections if not c.muxed_conn.is_closed
+                    ]
+                    if active_connections:
+                        logger.debug(
+                            "Peer %s still has %d active connections, not removing",
+                            peer_id,
+                            len(active_connections),
+                        )
+                        continue
+                # Remove Peer - no more active connections
                 self._handle_dead_peer(peer_id)
 
     def handle_subscription(

libp2p/pubsub/validators.py

@@ -28,10 +28,40 @@ def signature_validator(msg: rpc_pb2.Message) -> bool:
         logger.debug("Reject because no signature attached for msg: %s", msg)
         return False
 
-    # Validate if message sender matches message signer,
-    # i.e., check if `msg.key` matches `msg.from_id`
-    msg_pubkey = deserialize_public_key(msg.key)
-    if ID.from_pubkey(msg_pubkey) != msg.from_id:
+    # Get the public key - either from msg.key or extracted from peer ID
+    # For Ed25519 and other small keys, the key is embedded in the peer ID
+    # For RSA keys (which are too large), the key is sent in msg.key
+    msg_pubkey = None
+    sender_id = ID(msg.from_id)
+
+    # First, try to extract public key from the peer ID (from_id)
+    try:
+        msg_pubkey = sender_id.extract_public_key()
+        if msg_pubkey:
+            logger.debug(
+                "Extracted public key from peer ID, type: %s", msg_pubkey.get_type()
+            )
+    except Exception as e:
+        logger.debug("Could not extract key from peer ID: %s", e)
+
+    # If we couldn't extract from peer ID, use msg.key
+    if msg_pubkey is None:
+        if not msg.key or len(msg.key) == 0:
+            logger.debug(
+                "Reject because no key attached and cannot extract from peer "
+                "ID for msg: %s",
+                msg,
+            )
+            return False
+        logger.debug(
+            "Using msg.key for signature verification, length: %d", len(msg.key)
+        )
+        try:
+            msg_pubkey = deserialize_public_key(msg.key)
+        except Exception as e:
+            logger.debug("Failed to deserialize public key: %s", e)
+            raise
+    if ID.from_pubkey(msg_pubkey) != sender_id:
         logger.debug(
             "Reject because signing key does not match sender ID for msg: %s", msg
         )

newsfragments/1106.feature.rst

@@ -0,0 +1 @@
+Improved RSA key compatibility with other libp2p implementations, added public key extraction from peer IDs for Ed25519/Secp256k1 keys, and enhanced pubsub connection management to prevent premature peer removal and service crashes.