added record validators in kad-dht

Author: lla-dane

libp2p/kad_dht/common.py

@@ -9,6 +9,13 @@
 # Constants for the Kademlia algorithm
 ALPHA = 3  # Concurrency parameter
 PROTOCOL_ID = TProtocol("/ipfs/kad/1.0.0")
+PROTOCOL_PREFIX = TProtocol("/ipfs")
 QUERY_TIMEOUT = 10
 
 TTL = DEFAULT_TTL = 24 * 60 * 60  # 24 hours in seconds
+
+# Default parameters
+BUCKET_SIZE = 20  # k in the Kademlia paper
+MAXIMUM_BUCKETS = 256  # Maximum number of buckets (for 256-bit keys)
+PEER_REFRESH_INTERVAL = 60  # Interval to refresh peers in seconds
+STALE_PEER_THRESHOLD = 3600  # Time in seconds after which a peer is considered stale

libp2p/kad_dht/kad_dht.py

@@ -22,6 +22,7 @@
     IHost,
 )
 from libp2p.discovery.random_walk.rt_refresh_manager import RTRefreshManager
+from libp2p.custom_types import TProtocol
 from libp2p.network.stream.net_stream import (
     INetStream,
 )
@@ -31,13 +32,17 @@
 from libp2p.peer.peerinfo import (
     PeerInfo,
 )
+from libp2p.records.pubkey import PublicKeyValidator
+from libp2p.records.validator import NamespacedValidator, Validator
 from libp2p.tools.async_service import (
     Service,
 )
 
 from .common import (
     ALPHA,
+    BUCKET_SIZE,
     PROTOCOL_ID,
+    PROTOCOL_PREFIX,
     QUERY_TIMEOUT,
 )
 from .pb.kademlia_pb2 import (
@@ -89,7 +94,16 @@ class KadDHT(Service):
 
     """
 
-    def __init__(self, host: IHost, mode: DHTMode, enable_random_walk: bool = False):
+    def __init__(
+        self,
+        host: IHost,
+        mode: DHTMode, enable_random_walk: bool = False,
+        validator: NamespacedValidator | None = None,
+        validator_changed: bool = False,
+        protocol_prefix: TProtocol = PROTOCOL_PREFIX,
+        enable_providers: bool = True,
+        enable_values: bool = True,
+    ):
         """
         Initialize a new Kademlia DHT node.
 
@@ -112,6 +126,16 @@ def __init__(self, host: IHost, mode: DHTMode, enable_random_walk: bool = False)
         # Initialize the routing table
         self.routing_table = RoutingTable(self.local_peer_id, self.host)
 
+        self.protocol_prefix = protocol_prefix
+        self.enable_providers = enable_providers
+        self.enable_values = enable_values
+
+        self.validator = validator
+
+        # If true implies that the validator has been changed and that
+        # Defaults should not be used
+        self.validator_changed = validator_changed
+
         # Initialize peer routing
         self.peer_routing = PeerRouting(host, self.routing_table)
 
@@ -205,6 +229,84 @@ async def stop(self) -> None:
         else:
             logger.info("RT Refresh Manager was not running (Random Walk disabled)")
 
+    async def apply_fallbacks(self, host: IHost) -> None:
+        """
+        Apply fallback validators if not explicitely changed by the user
+
+        This sets default validators like 'pk' and 'ipns' if they are missing and
+        the default validator set hasn't been overridden.
+        """
+        if not self.validator_changed:
+            if not isinstance(self.validator, NamespacedValidator):
+                raise ValueError(
+                    "Default validator was changed without marking it True"
+                )
+
+            if "pk" not in self.validator._validators:
+                self.validator._validators["pk"] = PublicKeyValidator()
+
+            # TODO: Do the same thing for ipns, but need to implement first.
+
+    def validate_config(self) -> None:
+        """
+        Validate the DHT config.
+        """
+        if self.protocol_prefix != PROTOCOL_PREFIX:
+            return  # Skip validation for non-standart prefixes
+
+        for bucket in self.routing_table.buckets:
+            if bucket.bucket_size != BUCKET_SIZE:
+                raise ValueError(
+                    f"{PROTOCOL_PREFIX} prefix must use bucket size {BUCKET_SIZE}"
+                )
+
+        if not self.enable_providers:
+            raise ValueError(f"{PROTOCOL_PREFIX} prefix must have providers enabled")
+
+        if not self.enable_values:
+            raise ValueError(f"{PROTOCOL_PREFIX} prefix must have values enabled")
+
+        if not isinstance(self.validator, NamespacedValidator):
+            raise ValueError(
+                f"{PROTOCOL_PREFIX} prefix must use a namespace type validator"
+            )
+
+        vmap = self.validator._validators
+
+        # TODO: Need to add ipns also in the check
+        if set(vmap.keys()) != {"pk"}:
+            raise ValueError(f"{PROTOCOL_PREFIX} must have 'pk' and 'ipns' validators")
+
+        pk_validator = vmap.get("pk")
+        if not isinstance(pk_validator, PublicKeyValidator):
+            raise TypeError("'pk' namesapce must use PublicKeyValidator")
+
+        # TODO: ipns checks
+
+    def set_validator(self, val: NamespacedValidator) -> None:
+        """
+        Set a custom validator for the DHT config.
+
+        This marks the validator as explicitly changed, so the default
+        validators (pk and ipns) will not be automatically applied later.
+        """
+        self.validator = val
+        self.validator_changed = True
+        return
+
+    def set_namespace_validator(self, ns: str, val: Validator) -> None:
+        """
+        Adds a validator under a specofic namespace to the current DHT config.
+
+        Raises an error if the current validator is not a NamespacedValidator
+        """
+        if not isinstance(self.validator, NamespacedValidator):
+            raise TypeError(
+                "Can only add namespaced validators to a NamespacedValidator"
+            )
+
+        self.validator._validators["ns"] = val
+
     async def switch_mode(self, new_mode: DHTMode) -> DHTMode:
         """
         Switch the DHT mode.
@@ -539,6 +641,23 @@ async def put_value(self, key: bytes, value: bytes) -> None:
         """
         logger.debug(f"Storing value for key {key.hex()}")
 
+        if self.validator is not None:
+            # Dont allow local users to put bad values
+            self.validator.validate(key.decode("utf-8"), value)
+
+            old = self.value_store.get(key)
+            if old is not None and old != value:
+                # Select which value is better
+                try:
+                    index = self.validator.select(key.decode("utf-8"), [value, old])
+                    if index != 0:
+                        raise ValueError(
+                            "Refusing to replace newer value with the older one"
+                        )
+                except Exception as e:
+                    logger.debug(f"Validation select error for key {key.hex()}: {e}")
+                    raise
+
         # 1. Store locally first
         self.value_store.put(key, value)
         try:

libp2p/kad_dht/pb/kademlia.proto

@@ -36,3 +36,4 @@ message Message {
     repeated Peer closerPeers = 8;
     repeated Peer providerPeers = 9;
 }
+`

libp2p/kad_dht/routing_table.py

@@ -25,7 +25,10 @@
 )
 
 from .common import (
+    BUCKET_SIZE,
+    PEER_REFRESH_INTERVAL,
     PROTOCOL_ID,
+    STALE_PEER_THRESHOLD,
 )
 from .pb.kademlia_pb2 import (
     Message,
@@ -34,12 +37,6 @@
 # logger = logging.getLogger("libp2p.kademlia.routing_table")
 logger = logging.getLogger("kademlia-example.routing_table")
 
-# Default parameters
-BUCKET_SIZE = 20  # k in the Kademlia paper
-MAXIMUM_BUCKETS = 256  # Maximum number of buckets (for 256-bit keys)
-PEER_REFRESH_INTERVAL = 60  # Interval to refresh peers in seconds
-STALE_PEER_THRESHOLD = 3600  # Time in seconds after which a peer is considered stale
-
 
 def peer_id_to_key(peer_id: ID) -> bytes:
     """

libp2p/records/pb/record.proto

@@ -1,4 +1,4 @@
-from libp2p.records.pb import record_pb2
+from libp2p.kad_dht.pb import kademlia_pb2 as record_pb2
 
 
 def make_put_record(key: str, value: bytes) -> record_pb2.Record:

libp2p/records/pb/record_pb2.py

@@ -24,6 +24,7 @@
 from libp2p.peer.peerinfo import (
     PeerInfo,
 )
+from libp2p.records.validator import Validator
 from libp2p.tools.async_service import (
     background_trio_service,
 )
@@ -38,6 +39,14 @@
 TEST_TIMEOUT = 5  # Timeout in seconds
 
 
+class BlankValidator(Validator):
+    def validate(self, key: str, value: bytes) -> None:
+        return
+
+    def select(self, key: str, values: list[bytes]) -> int:
+        return 0
+
+
 @pytest.fixture
 async def dht_pair(security_protocol):
     """Create a pair of connected DHT nodes for testing."""