diff options
| -rw-r--r-- | snekbox/memfs.py | 20 | ||||
| -rw-r--r-- | snekbox/nsjail.py | 13 |
2 files changed, 27 insertions, 6 deletions
diff --git a/snekbox/memfs.py b/snekbox/memfs.py index d02c029..259f6f7 100644 --- a/snekbox/memfs.py +++ b/snekbox/memfs.py @@ -22,7 +22,13 @@ __all__ = ("MemFS",) class MemFS: """An in-memory temporary file system.""" - def __init__(self, instance_size: int, root_dir: str | Path = "/memfs") -> None: + def __init__( + self, + instance_size: int, + root_dir: str | Path = "/memfs", + home: str = "home", + output: str = "output", + ) -> None: """ Initialize an in-memory temporary file system. @@ -33,10 +39,14 @@ class MemFS: Args: instance_size: Size limit of each tmpfs instance in bytes. root_dir: Root directory to mount instances in. + home: Name of the home directory. + output: Name of the output directory within home. If empty, uses home. """ self.instance_size = instance_size self.root_dir = Path(root_dir) self.root_dir.mkdir(exist_ok=True, parents=True) + self._home_name = home + self._output_name = output for _ in range(10): name = str(uuid4()) @@ -82,12 +92,12 @@ class MemFS: @property def home(self) -> Path: """Path to home directory.""" - return self.path / "home" + return self.path / self._home_name @property def output(self) -> Path: """Path to output directory.""" - return self.home / "output" + return self.home / self._output_name def __enter__(self) -> MemFS: return self @@ -112,7 +122,7 @@ class MemFS: def files(self, limit: int, pattern: str = "**/*") -> Generator[FileAttachment, None, None]: """ - Yields FileAttachments for files in the MemFS. + Yields FileAttachments for files found in the output directory. Args: limit: The maximum number of files to parse. @@ -134,7 +144,7 @@ class MemFS: preload_dict: bool = False, ) -> list[FileAttachment]: """ - Return a sorted list of output files found in the MemFS. + Return a sorted list of file paths within the output directory. Args: limit: The maximum number of files to parse. diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py index 4933a0a..c5f2f80 100644 --- a/snekbox/nsjail.py +++ b/snekbox/nsjail.py @@ -52,6 +52,8 @@ class NsJail: max_output_size: int = 1_000_000, read_chunk_size: int = 10_000, memfs_instance_size: int = 48 * Size.MiB, + memfs_home: str = "home", + memfs_output: str = "output", files_limit: int | None = 100, files_timeout: float | None = 8, files_pattern: str = "**/*", @@ -65,6 +67,9 @@ class NsJail: max_output_size: Maximum size of the output in bytes. read_chunk_size: Size of the read buffer in bytes. memfs_instance_size: Size of the tmpfs instance in bytes. + memfs_home: Name of the mounted home directory. + memfs_output: Name of the output directory within home, + can be empty to use home as output. files_limit: Maximum number of output files to parse. files_timeout: Maximum time in seconds to wait for output files to be read. files_pattern: Pattern to match files to attach. @@ -74,6 +79,8 @@ class NsJail: self.max_output_size = max_output_size self.read_chunk_size = read_chunk_size self.memfs_instance_size = memfs_instance_size + self.memfs_home = memfs_home + self.memfs_output = memfs_output self.files_limit = files_limit self.files_timeout = files_timeout self.files_pattern = files_pattern @@ -191,7 +198,11 @@ class NsJail: *nsjail_args, ) - with NamedTemporaryFile() as nsj_log, MemFS(self.memfs_instance_size) as fs: + with NamedTemporaryFile() as nsj_log, MemFS( + instance_size=self.memfs_instance_size, + home=self.memfs_home, + output=self.memfs_output, + ) as fs: nsjail_args = ( # Mount `home` with Read/Write access "--bindmount", |