From ce7d85ae441c12e230f13b60fdb53cc1b2dc2137 Mon Sep 17 00:00:00 2001
From: MarkKoz <KozlovMark@gmail.com>
Date: Mon, 20 Dec 2021 12:51:50 -0800
Subject: [PATCH] Add cgroupv2 initialisation

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

Fix #126
Fix #102
---
 snekbox/nsjail.py | 48 ++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 45 insertions(+), 3 deletions(-)

diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 45603592..61d7f54c 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -43,6 +43,9 @@ def __init__(self, nsjail_binary: str = NSJAIL_PATH):
         self.config = self._read_config()
         self.cgroup_version = self._get_cgroup_version()
 
+        if self.cgroup_version == 2:
+            self._init_cgroupv2()
+
         log.info(f"Assuming cgroup version {self.cgroup_version}.")
 
     def _get_cgroup_version(self) -> int:
@@ -147,6 +150,35 @@ def _create_dynamic_cgroups(self) -> str:
 
         return cgroup
 
+    def _init_cgroupv2(self) -> None:
+        """Ensure cgroupv2 children have controllers enabled and memory swapping is disabled."""
+        cgroup_mount = Path(self.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}")
+
     @staticmethod
     def _parse_log(log_lines: Iterable[str]) -> None:
         """Parse and log NsJail's log messages."""
@@ -226,7 +258,16 @@ def python3(
         `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 = None
+        if self.cgroup_version == 1:
+            cgroup = self._create_dynamic_cgroups()
+            nsjail_args = (
+                "--cgroup_mem_parent", cgroup,
+                "--cgroup_pids_parent", cgroup,
+                *nsjail_args,
+            )
+        else:
+            nsjail_args = ("--use_cgroupv2", *nsjail_args)
 
         with NamedTemporaryFile() as nsj_log:
             args = (
@@ -278,7 +319,8 @@ def python3(
         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)