From 30b5811d3983a366d0821135d08e5d83b9e86e31 Mon Sep 17 00:00:00 2001 From: Winter-Soren Date: Thu, 29 May 2025 20:07:48 +0530 Subject: [PATCH 1/4] feat: store pubkey and peerid in peerstore --- .../example_encryption_insecure.py | 3 ++ libp2p/__init__.py | 4 ++- libp2p/security/insecure/transport.py | 29 +++++++++++++++++-- tests/core/stream_muxer/test_yamux.py | 2 +- tests/utils/factories.py | 4 +-- 5 files changed, 35 insertions(+), 7 deletions(-) diff --git a/examples/doc-examples/example_encryption_insecure.py b/examples/doc-examples/example_encryption_insecure.py index acd947e7..c1536808 100644 --- a/examples/doc-examples/example_encryption_insecure.py +++ b/examples/doc-examples/example_encryption_insecure.py @@ -27,6 +27,9 @@ async def main(): # 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 diff --git a/libp2p/__init__.py b/libp2p/__init__.py index c05d05e5..8a84ab05 100644 --- a/libp2p/__init__.py +++ b/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 diff --git a/libp2p/security/insecure/transport.py b/libp2p/security/insecure/transport.py index 4666cc78..8f7f9a70 100644 --- a/libp2p/security/insecure/transport.py +++ b/libp2p/security/insecure/transport.py @@ -1,8 +1,10 @@ from typing import ( + Callable, Optional, ) from libp2p.abc import ( + IPeerStore, IRawConnection, ISecureConn, ) @@ -10,6 +12,7 @@ from libp2p.crypto.exceptions import ( MissingDeserializerError, ) from libp2p.crypto.keys import ( + KeyPair, PrivateKey, PublicKey, ) @@ -34,6 +37,9 @@ 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, ) @@ -106,6 +112,7 @@ async def run_handshake( conn: IRawConnection, is_initiator: bool, remote_peer_id: ID, + peerstore: Optional[IPeerStore] = None, ) -> ISecureConn: """Raise `HandshakeFailure` when handshake failed.""" msg = make_exchange_message(local_private_key.get_public_key()) @@ -159,7 +166,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 @@ -170,6 +184,15 @@ class InsecureTransport(BaseSecureTransport): transport does not add any additional security. """ + def __init__( + self, + local_key_pair: KeyPair, + secure_bytes_provider: Optional[Callable[[int], bytes]] = None, + peerstore: Optional[IPeerStore] = None, + ) -> None: + 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 @@ -179,7 +202,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, 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: @@ -190,7 +213,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 ) diff --git a/tests/core/stream_muxer/test_yamux.py b/tests/core/stream_muxer/test_yamux.py index fa25af9f..4bfa0199 100644 --- a/tests/core/stream_muxer/test_yamux.py +++ b/tests/core/stream_muxer/test_yamux.py @@ -71,7 +71,7 @@ async def secure_conn_pair(key_pair, peer_id): client_rw = TrioStreamAdapter(client_send, client_receive) server_rw = TrioStreamAdapter(server_send, server_receive) - insecure_transport = InsecureTransport(key_pair) + insecure_transport = InsecureTransport(key_pair, peerstore=None) async def run_outbound(nursery_results): with trio.move_on_after(5): diff --git a/tests/utils/factories.py b/tests/utils/factories.py index 1fe32344..f661ed6e 100644 --- a/tests/utils/factories.py +++ b/tests/utils/factories.py @@ -159,8 +159,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: From 2248108b54ec3efec4d9524505f3aa8f3f089a66 Mon Sep 17 00:00:00 2001 From: Winter-Soren Date: Wed, 11 Jun 2025 23:51:57 +0530 Subject: [PATCH 2/4] fixed types and improved code --- libp2p/security/insecure/transport.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/libp2p/security/insecure/transport.py b/libp2p/security/insecure/transport.py index bd01004b..cf92e8d9 100644 --- a/libp2p/security/insecure/transport.py +++ b/libp2p/security/insecure/transport.py @@ -1,7 +1,4 @@ -from typing import ( - Callable, - Optional, -) +from collections.abc import Callable from libp2p.abc import ( IPeerStore, @@ -45,6 +42,7 @@ from libp2p.security.base_session import ( ) from libp2p.security.base_transport import ( BaseSecureTransport, + default_secure_bytes_provider, ) from libp2p.security.exceptions import ( HandshakeFailure, @@ -111,8 +109,8 @@ async def run_handshake( local_private_key: PrivateKey, conn: IRawConnection, is_initiator: bool, - remote_peer_id: ID, - peerstore: Optional[IPeerStore] = None, + 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()) @@ -196,9 +194,12 @@ class InsecureTransport(BaseSecureTransport): def __init__( self, local_key_pair: KeyPair, - secure_bytes_provider: Optional[Callable[[int], bytes]] = None, - peerstore: Optional[IPeerStore] = None, + 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 @@ -210,6 +211,7 @@ 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.peerstore ) From 8c16b316ac6a74fa5c2eaefd54fa7fc682e3855e Mon Sep 17 00:00:00 2001 From: Winter-Soren Date: Wed, 18 Jun 2025 23:32:48 +0530 Subject: [PATCH 3/4] added newsfragement and tests that would fail without these changes but pass with them --- newsfragments/631.feature.rst | 3 + .../test_insecure_peerstore_integration.py | 147 ++++++++++++++++++ 2 files changed, 150 insertions(+) create mode 100644 newsfragments/631.feature.rst create mode 100644 tests/core/security/test_insecure_peerstore_integration.py diff --git a/newsfragments/631.feature.rst b/newsfragments/631.feature.rst new file mode 100644 index 00000000..f8844ebd --- /dev/null +++ b/newsfragments/631.feature.rst @@ -0,0 +1,3 @@ +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. diff --git a/tests/core/security/test_insecure_peerstore_integration.py b/tests/core/security/test_insecure_peerstore_integration.py new file mode 100644 index 00000000..302c9470 --- /dev/null +++ b/tests/core/security/test_insecure_peerstore_integration.py @@ -0,0 +1,147 @@ +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.peerstore import PeerStore +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" From fd818d91026edc01ed70c26073bedea5a73b0017 Mon Sep 17 00:00:00 2001 From: Winter-Soren Date: Sun, 22 Jun 2025 16:01:02 +0530 Subject: [PATCH 4/4] test: added tests to ensure handshake adds pubkey to existing peer ID without one; peerstore unchanged on ID mismatch --- newsfragments/631.feature.rst | 4 + .../test_insecure_peerstore_integration.py | 167 ++++++++++++++++++ 2 files changed, 171 insertions(+) diff --git a/newsfragments/631.feature.rst b/newsfragments/631.feature.rst index f8844ebd..1d88c779 100644 --- a/newsfragments/631.feature.rst +++ b/newsfragments/631.feature.rst @@ -1,3 +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 diff --git a/tests/core/security/test_insecure_peerstore_integration.py b/tests/core/security/test_insecure_peerstore_integration.py index 302c9470..9634fc69 100644 --- a/tests/core/security/test_insecure_peerstore_integration.py +++ b/tests/core/security/test_insecure_peerstore_integration.py @@ -5,7 +5,9 @@ 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 @@ -145,3 +147,168 @@ async def test_insecure_transport_without_peerstore(): # 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()