From 91375cdff28e96a97d334547322771baad0741a5 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Thu, 29 Sep 2022 12:21:28 -0400 Subject: [PATCH] Island: Register a thread-safe IAgentEventQueue --- monkey/monkey_island/cc/services/initialize.py | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) diff --git a/monkey/monkey_island/cc/services/initialize.py b/monkey/monkey_island/cc/services/initialize.py index f1b86700e..79b29940f 100644 --- a/monkey/monkey_island/cc/services/initialize.py +++ b/monkey/monkey_island/cc/services/initialize.py @@ -15,7 +15,11 @@ from common.agent_event_serializers import ( register_common_agent_event_serializers, ) from common.aws import AWSInstance -from common.event_queue import IAgentEventQueue, PyPubSubAgentEventQueue +from common.event_queue import ( + IAgentEventQueue, + LockingAgentEventQueueDecorator, + PyPubSubAgentEventQueue, +) from common.utils.file_utils import get_binary_io_sha256_hash from monkey_island.cc.event_queue import ( IIslandEventQueue, @@ -104,11 +108,17 @@ def _register_conventions(container: DIContainer): def _register_event_queues(container: DIContainer): - container.register_instance(IAgentEventQueue, container.resolve(PyPubSubAgentEventQueue)) + agent_event_queue = container.resolve(PyPubSubAgentEventQueue) + container.register_instance(IAgentEventQueue, _decorate_agent_event_queue(agent_event_queue)) + island_event_queue = container.resolve(PyPubSubIslandEventQueue) container.register_instance(IIslandEventQueue, _decorate_island_event_queue(island_event_queue)) +def _decorate_agent_event_queue(agent_event_queue: IAgentEventQueue): + return LockingAgentEventQueueDecorator(agent_event_queue) + + def _decorate_island_event_queue(island_event_queue: IIslandEventQueue): return LockingIslandEventQueueDecorator(island_event_queue)