From 30b5811d3983a366d0821135d08e5d83b9e86e31 Mon Sep 17 00:00:00 2001 From: Winter-Soren Date: Thu, 29 May 2025 20:07:48 +0530 Subject: [PATCH 01/11] 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 02/11] 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 03/11] 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 dfc0bb4ec8a7e208ac0c757f97f0b07d5cc95b64 Mon Sep 17 00:00:00 2001 From: Luca Vivona Date: Thu, 19 Jun 2025 21:24:39 -0400 Subject: [PATCH 04/11] chore(kad_dht): centralize shared values in common.py --- libp2p/kad_dht/common.py | 10 ++++++++++ libp2p/kad_dht/kad_dht.py | 12 ++++++------ libp2p/kad_dht/peer_routing.py | 8 +++++--- libp2p/kad_dht/provider_store.py | 9 ++++++--- libp2p/kad_dht/value_store.py | 10 +++++----- 5 files changed, 32 insertions(+), 17 deletions(-) create mode 100644 libp2p/kad_dht/common.py diff --git a/libp2p/kad_dht/common.py b/libp2p/kad_dht/common.py new file mode 100644 index 00000000..ccabb002 --- /dev/null +++ b/libp2p/kad_dht/common.py @@ -0,0 +1,10 @@ +from libp2p.custom_types import ( + TProtocol, +) + +# Constants for the Kademlia algorithm +ALPHA = 3 # Concurrency parameter +PROTOCOL_ID = TProtocol("/ipfs/kad/1.0.0") +QUERY_TIMEOUT = 10 + +TTL = DEFAULT_TTL = 24 * 60 * 60 # 24 hours in seconds \ No newline at end of file diff --git a/libp2p/kad_dht/kad_dht.py b/libp2p/kad_dht/kad_dht.py index 7daad4cb..8a917bda 100644 --- a/libp2p/kad_dht/kad_dht.py +++ b/libp2p/kad_dht/kad_dht.py @@ -49,16 +49,16 @@ from .routing_table import ( from .value_store import ( ValueStore, ) +from .common import ( + PROTOCOL_ID, + ALPHA, + QUERY_TIMEOUT +) logger = logging.getLogger("kademlia-example.kad_dht") # logger = logging.getLogger("libp2p.kademlia") # Default parameters -PROTOCOL_ID = TProtocol("/ipfs/kad/1.0.0") -ROUTING_TABLE_REFRESH_INTERVAL = 1 * 60 # 1 min in seconds for testing -TTL = 24 * 60 * 60 # 24 hours in seconds -ALPHA = 3 -QUERY_TIMEOUT = 10 # seconds - +ROUTING_TABLE_REFRESH_INTERVAL = 60 # 1 min in seconds for testing class DHTMode(Enum): """DHT operation modes.""" diff --git a/libp2p/kad_dht/peer_routing.py b/libp2p/kad_dht/peer_routing.py index f3689e11..2e2107a0 100644 --- a/libp2p/kad_dht/peer_routing.py +++ b/libp2p/kad_dht/peer_routing.py @@ -25,12 +25,17 @@ from libp2p.peer.peerinfo import ( PeerInfo, ) + from .pb.kademlia_pb2 import ( Message, ) from .routing_table import ( RoutingTable, ) +from .common import ( + PROTOCOL_ID, + ALPHA +) from .utils import ( sort_peer_ids_by_distance, ) @@ -38,10 +43,7 @@ from .utils import ( # logger = logging.getLogger("libp2p.kademlia.peer_routing") logger = logging.getLogger("kademlia-example.peer_routing") -# Constants for the Kademlia algorithm -ALPHA = 3 # Concurrency parameter MAX_PEER_LOOKUP_ROUNDS = 20 # Maximum number of rounds in peer lookup -PROTOCOL_ID = TProtocol("/ipfs/kad/1.0.0") class PeerRouting(IPeerRouting): diff --git a/libp2p/kad_dht/provider_store.py b/libp2p/kad_dht/provider_store.py index 00ac6010..4938be2e 100644 --- a/libp2p/kad_dht/provider_store.py +++ b/libp2p/kad_dht/provider_store.py @@ -33,6 +33,12 @@ from .pb.kademlia_pb2 import ( Message, ) +from .common import ( + PROTOCOL_ID, + ALPHA, + QUERY_TIMEOUT +) + # logger = logging.getLogger("libp2p.kademlia.provider_store") logger = logging.getLogger("kademlia-example.provider_store") @@ -40,9 +46,6 @@ logger = logging.getLogger("kademlia-example.provider_store") PROVIDER_RECORD_REPUBLISH_INTERVAL = 22 * 60 * 60 # 22 hours in seconds PROVIDER_RECORD_EXPIRATION_INTERVAL = 48 * 60 * 60 # 48 hours in seconds PROVIDER_ADDRESS_TTL = 30 * 60 # 30 minutes in seconds -PROTOCOL_ID = TProtocol("/ipfs/kad/1.0.0") -ALPHA = 3 # Number of parallel queries/advertisements -QUERY_TIMEOUT = 10 # Timeout for each query in seconds class ProviderRecord: diff --git a/libp2p/kad_dht/value_store.py b/libp2p/kad_dht/value_store.py index a2e54776..c133fd00 100644 --- a/libp2p/kad_dht/value_store.py +++ b/libp2p/kad_dht/value_store.py @@ -23,14 +23,14 @@ from .pb.kademlia_pb2 import ( Message, ) +from .common import ( + PROTOCOL_ID, + DEFAULT_TTL +) + # logger = logging.getLogger("libp2p.kademlia.value_store") logger = logging.getLogger("kademlia-example.value_store") -# Default time to live for values in seconds (24 hours) -DEFAULT_TTL = 24 * 60 * 60 -PROTOCOL_ID = TProtocol("/ipfs/kad/1.0.0") - - class ValueStore: """ Store for key-value pairs in a Kademlia DHT. From 79ac01308c492169046bc04bcd90e1f25bf4037f Mon Sep 17 00:00:00 2001 From: Luca Vivona Date: Thu, 19 Jun 2025 21:38:02 -0400 Subject: [PATCH 05/11] remove: unused custom_types TProtocol import --- libp2p/kad_dht/peer_routing.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/libp2p/kad_dht/peer_routing.py b/libp2p/kad_dht/peer_routing.py index 2e2107a0..fa938707 100644 --- a/libp2p/kad_dht/peer_routing.py +++ b/libp2p/kad_dht/peer_routing.py @@ -15,9 +15,6 @@ from libp2p.abc import ( INetStream, IPeerRouting, ) -from libp2p.custom_types import ( - TProtocol, -) from libp2p.peer.id import ( ID, ) From d03ca45bd8db31580348455b022f8fa76edb044e Mon Sep 17 00:00:00 2001 From: Luca Vivona Date: Fri, 20 Jun 2025 11:57:50 -0400 Subject: [PATCH 06/11] style: fix flake8 linting errors --- libp2p/kad_dht/common.py | 2 +- libp2p/kad_dht/kad_dht.py | 4 +--- libp2p/kad_dht/value_store.py | 1 + 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/libp2p/kad_dht/common.py b/libp2p/kad_dht/common.py index ccabb002..8b82ea97 100644 --- a/libp2p/kad_dht/common.py +++ b/libp2p/kad_dht/common.py @@ -7,4 +7,4 @@ ALPHA = 3 # Concurrency parameter PROTOCOL_ID = TProtocol("/ipfs/kad/1.0.0") QUERY_TIMEOUT = 10 -TTL = DEFAULT_TTL = 24 * 60 * 60 # 24 hours in seconds \ No newline at end of file +TTL = DEFAULT_TTL = 24 * 60 * 60 # 24 hours in seconds diff --git a/libp2p/kad_dht/kad_dht.py b/libp2p/kad_dht/kad_dht.py index 8a917bda..f82c0063 100644 --- a/libp2p/kad_dht/kad_dht.py +++ b/libp2p/kad_dht/kad_dht.py @@ -18,9 +18,6 @@ import varint from libp2p.abc import ( IHost, ) -from libp2p.custom_types import ( - TProtocol, -) from libp2p.network.stream.net_stream import ( INetStream, ) @@ -60,6 +57,7 @@ logger = logging.getLogger("kademlia-example.kad_dht") # Default parameters ROUTING_TABLE_REFRESH_INTERVAL = 60 # 1 min in seconds for testing + class DHTMode(Enum): """DHT operation modes.""" diff --git a/libp2p/kad_dht/value_store.py b/libp2p/kad_dht/value_store.py index c133fd00..488735e6 100644 --- a/libp2p/kad_dht/value_store.py +++ b/libp2p/kad_dht/value_store.py @@ -31,6 +31,7 @@ from .common import ( # logger = logging.getLogger("libp2p.kademlia.value_store") logger = logging.getLogger("kademlia-example.value_store") + class ValueStore: """ Store for key-value pairs in a Kademlia DHT. From 811c217ee697cc032764e9ac722890ab6fd77e1d Mon Sep 17 00:00:00 2001 From: Luca Vivona Date: Fri, 20 Jun 2025 16:01:11 -0400 Subject: [PATCH 07/11] style: isort fix ording of imports --- libp2p/kad_dht/common.py | 4 ++++ libp2p/kad_dht/kad_dht.py | 10 +++++----- libp2p/kad_dht/peer_routing.py | 9 ++++----- libp2p/kad_dht/provider_store.py | 11 +++++------ libp2p/kad_dht/routing_table.py | 20 ++++++++------------ libp2p/kad_dht/value_store.py | 9 ++++----- 6 files changed, 30 insertions(+), 33 deletions(-) diff --git a/libp2p/kad_dht/common.py b/libp2p/kad_dht/common.py index 8b82ea97..d7fb6c76 100644 --- a/libp2p/kad_dht/common.py +++ b/libp2p/kad_dht/common.py @@ -1,3 +1,7 @@ +""" +Common Objects for Kademlia DHT implementation. +""" + from libp2p.custom_types import ( TProtocol, ) diff --git a/libp2p/kad_dht/kad_dht.py b/libp2p/kad_dht/kad_dht.py index f82c0063..b47d59e7 100644 --- a/libp2p/kad_dht/kad_dht.py +++ b/libp2p/kad_dht/kad_dht.py @@ -31,6 +31,11 @@ from libp2p.tools.async_service import ( Service, ) +from .common import ( + ALPHA, + PROTOCOL_ID, + QUERY_TIMEOUT, +) from .pb.kademlia_pb2 import ( Message, ) @@ -46,11 +51,6 @@ from .routing_table import ( from .value_store import ( ValueStore, ) -from .common import ( - PROTOCOL_ID, - ALPHA, - QUERY_TIMEOUT -) logger = logging.getLogger("kademlia-example.kad_dht") # logger = logging.getLogger("libp2p.kademlia") diff --git a/libp2p/kad_dht/peer_routing.py b/libp2p/kad_dht/peer_routing.py index fa938707..8e40ccf5 100644 --- a/libp2p/kad_dht/peer_routing.py +++ b/libp2p/kad_dht/peer_routing.py @@ -22,17 +22,16 @@ from libp2p.peer.peerinfo import ( PeerInfo, ) - +from .common import ( + ALPHA, + PROTOCOL_ID, +) from .pb.kademlia_pb2 import ( Message, ) from .routing_table import ( RoutingTable, ) -from .common import ( - PROTOCOL_ID, - ALPHA -) from .utils import ( sort_peer_ids_by_distance, ) diff --git a/libp2p/kad_dht/provider_store.py b/libp2p/kad_dht/provider_store.py index 4938be2e..5c34f0c7 100644 --- a/libp2p/kad_dht/provider_store.py +++ b/libp2p/kad_dht/provider_store.py @@ -29,16 +29,15 @@ from libp2p.peer.peerinfo import ( PeerInfo, ) +from .common import ( + ALPHA, + PROTOCOL_ID, + QUERY_TIMEOUT, +) from .pb.kademlia_pb2 import ( Message, ) -from .common import ( - PROTOCOL_ID, - ALPHA, - QUERY_TIMEOUT -) - # logger = logging.getLogger("libp2p.kademlia.provider_store") logger = logging.getLogger("kademlia-example.provider_store") diff --git a/libp2p/kad_dht/routing_table.py b/libp2p/kad_dht/routing_table.py index 4377c591..af4c7439 100644 --- a/libp2p/kad_dht/routing_table.py +++ b/libp2p/kad_dht/routing_table.py @@ -2,21 +2,16 @@ Kademlia DHT routing table implementation. """ -from collections import ( - OrderedDict, -) +from collections import OrderedDict import logging import time import trio -from libp2p.abc import ( - IHost, +from libp2p.abc import IHost +from libp2p.kad_dht.utils import ( + xor_distance, ) -from libp2p.custom_types import ( - TProtocol, -) -from libp2p.kad_dht.utils import xor_distance from libp2p.peer.id import ( ID, ) @@ -24,9 +19,10 @@ from libp2p.peer.peerinfo import ( PeerInfo, ) -from .pb.kademlia_pb2 import ( - Message, +from .common import ( + PROTOCOL_ID, ) +from .pb.kademlia_pb2 import Message # logger = logging.getLogger("libp2p.kademlia.routing_table") logger = logging.getLogger("kademlia-example.routing_table") @@ -243,7 +239,7 @@ class KBucket: raise ValueError(f"Peer {peer_id} not in bucket") # Default protocol ID for Kademlia DHT - protocol_id = TProtocol("/ipfs/kad/1.0.0") + protocol_id = PROTOCOL_ID try: # Open a stream to the peer with the DHT protocol diff --git a/libp2p/kad_dht/value_store.py b/libp2p/kad_dht/value_store.py index 488735e6..b79425fd 100644 --- a/libp2p/kad_dht/value_store.py +++ b/libp2p/kad_dht/value_store.py @@ -19,15 +19,14 @@ from libp2p.peer.id import ( ID, ) +from .common import ( + DEFAULT_TTL, + PROTOCOL_ID, +) from .pb.kademlia_pb2 import ( Message, ) -from .common import ( - PROTOCOL_ID, - DEFAULT_TTL -) - # logger = logging.getLogger("libp2p.kademlia.value_store") logger = logging.getLogger("kademlia-example.value_store") From ebdde7b5aa99bc9051756bd78dd7fccd432b0d86 Mon Sep 17 00:00:00 2001 From: Luca Vivona Date: Sat, 21 Jun 2025 15:08:11 -0400 Subject: [PATCH 08/11] style: enforce multiline import style for consistency --- libp2p/kad_dht/common.py | 2 +- libp2p/kad_dht/kad_dht.py | 4 +++- libp2p/kad_dht/routing_table.py | 13 +++++++------ 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/libp2p/kad_dht/common.py b/libp2p/kad_dht/common.py index d7fb6c76..392bffd0 100644 --- a/libp2p/kad_dht/common.py +++ b/libp2p/kad_dht/common.py @@ -1,5 +1,5 @@ """ -Common Objects for Kademlia DHT implementation. +Shared constants and protocol parameters for the Kademlia DHT. """ from libp2p.custom_types import ( diff --git a/libp2p/kad_dht/kad_dht.py b/libp2p/kad_dht/kad_dht.py index b47d59e7..dcf323ba 100644 --- a/libp2p/kad_dht/kad_dht.py +++ b/libp2p/kad_dht/kad_dht.py @@ -5,7 +5,9 @@ This module provides a complete Distributed Hash Table (DHT) implementation based on the Kademlia algorithm and protocol. """ -from enum import Enum +from enum import ( + Enum, +) import logging import time diff --git a/libp2p/kad_dht/routing_table.py b/libp2p/kad_dht/routing_table.py index af4c7439..9ca59031 100644 --- a/libp2p/kad_dht/routing_table.py +++ b/libp2p/kad_dht/routing_table.py @@ -2,13 +2,17 @@ Kademlia DHT routing table implementation. """ -from collections import OrderedDict +from collections import ( + OrderedDict, +) import logging import time import trio -from libp2p.abc import IHost +from libp2p.abc import ( + IHost, +) from libp2p.kad_dht.utils import ( xor_distance, ) @@ -238,12 +242,9 @@ class KBucket: if not peer_info: raise ValueError(f"Peer {peer_id} not in bucket") - # Default protocol ID for Kademlia DHT - protocol_id = PROTOCOL_ID - try: # Open a stream to the peer with the DHT protocol - stream = await self.host.new_stream(peer_id, [protocol_id]) + stream = await self.host.new_stream(peer_id, [PROTOCOL_ID]) try: # Create ping protobuf message From feb8db6655bbbd36e876ae75620d2d9bb6c6b554 Mon Sep 17 00:00:00 2001 From: Luca Vivona Date: Sun, 22 Jun 2025 00:15:44 -0400 Subject: [PATCH 09/11] style: enforce multiline import style --- libp2p/kad_dht/routing_table.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/libp2p/kad_dht/routing_table.py b/libp2p/kad_dht/routing_table.py index 9ca59031..15b6721e 100644 --- a/libp2p/kad_dht/routing_table.py +++ b/libp2p/kad_dht/routing_table.py @@ -26,7 +26,9 @@ from libp2p.peer.peerinfo import ( from .common import ( PROTOCOL_ID, ) -from .pb.kademlia_pb2 import Message +from .pb.kademlia_pb2 import ( + Message, +) # logger = logging.getLogger("libp2p.kademlia.routing_table") logger = logging.getLogger("kademlia-example.routing_table") From 3a4338e1df28842ab3bc7fd113b6acf9f0cfadef Mon Sep 17 00:00:00 2001 From: Luca Vivona Date: Sun, 22 Jun 2025 00:25:48 -0400 Subject: [PATCH 10/11] chore: eliminate self.protocol_id attribute \w in PeerRouting --- libp2p/kad_dht/peer_routing.py | 3 +-- tests/core/kad_dht/test_unit_peer_routing.py | 1 - 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/libp2p/kad_dht/peer_routing.py b/libp2p/kad_dht/peer_routing.py index 8e40ccf5..4bcdb647 100644 --- a/libp2p/kad_dht/peer_routing.py +++ b/libp2p/kad_dht/peer_routing.py @@ -60,7 +60,6 @@ class PeerRouting(IPeerRouting): """ self.host = host self.routing_table = routing_table - self.protocol_id = PROTOCOL_ID async def find_peer(self, peer_id: ID) -> PeerInfo | None: """ @@ -245,7 +244,7 @@ class PeerRouting(IPeerRouting): # Open a stream to the peer using the Kademlia protocol logger.debug(f"Opening stream to {peer} for closest peers query") try: - stream = await self.host.new_stream(peer, [self.protocol_id]) + stream = await self.host.new_stream(peer, [PROTOCOL_ID]) logger.debug(f"Stream opened to {peer}") except Exception as e: logger.warning(f"Failed to open stream to {peer}: {e}") diff --git a/tests/core/kad_dht/test_unit_peer_routing.py b/tests/core/kad_dht/test_unit_peer_routing.py index 72320b73..ffe20655 100644 --- a/tests/core/kad_dht/test_unit_peer_routing.py +++ b/tests/core/kad_dht/test_unit_peer_routing.py @@ -89,7 +89,6 @@ class TestPeerRouting: assert peer_routing.host == mock_host assert peer_routing.routing_table == mock_routing_table - assert peer_routing.protocol_id == PROTOCOL_ID @pytest.mark.trio async def test_find_peer_local_host(self, peer_routing, mock_host): From fd818d91026edc01ed70c26073bedea5a73b0017 Mon Sep 17 00:00:00 2001 From: Winter-Soren Date: Sun, 22 Jun 2025 16:01:02 +0530 Subject: [PATCH 11/11] 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()