Merge pull request #815 from lla-dane/kad-record

Signed-Peer-Record support in KAD-DHT message transfer mechanism.

Author: seetadev

libp2p/__init__.py

@@ -49,6 +49,7 @@
 )
 from libp2p.peer.peerstore import (
     PeerStore,
+    create_signed_peer_record,
 )
 from libp2p.security.insecure.transport import (
     PLAINTEXT_PROTOCOL_ID,
@@ -155,7 +156,6 @@ def get_default_muxer_options() -> TMuxerOptions:
     else:  # YAMUX is default
         return create_yamux_muxer_option()
 
-
 def new_swarm(
     key_pair: KeyPair | None = None,
     muxer_opt: TMuxerOptions | None = None,

libp2p/abc.py

@@ -970,6 +970,14 @@ def peers_with_addrs(self) -> list[ID]:
 
     # --------CERTIFIED-ADDR-BOOK----------
 
+    @abstractmethod
+    def get_local_record(self) -> Optional["Envelope"]:
+        """Get the local-peer-record wrapped in Envelope"""
+
+    @abstractmethod
+    def set_local_record(self, envelope: "Envelope") -> None:
+        """Set the local-peer-record wrapped in Envelope"""
+
     @abstractmethod
     def consume_peer_record(self, envelope: "Envelope", ttl: int) -> bool:
         """

libp2p/host/basic_host.py

@@ -43,6 +43,7 @@
 from libp2p.peer.peerinfo import (
     PeerInfo,
 )
+from libp2p.peer.peerstore import create_signed_peer_record
 from libp2p.protocol_muxer.exceptions import (
     MultiselectClientError,
     MultiselectError,
@@ -110,6 +111,14 @@ def __init__(
         if bootstrap:
             self.bootstrap = BootstrapDiscovery(network, bootstrap)
 
+        # Cache a signed-record if the local-node in the PeerStore
+        envelope = create_signed_peer_record(
+            self.get_id(),
+            self.get_addrs(),
+            self.get_private_key(),
+        )
+        self.get_peerstore().set_local_record(envelope)
+
     def get_id(self) -> ID:
         """
         :return: peer_id of host

libp2p/identity/identify/identify.py

@@ -15,8 +15,7 @@
 from libp2p.network.stream.exceptions import (
     StreamClosed,
 )
-from libp2p.peer.envelope import seal_record
-from libp2p.peer.peer_record import PeerRecord
+from libp2p.peer.peerstore import env_to_send_in_RPC
 from libp2p.utils import (
     decode_varint_with_size,
     get_agent_version,
@@ -66,9 +65,7 @@ def _mk_identify_protobuf(
     protocols = tuple(str(p) for p in host.get_mux().get_protocols() if p is not None)
 
     # Create a signed peer-record for the remote peer
-    record = PeerRecord(host.get_id(), host.get_addrs())
-    envelope = seal_record(record, host.get_private_key())
-    protobuf = envelope.marshal_envelope()
+    envelope_bytes, _ = env_to_send_in_RPC(host)
 
     observed_addr = observed_multiaddr.to_bytes() if observed_multiaddr else b""
     return Identify(
@@ -78,7 +75,7 @@ def _mk_identify_protobuf(
         listen_addrs=map(_multiaddr_to_bytes, laddrs),
         observed_addr=observed_addr,
         protocols=protocols,
-        signedPeerRecord=protobuf,
+        signedPeerRecord=envelope_bytes,
     )
 
 

libp2p/kad_dht/kad_dht.py

@@ -22,15 +22,18 @@
     IHost,
 )
 from libp2p.discovery.random_walk.rt_refresh_manager import RTRefreshManager
+from libp2p.kad_dht.utils import maybe_consume_signed_record
 from libp2p.network.stream.net_stream import (
     INetStream,
 )
+from libp2p.peer.envelope import Envelope
 from libp2p.peer.id import (
     ID,
 )
 from libp2p.peer.peerinfo import (
     PeerInfo,
 )
+from libp2p.peer.peerstore import env_to_send_in_RPC
 from libp2p.tools.async_service import (
     Service,
 )
@@ -234,6 +237,9 @@ async def handle_stream(self, stream: INetStream) -> None:
         await self.add_peer(peer_id)
         logger.debug(f"Added peer {peer_id} to routing table")
 
+        closer_peer_envelope: Envelope | None = None
+        provider_peer_envelope: Envelope | None = None
+
         try:
             # Read varint-prefixed length for the message
             length_prefix = b""
@@ -274,6 +280,14 @@ async def handle_stream(self, stream: INetStream) -> None:
                     )
                     logger.debug(f"Found {len(closest_peers)} peers close to target")
 
+                    # Consume the source signed_peer_record if sent
+                    if not maybe_consume_signed_record(message, self.host, peer_id):
+                        logger.error(
+                            "Received an invalid-signed-record, dropping the stream"
+                        )
+                        await stream.close()
+                        return
+
                     # Build response message with protobuf
                     response = Message()
                     response.type = Message.MessageType.FIND_NODE
@@ -298,6 +312,21 @@ async def handle_stream(self, stream: INetStream) -> None:
                         except Exception:
                             pass
 
+                        # Add the signed-peer-record for each peer in the peer-proto
+                        # if cached in the peerstore
+                        closer_peer_envelope = (
+                            self.host.get_peerstore().get_peer_record(peer)
+                        )
+
+                        if closer_peer_envelope is not None:
+                            peer_proto.signedRecord = (
+                                closer_peer_envelope.marshal_envelope()
+                            )
+
+                    # Create sender_signed_peer_record
+                    envelope_bytes, _ = env_to_send_in_RPC(self.host)
+                    response.senderRecord = envelope_bytes
+
                     # Serialize and send response
                     response_bytes = response.SerializeToString()
                     await stream.write(varint.encode(len(response_bytes)))
@@ -312,6 +341,14 @@ async def handle_stream(self, stream: INetStream) -> None:
                     key = message.key
                     logger.debug(f"Received ADD_PROVIDER for key {key.hex()}")
 
+                    # Consume the source signed-peer-record if sent
+                    if not maybe_consume_signed_record(message, self.host, peer_id):
+                        logger.error(
+                            "Received an invalid-signed-record, dropping the stream"
+                        )
+                        await stream.close()
+                        return
+
                     # Extract provider information
                     for provider_proto in message.providerPeers:
                         try:
@@ -338,6 +375,17 @@ async def handle_stream(self, stream: INetStream) -> None:
                             logger.debug(
                                 f"Added provider {provider_id} for key {key.hex()}"
                             )
+
+                            # Process the signed-records of provider if sent
+                            if not maybe_consume_signed_record(
+                                provider_proto, self.host
+                            ):
+                                logger.error(
+                                    "Received an invalid-signed-record,"
+                                    "dropping the stream"
+                                )
+                                await stream.close()
+                                return
                         except Exception as e:
                             logger.warning(f"Failed to process provider info: {e}")
 
@@ -346,6 +394,10 @@ async def handle_stream(self, stream: INetStream) -> None:
                     response.type = Message.MessageType.ADD_PROVIDER
                     response.key = key
 
+                    # Add sender's signed-peer-record
+                    envelope_bytes, _ = env_to_send_in_RPC(self.host)
+                    response.senderRecord = envelope_bytes
+
                     response_bytes = response.SerializeToString()
                     await stream.write(varint.encode(len(response_bytes)))
                     await stream.write(response_bytes)
@@ -357,6 +409,14 @@ async def handle_stream(self, stream: INetStream) -> None:
                     key = message.key
                     logger.debug(f"Received GET_PROVIDERS request for key {key.hex()}")
 
+                    # Consume the source signed_peer_record if sent
+                    if not maybe_consume_signed_record(message, self.host, peer_id):
+                        logger.error(
+                            "Received an invalid-signed-record, dropping the stream"
+                        )
+                        await stream.close()
+                        return
+
                     # Find providers for the key
                     providers = self.provider_store.get_providers(key)
                     logger.debug(
@@ -368,12 +428,28 @@ async def handle_stream(self, stream: INetStream) -> None:
                     response.type = Message.MessageType.GET_PROVIDERS
                     response.key = key
 
+                    # Create sender_signed_peer_record for the response
+                    envelope_bytes, _ = env_to_send_in_RPC(self.host)
+                    response.senderRecord = envelope_bytes
+
                     # Add provider information to response
                     for provider_info in providers:
                         provider_proto = response.providerPeers.add()
                         provider_proto.id = provider_info.peer_id.to_bytes()
                         provider_proto.connection = Message.ConnectionType.CAN_CONNECT
 
+                        # Add provider signed-records if cached
+                        provider_peer_envelope = (
+                            self.host.get_peerstore().get_peer_record(
+                                provider_info.peer_id
+                            )
+                        )
+
+                        if provider_peer_envelope is not None:
+                            provider_proto.signedRecord = (
+                                provider_peer_envelope.marshal_envelope()
+                            )
+
                         # Add addresses if available
                         for addr in provider_info.addrs:
                             provider_proto.addrs.append(addr.to_bytes())
@@ -397,6 +473,16 @@ async def handle_stream(self, stream: INetStream) -> None:
                             peer_proto.id = peer.to_bytes()
                             peer_proto.connection = Message.ConnectionType.CAN_CONNECT
 
+                            # Add the signed-records of closest_peers if cached
+                            closer_peer_envelope = (
+                                self.host.get_peerstore().get_peer_record(peer)
+                            )
+
+                            if closer_peer_envelope is not None:
+                                peer_proto.signedRecord = (
+                                    closer_peer_envelope.marshal_envelope()
+                                )
+
                             # Add addresses if available
                             try:
                                 addrs = self.host.get_peerstore().addrs(peer)
@@ -417,6 +503,14 @@ async def handle_stream(self, stream: INetStream) -> None:
                     key = message.key
                     logger.debug(f"Received GET_VALUE request for key {key.hex()}")
 
+                    # Consume the sender_signed_peer_record
+                    if not maybe_consume_signed_record(message, self.host, peer_id):
+                        logger.error(
+                            "Received an invalid-signed-record, dropping the stream"
+                        )
+                        await stream.close()
+                        return
+
                     value = self.value_store.get(key)
                     if value:
                         logger.debug(f"Found value for key {key.hex()}")
@@ -431,6 +525,10 @@ async def handle_stream(self, stream: INetStream) -> None:
                         response.record.value = value
                         response.record.timeReceived = str(time.time())
 
+                        # Create sender_signed_peer_record
+                        envelope_bytes, _ = env_to_send_in_RPC(self.host)
+                        response.senderRecord = envelope_bytes
+
                         # Serialize and send response
                         response_bytes = response.SerializeToString()
                         await stream.write(varint.encode(len(response_bytes)))
@@ -444,6 +542,10 @@ async def handle_stream(self, stream: INetStream) -> None:
                         response.type = Message.MessageType.GET_VALUE
                         response.key = key
 
+                        # Create sender_signed_peer_record for the response
+                        envelope_bytes, _ = env_to_send_in_RPC(self.host)
+                        response.senderRecord = envelope_bytes
+
                         # Add closest peers to key
                         closest_peers = self.routing_table.find_local_closest_peers(
                             key, 20
@@ -462,6 +564,16 @@ async def handle_stream(self, stream: INetStream) -> None:
                             peer_proto.id = peer.to_bytes()
                             peer_proto.connection = Message.ConnectionType.CAN_CONNECT
 
+                            # Add signed-records of closer-peers if cached
+                            closer_peer_envelope = (
+                                self.host.get_peerstore().get_peer_record(peer)
+                            )
+
+                            if closer_peer_envelope is not None:
+                                peer_proto.signedRecord = (
+                                    closer_peer_envelope.marshal_envelope()
+                                )
+
                             # Add addresses if available
                             try:
                                 addrs = self.host.get_peerstore().addrs(peer)
@@ -484,6 +596,15 @@ async def handle_stream(self, stream: INetStream) -> None:
                     key = message.record.key
                     value = message.record.value
                     success = False
+
+                    # Consume the source signed_peer_record if sent
+                    if not maybe_consume_signed_record(message, self.host, peer_id):
+                        logger.error(
+                            "Received an invalid-signed-record, dropping the stream"
+                        )
+                        await stream.close()
+                        return
+
                     try:
                         if not (key and value):
                             raise ValueError(
@@ -504,6 +625,12 @@ async def handle_stream(self, stream: INetStream) -> None:
                         response.type = Message.MessageType.PUT_VALUE
                         if success:
                             response.key = key
+
+                        # Create sender_signed_peer_record for the response
+                        envelope_bytes, _ = env_to_send_in_RPC(self.host)
+                        response.senderRecord = envelope_bytes
+
+                        # Serialize and send response
                         response_bytes = response.SerializeToString()
                         await stream.write(varint.encode(len(response_bytes)))
                         await stream.write(response_bytes)

libp2p/kad_dht/pb/kademlia.proto

@@ -27,6 +27,7 @@ message Message {
         bytes id = 1;
         repeated bytes addrs = 2;
         ConnectionType connection = 3;
+        optional bytes signedRecord = 4;    // Envelope(PeerRecord) encoded
     }
 
     MessageType type = 1;
@@ -35,4 +36,6 @@ message Message {
     Record record = 3;
     repeated Peer closerPeers = 8;
     repeated Peer providerPeers = 9;
+
+    optional bytes senderRecord = 11;   // Envelope(PeerRecord) encoded
 }