diff --git a/chia/_tests/timelord/test_timelord.py b/chia/_tests/timelord/test_timelord.py index 80ca02ed05..a0fec78595 100644 --- a/chia/_tests/timelord/test_timelord.py +++ b/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() diff --git a/chia/timelord/timelord.py b/chia/timelord/timelord.py index 11143a009d..526d3e11b8 100644 --- a/chia/timelord/timelord.py +++ b/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: