diff --git a/server/__init__.py b/server/__init__.py index 447b81614..bf6a5c40a 100644 --- a/server/__init__.py +++ b/server/__init__.py @@ -118,6 +118,7 @@ from .asyncio_extensions import map_suppress, synchronizedmethod from .broadcast_service import BroadcastService +from .client_message_queue_service import ClientMessageQueueService from .config import TRACE, config from .configuration_service import ConfigurationService from .core import Service, create_services @@ -144,6 +145,7 @@ __all__ = ( "BroadcastService", + "ClientMessageQueueService", "ConfigurationService", "GameConnection", "GameService", diff --git a/server/client_message_queue_service.py b/server/client_message_queue_service.py new file mode 100644 index 000000000..9f4bfb607 --- /dev/null +++ b/server/client_message_queue_service.py @@ -0,0 +1,133 @@ +""" +Forward RabbitMQ messages from trusted microservices to connected clients. + +# Wire contract +Publishers post to the `MQ_EXCHANGE_NAME` topic exchange with routing key +`request.client.notify`. Addressing lives in AMQP message headers: + +- `user-id` (int, optional): forward the body to the player with this id, if + connected to this lobby instance. If not connected, the message is logged + and acked. +- `channel` (str, optional, reserved): future per-channel pub/sub. Currently + recognised but not yet implemented. +- If neither header is set, the message is broadcast to every authenticated + client connected to this instance. + +The message body is a UTF-8 JSON object and is forwarded to the client +verbatim. The lobby server does not validate or rewrap it; producers are +trusted because the broker is reachable only from internal services. +""" + +import json +import logging +import socket +from typing import TYPE_CHECKING, Any, ClassVar, Optional + +from aio_pika.abc import AbstractIncomingMessage, AbstractQueue + +from .config import config +from .core import Service +from .decorators import with_logger +from .message_queue_service import MessageQueueService +from .player_service import PlayerService + +if TYPE_CHECKING: + from server import ServerInstance + + +CLIENT_NOTIFY_ROUTING_KEY = "request.client.notify" + + +@with_logger +class ClientMessageQueueService(Service): + """Consume `request.client.notify` messages and forward to local clients.""" + + _logger: ClassVar[logging.Logger] + + def __init__( + self, + server: "ServerInstance", + message_queue_service: MessageQueueService, + player_service: PlayerService, + ): + """Wire dependencies; consumer is started in `initialize`.""" + self.server = server + self.message_queue_service = message_queue_service + self.player_service = player_service + self._queue: Optional[AbstractQueue] = None + self._consumer_tag: Optional[str] = None + + async def initialize(self) -> None: + # Queue naming follows `..` plus a + # per-instance suffix because each lobby pod has its own queue (vs. + # the API's shared queues like `faf-lobby.api.event.update`). On k8s + # `socket.gethostname()` resolves to the pod name (e.g. + # `faf-lobby-server-6d9c4588ff-lzdcr`); locally it's the dev's host. + queue_name = ( + f"{config.MQ_EXCHANGE_NAME}.lobby.client.notify" + f".{socket.gethostname()}" + ) + result = await self.message_queue_service.declare_queue_and_consume( + exchange_name=config.MQ_EXCHANGE_NAME, + routing_key=CLIENT_NOTIFY_ROUTING_KEY, + callback=self._on_message, + queue_name=queue_name, + ) + if result is not None: + self._queue, self._consumer_tag = result + + async def shutdown(self) -> None: + if self._queue is not None and self._consumer_tag is not None: + await self._queue.cancel(self._consumer_tag) + self._queue = None + self._consumer_tag = None + + async def _on_message(self, message: AbstractIncomingMessage) -> None: + async with message.process(requeue=False): + try: + payload = json.loads(message.body) + except (ValueError, UnicodeDecodeError): + self._logger.warning( + "Dropping client-notify message with non-JSON body" + ) + return + + if not isinstance(payload, dict): + self._logger.warning( + "Dropping client-notify message: payload is not a JSON object" + ) + return + + headers = message.headers or {} + user_id = headers.get("user-id") + channel = headers.get("channel") + + if user_id is not None: + self._dispatch_to_user(user_id, payload) + elif channel is not None: + self._logger.info( + "client-notify channel %r received but channel routing is " + "not yet implemented; dropping", + channel, + ) + else: + self.server.write_broadcast(payload) + + def _dispatch_to_user(self, user_id: Any, payload: dict) -> None: + try: + player_id = int(user_id) + except (TypeError, ValueError): + self._logger.warning( + "Dropping client-notify message: invalid user-id %r", user_id + ) + return + + player = self.player_service[player_id] + if player is None: + self._logger.warning( + "Dropping client-notify message: user %s not connected here", + player_id, + ) + return + + player.write_message(payload) diff --git a/server/message_queue_service.py b/server/message_queue_service.py index 458e164cd..f346e4c19 100644 --- a/server/message_queue_service.py +++ b/server/message_queue_service.py @@ -5,11 +5,17 @@ import asyncio import json import logging -from typing import ClassVar, Iterable, Optional +from typing import Awaitable, Callable, ClassVar, Iterable, Optional import aio_pika from aio_pika import DeliveryMode, ExchangeType -from aio_pika.abc import AbstractChannel, AbstractConnection, AbstractExchange +from aio_pika.abc import ( + AbstractChannel, + AbstractConnection, + AbstractExchange, + AbstractIncomingMessage, + AbstractQueue +) from aio_pika.exceptions import ProbableAuthenticationError from .asyncio_extensions import synchronizedmethod @@ -187,6 +193,49 @@ async def publish_many( routing ) + async def declare_queue_and_consume( + self, + exchange_name: str, + routing_key: str, + callback: Callable[[AbstractIncomingMessage], Awaitable[None]], + queue_name: str = "", + exclusive: bool = True, + auto_delete: bool = True, + durable: bool = False, + ) -> Optional[tuple[AbstractQueue, str]]: + """ + Declare a queue, bind it to an exchange with the given routing key, and + start consuming. Returns `(queue, consumer_tag)` so the caller can + cancel on shutdown. Returns None if the broker connection is not ready. + """ + await self.initialize() + if not self._is_ready: + self._logger.warning( + "Not connected to RabbitMQ, unable to declare consumer queue." + ) + return None + + assert self._channel is not None + + exchange = self._exchanges.get(exchange_name) + if exchange is None: + raise KeyError(f"Unknown exchange {exchange_name}.") + + queue = await self._channel.declare_queue( + queue_name, + exclusive=exclusive, + auto_delete=auto_delete, + durable=durable, + ) + await queue.bind(exchange, routing_key=routing_key) + consumer_tag = await queue.consume(callback) + + self._logger.debug( + "Consuming from queue %r bound to %s/%s", + queue.name, exchange_name, routing_key, + ) + return queue, consumer_tag + @synchronizedmethod("initialization_lock") async def reconnect(self) -> None: self._is_ready = False diff --git a/tests/unit_tests/test_client_message_queue_service.py b/tests/unit_tests/test_client_message_queue_service.py new file mode 100644 index 000000000..23141007c --- /dev/null +++ b/tests/unit_tests/test_client_message_queue_service.py @@ -0,0 +1,202 @@ +import json +from unittest import mock + +import pytest + +from server import ServerInstance +from server.client_message_queue_service import ( + CLIENT_NOTIFY_ROUTING_KEY, + ClientMessageQueueService +) +from server.config import config + + +def make_incoming_message(body: bytes, headers: dict | None = None): + """Build a stand-in for aio_pika's IncomingMessage.""" + message = mock.Mock() + message.body = body + message.headers = headers + + process_cm = mock.MagicMock() + process_cm.__aenter__ = mock.AsyncMock(return_value=None) + process_cm.__aexit__ = mock.AsyncMock(return_value=False) + message.process = mock.Mock(return_value=process_cm) + return message + + +@pytest.fixture +def server_instance(): + return mock.create_autospec(ServerInstance) + + +@pytest.fixture +def fake_player_service(): + """Stand in for PlayerService supporting __getitem__/__setitem__.""" + class _FakePlayerService: + def __init__(self): + self._players = {} + + def __getitem__(self, player_id): + return self._players.get(player_id) + + def __setitem__(self, player_id, player): + self._players[player_id] = player + + return _FakePlayerService() + + +@pytest.fixture +async def client_message_queue_service(server_instance, fake_player_service): + queue = mock.Mock() + queue.cancel = mock.AsyncMock() + mq_service = mock.Mock() + mq_service.declare_queue_and_consume = mock.AsyncMock( + return_value=(queue, "consumer-tag-123") + ) + service = ClientMessageQueueService( + server=server_instance, + message_queue_service=mq_service, + player_service=fake_player_service, + ) + await service.initialize() + yield service + await service.shutdown() + + +async def test_shutdown_cancels_consumer( + server_instance, fake_player_service +): + queue = mock.Mock() + queue.cancel = mock.AsyncMock() + mq_service = mock.Mock() + mq_service.declare_queue_and_consume = mock.AsyncMock( + return_value=(queue, "consumer-tag-xyz") + ) + service = ClientMessageQueueService( + server=server_instance, + message_queue_service=mq_service, + player_service=fake_player_service, + ) + await service.initialize() + await service.shutdown() + + queue.cancel.assert_awaited_once_with("consumer-tag-xyz") + assert service._queue is None + assert service._consumer_tag is None + + +async def test_shutdown_noop_when_broker_unavailable( + server_instance, fake_player_service +): + mq_service = mock.Mock() + mq_service.declare_queue_and_consume = mock.AsyncMock(return_value=None) + service = ClientMessageQueueService( + server=server_instance, + message_queue_service=mq_service, + player_service=fake_player_service, + ) + await service.initialize() + # Should not raise even though no queue was ever created. + await service.shutdown() + + +async def test_initialize_declares_consumer(client_message_queue_service): + mq = client_message_queue_service.message_queue_service + mq.declare_queue_and_consume.assert_awaited_once() + kwargs = mq.declare_queue_and_consume.await_args.kwargs + assert kwargs["exchange_name"] == config.MQ_EXCHANGE_NAME + assert kwargs["routing_key"] == CLIENT_NOTIFY_ROUTING_KEY + assert kwargs["callback"] == client_message_queue_service._on_message + + +async def test_dispatch_to_connected_user( + client_message_queue_service, fake_player_service +): + player = mock.Mock() + player.write_message = mock.Mock() + fake_player_service[42] = player + + payload = {"command": "notice", "text": "hi"} + msg = make_incoming_message(json.dumps(payload).encode(), {"user-id": 42}) + + await client_message_queue_service._on_message(msg) + + player.write_message.assert_called_once_with(payload) + client_message_queue_service.server.write_broadcast.assert_not_called() + + +async def test_dispatch_to_disconnected_user_is_dropped( + client_message_queue_service, caplog +): + payload = {"command": "notice"} + msg = make_incoming_message(json.dumps(payload).encode(), {"user-id": 999}) + + await client_message_queue_service._on_message(msg) + + client_message_queue_service.server.write_broadcast.assert_not_called() + assert any("not connected here" in m for m in caplog.messages) + + +async def test_broadcast_when_no_user_id_header(client_message_queue_service): + payload = {"command": "announcement", "text": "hello world"} + msg = make_incoming_message(json.dumps(payload).encode(), headers=None) + + await client_message_queue_service._on_message(msg) + + client_message_queue_service.server.write_broadcast.assert_called_once_with(payload) + + +async def test_broadcast_when_headers_empty(client_message_queue_service): + payload = {"command": "announcement"} + msg = make_incoming_message(json.dumps(payload).encode(), headers={}) + + await client_message_queue_service._on_message(msg) + + client_message_queue_service.server.write_broadcast.assert_called_once_with(payload) + + +async def test_malformed_json_body_is_dropped( + client_message_queue_service, caplog +): + msg = make_incoming_message(b"not json at all", headers=None) + + await client_message_queue_service._on_message(msg) + + client_message_queue_service.server.write_broadcast.assert_not_called() + assert any("non-JSON body" in m for m in caplog.messages) + + +async def test_non_object_json_body_is_dropped(client_message_queue_service): + msg = make_incoming_message(b"[1, 2, 3]", headers=None) + + await client_message_queue_service._on_message(msg) + + client_message_queue_service.server.write_broadcast.assert_not_called() + + +async def test_invalid_user_id_header_is_dropped( + client_message_queue_service, caplog +): + msg = make_incoming_message( + json.dumps({"command": "x"}).encode(), + headers={"user-id": "not-an-int"}, + ) + + await client_message_queue_service._on_message(msg) + + client_message_queue_service.server.write_broadcast.assert_not_called() + assert any("invalid user-id" in m for m in caplog.messages) + + +async def test_channel_header_currently_drops( + client_message_queue_service, caplog +): + msg = make_incoming_message( + json.dumps({"command": "x"}).encode(), + headers={"channel": "matchmaker"}, + ) + + await client_message_queue_service._on_message(msg) + + client_message_queue_service.server.write_broadcast.assert_not_called() + assert any("channel routing is not yet implemented" in m for m in caplog.messages) diff --git a/tests/unit_tests/test_message_queue_service.py b/tests/unit_tests/test_message_queue_service.py index b285771eb..edd2aadc8 100644 --- a/tests/unit_tests/test_message_queue_service.py +++ b/tests/unit_tests/test_message_queue_service.py @@ -44,6 +44,67 @@ async def test_incorrect_port(mocker, caplog): assert expected_warning in caplog.messages +async def test_declare_queue_and_consume(mq_service): + received = asyncio.Event() + captured: dict = {} + + async def callback(message): + async with message.process(): + captured["body"] = message.body + captured["headers"] = dict(message.headers or {}) + received.set() + + result = await mq_service.declare_queue_and_consume( + exchange_name="test_exchange", + routing_key="consume.test", + callback=callback, + ) + assert result is not None + queue, consumer_tag = result + assert consumer_tag + + await mq_service.publish( + "test_exchange", + "consume.test", + {"hello": "world"}, + delivery_mode=aio_pika.DeliveryMode.NOT_PERSISTENT, + ) + + await asyncio.wait_for(received.wait(), timeout=5) + assert captured["body"] == b'{"hello": "world"}' + + await queue.cancel(consumer_tag) + + +async def test_declare_queue_and_consume_unknown_exchange(mq_service): + async def callback(_message): + pass + + with pytest.raises(KeyError): + await mq_service.declare_queue_and_consume( + exchange_name="not_declared", + routing_key="anything", + callback=callback, + ) + + +async def test_declare_queue_and_consume_not_ready(mocker, caplog): + from server.message_queue_service import ConnectionAttemptFailed + + service = MessageQueueService() + service._connect = mock.AsyncMock(side_effect=ConnectionAttemptFailed) + + async def callback(_message): + pass + + result = await service.declare_queue_and_consume( + exchange_name="test_exchange", + routing_key="x", + callback=callback, + ) + assert result is None + + async def test_several_initializations_connect_only_once(): service = MessageQueueService()