Account for NsJail's use_cgroupv2 setting when detecting cgroup version

2 files changed, 38 insertions, 15 deletions

diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index bac8af3..6253c48 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -25,7 +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")
 
-
 # Limit of stdout bytes we consume before terminating nsjail
 OUTPUT_MAX = 1_000_000  # 1 MB
 READ_CHUNK_SIZE = 10_000  # chars
@@ -41,8 +40,9 @@ class NsJail:
     def __init__(self, nsjail_binary: str = NSJAIL_PATH):
         self.nsjail_binary = nsjail_binary
         self.config = self._read_config()
+        self.cgroup_version = utils.cgroup.get_version(self.config)
 
-        log.info(f"Cgroups version: {utils.cgroup.probe_version()}")
+        log.info(f"Assuming cgroup version {self.cgroup_version}.")
 
     @staticmethod
     def _read_config() -> NsJailConfig:
@@ -149,16 +149,13 @@ class NsJail:
         cgroup = utils.cgroup.create_dynamic(self.config)
 
         with NamedTemporaryFile() as nsj_log:
-            if utils.cgroup.probe_version() == 2:
+            if self.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
diff --git a/snekbox/utils/cgroup.py b/snekbox/utils/cgroup.py
index 018fbaa..b58a7e5 100644
--- a/snekbox/utils/cgroup.py
+++ b/snekbox/utils/cgroup.py
@@ -6,9 +6,6 @@ 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:
     """
@@ -54,9 +51,38 @@ def create_dynamic(config: NsJailConfig) -> str:
     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
+def get_version(config: NsJailConfig) -> int:
+    """
+    Examine the filesystem and return the guessed cgroup 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 = (
+        config.cgroup_mem_mount,
+        config.cgroup_pids_mount,
+        config.cgroup_net_cls_mount,
+        config.cgroup_cpu_mount
+    )
+    v1_exists = any(Path(mount).exists() for mount in cgroup_mounts)
+
+    controllers_path = Path(config.cgroupv2_mount, "cgroup.controllers")
+    v2_exists = controllers_path.exists()
+
+    config_version = 2 if 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