Account for NsJail's use_cgroupv2 setting when detecting cgroup version

snekbox/nsjail.py

@@ -26,9 +26,6 @@
 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
 READ_CHUNK_SIZE = 10_000  # chars
@@ -44,19 +41,45 @@ class NsJail:
     def __init__(self, nsjail_binary: str = NSJAIL_PATH):
         self.nsjail_binary = nsjail_binary
         self.config = self._read_config()
+        self.cgroup_version = self._get_cgroup_version()
 
-        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
+        log.info(f"Assuming cgroup version {self.cgroup_version}.")
 
-        if DEBUG:
-            log.info(f"Guessed cgroups version: {version}")
+    def _get_cgroup_version(self) -> int:
+        """
+        Examine the filesystem and return the guessed cgroup version.
 
-        return version
+        Fall back to use_cgroupv2 in the NsJail config if either both v1 and v2 seem to be enabled,
+        or neither seem to be enabled.
+        """
+        cgroup_mounts = (
+            self.config.cgroup_mem_mount,
+            self.config.cgroup_pids_mount,
+            self.config.cgroup_net_cls_mount,
+            self.config.cgroup_cpu_mount
+        )
+        v1_exists = any(Path(mount).exists() for mount in cgroup_mounts)
+
+        controllers_path = Path(self.config.cgroupv2_mount, "cgroup.controllers")
+        v2_exists = controllers_path.exists()
+
+        config_version = 2 if self.config.use_cgroupv2 else 1
+
+        if v1_exists and v2_exists:
+            # Probably hybrid mode. Use whatever is set in the config.
+            return config_version
+        elif v1_exists:
+            return 1
+        elif v2_exists:
+            return 2
+        else:
+            log.warning(
+                f"Neither the cgroupv1 controller mounts, nor {str(controllers_path)!r} exists. "
+                "Either cgroup_xxx_mount and cgroupv2_mount are misconfigured, or all "
+                "corresponding v1 controllers are disabled on the system. "
+                "Falling back to the use_cgroupv2 NsJail setting."
+            )
+            return config_version
 
     @staticmethod
     def _read_config() -> NsJailConfig:
@@ -206,16 +229,10 @@ def python3(
         cgroup = self._create_dynamic_cgroups()
 
         with NamedTemporaryFile() as nsj_log:
-            if self._probe_cgroup_version() == 2:
-                nsjail_args = (["--use_cgroupv2"]).extend(nsjail_args)
-
             args = (
                 self.nsjail_binary,
                 "--config", NSJAIL_CFG,
                 "--log", nsj_log.name,
-                # Set our dynamically created parent cgroups
-                "--cgroup_mem_parent", cgroup,
-                "--cgroup_pids_parent", cgroup,
                 *nsjail_args,
                 "--",
                 self.config.exec_bin.path, *self.config.exec_bin.arg, *py_args, code