From 48cded4c7cb36c53cb4de244eaa3a931206dd161 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Tue, 15 Mar 2022 12:25:48 -0400 Subject: [PATCH] Agent: Make CachingAgentRepository fully thread-safe --- .../infection_monkey/exploit/caching_agent_repository.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/monkey/infection_monkey/exploit/caching_agent_repository.py b/monkey/infection_monkey/exploit/caching_agent_repository.py index 2e52990b9..8a55b3f63 100644 --- a/monkey/infection_monkey/exploit/caching_agent_repository.py +++ b/monkey/infection_monkey/exploit/caching_agent_repository.py @@ -1,4 +1,5 @@ import io +import threading from functools import lru_cache from typing import Mapping @@ -19,9 +20,15 @@ class CachingAgentRepository(IAgentRepository): def __init__(self, island_url: str, proxies: Mapping[str, str]): self._island_url = island_url self._proxies = proxies + self._lock = threading.Lock() def get_agent_binary(self, os: str, _: str = None) -> io.BytesIO: - return io.BytesIO(self._download_binary_from_island(os)) + # If multiple calls to get_agent_binary() are made simultaneously before the result of + # _download_binary_from_island() is cached, then multiple requests will be sent to the + # island. Add a mutex in front of the call to _download_agent_binary_from_island() so + # that only one request per OS will be sent to the island. + with self._lock: + return io.BytesIO(self._download_binary_from_island(os)) @lru_cache(maxsize=None) def _download_binary_from_island(self, os: str) -> bytes: