From 30b5811d3983a366d0821135d08e5d83b9e86e31 Mon Sep 17 00:00:00 2001 From: Winter-Soren Date: Thu, 29 May 2025 20:07:48 +0530 Subject: [PATCH] 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: