chia-blockchain/tests/core/server/test_dos.py

# flake8: noqa: F811, F401
import asyncio
import logging

import pytest
from aiohttp import ClientSession, ClientTimeout, ServerDisconnectedError, WSCloseCode, WSMessage, WSMsgType

from chia.full_node.full_node_api import FullNodeAPI
from chia.protocols import full_node_protocol
from chia.protocols.protocol_message_types import ProtocolMessageTypes
from chia.protocols.shared_protocol import Handshake
from chia.server.outbound_message import make_msg, Message
from chia.server.rate_limits import RateLimiter
from chia.server.server import ssl_context_for_client
from chia.server.ws_connection import WSChiaConnection
from chia.types.peer_info import PeerInfo
from chia.util.errors import Err
from chia.util.ints import uint16, uint64
from tests.time_out_assert import time_out_assert

log = logging.getLogger(__name__)


async def get_block_path(full_node: FullNodeAPI):
    blocks_list = [await full_node.full_node.blockchain.get_full_peak()]
    assert blocks_list[0] is not None
    while blocks_list[0].height != 0:
        b = await full_node.full_node.block_store.get_full_block(blocks_list[0].prev_header_hash)
        assert b is not None
        blocks_list.insert(0, b)
    return blocks_list


class FakeRateLimiter:
    def process_msg_and_check(self, msg, capa, capb):
        return True


class TestDos:
    @pytest.mark.asyncio
    async def test_large_message_disconnect_and_ban(self, setup_two_nodes_fixture, self_hostname):
        nodes, _ = setup_two_nodes_fixture
        server_1 = nodes[0].full_node.server
        server_2 = nodes[1].full_node.server

        # Use the server_2 ssl information to connect to server_1, and send a huge message
        timeout = ClientTimeout(total=10)
        session = ClientSession(timeout=timeout)
        url = f"wss://{self_hostname}:{server_1._port}/ws"

        ssl_context = ssl_context_for_client(
            server_2.chia_ca_crt_path, server_2.chia_ca_key_path, server_2.p2p_crt_path, server_2.p2p_key_path
        )
        ws = await session.ws_connect(
            url, autoclose=True, autoping=True, heartbeat=60, ssl=ssl_context, max_msg_size=100 * 1024 * 1024
        )
        assert not ws.closed
        await ws.close()
        assert ws.closed

        ws = await session.ws_connect(
            url, autoclose=True, autoping=True, heartbeat=60, ssl=ssl_context, max_msg_size=100 * 1024 * 1024
        )
        assert not ws.closed

        large_msg: bytes = bytes([0] * (60 * 1024 * 1024))
        await ws.send_bytes(large_msg)

        response: WSMessage = await ws.receive()
        print(response)
        assert response.type == WSMsgType.CLOSE
        assert response.data == WSCloseCode.MESSAGE_TOO_BIG
        await ws.close()

        # Now test that the ban is active
        await asyncio.sleep(5)
        assert ws.closed
        try:
            ws = await session.ws_connect(
                url, autoclose=True, autoping=True, heartbeat=60, ssl=ssl_context, max_msg_size=100 * 1024 * 1024
            )
            response: WSMessage = await ws.receive()
            assert response.type == WSMsgType.CLOSE
        except ServerDisconnectedError:
            pass
        await session.close()

    @pytest.mark.asyncio
    async def test_bad_handshake_and_ban(self, setup_two_nodes_fixture, self_hostname):
        nodes, _ = setup_two_nodes_fixture
        server_1 = nodes[0].full_node.server
        server_2 = nodes[1].full_node.server

        server_1.invalid_protocol_ban_seconds = 10
        # Use the server_2 ssl information to connect to server_1, and send a huge message
        timeout = ClientTimeout(total=10)
        session = ClientSession(timeout=timeout)
        url = f"wss://{self_hostname}:{server_1._port}/ws"

        ssl_context = ssl_context_for_client(
            server_2.chia_ca_crt_path, server_2.chia_ca_key_path, server_2.p2p_crt_path, server_2.p2p_key_path
        )
        ws = await session.ws_connect(
            url, autoclose=True, autoping=True, heartbeat=60, ssl=ssl_context, max_msg_size=100 * 1024 * 1024
        )
        await ws.send_bytes(bytes([1] * 1024))

        response: WSMessage = await ws.receive()
        print(response)
        assert response.type == WSMsgType.CLOSE
        assert response.data == WSCloseCode.PROTOCOL_ERROR
        await ws.close()

        # Now test that the ban is active
        await asyncio.sleep(5)
        assert ws.closed
        try:
            ws = await session.ws_connect(
                url, autoclose=True, autoping=True, heartbeat=60, ssl=ssl_context, max_msg_size=100 * 1024 * 1024
            )
            response: WSMessage = await ws.receive()
            assert response.type == WSMsgType.CLOSE
        except ServerDisconnectedError:
            pass
        await asyncio.sleep(6)

        # Ban expired
        await session.ws_connect(
            url, autoclose=True, autoping=True, heartbeat=60, ssl=ssl_context, max_msg_size=100 * 1024 * 1024
        )

        await session.close()

    @pytest.mark.asyncio
    async def test_invalid_protocol_handshake(self, setup_two_nodes_fixture, self_hostname):
        nodes, _ = setup_two_nodes_fixture
        server_1 = nodes[0].full_node.server
        server_2 = nodes[1].full_node.server

        server_1.invalid_protocol_ban_seconds = 10
        # Use the server_2 ssl information to connect to server_1
        timeout = ClientTimeout(total=10)
        session = ClientSession(timeout=timeout)
        url = f"wss://{self_hostname}:{server_1._port}/ws"

        ssl_context = ssl_context_for_client(
            server_2.chia_ca_crt_path, server_2.chia_ca_key_path, server_2.p2p_crt_path, server_2.p2p_key_path
        )
        ws = await session.ws_connect(
            url, autoclose=True, autoping=True, heartbeat=60, ssl=ssl_context, max_msg_size=100 * 1024 * 1024
        )

        # Construct an otherwise valid handshake message
        handshake: Handshake = Handshake("test", "0.0.32", "1.0.0.0", 3456, 1, [(1, "1")])
        outbound_handshake: Message = Message(2, None, bytes(handshake))  # 2 is an invalid ProtocolType
        await ws.send_bytes(bytes(outbound_handshake))

        response: WSMessage = await ws.receive()
        print(response)
        assert response.type == WSMsgType.CLOSE
        assert response.data == WSCloseCode.PROTOCOL_ERROR
        assert response.extra == str(int(Err.INVALID_HANDSHAKE.value))  # We want INVALID_HANDSHAKE and not UNKNOWN
        await ws.close()
        await session.close()
        await asyncio.sleep(1)  # give some time for cleanup to work

    @pytest.mark.asyncio
    async def test_spam_tx(self, setup_two_nodes_fixture, self_hostname):
        nodes, _ = setup_two_nodes_fixture
        full_node_1, full_node_2 = nodes
        server_1 = nodes[0].full_node.server
        server_2 = nodes[1].full_node.server

        await server_2.start_client(PeerInfo(self_hostname, uint16(server_1._port)), full_node_2.full_node.on_connect)

        assert len(server_1.all_connections) == 1

        ws_con: WSChiaConnection = list(server_1.all_connections.values())[0]
        ws_con_2: WSChiaConnection = list(server_2.all_connections.values())[0]

        ws_con.peer_host = "1.2.3.4"
        ws_con_2.peer_host = "1.2.3.4"

        new_tx_message = make_msg(
            ProtocolMessageTypes.new_transaction,
            full_node_protocol.NewTransaction(bytes([9] * 32), uint64(0), uint64(0)),
        )
        for i in range(4000):
            await ws_con._send_message(new_tx_message)

        await asyncio.sleep(1)
        assert not ws_con.closed

        # Tests outbound rate limiting, we will not send too much data
        for i in range(2000):
            await ws_con._send_message(new_tx_message)

        await asyncio.sleep(1)
        assert not ws_con.closed

        # Remove outbound rate limiter to test inbound limits
        ws_con.outbound_rate_limiter = RateLimiter(incoming=True, percentage_of_limit=10000)

        with pytest.raises(ConnectionResetError):
            for i in range(6000):
                await ws_con._send_message(new_tx_message)
                await asyncio.sleep(0)
        await asyncio.sleep(1)

        def is_closed():
            return ws_con.closed

        await time_out_assert(15, is_closed)

        assert ws_con.closed

        def is_banned():
            return "1.2.3.4" in server_2.banned_peers

        await time_out_assert(15, is_banned)

    @pytest.mark.asyncio
    async def test_spam_message_non_tx(self, setup_two_nodes_fixture, self_hostname):
        nodes, _ = setup_two_nodes_fixture
        full_node_1, full_node_2 = nodes
        server_1 = nodes[0].full_node.server
        server_2 = nodes[1].full_node.server

        await server_2.start_client(PeerInfo(self_hostname, uint16(server_1._port)), full_node_2.full_node.on_connect)

        assert len(server_1.all_connections) == 1

        ws_con: WSChiaConnection = list(server_1.all_connections.values())[0]
        ws_con_2: WSChiaConnection = list(server_2.all_connections.values())[0]

        ws_con.peer_host = "1.2.3.4"
        ws_con_2.peer_host = "1.2.3.4"

        def is_closed():
            return ws_con.closed

        new_message = make_msg(
            ProtocolMessageTypes.request_mempool_transactions,
            full_node_protocol.RequestMempoolTransactions(bytes([])),
        )
        for i in range(2):
            await ws_con._send_message(new_message)
        await asyncio.sleep(1)
        assert not ws_con.closed

        # Tests outbound rate limiting, we will not send too much data
        for i in range(10):
            await ws_con._send_message(new_message)

        await asyncio.sleep(1)
        assert not ws_con.closed

        # Remove outbound rate limiter to test inbound limits
        ws_con.outbound_rate_limiter = RateLimiter(incoming=True, percentage_of_limit=10000)

        for i in range(6):
            await ws_con._send_message(new_message)
        await time_out_assert(15, is_closed)

        # Banned
        def is_banned():
            return "1.2.3.4" in server_2.banned_peers

        await time_out_assert(15, is_banned)

    @pytest.mark.asyncio
    async def test_spam_message_too_large(self, setup_two_nodes_fixture, self_hostname):
        nodes, _ = setup_two_nodes_fixture
        full_node_1, full_node_2 = nodes
        server_1 = nodes[0].full_node.server
        server_2 = nodes[1].full_node.server

        await server_2.start_client(PeerInfo(self_hostname, uint16(server_1._port)), full_node_2.full_node.on_connect)

        assert len(server_1.all_connections) == 1

        ws_con: WSChiaConnection = list(server_1.all_connections.values())[0]
        ws_con_2: WSChiaConnection = list(server_2.all_connections.values())[0]

        ws_con.peer_host = "1.2.3.4"
        ws_con_2.peer_host = "1.2.3.4"

        def is_closed():
            return ws_con.closed

        new_message = make_msg(
            ProtocolMessageTypes.request_mempool_transactions,
            full_node_protocol.RequestMempoolTransactions(bytes([0] * 5 * 1024 * 1024)),
        )
        # Tests outbound rate limiting, we will not send big messages
        await ws_con._send_message(new_message)

        await asyncio.sleep(1)
        assert not ws_con.closed

        # Remove outbound rate limiter to test inbound limits
        ws_con.outbound_rate_limiter = FakeRateLimiter()

        await ws_con._send_message(new_message)
        await time_out_assert(15, is_closed)

        # Banned
        def is_banned():
            return "1.2.3.4" in server_2.banned_peers

        await time_out_assert(15, is_banned)