Move cgroup functions to a new utility module

3 files changed, 69 insertions, 61 deletions

diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index a1092cf..bac8af3 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -4,7 +4,6 @@ import re
 import subprocess
 import sys
 import textwrap
-import uuid
 from pathlib import Path
 from subprocess import CompletedProcess
 from tempfile import NamedTemporaryFile
@@ -12,7 +11,7 @@ from typing import Iterable
 
 from google.protobuf import text_format
 
-from snekbox import DEBUG
+from snekbox import DEBUG, utils
 from snekbox.config_pb2 import NsJailConfig
 
 log = logging.getLogger(__name__)
@@ -26,8 +25,6 @@ LOG_BLACKLIST = ("Process will be ",)
 NSJAIL_PATH = os.getenv("NSJAIL_PATH", "/usr/sbin/nsjail")
 NSJAIL_CFG = os.getenv("NSJAIL_CFG", "./config/snekbox.cfg")
 
-# If this file is present, cgroupv2 should be enabled
-CGROUPV2_PROBE_PATH = Path("/sys/fs/cgroup/cgroup.controllers")
 
 # Limit of stdout bytes we consume before terminating nsjail
 OUTPUT_MAX = 1_000_000  # 1 MB
@@ -45,18 +42,7 @@ class NsJail:
         self.nsjail_binary = nsjail_binary
         self.config = self._read_config()
 
-        log.info(f"Cgroups version: {self._probe_cgroup_version()}")
-
-    @staticmethod
-    def _probe_cgroup_version() -> int:
-        """Poll the filesystem and return the guessed cgroup version."""
-        # Right now we check whenever the controller path exists
-        version = 2 if CGROUPV2_PROBE_PATH.exists() else 1
-
-        if DEBUG:
-            log.info(f"Guessed cgroups version: {version}")
-
-        return version
+        log.info(f"Cgroups version: {utils.cgroup.probe_version()}")
 
     @staticmethod
     def _read_config() -> NsJailConfig:
@@ -81,49 +67,6 @@ class NsJail:
 
         return config
 
-    def _create_dynamic_cgroups(self) -> str:
-        """
-        Create a PID and memory cgroup for NsJail to use as the parent cgroup.
-
-        Returns the name of the cgroup, located in the cgroup root.
-
-        NsJail doesn't do this automatically because it requires privileges NsJail usually doesn't
-        have.
-
-        Disables memory swapping.
-        """
-        # Pick a name for the cgroup
-        cgroup = "snekbox-" + str(uuid.uuid4())
-
-        pids = Path(self.config.cgroup_pids_mount, cgroup)
-        mem = Path(self.config.cgroup_mem_mount, cgroup)
-        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 before the swap limit.
-        #
-        # Since child cgroups are dynamically created, the swap limit has to be set on the parent
-        # instead so that children inherit it. Given the swap's dependency on the memory limit,
-        # the memory limit must also be set on the parent. NsJail only sets the memory limit for
-        # child cgroups, not the parent.
-        (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.
-            # Therefore, setting it equal to the memory limit effectively disables swapping.
-            (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. "
-                "This is probably because CONFIG_MEMCG_SWAP or CONFIG_MEMCG_SWAP_ENABLED is unset. "
-                "Please ensure swap memory is disabled on the system."
-            )
-
-        return cgroup
-
     @staticmethod
     def _parse_log(log_lines: Iterable[str]) -> None:
         """Parse and log NsJail's log messages."""
@@ -203,10 +146,10 @@ class NsJail:
         `py_args` are arguments to pass to the Python subprocess before the code,
         which is the last argument. By default, it's "-c", which executes the code given.
         """
-        cgroup = self._create_dynamic_cgroups()
+        cgroup = utils.cgroup.create_dynamic(self.config)
 
         with NamedTemporaryFile() as nsj_log:
-            if self._probe_cgroup_version() == 2:
+            if utils.cgroup.probe_version() == 2:
                 nsjail_args = (["--use_cgroupv2"]).extend(nsjail_args)
 
             args = (
diff --git a/snekbox/utils/__init__.py b/snekbox/utils/__init__.py
new file mode 100644
index 0000000..22ece40
--- /dev/null
+++ b/snekbox/utils/__init__.py
@@ -0,0 +1,3 @@
+from . import cgroup
+
+__all__ = ("cgroup",)
diff --git a/snekbox/utils/cgroup.py b/snekbox/utils/cgroup.py
new file mode 100644
index 0000000..018fbaa
--- /dev/null
+++ b/snekbox/utils/cgroup.py
@@ -0,0 +1,62 @@
+import logging
+import uuid
+from pathlib import Path
+
+from snekbox.config_pb2 import NsJailConfig
+
+log = logging.getLogger(__name__)
+
+# If this file is present, cgroupv2 should be enabled
+CGROUPV2_PROBE_PATH = Path("/sys/fs/cgroup/cgroup.controllers")
+
+
+def create_dynamic(config: NsJailConfig) -> str:
+    """
+    Create a PID and memory cgroup for NsJail to use as the parent cgroup.
+
+    Returns the name of the cgroup, located in the cgroup root.
+
+    NsJail doesn't do this automatically because it requires privileges NsJail usually doesn't
+    have.
+
+    Disables memory swapping.
+    """
+    # Pick a name for the cgroup
+    cgroup = "snekbox-" + str(uuid.uuid4())
+
+    pids = Path(config.cgroup_pids_mount, cgroup)
+    mem = Path(config.cgroup_mem_mount, cgroup)
+    mem_max = str(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 before the swap limit.
+    #
+    # Since child cgroups are dynamically created, the swap limit has to be set on the parent
+    # instead so that children inherit it. Given the swap's dependency on the memory limit,
+    # the memory limit must also be set on the parent. NsJail only sets the memory limit for
+    # child cgroups, not the parent.
+    (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.
+        # Therefore, setting it equal to the memory limit effectively disables swapping.
+        (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. "
+            "This is probably because CONFIG_MEMCG_SWAP or CONFIG_MEMCG_SWAP_ENABLED is unset. "
+            "Please ensure swap memory is disabled on the system."
+        )
+
+    return cgroup
+
+
+def probe_version() -> int:
+    """Poll the filesystem and return the guessed cgroup version."""
+    # Right now we check whenever the controller path exists
+    version = 2 if CGROUPV2_PROBE_PATH.exists() else 1
+    log.debug(f"Guessed cgroups version: {version}")
+    return version