feat: store pubkey and peerid in peerstore

This commit is contained in:
Winter-Soren
2025-05-29 20:07:48 +05:30
parent 67ca1d7769
commit 30b5811d39
5 changed files with 35 additions and 7 deletions

View File

@ -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

View File

@ -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

View File

@ -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
)

View File

@ -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):

View File

@ -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: