Read cgroup values from the config file

1 files changed, 7 insertions, 18 deletions

diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 0b6d1c6..f5df64c 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -22,13 +22,8 @@ LOG_PATTERN = re.compile(
 )
 LOG_BLACKLIST = ("Process will be ",)
 
-# Explicitly define constants for NsJail's default values.
-CGROUP_PIDS_PARENT = Path("/sys/fs/cgroup/pids/NSJAIL")
-CGROUP_MEMORY_PARENT = Path("/sys/fs/cgroup/memory/NSJAIL")
-
 NSJAIL_PATH = os.getenv("NSJAIL_PATH", "/usr/sbin/nsjail")
 NSJAIL_CFG = os.getenv("NSJAIL_CFG", "./config/snekbox.cfg")
-MEM_MAX = 52428800
 
 # Limit of stdout bytes we consume before terminating nsjail
 OUTPUT_MAX = 1_000_000  # 1 MB
@@ -58,11 +53,7 @@ class NsJail:
 
         return config
 
-    @staticmethod
-    def _create_parent_cgroups(
-        pids: Path = CGROUP_PIDS_PARENT,
-        mem: Path = CGROUP_MEMORY_PARENT
-    ) -> None:
+    def _create_parent_cgroups(self) -> None:
         """
         Create the PIDs and memory cgroups which NsJail will use as its parent cgroups.
 
@@ -71,16 +62,20 @@ class NsJail:
 
         Disables memory swapping.
         """
+        pids = Path(self.config.cgroup_pids_mount, self.config.cgroup_pids_parent)
+        mem = Path(self.config.cgroup_mem_mount, self.config.cgroup_mem_parent)
+        mem_max = str(self.config.cgroup_mem_max)
+
         pids.mkdir(parents=True, exist_ok=True)
         mem.mkdir(parents=True, exist_ok=True)
 
         # Swap limit cannot be set to a value lower than memory.limit_in_bytes.
         # Therefore, this must be set first.
-        (mem / "memory.limit_in_bytes").write_text(str(MEM_MAX), encoding="utf-8")
+        (mem / "memory.limit_in_bytes").write_text(mem_max, encoding="utf-8")
 
         try:
             # Swap limit is specified as the sum of the memory and swap limits.
-            (mem / "memory.memsw.limit_in_bytes").write_text(str(MEM_MAX), encoding="utf-8")
+            (mem / "memory.memsw.limit_in_bytes").write_text(mem_max, encoding="utf-8")
         except PermissionError:
             log.warning(
                 "Failed to set the memory swap limit for the cgroup. "
@@ -158,12 +153,6 @@ class NsJail:
                 self.nsjail_binary,
                 "--config", NSJAIL_CFG,
                 "--log", nsj_log.name,
-                f"--cgroup_mem_max={MEM_MAX}",
-                "--cgroup_mem_mount", str(CGROUP_MEMORY_PARENT.parent),
-                "--cgroup_mem_parent", CGROUP_MEMORY_PARENT.name,
-                "--cgroup_pids_max=1",
-                "--cgroup_pids_mount", str(CGROUP_PIDS_PARENT.parent),
-                "--cgroup_pids_parent", CGROUP_PIDS_PARENT.name,
                 "--",
                 self.python_binary, "-Squ", "-c", code
             )