Merge pull request #1060 from sumanjeet0012/fix/kademlia_key_decoding

Migrate DHT API to accept string keys.

Author: seetadev

examples/kademlia/kademlia.py

@@ -13,7 +13,6 @@
 import secrets
 import sys
 
-import base58
 from multiaddr import (
     Multiaddr,
 )
@@ -32,9 +31,6 @@
     DHTMode,
     KadDHT,
 )
-from libp2p.kad_dht.utils import (
-    create_key_from_binary,
-)
 from libp2p.tools.async_service import (
     background_trio_service,
 )
@@ -191,57 +187,46 @@ async def run_node(
             # Start the DHT service
             async with background_trio_service(dht):
                 logger.info(f"DHT service started in {dht_mode.value} mode")
-                val_key = create_key_from_binary(b"py-libp2p kademlia example value")
-                content = b"Hello from python node "
-                content_key = create_key_from_binary(content)
+
+                # Example 1: Simple Key-Value Storage
+                # Just use a string key directly - DHT API accepts strings!
+                key = "my-example-key"
+                value = b"Hello from py-libp2p!"
+
+                # Example 2: Content Provider Advertisement
+                content_id = "my-content-identifier"
 
                 if dht_mode == DHTMode.SERVER:
-                    # Store a value in the DHT
-                    msg = "Hello message from Sumanjeet"
-                    val_data = msg.encode()
-                    await dht.put_value(val_key, val_data)
-                    logger.info(
-                        f"Stored value '{val_data.decode()}'"
-                        f"with key: {base58.b58encode(val_key).decode()}"
-                    )
+                    # Store key-value pair in the DHT
+                    await dht.put_value(key, value)
+                    logger.info(f"Stored value: {value.decode()} with key: {key}")
 
-                    # Advertise as content server
-                    success = await dht.provider_store.provide(content_key)
+                    # Advertise as a provider for content
+                    success = await dht.provide(content_id)
                     if success:
-                        logger.info(
-                            "Successfully advertised as server"
-                            f"for content: {content_key.hex()}"
-                        )
+                        logger.info(f"Advertised as provider for content: {content_id}")
                     else:
-                        logger.warning("Failed to advertise as content server")
+                        logger.warning("Failed to advertise as provider")
 
                 else:
-                    # retrieve the value
-                    logger.info(
-                        "Looking up key: %s", base58.b58encode(val_key).decode()
-                    )
-                    val_data = await dht.get_value(val_key)
-                    if val_data:
-                        try:
-                            logger.info(f"Retrieved value: {val_data.decode()}")
-                        except UnicodeDecodeError:
-                            logger.info(f"Retrieved value (bytes): {val_data!r}")
+                    # Retrieve value from DHT using the same key
+                    logger.info(f"Looking up key: {key}")
+                    retrieved_value = await dht.get_value(key)
+                    if retrieved_value:
+                        logger.info(f"Retrieved value: {retrieved_value.decode()}")
                     else:
                         logger.warning("Failed to retrieve value")
 
-                    # Also check if we can find servers for our own content
-                    logger.info("Looking for servers of content: %s", content_key.hex())
-                    providers = await dht.provider_store.find_providers(content_key)
+                    # Find providers for content
+                    logger.info(f"Looking for providers of content: {content_id}")
+                    providers = await dht.find_providers(content_id)
                     if providers:
                         logger.info(
-                            "Found %d servers for content: %s",
-                            len(providers),
-                            [p.peer_id.pretty() for p in providers],
+                            f"Found {len(providers)} providers: "
+                            f"{[p.peer_id.pretty() for p in providers]}"
                         )
                     else:
-                        logger.warning(
-                            "No servers found for content %s", content_key.hex()
-                        )
+                        logger.warning("No providers found")
 
                 # Keep the node running
                 while True:

libp2p/kad_dht/kad_dht.py

@@ -763,46 +763,45 @@ async def find_peer(self, peer_id: ID) -> PeerInfo | None:
 
     # Value storage and retrieval methods
 
-    async def put_value(self, key: bytes, value: bytes) -> None:
+    async def put_value(self, key: str, value: bytes) -> None:
         """
         Store a value in the DHT.
+
+        Args:
+            key: String key (will be converted to bytes for storage)
+            value: Binary value to store
+
         """
-        logger.debug(f"Storing value for key {key.hex()}")
+        logger.debug(f"Storing value for key {key}")
 
-        if key.decode("utf-8").startswith("/"):
-            if self.validator is not None:
-                # Dont allow local users to put bad values
-                self.validator.validate(key.decode("utf-8"), value)
+        # Validate if key starts with "/" (namespaced keys like /pk/...)
+        if self.validator is not None and key.startswith("/"):
+            self.validator.validate(key, value)
 
-                old_value_record = self.value_store.get(key)
-                if old_value_record is not None and old_value_record.value != value:
-                    # Select which value is better
-                    try:
-                        index = self.validator.select(
-                            key.decode("utf-8"), [value, old_value_record.value]
-                        )
-                        if index != 0:
-                            raise ValueError(
-                                "Refusing to replace newer value with the older one"
-                            )
-                    except Exception as e:
-                        logger.debug(f"Validation error for key {key.hex()}: {e}")
-                        raise
+            key_bytes = key.encode("utf-8")
+            old_value_record = self.value_store.get(key_bytes)
+            if old_value_record is not None and old_value_record.value != value:
+                index = self.validator.select(key, [value, old_value_record.value])
+                if index != 0:
+                    raise ValueError(
+                        "Refusing to replace newer value with the older one"
+                    )
+
+        # Convert string key to bytes for storage
+        key_bytes = key.encode("utf-8")
 
         # 1. Store locally first
-        self.value_store.put(key, value)
+        self.value_store.put(key_bytes, value)
         try:
             decoded_value = value.decode("utf-8")
         except UnicodeDecodeError:
             decoded_value = value.hex()
-        logger.debug(
-            f"Stored value locally for key {key.hex()} with value {decoded_value}"
-        )
+        logger.debug(f"Stored value locally for key {key} with value {decoded_value}")
 
         # 2. Get closest peers, excluding self
         closest_peers = [
             peer
-            for peer in self.routing_table.find_local_closest_peers(key)
+            for peer in self.routing_table.find_local_closest_peers(key_bytes)
             if peer != self.local_peer_id
         ]
         logger.debug(f"Found {len(closest_peers)} peers to store value at")
@@ -817,7 +816,7 @@ async def store_one(idx: int, peer: ID) -> None:
                 try:
                     with trio.move_on_after(QUERY_TIMEOUT):
                         success = await self.value_store._store_at_peer(
-                            peer, key, value
+                            peer, key_bytes, value
                         )
                         batch_results[idx] = success
                         if success:
@@ -835,19 +834,32 @@ async def store_one(idx: int, peer: ID) -> None:
 
         logger.info(f"Successfully stored value at {stored_count} peers")
 
-    async def get_value(self, key: bytes) -> bytes | None:
-        logger.debug(f"Getting value for key: {key.hex()}")
+    async def get_value(self, key: str) -> bytes | None:
+        """
+        Retrieve a value from the DHT.
+
+        Args:
+            key: String key (will be converted to bytes for lookup)
+
+        Returns:
+            The value if found, None otherwise
+
+        """
+        logger.debug(f"Getting value for key: {key}")
+
+        # Convert string key to bytes for lookup
+        key_bytes = key.encode("utf-8")
 
         # 1. Check local store first
-        value_record = self.value_store.get(key)
+        value_record = self.value_store.get(key_bytes)
         if value_record:
             logger.debug("Found value locally")
             return value_record.value
 
         # 2. Get closest peers, excluding self
         closest_peers = [
             peer
-            for peer in self.routing_table.find_local_closest_peers(key)
+            for peer in self.routing_table.find_local_closest_peers(key_bytes)
             if peer != self.local_peer_id
         ]
         logger.debug(f"Searching {len(closest_peers)} peers for value")
@@ -861,7 +873,7 @@ async def query_one(peer: ID) -> None:
                 nonlocal found_value
                 try:
                     with trio.move_on_after(QUERY_TIMEOUT):
-                        value = await self.value_store._get_from_peer(peer, key)
+                        value = await self.value_store._get_from_peer(peer, key_bytes)
                         if value is not None and found_value is None:
                             found_value = value
                             logger.debug(f"Found value at peer {peer}")
@@ -873,12 +885,12 @@ async def query_one(peer: ID) -> None:
                     nursery.start_soon(query_one, peer)
 
             if found_value is not None:
-                self.value_store.put(key, found_value)
+                self.value_store.put(key_bytes, found_value)
                 logger.info("Successfully retrieved value from network")
                 return found_value
 
         # 4. Not found
-        logger.warning(f"Value not found for key {key.hex()}")
+        logger.warning(f"Value not found for key {key}")
         return None
 
     # Add these methods in the Utility methods section
@@ -899,17 +911,19 @@ async def add_peer(self, peer_id: ID) -> bool:
         """
         return await self.routing_table.add_peer(peer_id)
 
-    async def provide(self, key: bytes) -> bool:
+    async def provide(self, key: str) -> bool:
         """
         Reference to provider_store.provide for convenience.
         """
-        return await self.provider_store.provide(key)
+        key_bytes = key.encode("utf-8")
+        return await self.provider_store.provide(key_bytes)
 
-    async def find_providers(self, key: bytes, count: int = 20) -> list[PeerInfo]:
+    async def find_providers(self, key: str, count: int = 20) -> list[PeerInfo]:
         """
         Reference to provider_store.find_providers for convenience.
         """
-        return await self.provider_store.find_providers(key, count)
+        key_bytes = key.encode("utf-8")
+        return await self.provider_store.find_providers(key_bytes, count)
 
     def get_routing_table_size(self) -> int:
         """

libp2p/kad_dht/value_store.py

@@ -72,7 +72,7 @@ def put(self, key: bytes, value: bytes, validity: float = 0.0) -> None:
         logger.debug(
             "Storing value for key %s... with validity %s", key.hex(), validity
         )
-        record = make_put_record(key.decode("utf-8"), value)
+        record = make_put_record(key, value)
         record.timeReceived = str(time.time)
 
         self.store[key] = (record, validity)

libp2p/records/record.py

@@ -1,19 +1,19 @@
 from libp2p.kad_dht.pb import kademlia_pb2 as record_pb2
 
 
-def make_put_record(key: str, value: bytes) -> record_pb2.Record:
+def make_put_record(key: bytes, value: bytes) -> record_pb2.Record:
     """
     Create a new Record object with the specified key and value.
 
     Args:
-        key (str): The key for the record, which will be encoded as bytes.
+        key (bytes): The key for the record.
         value (bytes): The value to associate with the key in the record.
 
     Returns:
         record_pb2.Record: A Record object containing the provided key and value.
 
     """
     record = record_pb2.Record()
-    record.key = key.encode()
+    record.key = key
     record.value = value
     return record

newsfragments/1059.bugfix.rst

@@ -0,0 +1 @@
+Kademlia DHT API now accepts string keys instead of bytes (``put_value(key: str, ...)``). Fixes UnicodeDecodeError with binary multihash keys.

tests/core/kad_dht/test_kad_dht.py

@@ -255,14 +255,15 @@ async def test_put_and_get_value(dht_pair: tuple[KadDHT, KadDHT]):
     dht_a, dht_b = dht_pair
     # dht_a.peer_routing.routing_table.add_peer(dht_b.pe)
     peer_b_info = PeerInfo(dht_b.host.get_id(), dht_b.host.get_addrs())
-    # Generate a random key and value
-    key = b"rendom_key"
+    # Generate a random key and value (use string key for API)
+    key = "random_key"
+    key_bytes = key.encode("utf-8")
     value = b"test-value"
 
     # First add the value directly to node A's store to verify storage works
-    dht_a.value_store.put(key, value)
+    dht_a.value_store.put(key_bytes, value)
     logger.debug("Local value store: %s", dht_a.value_store.store)
-    local_value_record = dht_a.value_store.get(key)
+    local_value_record = dht_a.value_store.get(key_bytes)
     assert local_value_record is not None
     assert local_value_record.value == value, "Local value storage failed"
     print("number of nodes in peer store", dht_a.host.get_peerstore().peer_ids())
@@ -305,7 +306,7 @@ async def test_put_and_get_value(dht_pair: tuple[KadDHT, KadDHT]):
     assert record_b.seq == record_b_put_value.seq
 
     # # Log debugging information
-    logger.debug("Put value with key %s...", key.hex()[:10])
+    logger.debug("Put value with key %s...", key[:10])
     logger.debug("Node A value store: %s", dht_a.value_store.store)
 
     # # Allow more time for the value to propagate
@@ -353,15 +354,15 @@ async def test_put_and_get_value(dht_pair: tuple[KadDHT, KadDHT]):
     value = keypair.public_key.serialize()
 
     with trio.fail_after(TEST_TIMEOUT):
-        await dht_a.put_value(key.encode(), value)
+        await dht_a.put_value(key, value)  # Now accepts string directly
 
     # INVALID KEY PAIR
     key = "/pk/abcdef1234567890"  # Not a valid multihash
     value = b"not-a-real-key"
 
     with trio.fail_after(TEST_TIMEOUT):
         with pytest.raises(InvalidRecordType, match="valid multihash"):
-            await dht_a.put_value(key.encode(), value)
+            await dht_a.put_value(key, value)  # Now accepts string directly
 
 
 @pytest.mark.trio
@@ -372,10 +373,11 @@ async def test_provide_and_find_providers(dht_pair: tuple[KadDHT, KadDHT]):
 
     # Generate a random content ID
     content = f"test-content-{uuid.uuid4()}".encode()
-    content_id = b"randome_content"
+    content_id = "randome_content"  # String for API
+    content_id_bytes = content_id.encode("utf-8")  # Bytes for internal storage
 
     # Store content on the first node
-    dht_a.value_store.put(content_id, content)
+    dht_a.value_store.put(content_id_bytes, content)
 
     # An extra FIND_NODE req is sent between the 2 nodes while dht creation,
     # so both the nodes will have records of each other before PUT_VALUE req is sent
@@ -562,12 +564,13 @@ async def test_dht_req_fail_with_invalid_record_transfer(
     peer_b_info = PeerInfo(dht_b.host.get_id(), dht_b.host.get_addrs())
 
     # Generate a random key and value
-    key = b"rendom_key"
+    key = "random_key"  # String for API
+    key_bytes = key.encode("utf-8")  # Bytes for internal storage
     value = b"test-value"
 
     # First add the value directly to node A's store to verify storage works
-    dht_a.value_store.put(key, value)
-    local_value = dht_a.value_store.get(key)
+    dht_a.value_store.put(key_bytes, value)
+    local_value = dht_a.value_store.get(key_bytes)
     assert local_value is not None
     assert local_value.value == value, "Local value storage failed"
     await dht_a.routing_table.add_peer(peer_b_info)
@@ -583,7 +586,7 @@ async def test_dht_req_fail_with_invalid_record_transfer(
         dht_a.host.get_peerstore().set_local_record(envelope)
 
     await dht_a.put_value(key, value)
-    retrieved_value_record = dht_b.value_store.get(key)
+    retrieved_value_record = dht_b.value_store.get(key_bytes)
 
     # This proves that DHT_B rejected DHT_A PUT_RECORD req upon receiving
     # the corrupted invalid record
@@ -596,7 +599,7 @@ async def test_dht_req_fail_with_invalid_record_transfer(
     dht_a.host.get_peerstore().set_local_record(false_envelope)
 
     await dht_a.put_value(key, value)
-    retrieved_value_record = dht_b.value_store.get(key)
+    retrieved_value_record = dht_b.value_store.get(key_bytes)
 
     # This proves that DHT_B rejected DHT_A PUT_RECORD req upon receving
     # the record with a different peer_id regardless of a valid signature