Add redirect_policy to ClusterClient (leader-trust allowlist)

ISSUE-04 — a compromised peer could return any address as the
current leader; the client blindly opened a TCP connection to that
address (SSRF-style attack surface). ClusterClient now accepts an
optional ``redirect_policy`` callable that authorizes each redirect
target. Convenience ``allowlist_policy`` helper builds one from a
static address list — the common pattern "only redirect to hosts
I seed-listed."

Backwards-compatible: policy defaults to None which preserves the
legacy "accept any redirect" behavior. Users should set it in
production where network trust isn't absolute.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: antoineleclair

src/dqliteclient/cluster.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import random
+from collections.abc import Callable
 
 from dqliteclient.connection import DqliteConnection, _parse_address
 from dqliteclient.exceptions import (
@@ -15,6 +16,11 @@
 from dqliteclient.retry import retry_with_backoff
 
 
+# Type alias for a redirect-target policy. Returns True if the address
+# should be accepted, False to reject with a ClusterError.
+RedirectPolicy = Callable[[str], bool]
+
+
 class ClusterClient:
     """Client with automatic leader detection and failover."""
 
@@ -23,23 +29,47 @@ def __init__(
         node_store: NodeStore,
         *,
         timeout: float = 10.0,
+        redirect_policy: RedirectPolicy | None = None,
     ) -> None:
         """Initialize cluster client.
 
         Args:
             node_store: Store for cluster node information
             timeout: Connection timeout in seconds
+            redirect_policy: Optional callable ``(address) -> bool`` that
+                authorizes each leader redirect target. If None (default),
+                redirects are accepted — this preserves backward
+                compatibility but permits a compromised peer to redirect
+                clients to arbitrary hosts (SSRF-style). Supply a
+                callable or the ``only_nodes_in_store`` helper to
+                constrain where redirects can go.
         """
         if timeout <= 0:
             raise ValueError(f"timeout must be positive, got {timeout}")
         self._node_store = node_store
         self._timeout = timeout
+        self._redirect_policy = redirect_policy
 
     @classmethod
-    def from_addresses(cls, addresses: list[str], timeout: float = 10.0) -> "ClusterClient":
+    def from_addresses(
+        cls,
+        addresses: list[str],
+        timeout: float = 10.0,
+        *,
+        redirect_policy: RedirectPolicy | None = None,
+    ) -> "ClusterClient":
         """Create cluster client from list of addresses."""
         store = MemoryNodeStore(addresses)
-        return cls(store, timeout=timeout)
+        return cls(store, timeout=timeout, redirect_policy=redirect_policy)
+
+    def _check_redirect(self, address: str) -> None:
+        """Reject leader-redirect targets that fail the configured policy."""
+        if self._redirect_policy is None:
+            return
+        if not self._redirect_policy(address):
+            raise ClusterError(
+                f"Leader redirect to {address!r} rejected by redirect_policy"
+            )
 
     async def find_leader(self) -> str:
         """Find the current cluster leader.
@@ -68,6 +98,12 @@ async def find_leader(self) -> str:
                     self._query_leader(node.address), timeout=self._timeout
                 )
                 if leader_address:
+                    # Only leader_address values that did NOT come from
+                    # node.address itself need authorizing — those are
+                    # real redirects. If the server returned its own
+                    # address, it's the leader and already in the store.
+                    if leader_address != node.address:
+                        self._check_redirect(leader_address)
                     return leader_address
             except TimeoutError as e:
                 errors.append(f"{node.address}: timed out")
@@ -143,3 +179,19 @@ async def try_connect() -> DqliteConnection:
     async def update_nodes(self, nodes: list[NodeInfo]) -> None:
         """Update the node store with new node information."""
         await self._node_store.set_nodes(nodes)
+
+
+def allowlist_policy(addresses: list[str] | set[str]) -> RedirectPolicy:
+    """Build a redirect policy that accepts only the given addresses.
+
+    Useful for the common case: "only allow redirects to hosts I've
+    explicitly seed-listed." Addresses are matched by exact string
+    equality — callers that need CIDR / DNS / wildcard matching should
+    supply their own callable.
+    """
+    allowed = set(addresses)
+
+    def policy(addr: str) -> bool:
+        return addr in allowed
+
+    return policy

tests/test_redirect_policy.py

@@ -0,0 +1,67 @@
+"""Leader-redirect allowlist policy (ISSUE-04).
+
+A compromised peer can return any address as the leader; the client
+used to open a TCP connection to whatever it was told. ClusterClient
+now takes an optional ``redirect_policy`` callable that authorizes
+each redirect target; a convenience ``allowlist_policy`` helper builds
+one from a static list of addresses.
+"""
+
+import asyncio
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+from dqliteclient.cluster import ClusterClient, allowlist_policy
+from dqliteclient.exceptions import ClusterError
+from dqliteclient.node_store import MemoryNodeStore
+
+
+class TestRedirectPolicy:
+    def test_accepts_seeded_redirect(self) -> None:
+        store = MemoryNodeStore(["10.0.0.1:9001", "10.0.0.2:9001"])
+        cc = ClusterClient(
+            store,
+            timeout=5.0,
+            redirect_policy=allowlist_policy(["10.0.0.1:9001", "10.0.0.2:9001"]),
+        )
+        # Simulate: probing 10.0.0.1 reports 10.0.0.2 as leader.
+        with patch.object(cc, "_query_leader", new=AsyncMock(return_value="10.0.0.2:9001")):
+            result = asyncio.run(cc.find_leader())
+        assert result == "10.0.0.2:9001"
+
+    def test_rejects_unknown_redirect(self) -> None:
+        store = MemoryNodeStore(["10.0.0.1:9001"])
+        cc = ClusterClient(
+            store,
+            timeout=5.0,
+            redirect_policy=allowlist_policy(["10.0.0.1:9001"]),
+        )
+        # A compromised peer redirects to an attacker-controlled host.
+        with patch.object(cc, "_query_leader", new=AsyncMock(return_value="attacker.com:9001")):
+            with pytest.raises(ClusterError, match="rejected"):
+                asyncio.run(cc.find_leader())
+
+    def test_no_policy_means_any_redirect_accepted(self) -> None:
+        """Default (None) policy preserves legacy behavior."""
+        store = MemoryNodeStore(["10.0.0.1:9001"])
+        cc = ClusterClient(store, timeout=5.0)
+        with patch.object(cc, "_query_leader", new=AsyncMock(return_value="anywhere.invalid:9001")):
+            result = asyncio.run(cc.find_leader())
+        assert result == "anywhere.invalid:9001"
+
+    def test_self_leader_bypasses_policy(self) -> None:
+        """If the queried node is the leader (returns its own address),
+        the redirect policy doesn't apply — the address is already in the
+        seed list by definition."""
+        store = MemoryNodeStore(["10.0.0.1:9001"])
+        # Policy that rejects everything — but the node returning its own
+        # address isn't a real redirect, so it's accepted.
+        cc = ClusterClient(
+            store,
+            timeout=5.0,
+            redirect_policy=lambda _a: False,
+        )
+        with patch.object(cc, "_query_leader", new=AsyncMock(return_value="10.0.0.1:9001")):
+            result = asyncio.run(cc.find_leader())
+        assert result == "10.0.0.1:9001"