Merge pull request #631 from Winter-Soren/feat/619-store-pubkey-peerid-peerstore

feat: store pubkey and peerid in peerstore

Author: seetadev

examples/doc-examples/example_encryption_insecure.py

@@ -24,6 +24,12 @@ async def main():
     insecure_transport = InsecureTransport(
         # local_key_pair: The key pair used for libp2p identity
         local_key_pair=key_pair,
+        # secure_bytes_provider: Optional function to generate secure random bytes
+        # (defaults to secrets.token_bytes)
+        secure_bytes_provider=None,  # Use default implementation
+        # peerstore: Optional peerstore to store peer IDs and public keys
+        # (defaults to None)
+        peerstore=None,
     )
 
     # Create a security options dictionary mapping protocol ID to transport

libp2p/__init__.py

@@ -200,7 +200,9 @@ def new_swarm(
             key_pair, noise_privkey=noise_key_pair.private_key
         ),
         TProtocol(secio.ID): secio.Transport(key_pair),
-        TProtocol(PLAINTEXT_PROTOCOL_ID): InsecureTransport(key_pair),
+        TProtocol(PLAINTEXT_PROTOCOL_ID): InsecureTransport(
+            key_pair, peerstore=peerstore_opt
+        ),
     }
 
     # Use given muxer preference if provided, otherwise use global default

libp2p/security/insecure/transport.py

@@ -1,11 +1,15 @@
+from collections.abc import Callable
+
 from libp2p.abc import (
+    IPeerStore,
     IRawConnection,
     ISecureConn,
 )
 from libp2p.crypto.exceptions import (
     MissingDeserializerError,
 )
 from libp2p.crypto.keys import (
+    KeyPair,
     PrivateKey,
     PublicKey,
 )
@@ -30,11 +34,15 @@
 from libp2p.peer.id import (
     ID,
 )
+from libp2p.peer.peerstore import (
+    PeerStoreError,
+)
 from libp2p.security.base_session import (
     BaseSession,
 )
 from libp2p.security.base_transport import (
     BaseSecureTransport,
+    default_secure_bytes_provider,
 )
 from libp2p.security.exceptions import (
     HandshakeFailure,
@@ -102,6 +110,7 @@ async def run_handshake(
     conn: IRawConnection,
     is_initiator: bool,
     remote_peer_id: ID | None,
+    peerstore: IPeerStore | None = None,
 ) -> ISecureConn:
     """Raise `HandshakeFailure` when handshake failed."""
     msg = make_exchange_message(local_private_key.get_public_key())
@@ -164,7 +173,14 @@ async def run_handshake(
         conn=conn,
     )
 
-    # TODO: Store `pubkey` and `peer_id` to `PeerStore`
+    # Store `pubkey` and `peer_id` to `PeerStore`
+    if peerstore is not None:
+        try:
+            peerstore.add_pubkey(received_peer_id, received_pubkey)
+        except PeerStoreError:
+            # If peer ID and pubkey don't match, it would have already been caught above
+            # This might happen if the peer is already in the store
+            pass
 
     return secure_conn
 
@@ -175,6 +191,18 @@ class InsecureTransport(BaseSecureTransport):
     transport does not add any additional security.
     """
 
+    def __init__(
+        self,
+        local_key_pair: KeyPair,
+        secure_bytes_provider: Callable[[int], bytes] | None = None,
+        peerstore: IPeerStore | None = None,
+    ) -> None:
+        # If secure_bytes_provider is None, use the default one
+        if secure_bytes_provider is None:
+            secure_bytes_provider = default_secure_bytes_provider
+        super().__init__(local_key_pair, secure_bytes_provider)
+        self.peerstore = peerstore
+
     async def secure_inbound(self, conn: IRawConnection) -> ISecureConn:
         """
         Secure the connection, either locally or by communicating with opposing
@@ -183,8 +211,9 @@ async def secure_inbound(self, conn: IRawConnection) -> ISecureConn:
 
         :return: secure connection object (that implements secure_conn_interface)
         """
+        # For inbound connections, we don't know the remote peer ID yet
         return await run_handshake(
-            self.local_peer, self.local_private_key, conn, False, None
+            self.local_peer, self.local_private_key, conn, False, None, self.peerstore
         )
 
     async def secure_outbound(self, conn: IRawConnection, peer_id: ID) -> ISecureConn:
@@ -195,7 +224,7 @@ async def secure_outbound(self, conn: IRawConnection, peer_id: ID) -> ISecureCon
         :return: secure connection object (that implements secure_conn_interface)
         """
         return await run_handshake(
-            self.local_peer, self.local_private_key, conn, True, peer_id
+            self.local_peer, self.local_private_key, conn, True, peer_id, self.peerstore
         )
 
 

newsfragments/631.feature.rst

@@ -0,0 +1,7 @@
+Store public key and peer ID in peerstore during handshake
+
+Modified the InsecureTransport class to accept an optional peerstore parameter and updated the handshake process to store the received public key and peer ID in the peerstore when available.
+
+Added test cases to verify:
+1. The peerstore remains unchanged when handshake fails due to peer ID mismatch
+2. The handshake correctly adds a public key to a peer ID that already exists in the peerstore but doesn't have a public key yet