Adds a mount for /dev/shm

2 files changed, 31 insertions, 9 deletions

diff --git a/snekbox/memfs.py b/snekbox/memfs.py
index f6e2ca9..9401e43 100644
--- a/snekbox/memfs.py
+++ b/snekbox/memfs.py
@@ -34,7 +34,7 @@ def mount_tmpfs(name: str) -> Path:
             "-t",
             "tmpfs",
             "-o",
-            f"size={MemFSOptions.MEMFS_SIZE}",
+            f"size={MemFSOptions.MEMFS_SIZE_STR}",
             "tmpfs",
             str(tmp),
         ]
@@ -53,11 +53,14 @@ class MemFSOptions:
     """Options for memory file system."""
 
     # Size of the memory filesystem (per instance)
-    MEMFS_SIZE = "32M"
+    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 (32 MB)
     MAX_FILE_SIZE = 32 * 1024 * 1024
+    # Size of /dev/shm (16 MB)
+    SHM_SIZE = 16 * 1024 * 1024
 
 
 class MemoryTempDir:
@@ -79,6 +82,11 @@ class MemoryTempDir:
         """Path to home directory."""
         return Path(self.path, "home") if self.path else None
 
+    @property
+    def shm(self) -> Path | None:
+        """Path to /dev/shm."""
+        return Path(self.path, "dev", "shm") if self.path else None
+
     def __enter__(self) -> MemoryTempDir:
         # Generates a uuid tempdir
         with self.assignment_lock:
@@ -86,10 +94,17 @@ class MemoryTempDir:
                 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)  # Allow all access
+                    home.chmod(0o777)
+
+                    # Create a /dev/shm folder
+                    shm = self.path / "dev" / "shm"
+                    shm.mkdir(parents=True)
+                    shm.chmod(0o777)
+
                     self.assigned_names.add(name)
                     return self
             else:
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 816b490..47a5860 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 MemoryTempDir
+from snekbox.memfs import MemFSOptions, MemoryTempDir
 
 __all__ = ("NsJail",)
 
@@ -159,11 +159,18 @@ class NsJail:
         with NamedTemporaryFile() as nsj_log, MemoryTempDir() as temp_dir:
             # Add the temp dir to be mounted as cwd
             nsjail_args = (
-                "--bindmount",  # Mount temp dir in R/W mode
+                # 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}",
+                # Mount `home` in R/W mode
+                "--bindmount",
                 f"{temp_dir.home}:home",
-                "--cwd",  # Set cwd to temp dir
+                # Set cwd to temp dir
+                "--cwd",
                 "home",
-                "--env",  # Set $HOME to temp dir
+                # Set $HOME to temp dir
+                "--env",
                 "HOME=home",
                 *nsjail_args,
             )
@@ -185,7 +192,7 @@ class NsJail:
             if any(py_args):
                 args.extend(py_args)
 
-            # For cases with no `code` and `-m` in args, such as "-m timeit", don't add `code`
+            # For cases `-m` in args, such as "-m timeit", use `code` directly without file
             if "-m" not in args:
                 args.append(code if c_mode else "main.py")
                 log.info(f"args: {args}")
@@ -226,7 +233,7 @@ class NsJail:
                 attachments = sorted(temp_dir.attachments(), key=lambda a: a.name)
                 log.info(f"Found {len(attachments)} attachments.")
             except AttachmentError as err:
-                log.error(f"Failed to parse attachments: {err}")
+                log.warning(f"Failed to parse attachments: {err}")
                 return EvalResult(args, returncode, f"AttachmentError: {err}")
 
             log_lines = nsj_log.read().decode("utf-8").splitlines()