Agent: Make CachingAgentRepository fully thread-safe

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: