forked from p15670423/monkey
Agent: Make CachingAgentRepository fully thread-safe
This commit is contained in:
parent
153d65eca0
commit
48cded4c7c
|
@ -1,4 +1,5 @@
|
||||||
import io
|
import io
|
||||||
|
import threading
|
||||||
from functools import lru_cache
|
from functools import lru_cache
|
||||||
from typing import Mapping
|
from typing import Mapping
|
||||||
|
|
||||||
|
@ -19,9 +20,15 @@ class CachingAgentRepository(IAgentRepository):
|
||||||
def __init__(self, island_url: str, proxies: Mapping[str, str]):
|
def __init__(self, island_url: str, proxies: Mapping[str, str]):
|
||||||
self._island_url = island_url
|
self._island_url = island_url
|
||||||
self._proxies = proxies
|
self._proxies = proxies
|
||||||
|
self._lock = threading.Lock()
|
||||||
|
|
||||||
def get_agent_binary(self, os: str, _: str = None) -> io.BytesIO:
|
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)
|
@lru_cache(maxsize=None)
|
||||||
def _download_binary_from_island(self, os: str) -> bytes:
|
def _download_binary_from_island(self, os: str) -> bytes:
|
||||||
|
|
Loading…
Reference in New Issue