Refactored max size handling

3 files changed, 52 insertions, 65 deletions

diff --git a/snekbox/memfs.py b/snekbox/memfs.py
index 607aa86..035efed 100644
--- a/snekbox/memfs.py
+++ b/snekbox/memfs.py
@@ -16,13 +16,12 @@ from snekbox.snekio import FileAttachment
 
 log = logging.getLogger(__name__)
 
-
 NAMESPACE_DIR = Path("/memfs")
 NAMESPACE_DIR.mkdir(exist_ok=True)
 NAMESPACE_DIR.chmod(0o711)  # Execute only access for other users
 
 
-def mount_tmpfs(name: str) -> Path:
+def mount_tmpfs(name: str, size: int | str) -> Path:
     """Create and mount a tmpfs directory."""
     tmp = NAMESPACE_DIR / name
     tmp.mkdir()
@@ -34,7 +33,7 @@ def mount_tmpfs(name: str) -> Path:
             "-t",
             "tmpfs",
             "-o",
-            f"size={MemFSOptions.MEMFS_SIZE_STR}",
+            f"size={size}",
             "tmpfs",
             str(tmp),
         ]
@@ -49,29 +48,20 @@ def unmount_tmpfs(name: str) -> None:
     rmtree(tmp, ignore_errors=True)
 
 
-class MemFSOptions:
-    """Options for memory file system."""
-
-    # Size of the memory filesystem (per instance)
-    MEMFS_SIZE = 48 * 1024 * 1024
-    MEMFS_SIZE_STR = "48M"
-    # Maximum number of files attachments will be scanned for
-    MAX_FILES = 6
-    # Maximum size of a file attachment (8 MiB)
-    # 8 MB is also the discord bot upload limit
-    MAX_FILE_SIZE = 8 * 1024 * 1024
-    # Size of /dev/shm (16 MiB)
-    SHM_SIZE = 16 * 1024 * 1024
-
-
-class MemoryTempDir:
+class MemFS:
     """A temporary directory using tmpfs."""
 
     assignment_lock = BoundedSemaphore(1)
     assigned_names: set[str] = set()  # Pool of tempdir names in use
 
-    def __init__(self) -> None:
+    def __init__(self, instance_size: int) -> None:
+        """
+        Create a temporary directory using tmpfs.
+
+        size: Size limit of each tmpfs instance in bytes
+        """
         self.path: Path | None = None
+        self.instance_size = instance_size
 
     @property
     def name(self) -> str | None:
@@ -88,24 +78,20 @@ class MemoryTempDir:
         """Path to /dev/shm."""
         return Path(self.path, "dev", "shm") if self.path else None
 
-    def __enter__(self) -> MemoryTempDir:
+    def mkdir(self, path: str, chmod: int = 0o777) -> Path:
+        """Create a directory in the tempdir."""
+        f = Path(self.path, path)
+        f.mkdir(parents=True, exist_ok=True)
+        f.chmod(chmod)
+        return f
+
+    def __enter__(self) -> MemFS:
         # Generates a uuid tempdir
         with self.assignment_lock:
             for _ in range(10):
                 name = str(uuid4())
                 if name not in self.assigned_names:
-                    self.path = mount_tmpfs(name)
-
-                    # Create a home folder
-                    home = self.path / "home"
-                    home.mkdir()
-                    home.chmod(0o777)
-
-                    # Create a /dev/shm folder
-                    shm = self.path / "dev" / "shm"
-                    shm.mkdir(parents=True)
-                    shm.chmod(0o777)
-
+                    self.path = mount_tmpfs(name, self.instance_size)
                     self.assigned_names.add(name)
                     return self
             else:
@@ -122,24 +108,27 @@ class MemoryTempDir:
     @contextmanager
     def allow_write(self) -> None:
         """Temporarily allow writes to the root tempdir."""
+        backup = self.path.stat().st_mode
         self.path.chmod(0o777)
         yield
-        self.path.chmod(0o711)
+        self.path.chmod(backup)
 
-    def attachments(self) -> Generator[FileAttachment, None, None]:
+    def attachments(
+        self, max_count: int, max_size: int | None = None
+    ) -> Generator[FileAttachment, None, None]:
         """Return a list of attachments in the tempdir."""
-        # Look for any file starting with `output`
         count = 0
+        # Look for any file starting with `output`
         for file in self.home.glob("output*"):
-            if count >= MemFSOptions.MAX_FILES:
+            if count > max_count:
                 log.warning("Maximum number of attachments reached, skipping remaining files")
                 break
             if file.is_file():
                 count += 1
-                yield FileAttachment.from_path(file, MemFSOptions.MAX_FILE_SIZE)
+                yield FileAttachment.from_path(file, max_size)
 
     def cleanup(self) -> None:
-        """Remove files in temp dir, releases name."""
+        """Unmounts tmpfs, releases name."""
         if self.path is None:
             return
         # Remove the path folder
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 47a5860..8b3f6cd 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -11,7 +11,7 @@ from google.protobuf import text_format
 # noinspection PyProtectedMember
 from snekbox import DEBUG, utils
 from snekbox.config_pb2 import NsJailConfig
-from snekbox.memfs import MemFSOptions, MemoryTempDir
+from snekbox.memfs import MemFS
 
 __all__ = ("NsJail",)
 
@@ -39,11 +39,17 @@ 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,
+        max_attachments: int = 100,
+        max_attachment_size: int | None = None,
     ):
         self.nsjail_path = nsjail_path
         self.config_path = config_path
         self.max_output_size = max_output_size
         self.read_chunk_size = read_chunk_size
+        self.memfs_instance_size = memfs_instance_size
+        self.max_attachments = max_attachments
+        self.max_attachment_size = max_attachment_size
 
         self.config = self._read_config(config_path)
         self.cgroup_version = utils.cgroup.init(self.config)
@@ -133,7 +139,11 @@ class NsJail:
         return "".join(output)
 
     def python3(
-        self, code: str, *, nsjail_args: Iterable[str] = (), py_args: Iterable[str] = ("",)
+        self,
+        code: str,
+        *,
+        nsjail_args: Iterable[str] = (),
+        py_args: Iterable[str] = ("",),
     ) -> EvalResult:
         """
         Execute Python 3 code in an isolated environment and return the completed process.
@@ -156,16 +166,16 @@ class NsJail:
                 *nsjail_args,
             )
 
-        with NamedTemporaryFile() as nsj_log, MemoryTempDir() as temp_dir:
+        with NamedTemporaryFile() as nsj_log, MemFS(self.memfs_instance_size) as fs:
             # Add the temp dir to be mounted as cwd
             nsjail_args = (
                 # Mount a tmpfs at /dev/shm to support multiprocessing
                 "--mount",
                 # src:dst:fs_type:options
-                f"{temp_dir.shm}:/dev/shm:tmpfs:size={MemFSOptions.SHM_SIZE}",
+                f"{fs.shm}:/dev/shm:tmpfs:size={fs.instance_size}",
                 # Mount `home` in R/W mode
                 "--bindmount",
-                f"{temp_dir.home}:home",
+                f"{fs.home}:home",
                 # Set cwd to temp dir
                 "--cwd",
                 "home",
@@ -198,8 +208,8 @@ class NsJail:
                 log.info(f"args: {args}")
                 # Write the code to a file
                 if not c_mode:
-                    with temp_dir.allow_write():
-                        code_path = temp_dir.home / "main.py"
+                    with fs.allow_write():
+                        code_path = fs.home / "main.py"
                         code_path.write_text(code)
                     log.info(f"Created code file at [{code_path!r}].")
             else:
@@ -230,7 +240,10 @@ class NsJail:
             # Parse attachments
             try:
                 # Sort attachments by name lexically
-                attachments = sorted(temp_dir.attachments(), key=lambda a: a.name)
+                attachments = sorted(
+                    fs.attachments(self.max_attachments, self.max_attachment_size),
+                    key=lambda a: a.name,
+                )
                 log.info(f"Found {len(attachments)} attachments.")
             except AttachmentError as err:
                 log.warning(f"Failed to parse attachments: {err}")
diff --git a/snekbox/snekio.py b/snekbox/snekio.py
index 12645cb..eac6438 100644
--- a/snekbox/snekio.py
+++ b/snekbox/snekio.py
@@ -33,36 +33,21 @@ class FileAttachment:
     content: bytes
 
     @classmethod
-    def from_path(cls, file: Path, max_size: int) -> FileAttachment:
+    def from_path(cls, file: Path, max_size: int | None = None) -> FileAttachment:
         """Create an attachment from a path."""
         size = file.stat().st_size
-        if size > max_size:
+        if max_size is not None and size > max_size:
             raise AttachmentError(
                 f"File {file.name} too large: {sizeof_fmt(size)} "
                 f"exceeds the limit of {sizeof_fmt(max_size)}"
             )
-
-        with file.open("rb") as f:
-            content = f.read(max_size + 1)
-            size = len(content)
-            if len(content) > max_size:
-                raise AttachmentError(
-                    f"File {file.name} too large: {sizeof_fmt(len(content))} "
-                    f"exceeds the limit of {sizeof_fmt(max_size)}"
-                )
-        return cls(file.name, content)
+        return cls(file.name, file.read_bytes())
 
     @property
     def mime(self) -> str:
         """MIME type of the attachment."""
         return mimetypes.guess_type(self.name)[0]
 
-    def is_supported(self) -> bool:
-        """Return whether the attachment is supported."""
-        if self.mime.startswith("text/"):
-            return True
-        return self.mime in SUPPORTED_MIME_TYPES
-
     def to_dict(self) -> dict[str, str]:
         """Convert the attachment to a dict."""
         cmp = zlib.compress(self.content)