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

feat: store pubkey and peerid in peerstore

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,4 +1,7 @@
+from collections.abc import Callable
+
 from libp2p.abc import (
+    IPeerStore,
     IRawConnection,
     ISecureConn,
 )
@@ -6,6 +9,7 @@ from libp2p.crypto.exceptions import (
     MissingDeserializerError,
 )
 from libp2p.crypto.keys import (
+    KeyPair,
     PrivateKey,
     PublicKey,
 )
@@ -30,11 +34,15 @@ from libp2p.network.connection.exceptions import (
 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 @@ class InsecureTransport(BaseSecureTransport):
 
         :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 @@ class InsecureTransport(BaseSecureTransport):
         :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

tests/core/security/test_insecure_peerstore_integration.py

@@ -0,0 +1,314 @@
+import pytest
+import trio
+from trio.testing import memory_stream_pair
+
+from libp2p.abc import IRawConnection
+from libp2p.crypto.rsa import create_new_key_pair
+from libp2p.peer.id import ID
+from libp2p.peer.peerdata import PeerData
+from libp2p.peer.peerstore import PeerStore
+from libp2p.security.exceptions import HandshakeFailure
+from libp2p.security.insecure.transport import InsecureTransport
+
+
+# Adapter class to bridge between trio streams and libp2p raw connections
+class TrioStreamAdapter(IRawConnection):
+    def __init__(self, send_stream, receive_stream, is_initiator: bool = False):
+        self.send_stream = send_stream
+        self.receive_stream = receive_stream
+        self.is_initiator = is_initiator
+
+    async def write(self, data: bytes) -> None:
+        await self.send_stream.send_all(data)
+
+    async def read(self, n: int | None = None) -> bytes:
+        if n is None or n == -1:
+            raise ValueError("Reading unbounded not supported")
+        return await self.receive_stream.receive_some(n)
+
+    async def close(self) -> None:
+        await self.send_stream.aclose()
+        await self.receive_stream.aclose()
+
+    def get_remote_address(self) -> tuple[str, int] | None:
+        # Return None since this is a test adapter without real network info
+        return None
+
+
+@pytest.mark.trio
+async def test_insecure_transport_stores_pubkey_in_peerstore():
+    """
+    Test that InsecureTransport stores the pubkey and peerid in
+    peerstore during handshake.
+    """
+    # Create key pairs for both sides
+    local_key_pair = create_new_key_pair()
+    remote_key_pair = create_new_key_pair()
+
+    # Create peer IDs
+    remote_peer_id = ID.from_pubkey(remote_key_pair.public_key)
+
+    # Create peerstore
+    peerstore = PeerStore()
+
+    # Create memory streams for communication
+    local_send, remote_receive = memory_stream_pair()
+    remote_send, local_receive = memory_stream_pair()
+
+    # Create adapters
+    local_stream = TrioStreamAdapter(local_send, local_receive, is_initiator=True)
+    remote_stream = TrioStreamAdapter(remote_send, remote_receive, is_initiator=False)
+
+    # Create transports
+    local_transport = InsecureTransport(local_key_pair, peerstore=peerstore)
+    remote_transport = InsecureTransport(remote_key_pair, peerstore=None)
+
+    # Run handshake
+    async def run_local_handshake(nursery_results):
+        with trio.move_on_after(5):
+            local_conn = await local_transport.secure_outbound(
+                local_stream, remote_peer_id
+            )
+            nursery_results["local"] = local_conn
+
+    async def run_remote_handshake(nursery_results):
+        with trio.move_on_after(5):
+            remote_conn = await remote_transport.secure_inbound(remote_stream)
+            nursery_results["remote"] = remote_conn
+
+    nursery_results = {}
+    async with trio.open_nursery() as nursery:
+        nursery.start_soon(run_local_handshake, nursery_results)
+        nursery.start_soon(run_remote_handshake, nursery_results)
+        await trio.sleep(0.1)  # Give tasks a chance to finish
+
+    local_conn = nursery_results.get("local")
+    remote_conn = nursery_results.get("remote")
+
+    assert local_conn is not None, "Local handshake failed"
+    assert remote_conn is not None, "Remote handshake failed"
+
+    # Verify that the remote peer ID is in the peerstore
+    assert remote_peer_id in peerstore.peer_ids()
+
+    # Verify that the public key was stored and matches
+    stored_pubkey = peerstore.pubkey(remote_peer_id)
+    assert stored_pubkey is not None
+    assert stored_pubkey.serialize() == remote_key_pair.public_key.serialize()
+
+
+@pytest.mark.trio
+async def test_insecure_transport_without_peerstore():
+    """
+    Test that InsecureTransport works correctly
+    without a peerstore.
+    """
+    # Create key pairs for both sides
+    local_key_pair = create_new_key_pair()
+    remote_key_pair = create_new_key_pair()
+
+    # Create peer IDs
+    remote_peer_id = ID.from_pubkey(remote_key_pair.public_key)
+
+    # Create memory streams for communication
+    local_send, remote_receive = memory_stream_pair()
+    remote_send, local_receive = memory_stream_pair()
+
+    # Create adapters
+    local_stream = TrioStreamAdapter(local_send, local_receive, is_initiator=True)
+    remote_stream = TrioStreamAdapter(remote_send, remote_receive, is_initiator=False)
+
+    # Create transports without peerstore
+    local_transport = InsecureTransport(local_key_pair, peerstore=None)
+    remote_transport = InsecureTransport(remote_key_pair, peerstore=None)
+
+    # Run handshake
+    async def run_local_handshake(nursery_results):
+        with trio.move_on_after(5):
+            local_conn = await local_transport.secure_outbound(
+                local_stream, remote_peer_id
+            )
+            nursery_results["local"] = local_conn
+
+    async def run_remote_handshake(nursery_results):
+        with trio.move_on_after(5):
+            remote_conn = await remote_transport.secure_inbound(remote_stream)
+            nursery_results["remote"] = remote_conn
+
+    nursery_results = {}
+    async with trio.open_nursery() as nursery:
+        nursery.start_soon(run_local_handshake, nursery_results)
+        nursery.start_soon(run_remote_handshake, nursery_results)
+        await trio.sleep(0.1)  # Give tasks a chance to finish
+
+    local_conn = nursery_results.get("local")
+    remote_conn = nursery_results.get("remote")
+
+    # Verify that handshake still works without a peerstore
+    assert local_conn is not None, "Local handshake failed"
+    assert remote_conn is not None, "Remote handshake failed"
+
+
+@pytest.mark.trio
+async def test_peerstore_unchanged_when_handshake_fails():
+    """
+    Test that the peerstore remains unchanged if the handshake fails
+    due to a peer ID mismatch.
+    """
+    # Create key pairs for both sides
+    local_key_pair = create_new_key_pair()
+    remote_key_pair = create_new_key_pair()
+
+    # Create a third key pair to cause a mismatch
+    mismatch_key_pair = create_new_key_pair()
+
+    # Create peer IDs
+    remote_peer_id = ID.from_pubkey(remote_key_pair.public_key)
+    mismatch_peer_id = ID.from_pubkey(mismatch_key_pair.public_key)
+
+    # Create peerstore and add some initial data to verify it stays unchanged
+    peerstore = PeerStore()
+
+    # Store some initial data in peerstore to verify it remains unchanged
+    initial_key_pair = create_new_key_pair()
+    initial_peer_id = ID.from_pubkey(initial_key_pair.public_key)
+    peerstore.add_pubkey(initial_peer_id, initial_key_pair.public_key)
+
+    # Remember the initial state of the peerstore
+    initial_peer_ids = set(peerstore.peer_ids())
+
+    # Create memory streams for communication
+    local_send, remote_receive = memory_stream_pair()
+    remote_send, local_receive = memory_stream_pair()
+
+    # Create adapters
+    local_stream = TrioStreamAdapter(local_send, local_receive, is_initiator=True)
+    remote_stream = TrioStreamAdapter(remote_send, remote_receive, is_initiator=False)
+
+    # Create transports
+    local_transport = InsecureTransport(local_key_pair, peerstore=peerstore)
+    remote_transport = InsecureTransport(remote_key_pair, peerstore=None)
+
+    # Run handshake with mismatched peer_id
+    # (expecting remote_peer_id but sending mismatch_peer_id to cause a failure)
+    async def run_local_handshake(nursery_results):
+        with trio.move_on_after(5):
+            try:
+                # Pass mismatch_peer_id instead of remote_peer_id
+                # to cause a handshake failure
+                local_conn = await local_transport.secure_outbound(
+                    local_stream, mismatch_peer_id
+                )
+                nursery_results["local"] = local_conn
+            except HandshakeFailure:
+                nursery_results["local_error"] = True
+
+    async def run_remote_handshake(nursery_results):
+        with trio.move_on_after(5):
+            try:
+                remote_conn = await remote_transport.secure_inbound(remote_stream)
+                nursery_results["remote"] = remote_conn
+            except HandshakeFailure:
+                nursery_results["remote_error"] = True
+
+    nursery_results = {}
+    async with trio.open_nursery() as nursery:
+        nursery.start_soon(run_local_handshake, nursery_results)
+        nursery.start_soon(run_remote_handshake, nursery_results)
+        await trio.sleep(0.1)
+
+    # Verify that at least one side encountered an error
+    assert "local_error" in nursery_results or "remote_error" in nursery_results, (
+        "Expected handshake to fail due to peer ID mismatch"
+    )
+
+    # Verify that the peerstore remains unchanged
+    current_peer_ids = set(peerstore.peer_ids())
+    assert current_peer_ids == initial_peer_ids, (
+        "Peerstore should remain unchanged when handshake fails"
+    )
+
+    # Verify that neither the remote_peer_id nor mismatch_peer_id was added
+    assert remote_peer_id not in peerstore.peer_ids(), (
+        "Remote peer ID should not be added on handshake failure"
+    )
+    assert mismatch_peer_id not in peerstore.peer_ids(), (
+        "Mismatch peer ID should not be added on handshake failure"
+    )
+
+
+@pytest.mark.trio
+async def test_handshake_adds_pubkey_to_existing_peer():
+    """
+    Test that when a peer ID already exists in the peerstore but without
+    a public key, the handshake correctly adds the public key.
+
+    This tests the case where we might have a peer ID from another source
+    (like a routing table) but don't yet have its public key.
+    """
+    # Create key pairs for both sides
+    local_key_pair = create_new_key_pair()
+    remote_key_pair = create_new_key_pair()
+
+    # Create peer IDs
+    remote_peer_id = ID.from_pubkey(remote_key_pair.public_key)
+
+    # Create peerstore and add the peer ID without a public key
+    peerstore = PeerStore()
+
+    # Add the peer ID to the peerstore without its public key
+    # (adding an address for the peer, which creates the peer entry)
+    # This simulates having discovered a peer through DHT or other means
+    # without having its public key yet
+    peerstore.peer_data_map[remote_peer_id] = PeerData()
+
+    # Verify initial state - the peer ID should exist but without a public key
+    assert remote_peer_id in peerstore.peer_ids()
+    with pytest.raises(Exception):
+        peerstore.pubkey(remote_peer_id)
+
+    # Create memory streams for communication
+    local_send, remote_receive = memory_stream_pair()
+    remote_send, local_receive = memory_stream_pair()
+
+    # Create adapters
+    local_stream = TrioStreamAdapter(local_send, local_receive, is_initiator=True)
+    remote_stream = TrioStreamAdapter(remote_send, remote_receive, is_initiator=False)
+
+    # Create transports
+    local_transport = InsecureTransport(local_key_pair, peerstore=peerstore)
+    remote_transport = InsecureTransport(remote_key_pair, peerstore=None)
+
+    # Run handshake
+    async def run_local_handshake(nursery_results):
+        with trio.move_on_after(5):
+            local_conn = await local_transport.secure_outbound(
+                local_stream, remote_peer_id
+            )
+            nursery_results["local"] = local_conn
+
+    async def run_remote_handshake(nursery_results):
+        with trio.move_on_after(5):
+            remote_conn = await remote_transport.secure_inbound(remote_stream)
+            nursery_results["remote"] = remote_conn
+
+    nursery_results = {}
+    async with trio.open_nursery() as nursery:
+        nursery.start_soon(run_local_handshake, nursery_results)
+        nursery.start_soon(run_remote_handshake, nursery_results)
+        await trio.sleep(0.1)  # Give tasks a chance to finish
+
+    local_conn = nursery_results.get("local")
+    remote_conn = nursery_results.get("remote")
+
+    # Verify that the handshake succeeded
+    assert local_conn is not None, "Local handshake failed"
+    assert remote_conn is not None, "Remote handshake failed"
+
+    # Verify that the peer ID is still in the peerstore
+    assert remote_peer_id in peerstore.peer_ids()
+
+    # Verify that the public key was added
+    stored_pubkey = peerstore.pubkey(remote_peer_id)
+    assert stored_pubkey is not None
+    assert stored_pubkey.serialize() == remote_key_pair.public_key.serialize()

tests/core/stream_muxer/test_yamux.py

@@ -79,7 +79,7 @@ async def secure_conn_pair(key_pair, peer_id):
     client_rw = TrioStreamAdapter(client_send, client_receive, is_initiator=True)
     server_rw = TrioStreamAdapter(server_send, server_receive, is_initiator=False)
 
-    insecure_transport = InsecureTransport(key_pair)
+    insecure_transport = InsecureTransport(key_pair, peerstore=None)
 
     async def run_outbound(nursery_results):
         with trio.move_on_after(5):

tests/utils/factories.py

@@ -161,8 +161,8 @@ def noise_handshake_payload_factory() -> NoiseHandshakePayload:
     )
 
 
-def plaintext_transport_factory(key_pair: KeyPair) -> ISecureTransport:
-    return InsecureTransport(key_pair)
+def plaintext_transport_factory(key_pair: KeyPair, peerstore=None) -> ISecureTransport:
+    return InsecureTransport(key_pair, peerstore=peerstore)
 
 
 def secio_transport_factory(key_pair: KeyPair) -> ISecureTransport: