mirror of
https://github.com/varun-r-mallya/py-libp2p.git
synced 2025-12-31 20:36:24 +00:00
feat: store pubkey and peerid in peerstore
This commit is contained in:
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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
|
||||
)
|
||||
|
||||
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -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:
|
||||
|
||||
Reference in New Issue
Block a user