Ignore swap limits if the swap controller is disabled

Fix NsJail failing to set the swap limit because it tries to write to
a file that doesn't exist.

Log a warning if swap is on, the swap controller is disabled, and the
NsJail config is attempting to limit swap memory.

3 files changed, 70 insertions, 2 deletions

diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py
index 79d33ba..ac36551 100644
--- a/snekbox/nsjail.py
+++ b/snekbox/nsjail.py
@@ -40,6 +40,7 @@ class NsJail:
         self.nsjail_binary = nsjail_binary
         self.config = self._read_config()
         self.cgroup_version = utils.cgroup.init(self.config)
+        self.ignore_swap_limits = utils.swap.should_ignore_limit(self.config, self.cgroup_version)
 
         log.info(f"Assuming cgroup version {self.cgroup_version}.")
 
@@ -148,6 +149,11 @@ class NsJail:
         if self.cgroup_version == 2:
             nsjail_args = ("--use_cgroupv2", *nsjail_args)
 
+        if self.ignore_swap_limits:
+            nsjail_args = (
+                "--cgroup_mem_memsw_max", "0", "--cgroup_mem_swap_max", "-1", *nsjail_args
+            )
+
         with NamedTemporaryFile() as nsj_log:
             args = (
                 self.nsjail_binary,
diff --git a/snekbox/utils/__init__.py b/snekbox/utils/__init__.py
index 22ece40..5a7b632 100644
--- a/snekbox/utils/__init__.py
+++ b/snekbox/utils/__init__.py
@@ -1,3 +1,3 @@
-from . import cgroup
+from . import cgroup, swap
 
-__all__ = ("cgroup",)
+__all__ = ("cgroup", "swap")
diff --git a/snekbox/utils/swap.py b/snekbox/utils/swap.py
new file mode 100644
index 0000000..3e0d0aa
--- /dev/null
+++ b/snekbox/utils/swap.py
@@ -0,0 +1,62 @@
+import logging
+import uuid
+from pathlib import Path
+
+from snekbox.config_pb2 import NsJailConfig
+
+log = logging.getLogger(__name__)
+
+
+def controller_exists(config: NsJailConfig, cgroup_version: int) -> bool:
+    """Return True if the swap memory cgroup controller is enabled."""
+    if cgroup_version == 1:
+        return Path(config.cgroup_mem_mount, "memory.memsw.max_usage_in_bytes").exists()
+    else:
+        # Create a child cgroup because memory.swap isn't available in the root cgroup.
+        child = Path(config.cgroupv2_mount, f"snekbox-temp-{uuid.uuid4()}")
+        child.mkdir()
+        swap_controller_exists = (child / "memory.swap.max").exists()
+        child.rmdir()
+
+        return swap_controller_exists
+
+
+def is_enabled() -> bool:
+    """Return True if the total size of swap memory is greater than 0."""
+    with open("/proc/meminfo", "rb") as f:
+        for line in f:
+            name, value, *_ = line.split()
+            if name == b"SwapTotal:":
+                return value != b"0"
+
+    log.warning("Couldn't determine if swap is on or off. Assuming it's off.")
+    return False
+
+
+def should_ignore_limit(config: NsJailConfig, cgroup_version: int) -> bool:
+    """
+    Return True if a swap limit should not be configured for NsJail.
+
+    If the swap controller doesn't exist, then NsJail would fail when trying to limit the swap.
+    It would attempt to write to a file that doesn't exist. In such case, the swap limit arguments
+    should be set to their default values, so NsJail will avoid setting a swap limit.
+
+    Log a warning if swap is enabled but the swap controller isn't enabled.
+    """
+    if config.cgroup_mem_max <= 0:
+        return False
+
+    if config.cgroup_mem_memsw_max <= 0 and config.cgroup_mem_swap_max < 0:
+        log.warning("Memory is being limited, but swap memory is unlimited.")
+        return False
+
+    controller_missing = not controller_exists(config, cgroup_version)
+    if is_enabled() and controller_missing:
+        log.warning(
+            "Swap memory is available, but the swap memory controller is not enabled. This is "
+            "probably due to the CONFIG_MEMCG_SWAP or CONFIG_MEMCG_SWAP_ENABLED kernel "
+            "parameters not being set. NsJail will not be able to limit memory effectively. "
+            "Please turn off swap memory on the system, or enable the swap memory controller."
+        )
+
+    return controller_missing