"""Tests for core multiplexer handler.""" import asyncio from contextlib import suppress import ipaddress import os from unittest.mock import patch import pytest import snitun from snitun.exceptions import MultiplexerTransportClose, MultiplexerTransportError from snitun.multiplexer import channel as channel_module, core as core_module from snitun.multiplexer.channel import MultiplexerChannel from snitun.multiplexer.core import Multiplexer from snitun.multiplexer.crypto import CryptoTransport from snitun.multiplexer.message import ( CHANNEL_FLOW_PAUSE, CHANNEL_FLOW_PING, HEADER_SIZE, MultiplexerChannelId, MultiplexerMessage, ) from ..conftest import Client IP_ADDR = ipaddress.ip_address("8.8.8.8") async def test_init_multiplexer_server( test_server: list[Client], test_client: Client, crypto_key_iv: tuple[bytes, bytes], ) -> None: """Test to create a new Multiplexer from server socket.""" client = test_server[0] multiplexer = Multiplexer( CryptoTransport(*crypto_key_iv), client.reader, client.writer, snitun.PROTOCOL_VERSION, ) assert multiplexer.is_connected assert multiplexer._throttling is None multiplexer.shutdown() client.close.set() async def test_init_multiplexer_client( test_client: Client, crypto_key_iv: tuple[bytes, bytes], ) -> None: """Test to create a new Multiplexer from client socket.""" multiplexer = Multiplexer( CryptoTransport(*crypto_key_iv), test_client.reader, test_client.writer, snitun.PROTOCOL_VERSION, ) assert multiplexer.is_connected assert multiplexer._throttling is None multiplexer.shutdown() async def test_init_multiplexer_server_throttling( test_server: list[Client], test_client: Client, crypto_key_iv: tuple[bytes, bytes], ) -> None: """Test to create a new Multiplexer from server socket.""" client = test_server[0] multiplexer = Multiplexer( CryptoTransport(*crypto_key_iv), client.reader, client.writer, snitun.PROTOCOL_VERSION, throttling=500, ) assert multiplexer.is_connected assert multiplexer._throttling == 0.002 multiplexer.shutdown() client.close.set() async def test_init_multiplexer_client_throttling( test_client: Client, crypto_key_iv: tuple[bytes, bytes], ) -> None: """Test to create a new Multiplexer from client socket.""" multiplexer = Multiplexer( CryptoTransport(*crypto_key_iv), test_client.reader, test_client.writer, snitun.PROTOCOL_VERSION, throttling=500, ) assert multiplexer.is_connected assert multiplexer._throttling == 0.002 multiplexer.shutdown() async def test_multiplexer_server_close( multiplexer_server: Multiplexer, multiplexer_client: Multiplexer, ) -> None: """Test a close from server peers.""" assert multiplexer_server.is_connected assert multiplexer_client.is_connected multiplexer_server.shutdown() await asyncio.sleep(0.1) assert not multiplexer_server.is_connected assert not multiplexer_client.is_connected async def test_multiplexer_client_close( multiplexer_server: Multiplexer, multiplexer_client: Multiplexer, ) -> None: """Test a close from client peers.""" assert multiplexer_server.is_connected assert multiplexer_client.is_connected multiplexer_client.shutdown() await asyncio.sleep(0.1) assert not multiplexer_server.is_connected assert not multiplexer_client.is_connected async def test_multiplexer_ping( test_server: list[Client], multiplexer_client: Multiplexer, ) -> None: """Test a ping between peers.""" loop = asyncio.get_running_loop() client = test_server[0] ping_task = loop.create_task(multiplexer_client.ping()) await asyncio.sleep(0.1) data = await client.reader.read(60) data = multiplexer_client._crypto.decrypt(data) assert data[16] == CHANNEL_FLOW_PING assert int.from_bytes(data[17:21], "big") == 0 assert data[21:25] == b"ping" ping_task.cancel() async def test_multiplexer_ping_error( test_server: list[Client], multiplexer_client: Multiplexer, ) -> None: """Test a ping between peers.""" from snitun.multiplexer import core as multi_core loop = asyncio.get_running_loop() multi_core.PEER_TCP_TIMEOUT = 0.2 client = test_server[0] ping_task = loop.create_task(multiplexer_client.ping()) await asyncio.sleep(0.3) data = await client.reader.read(60) data = multiplexer_client._crypto.decrypt(data) assert data[16] == CHANNEL_FLOW_PING assert int.from_bytes(data[17:21], "big") == 0 assert data[21:25] == b"ping" assert ping_task.done() with pytest.raises(MultiplexerTransportError): raise ping_task.exception() multi_core.PEER_TCP_TIMEOUT = 90 async def test_multiplexer_ping_pong( multiplexer_client: Multiplexer, multiplexer_server: Multiplexer, ) -> None: """Test that without new channel callback can't create new channels.""" await multiplexer_client.ping() assert multiplexer_client._healthy.is_set() async def test_multiplexer_cant_init_channel( multiplexer_client: Multiplexer, multiplexer_server: Multiplexer, ) -> None: """Test that without new channel callback can't create new channels.""" assert not multiplexer_client._channels assert not multiplexer_server._channels # Disable new channels multiplexer_server._new_connections = None await multiplexer_client.create_channel(IP_ADDR, lambda _: None) await asyncio.sleep(0.1) assert multiplexer_client._channels assert not multiplexer_server._channels async def test_multiplexer_init_channel( multiplexer_client: Multiplexer, multiplexer_server: Multiplexer, ) -> None: """Test that new channels are created.""" assert not multiplexer_client._channels assert not multiplexer_server._channels channel = await multiplexer_client.create_channel(IP_ADDR, lambda _: None) await asyncio.sleep(0.1) assert multiplexer_client._channels assert multiplexer_server._channels assert multiplexer_client._channels[channel.id] assert multiplexer_server._channels[channel.id] assert multiplexer_client._channels[channel.id].ip_address == IP_ADDR assert multiplexer_server._channels[channel.id].ip_address == IP_ADDR async def test_multiplexer_init_channel_full( multiplexer_client: Multiplexer, raise_timeout: None, ) -> None: """Test that new channels are created but peer error is available.""" assert not multiplexer_client._channels with pytest.raises(MultiplexerTransportError): await multiplexer_client.create_channel(IP_ADDR, lambda _: None) await asyncio.sleep(0.1) assert not multiplexer_client._channels async def test_multiplexer_close_channel( multiplexer_client: Multiplexer, multiplexer_server: Multiplexer, ) -> None: """Test that channels are nice removed.""" assert not multiplexer_client._channels assert not multiplexer_server._channels channel = await multiplexer_client.create_channel(IP_ADDR, lambda _: None) await asyncio.sleep(0.1) assert multiplexer_client._channels assert multiplexer_server._channels assert multiplexer_client._channels[channel.id] assert multiplexer_server._channels[channel.id] assert multiplexer_client._channels[channel.id].ip_address == IP_ADDR assert multiplexer_server._channels[channel.id].ip_address == IP_ADDR multiplexer_client.delete_channel(channel) await asyncio.sleep(0.1) assert not multiplexer_client._channels assert not multiplexer_server._channels async def test_multiplexer_delete_unknown_channel( multiplexer_client: Multiplexer, multiplexer_server: Multiplexer, caplog: pytest.LogCaptureFixture, ) -> None: """Test deleting an unknown channel.""" assert not multiplexer_client._channels assert not multiplexer_server._channels non_existant_channel = MultiplexerChannel( multiplexer_server._queue, ipaddress.IPv4Address("127.0.0.1"), snitun.PROTOCOL_VERSION, ) await multiplexer_server._queue.put( non_existant_channel.id, non_existant_channel.init_close(), ) await asyncio.sleep(0.1) assert not multiplexer_client._channels assert not multiplexer_server._channels assert ( f"Receive close from unknown channel: {non_existant_channel.id}" in caplog.text ) async def test_multiplexer_delete_channel_called_multiple_times( multiplexer_client: Multiplexer, multiplexer_server: Multiplexer, ) -> None: """Test that channels can be deleted twice.""" assert not multiplexer_client._channels assert not multiplexer_server._channels channel = await multiplexer_client.create_channel(IP_ADDR, lambda _: None) await asyncio.sleep(0.1) assert multiplexer_client._channels assert multiplexer_server._channels assert multiplexer_client._channels[channel.id] assert multiplexer_server._channels[channel.id] assert multiplexer_client._channels[channel.id].ip_address == IP_ADDR assert multiplexer_server._channels[channel.id].ip_address == IP_ADDR multiplexer_client.delete_channel(channel) assert not multiplexer_client._channels multiplexer_client.delete_channel(channel) assert not multiplexer_client._channels await asyncio.sleep(0.1) assert not multiplexer_server._channels async def test_multiplexer_close_channel_full(multiplexer_client: Multiplexer) -> None: """Test that channels are nice removed but peer error is available.""" assert not multiplexer_client._channels channel = await multiplexer_client.create_channel(IP_ADDR, lambda _: None) await asyncio.sleep(0.1) assert multiplexer_client._channels multiplexer_client.delete_channel(channel) await asyncio.sleep(0.1) assert not multiplexer_client._channels async def test_multiplexer_data_channel( multiplexer_client: Multiplexer, multiplexer_server: Multiplexer, ) -> None: """Test that new channels are created.""" assert not multiplexer_client._channels assert not multiplexer_server._channels channel_client = await multiplexer_client.create_channel(IP_ADDR, lambda _: None) await asyncio.sleep(0.1) channel_server = multiplexer_server._channels.get(channel_client.id) assert channel_client assert channel_server await channel_client.write(b"test 1") await asyncio.sleep(0.1) data = await channel_server.read() assert data == b"test 1" await channel_server.write(b"test 2") await asyncio.sleep(0.1) data = await channel_client.read() assert data == b"test 2" async def test_multiplexer_channel_shutdown( multiplexer_client: Multiplexer, multiplexer_server: Multiplexer, ) -> None: """Test that new channels are created and graceful shutdown.""" loop = asyncio.get_running_loop() assert not multiplexer_client._channels assert not multiplexer_server._channels channel_client = await multiplexer_client.create_channel(IP_ADDR, lambda _: None) await asyncio.sleep(0.1) channel_server = multiplexer_server._channels.get(channel_client.id) client_read = loop.create_task(channel_client.read()) server_read = loop.create_task(channel_server.read()) assert not client_read.done() assert not server_read.done() multiplexer_client.shutdown() await asyncio.sleep(0.1) assert not multiplexer_client._channels assert client_read.done() with pytest.raises(MultiplexerTransportClose): raise client_read.exception() assert not multiplexer_server._channels assert server_read.done() with pytest.raises(MultiplexerTransportClose): raise server_read.exception() @patch.object(channel_module, "INCOMING_QUEUE_MAX_BYTES_CHANNEL", 1) @patch.object(core_module, "OUTGOING_QUEUE_MAX_BYTES_CHANNEL", 1) async def test_multiplexer_data_channel_abort_full( multiplexer_client: Multiplexer, multiplexer_server: Multiplexer, ) -> None: """Test that new channels are created.""" assert not multiplexer_client._channels assert not multiplexer_server._channels channel_client = await multiplexer_client.create_channel(IP_ADDR, lambda _: None) await asyncio.sleep(0.1) channel_server = multiplexer_server._channels.get(channel_client.id) assert channel_client assert channel_server large_msg = b"test xxxx" * 1000 with pytest.raises(MultiplexerTransportClose): for _ in range(1, 50000): await channel_client.write(large_msg) with pytest.raises(MultiplexerTransportClose): for _ in range(1, 50000): await channel_server.read() await asyncio.sleep(0.1) assert not multiplexer_client._channels assert not multiplexer_server._channels async def test_multiplexer_throttling( multiplexer_client: Multiplexer, multiplexer_server: Multiplexer, ) -> None: """Test that new channels are created and graceful shutdown.""" loop = asyncio.get_running_loop() assert not multiplexer_client._channels assert not multiplexer_server._channels data_in = [] channel_client = await multiplexer_client.create_channel(IP_ADDR, lambda _: None) await asyncio.sleep(0.1) channel_server = multiplexer_server._channels.get(channel_client.id) multiplexer_server._throttling = 0.1 multiplexer_client._throttling = 0.1 async def _sender() -> None: """Send data much as possible.""" for _ in range(1, 100000): await channel_client.write(b"data") async def _receiver() -> None: """Receive data much as possible.""" for _ in range(1, 100000): data = await channel_server.read() data_in.append(data) receiver = loop.create_task(_receiver()) sender = loop.create_task(_sender()) await asyncio.sleep(0.8) assert not receiver.done() assert len(data_in) <= 8 receiver.cancel() sender.cancel() with suppress(asyncio.CancelledError): await receiver with suppress(asyncio.CancelledError): await sender async def test_multiplexer_core_peer_timeout( multiplexer_client: Multiplexer, multiplexer_server: Multiplexer, ) -> None: """Test that new channels are created and graceful shutdown.""" from snitun.multiplexer import core as multi_core loop = asyncio.get_running_loop() multi_core.PEER_TCP_TIMEOUT = 0.2 assert not multiplexer_client._channels assert not multiplexer_server._channels channel_client = await multiplexer_client.create_channel(IP_ADDR, lambda _: None) await asyncio.sleep(0.1) channel_server = multiplexer_server._channels.get(channel_client.id) client_read = loop.create_task(channel_client.read()) server_read = loop.create_task(channel_server.read()) assert not client_read.done() assert not server_read.done() await multiplexer_client.ping() await asyncio.sleep(0.3) assert not multiplexer_client._channels assert not multiplexer_server._channels assert server_read.done() assert client_read.done() with pytest.raises(MultiplexerTransportClose): raise server_read.exception() with pytest.raises(MultiplexerTransportClose): raise client_read.exception() multi_core.PEER_TCP_TIMEOUT = 90 @patch.object(channel_module, "INCOMING_QUEUE_LOW_WATERMARK", HEADER_SIZE * 2) @patch.object(channel_module, "INCOMING_QUEUE_HIGH_WATERMARK", HEADER_SIZE * 3) async def test_remote_input_queue_goes_under_water( multiplexer_client: Multiplexer, multiplexer_server: Multiplexer, ) -> None: """Test the remote input queue going under water.""" assert not multiplexer_client._channels assert not multiplexer_server._channels client_channel_under_water: list[bool] = [] server_channel_under_water: list[bool] = [] def _on_client_channel_under_water(under_water: bool) -> None: client_channel_under_water.append(under_water) def _on_server_channel_under_water(under_water: bool) -> None: server_channel_under_water.append(under_water) channel_client = await multiplexer_client.create_channel( IP_ADDR, _on_client_channel_under_water, ) await asyncio.sleep(0.1) channel_server = multiplexer_server._channels.get(channel_client.id) channel_server.set_pause_resume_reader_callback(_on_server_channel_under_water) assert channel_client assert channel_server sent_messages: list[bytes] = [] message_count = 255 for i in range(message_count): payload = str(i).encode() sent_messages.append(payload) await channel_client.write(payload) await asyncio.sleep(0.1) assert client_channel_under_water == [True] assert server_channel_under_water == [] for i in range(message_count): data = await channel_server.read() assert data == sent_messages[i] await asyncio.sleep(0.1) assert client_channel_under_water == [True, False] assert server_channel_under_water == [] @patch.object(channel_module, "INCOMING_QUEUE_LOW_WATERMARK", HEADER_SIZE * 2) @patch.object(channel_module, "INCOMING_QUEUE_HIGH_WATERMARK", HEADER_SIZE * 3) async def test_remote_input_queue_goes_under_water_protocol_version_0( multiplexer_client: Multiplexer, multiplexer_server_peer_protocol_0: Multiplexer, ) -> None: """Test the remote input queue going under water with client protocol 0. Protocol 0 has no flow control. """ assert not multiplexer_client._channels assert not multiplexer_server_peer_protocol_0._channels client_channel_under_water: list[bool] = [] server_channel_under_water: list[bool] = [] def _on_client_channel_under_water(under_water: bool) -> None: client_channel_under_water.append(under_water) def _on_server_channel_under_water(under_water: bool) -> None: server_channel_under_water.append(under_water) channel_client = await multiplexer_client.create_channel( IP_ADDR, _on_client_channel_under_water, ) await asyncio.sleep(0.1) channel_server = multiplexer_server_peer_protocol_0._channels.get(channel_client.id) channel_server.set_pause_resume_reader_callback(_on_server_channel_under_water) assert channel_client assert channel_server sent_messages: list[bytes] = [] message_count = 255 for i in range(message_count): payload = str(i).encode() sent_messages.append(payload) await channel_client.write(payload) await asyncio.sleep(0.1) # No flow control for protocol 0 assert client_channel_under_water == [] assert server_channel_under_water == [] for i in range(message_count): data = await channel_server.read() assert data == sent_messages[i] await asyncio.sleep(0.1) # No flow control for protocol 0 assert client_channel_under_water == [] assert server_channel_under_water == [] async def test_sending_unknown_message_type( multiplexer_client: Multiplexer, multiplexer_server: Multiplexer, caplog: pytest.LogCaptureFixture, ) -> None: """Test that new channels are created.""" assert not multiplexer_client._channels assert not multiplexer_server._channels channel_client = await multiplexer_client.create_channel( IP_ADDR, lambda _: None, ) await asyncio.sleep(0.1) channel_server = multiplexer_server._channels.get(channel_client.id) assert channel_client assert channel_server channel_client._output.put_nowait( channel_client.id, MultiplexerMessage(channel_client.id, 255), ) await asyncio.sleep(0.1) assert "Receive unknown message type: 255" in caplog.text async def test_sending_pause_for_unknown_channel( multiplexer_client: Multiplexer, multiplexer_server: Multiplexer, caplog: pytest.LogCaptureFixture, ) -> None: """Test sending pause for unknown channel is logged.""" assert not multiplexer_client._channels assert not multiplexer_server._channels channel_client = await multiplexer_client.create_channel( IP_ADDR, lambda _: None, ) await asyncio.sleep(0.1) channel_server = multiplexer_server._channels.get(channel_client.id) assert channel_client assert channel_server wrong_channel_id = MultiplexerChannelId(os.urandom(16)) channel_client._output.put_nowait( channel_client.id, MultiplexerMessage(wrong_channel_id, CHANNEL_FLOW_PAUSE), ) await asyncio.sleep(0.1) assert ( f"Receive pause from unknown channel: {wrong_channel_id.hex()}" in caplog.text )