diff --git a/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py b/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py index 18a559f78..0e7ee8a50 100644 --- a/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py +++ b/monkey/monkey_island/cc/repository/file_repository_locking_decorator.py @@ -1,30 +1,27 @@ -from multiprocessing import Lock from typing import BinaryIO -from . import IFileRepository +from readerwriterlock import rwlock -# TODO: Use a read-write lock. See: -# - https://pypi.org/project/readerwriterlock/ -# - https://www.oreilly.com/library/view/python-cookbook/0596001673/ch06s04.html +from . import IFileRepository class FileRepositoryLockingDecorator(IFileRepository): def __init__(self, file_repository: IFileRepository): self._file_repository = file_repository - self._lock = Lock() + self._rwlock = rwlock.RWLockFair() def save_file(self, unsafe_file_name: str, file_contents: BinaryIO): - with self._lock: + with self._rwlock.gen_wlock(): return self._file_repository.save_file(unsafe_file_name, file_contents) def open_file(self, unsafe_file_name: str) -> BinaryIO: - with self._lock: + with self._rwlock.gen_rlock(): return self._file_repository.open_file(unsafe_file_name) def delete_file(self, unsafe_file_name: str): - with self._lock: + with self._rwlock.gen_wlock(): return self._file_repository.delete_file(unsafe_file_name) def delete_all_files(self): - with self._lock: + with self._rwlock.gen_wlock(): return self._file_repository.delete_all_files()