Merge pull request libp2p#1095 from sumanjeet0012/feat/kademlia_validator

Implement custom validators and quorum based value retrieval

Author: seetadev

examples/kademlia/kademlia.py

@@ -4,6 +4,7 @@
 A basic example of using the Kademlia DHT implementation, with all setup logic inlined.
 This example demonstrates both value storage/retrieval and content server
 advertisement/discovery.
+It also shows how to use custom validators for namespaced keys.
 """
 
 import argparse
@@ -31,6 +32,7 @@
     DHTMode,
     KadDHT,
 )
+from libp2p.records.validator import Validator
 from libp2p.tools.async_service import (
     background_trio_service,
 )
@@ -39,6 +41,46 @@
 )
 from libp2p.utils.paths import get_script_dir, join_paths
 
+
+# Custom validator for the "example" namespace
+class ExampleValidator(Validator):
+    """
+    A simple validator for the 'example' namespace.
+
+    This validator accepts any value and always selects the first value
+    when comparing multiple values.
+    """
+
+    def validate(self, key: str, value: bytes) -> None:
+        """
+        Validate a key-value pair.
+
+        In a real application, you might check:
+        - Value format/schema
+        - Signatures
+        - Size limits
+        - etc.
+        """
+        # For this example, we accept any value
+        # You can add custom validation logic here
+        if not value:
+            raise ValueError("Value cannot be empty")
+
+    def select(self, key: str, values: list[bytes]) -> int:
+        """
+        Select the best value from a list of values.
+
+        Returns the index of the selected value.
+        In this example, we simply return the first value (index 0).
+
+        In a real application, you might:
+        - Compare timestamps
+        - Check version numbers
+        - Verify signatures and pick the most recent valid one
+        """
+        return 0
+
+
 # Configure logging
 logging.basicConfig(
     level=logging.INFO,
@@ -175,6 +217,12 @@ async def run_node(
 
             await connect_to_bootstrap_nodes(host, bootstrap_nodes)
             dht = KadDHT(host, dht_mode)
+
+            # Register a custom validator for the "example" namespace
+            # This allows us to store values with keys like "/example/my-key"
+            dht.register_validator("example", ExampleValidator())
+            logger.info("Registered custom 'example' namespace validator")
+
             # take all peer ids from the host and add them to the dht
             for peer_id in host.get_peerstore().peer_ids():
                 await dht.routing_table.add_peer(peer_id)
@@ -188,12 +236,15 @@ async def run_node(
             async with background_trio_service(dht):
                 logger.info(f"DHT service started in {dht_mode.value} mode")
 
-                # Example 1: Simple Key-Value Storage
-                # Just use a string key directly - DHT API accepts strings!
-                key = "my-example-key"
+                # Example 1: Simple Key-Value Storage with namespaced key
+                # Keys MUST be namespaced (e.g., /namespace/key) for validation
+                # The namespace must have a registered validator
+                key = "/example/my-example-key"
                 value = b"Hello from py-libp2p!"
 
                 # Example 2: Content Provider Advertisement
+                # Provider keys use a different storage mechanism (provider store)
+                # that doesn't go through the value validation path
                 content_id = "my-content-identifier"
 
                 if dht_mode == DHTMode.SERVER:

libp2p/kad_dht/kad_dht.py

@@ -50,6 +50,7 @@
 )
 from .pb.kademlia_pb2 import (
     Message,
+    Record,
 )
 from .peer_routing import (
     PeerRouting,
@@ -77,6 +78,32 @@ class DHTMode(Enum):
     SERVER = "SERVER"
 
 
+# Timestamp validation constants
+MAX_TIMESTAMP_AGE = 24 * 60 * 60  # 24 hours in seconds
+MAX_TIMESTAMP_FUTURE = 5 * 60  # 5 minutes in the future in seconds
+
+
+def is_valid_timestamp(ts: float) -> bool:
+    """
+    Validate if a timestamp is within acceptable bounds.
+
+    Args:
+        ts: The timestamp to validate (Unix timestamp in seconds)
+
+    Returns:
+        bool: True if timestamp is valid (not too old and not too far in future)
+
+    """
+    current_time = time.time()
+    # Check if timestamp is not in the future by more than MAX_TIMESTAMP_FUTURE
+    if ts > current_time + MAX_TIMESTAMP_FUTURE:
+        return False
+    # Check if timestamp is not too far in the past
+    if current_time - ts > MAX_TIMESTAMP_AGE:
+        return False
+    return True
+
+
 class KadDHT(Service):
     """
     Kademlia DHT implementation for libp2p.
@@ -714,6 +741,13 @@ async def handle_stream(self, stream: INetStream) -> None:
                                 "Missing key or value in PUT_VALUE message"
                             )
 
+                        # Always validate the key-value pair before storing
+                        # Reject keys without registered namespace validators
+                        key_str = key.decode("utf-8")
+                        if self.validator is None:
+                            raise ValueError("Validator is required for DHT operations")
+                        self.validator.validate(key_str, value)
+
                         self.value_store.put(key, value)
                         logger.debug(f"Stored value {value.hex()} for key {key.hex()}")
                         success = True
@@ -771,24 +805,28 @@ async def put_value(self, key: str, value: bytes) -> None:
             key: String key (will be converted to bytes for storage)
             value: Binary value to store
 
+        Raises:
+            InvalidRecordType: If no validator is registered for the key's namespace
+            ValueError: If trying to replace a newer value with an older one
+
         """
         logger.debug(f"Storing value for key {key}")
 
-        # Validate if key starts with "/" (namespaced keys like /pk/...)
-        if self.validator is not None and key.startswith("/"):
-            self.validator.validate(key, value)
-
-            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"
-                    )
+        # Always validate the key-value pair using the namespaced validator
+        # This will raise InvalidRecordType if:
+        # - The key is not namespaced (doesn't start with / or has no second /)
+        # - No validator is registered for the key's namespace
+        # Following Go libp2p behavior where only namespaced keys are allowed
+        if self.validator is None:
+            raise ValueError("Validator is required for DHT operations")
+        self.validator.validate(key, value)
 
-        # Convert string key to bytes for storage
         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")
 
         # 1. Store locally first
         self.value_store.put(key_bytes, value)
@@ -834,15 +872,21 @@ 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: str) -> bytes | None:
+    async def get_value(self, key: str, quorum: int = 0) -> bytes | None:
         """
         Retrieve a value from the DHT.
 
         Args:
             key: String key (will be converted to bytes for lookup)
+            quorum: Minimum number of valid peer responses required for confidence.
+            If quorum > 0 and not met, the function still returns the best value
+            found (if any) but logs a warning. Set to 0 to disable quorum checking.
 
         Returns:
-            The value if found, None otherwise
+            The value if found (best value even if quorum not met), None otherwise.
+            Note: When quorum is not met, a warning is logged but the best available
+            value is still returned. This allows graceful degradation when the network
+            has insufficient peers.
 
         """
         logger.debug(f"Getting value for key: {key}")
@@ -864,32 +908,110 @@ async def get_value(self, key: str) -> bytes | None:
         ]
         logger.debug(f"Searching {len(closest_peers)} peers for value")
 
+        # Collect valid records from peers: mapping peer -> Record
+        valid_records: list[tuple[ID, Record]] = []
+
         # 3. Query ALPHA peers at a time in parallel
+        # Use list to track mutable state (pyrefly requirement)
+        total_responses_list: list[int] = [0]
         for i in range(0, len(closest_peers), ALPHA):
             batch = closest_peers[i : i + ALPHA]
-            found_value = None
 
             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_bytes)
-                        if value is not None and found_value is None:
-                            found_value = value
-                            logger.debug(f"Found value at peer {peer}")
+                        # Fetch the record directly to get timeReceived
+                        rec = await self.value_store._get_from_peer(
+                            peer, key_bytes, return_record=True
+                        )
+                        if rec is not None:
+                            total_responses_list[0] += 1
+                            # Validate the record's value
+                            try:
+                                if self.validator is None:
+                                    raise ValueError("Validator is required")
+                                if not isinstance(rec, Record):
+                                    raise TypeError("Expected Record type")
+                                self.validator.validate(key, rec.value)
+                                valid_records.append((peer, rec))
+                                logger.debug(f"Found valid record at peer {peer}")
+                            except Exception as e:
+                                logger.debug(
+                                    f"Received invalid record from {peer}, "
+                                    f"discarding: {e}"
+                                )
                 except Exception as e:
                     logger.debug(f"Error querying peer {peer}: {e}")
 
             async with trio.open_nursery() as nursery:
                 for peer in batch:
                     nursery.start_soon(query_one, peer)
 
-            if found_value is not None:
-                self.value_store.put(key_bytes, found_value)
-                logger.info("Successfully retrieved value from network")
-                return found_value
+            # If quorum is set and we have enough valid records, we can stop early
+            # Note: quorum counts valid records, not all responses.
+            if quorum and len(valid_records) >= quorum:
+                logger.debug(f"Quorum reached ({len(valid_records)} valid records)")
+                break
+
+        # 4. Select the best record if any valid records were found
+        if valid_records:
+            # Check if quorum was met
+            if quorum > 0 and len(valid_records) < quorum:
+                logger.warning(
+                    f"Quorum not met: found {len(valid_records)} valid records, "
+                    f"required {quorum}. Returning best value found."
+                )
+
+            # Select the best record using the validator
+            # Note: Following Go libp2p's approach, we use validator.select() to choose
+            # the best value, not timestamps. The timeReceived field is for local
+            # bookkeeping only, not for distributed consensus on the "best" record.
+            if self.validator is None:
+                raise ValueError("Validator is required for record selection")
+
+            values = [rec.value for _p, rec in valid_records]
+            best_idx = self.validator.select(key, values)
+            logger.debug(
+                f"Selected best value at index {best_idx}using validator.select()"
+            )
+
+            best_peer, best_rec = valid_records[best_idx]
+            best_value = best_rec.value
+
+            # Propagate the best record to peers that have different values
+            # This ensures network consistency, following Go libp2p's approach
+            outdated_peers: list[ID] = []
+            for peer, rec in valid_records:
+                # Propagate if the peer has a different value than the best
+                if rec.value != best_value:
+                    outdated_peers.append(peer)
 
-        # 4. Not found
+            if outdated_peers:
+                logger.debug(
+                    f"Propagating best value to {len(outdated_peers)}"
+                    "peers with outdated values"
+                )
+
+                async def propagate(peer: ID) -> None:
+                    try:
+                        with trio.move_on_after(QUERY_TIMEOUT):
+                            await self.value_store._store_at_peer(
+                                peer, key_bytes, best_value
+                            )
+                            logger.debug(f"Propagated updated record to peer {peer}")
+                    except Exception as e:
+                        logger.debug(f"Failed to propagate to peer {peer}: {e}")
+
+                async with trio.open_nursery() as nursery:
+                    for p in outdated_peers:
+                        nursery.start_soon(propagate, p)
+
+            # Store the best value locally
+            self.value_store.put(key_bytes, best_value)
+            logger.info("Successfully retrieved value from network")
+            return best_value
+
+        # 5. Not found
         logger.warning(f"Value not found for key {key}")
         return None
 
@@ -949,6 +1071,36 @@ def get_value_store_size(self) -> int:
         """
         return self.value_store.size()
 
+    def register_validator(self, namespace: str, validator: Validator) -> None:
+        """
+        Register a custom validator for a specific namespace.
+
+        This allows storing and retrieving values with custom namespaced keys
+        (e.g., /myapp/key). The validator will be used to validate values
+        before storing and after retrieval.
+
+        Args:
+            namespace: The namespace string (e.g., "myapp" for keys like /myapp/key)
+            validator: A Validator instance with validate() and select() methods
+
+        Example:
+            class MyValidator(Validator):
+                def validate(self, key: str, value: bytes) -> None:
+                    # Custom validation logic
+                    pass
+
+                def select(self, key: str, values: list[bytes]) -> int:
+                    return 0  # Return index of best value
+
+            dht.register_validator("myapp", MyValidator())
+            await dht.put_value("/myapp/my-key", b"my-value")
+
+        """
+        if self.validator is None:
+            self.validator = NamespacedValidator({namespace: validator})
+        else:
+            self.validator.add_validator(namespace, validator)
+
     def is_random_walk_enabled(self) -> bool:
         """
         Check if random walk peer discovery is enabled.