diff options
| -rw-r--r-- | snekbox/filesystem.py | 12 | ||||
| -rw-r--r-- | snekbox/nsjail.py | 3 | ||||
| -rw-r--r-- | tests/test_nsjail.py | 5 |
3 files changed, 11 insertions, 9 deletions
diff --git a/snekbox/filesystem.py b/snekbox/filesystem.py index 0cdfbf5..9f0634a 100644 --- a/snekbox/filesystem.py +++ b/snekbox/filesystem.py @@ -20,13 +20,13 @@ libc.mount.argtypes = ( libc.umount2.argtypes = (ctypes.c_char_p, ctypes.c_int) -class Size(int): - """Size in bytes.""" +class Size(IntEnum): + """Size multipliers for bytes.""" - @classmethod - def from_mb(cls, mb: int) -> Size: - """Create a Size from mebibytes.""" - return cls(mb * 1024 * 1024) + KB = 1024 + MB = 1024**2 + GB = 1024**3 + TB = 1024**4 class UnmountFlags(IntEnum): diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py index 0dd49e5..887bc9a 100644 --- a/snekbox/nsjail.py +++ b/snekbox/nsjail.py @@ -10,6 +10,7 @@ from google.protobuf import text_format from snekbox import DEBUG, utils from snekbox.config_pb2 import NsJailConfig +from snekbox.filesystem import Size from snekbox.memfs import MemFS from snekbox.process import EvalResult from snekbox.snekio import FileAttachment @@ -50,7 +51,7 @@ class NsJail: config_path: str = "./config/snekbox.cfg", max_output_size: int = 1_000_000, read_chunk_size: int = 10_000, - memfs_instance_size: int = 48 * 1024 * 1024, + memfs_instance_size: int = 48 * Size.MB, files_limit: int | None = 100, files_timeout: float | None = 8, files_pattern: str = "**/*", diff --git a/tests/test_nsjail.py b/tests/test_nsjail.py index 839d3ec..927b8c0 100644 --- a/tests/test_nsjail.py +++ b/tests/test_nsjail.py @@ -9,6 +9,7 @@ from itertools import product from pathlib import Path from textwrap import dedent +from snekbox.filesystem import Size from snekbox.nsjail import NsJail from snekbox.snekio import FileAttachment @@ -18,7 +19,7 @@ class NsJailTests(unittest.TestCase): super().setUp() # Specify lower limits for unit tests to complete within time limits - self.nsjail = NsJail(memfs_instance_size=2 * 1024 * 1024) + self.nsjail = NsJail(memfs_instance_size=2 * Size.MB) self.logger = logging.getLogger("snekbox.nsjail") self.logger.setLevel(logging.WARNING) @@ -199,7 +200,7 @@ class NsJailTests(unittest.TestCase): """ ).strip() - nsjail = NsJail(memfs_instance_size=32 * 1024 * 1024, files_timeout=1) + nsjail = NsJail(memfs_instance_size=32 * Size.MB, files_timeout=1) result = nsjail.python3(["-c", code]) self.assertEqual(result.returncode, None) self.assertEqual( |