Add cgroupv2 initialisation

Ensure the cgroupv2 mount exists, subtree_control is not empty, and
swap is disabled.

Fix #126
Fix #102

2 files changed, 46 insertions, 6 deletions

diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 6253c48..2d84d3b 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -42,6 +42,9 @@ class NsJail:
         self.config = self._read_config()
         self.cgroup_version = utils.cgroup.get_version(self.config)
 
+        if self.cgroup_version == 2:
+            utils.cgroup.init_v2(self.config)
+
         log.info(f"Assuming cgroup version {self.cgroup_version}.")
 
     @staticmethod
@@ -146,12 +149,18 @@ 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 = utils.cgroup.create_dynamic(self.config)
+        cgroup = None
+        if self.cgroup_version == 1:
+            cgroup = utils.cgroup.create_dynamic(self.config)
+            nsjail_args = (
+                "--cgroup_mem_parent", cgroup,
+                "--cgroup_pids_parent", cgroup,
+                *nsjail_args,
+            )
+        else:
+            nsjail_args = ("--use_cgroupv2", *nsjail_args)
 
         with NamedTemporaryFile() as nsj_log:
-            if self.cgroup_version == 2:
-                nsjail_args = (["--use_cgroupv2"]).extend(nsjail_args)
-
             args = (
                 self.nsjail_binary,
                 "--config", NSJAIL_CFG,
@@ -201,7 +210,8 @@ class NsJail:
         log.info(f"nsjail return code: {returncode}")
 
         # Remove the dynamically created cgroups once we're done
-        Path(self.config.cgroup_mem_mount, cgroup).rmdir()
-        Path(self.config.cgroup_pids_mount, cgroup).rmdir()
+        if self.cgroup_version == 1:
+            Path(self.config.cgroup_mem_mount, cgroup).rmdir()
+            Path(self.config.cgroup_pids_mount, cgroup).rmdir()
 
         return CompletedProcess(args, returncode, output, None)
diff --git a/snekbox/utils/cgroup.py b/snekbox/utils/cgroup.py
index b58a7e5..4593412 100644
--- a/snekbox/utils/cgroup.py
+++ b/snekbox/utils/cgroup.py
@@ -86,3 +86,33 @@ def get_version(config: NsJailConfig) -> int:
             "Falling back to the use_cgroupv2 NsJail setting."
         )
         return config_version
+
+
+def init_v2(config: NsJailConfig) -> None:
+    """Ensure cgroupv2 children have controllers enabled and memory swapping is disabled."""
+    cgroup_mount = Path(config.cgroupv2_mount)
+
+    # Swap has to be disabled since NsJail doesn't do it.
+    (cgroup_mount / "memory.swap.max").write_text("0")
+
+    # If the root's subtree_control already has some controllers enabled,
+    # no further action is necessary.
+    if (cgroup_mount / "cgroup.subtree_control").read_text().strip():
+        return
+
+    # Move all processes from the cgroupv2 mount to a child cgroup.
+    # This is necessary to be able to write to subtree_control in the parent later.
+    # Otherwise, a write operation would yield a "device or resource busy" error.
+    init_cgroup = cgroup_mount / "init"
+    init_cgroup.mkdir(parents=True, exist_ok=True)
+
+    procs = (cgroup_mount / "cgroup.procs").read_text().split()
+    for proc in procs:
+        (init_cgroup / "cgroup.procs").write_text(proc)
+
+    # Enable all available controllers for child cgroups.
+    # This also retroactively enables controllers for children that already exist,
+    # including the "init" child created just before.
+    controllers = (cgroup_mount / "cgroup.controllers").read_text().split()
+    for controller in controllers:
+        (cgroup_mount / "cgroup.subtree_control").write_text(f"+{controller}")