Close VDF client TCP connections properly (#20602)

* Close VDF client TCP connections properly

_handle_client() did not close the StreamWriter for non-whitelisted
TCP connections, and free_clients had no upper bound once all 3
chains were mapped.

Fix:
- Explicitly close and await writer for non-whitelisted IPs
- Cap free_clients at max_free_clients (10)
- Close remaining free_clients connections during shutdown

Add unit tests for all three _handle_client paths (reject, accept,
cap overflow).

* Remove unrelated test to match private repo PR

Remove test_invalid_vdf_proof_is_ignored_in_process_communication and its
helpers (_NullTransport, _make_null_writer) so the public PR matches the
private security fix exactly.

* Move writer close outside lock in _handle_client

Avoid holding self.lock during writer.close()/wait_closed() to
prevent blocking timelord operations on a slow TCP close handshake.

chia/_tests/timelord/test_timelord.py

@@ -1,7 +1,11 @@
 from __future__ import annotations
 
+import asyncio
+from unittest.mock import AsyncMock, MagicMock, patch
+
 import pytest
 
+from chia.timelord.timelord import Timelord
 from chia.timelord.timelord_service import TimelordService
 
 
@@ -9,3 +13,67 @@ from chia.timelord.timelord_service import TimelordService
 async def test_timelord_has_no_server(timelord_service: TimelordService) -> None:
     timelord_server = timelord_service._node.server
     assert timelord_server.webserver is None
+
+
+def _make_mock_writer(ip: str = "127.0.0.1") -> MagicMock:
+    writer = MagicMock(spec=asyncio.StreamWriter)
+    writer.get_extra_info.return_value = (ip, 12345)
+    writer.close = MagicMock()
+    writer.wait_closed = AsyncMock()
+    return writer
+
+
+def _make_timelord_stub(ip_whitelist: list[str]) -> Timelord:
+    with patch.object(Timelord, "__init__", lambda self, *a, **kw: None):
+        tl = Timelord.__new__(Timelord)
+    tl.free_clients = []
+    tl.max_free_clients = 10
+    tl.ip_whitelist = ip_whitelist
+    tl.lock = asyncio.Lock()
+    return tl
+
+
+class TestHandleClient:
+    @pytest.mark.anyio
+    async def test_non_whitelisted_ip_is_rejected_and_closed(self) -> None:
+        tl = _make_timelord_stub(ip_whitelist=["127.0.0.1"])
+        reader = MagicMock(spec=asyncio.StreamReader)
+        writer = _make_mock_writer(ip="10.0.0.99")
+
+        await tl._handle_client(reader, writer)
+
+        assert len(tl.free_clients) == 0
+        writer.close.assert_called_once()
+        writer.wait_closed.assert_awaited_once()
+
+    @pytest.mark.anyio
+    async def test_whitelisted_ip_is_accepted(self) -> None:
+        tl = _make_timelord_stub(ip_whitelist=["127.0.0.1"])
+        reader = MagicMock(spec=asyncio.StreamReader)
+        writer = _make_mock_writer(ip="127.0.0.1")
+
+        await tl._handle_client(reader, writer)
+
+        assert len(tl.free_clients) == 1
+        assert tl.free_clients[0] == ("127.0.0.1", reader, writer)
+        writer.close.assert_not_called()
+
+    @pytest.mark.anyio
+    async def test_excess_clients_beyond_cap_are_rejected(self) -> None:
+        tl = _make_timelord_stub(ip_whitelist=["127.0.0.1"])
+        tl.max_free_clients = 3
+
+        for _ in range(3):
+            reader = MagicMock(spec=asyncio.StreamReader)
+            writer = _make_mock_writer(ip="127.0.0.1")
+            await tl._handle_client(reader, writer)
+
+        assert len(tl.free_clients) == 3
+
+        overflow_reader = MagicMock(spec=asyncio.StreamReader)
+        overflow_writer = _make_mock_writer(ip="127.0.0.1")
+        await tl._handle_client(overflow_reader, overflow_writer)
+
+        assert len(tl.free_clients) == 3
+        overflow_writer.close.assert_called_once()
+        overflow_writer.wait_closed.assert_awaited_once()

chia/timelord/timelord.py

@@ -98,6 +98,7 @@ class Timelord:
         self.constants = constants
         self._shut_down = False
         self.free_clients: list[tuple[str, asyncio.StreamReader, asyncio.StreamWriter]] = []
+        self.max_free_clients: int = 10
         self.ip_whitelist = self.config["vdf_clients"]["ip"]
         self._server: ChiaServer | None = None
         self.chain_type_to_stream: dict[Chain, tuple[str, asyncio.StreamReader, asyncio.StreamWriter]] = {}
@@ -189,6 +190,10 @@ class Timelord:
                 self.main_loop.cancel()
             if self.bluebox_pool is not None:
                 self.bluebox_pool.shutdown()
+            for _, _, writer in self.free_clients:
+                with contextlib.suppress(Exception):
+                    writer.close()
+            self.free_clients.clear()
 
     def get_connections(self, request_node_type: NodeType | None) -> list[dict[str, Any]]:
         return default_get_connections(server=self.server, request_node_type=request_node_type)
@@ -212,12 +217,22 @@ class Timelord:
         self._server = server
 
     async def _handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+        should_close = False
         async with self.lock:
             client_ip = writer.get_extra_info("peername")[0]
             log.debug(f"New timelord connection from client: {client_ip}.")
-            if client_ip in self.ip_whitelist:
+            if client_ip not in self.ip_whitelist:
+                log.warning(f"Rejected VDF client from non-whitelisted IP: {client_ip}")
+                should_close = True
+            elif len(self.free_clients) >= self.max_free_clients:
+                log.warning(f"Too many free VDF clients ({len(self.free_clients)}), rejecting {client_ip}")
+                should_close = True
+            else:
                 self.free_clients.append((client_ip, reader, writer))
                 log.debug(f"Added new VDF client {client_ip}.")
+        if should_close:
+            writer.close()
+            await writer.wait_closed()
 
     async def _stop_chain(self, chain: Chain) -> None:
         try: