diff options
| -rw-r--r-- | snekbox/nsjail.py | 65 | ||||
| -rw-r--r-- | snekbox/utils/__init__.py | 3 | ||||
| -rw-r--r-- | snekbox/utils/cgroup.py | 62 |
3 files changed, 69 insertions, 61 deletions
diff --git a/snekbox/nsjail.py b/snekbox/nsjail.py index a1092cf..bac8af3 100644 --- a/snekbox/nsjail.py +++ b/snekbox/nsjail.py @@ -4,7 +4,6 @@ import re import subprocess import sys import textwrap -import uuid from pathlib import Path from subprocess import CompletedProcess from tempfile import NamedTemporaryFile @@ -12,7 +11,7 @@ from typing import Iterable from google.protobuf import text_format -from snekbox import DEBUG +from snekbox import DEBUG, utils from snekbox.config_pb2 import NsJailConfig log = logging.getLogger(__name__) @@ -26,8 +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") -# 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 @@ -45,18 +42,7 @@ class NsJail: self.nsjail_binary = nsjail_binary self.config = self._read_config() - 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 - - if DEBUG: - log.info(f"Guessed cgroups version: {version}") - - return version + log.info(f"Cgroups version: {utils.cgroup.probe_version()}") @staticmethod def _read_config() -> NsJailConfig: @@ -81,49 +67,6 @@ class NsJail: return config - def _create_dynamic_cgroups(self) -> str: - """ - Create a PID and memory cgroup for NsJail to use as the parent cgroup. - - Returns the name of the cgroup, located in the cgroup root. - - NsJail doesn't do this automatically because it requires privileges NsJail usually doesn't - have. - - Disables memory swapping. - """ - # Pick a name for the cgroup - cgroup = "snekbox-" + str(uuid.uuid4()) - - pids = Path(self.config.cgroup_pids_mount, cgroup) - mem = Path(self.config.cgroup_mem_mount, cgroup) - mem_max = str(self.config.cgroup_mem_max) - - pids.mkdir(parents=True, exist_ok=True) - mem.mkdir(parents=True, exist_ok=True) - - # Swap limit cannot be set to a value lower than memory.limit_in_bytes. - # Therefore, this must be set before the swap limit. - # - # Since child cgroups are dynamically created, the swap limit has to be set on the parent - # instead so that children inherit it. Given the swap's dependency on the memory limit, - # the memory limit must also be set on the parent. NsJail only sets the memory limit for - # child cgroups, not the parent. - (mem / "memory.limit_in_bytes").write_text(mem_max, encoding="utf-8") - - try: - # Swap limit is specified as the sum of the memory and swap limits. - # Therefore, setting it equal to the memory limit effectively disables swapping. - (mem / "memory.memsw.limit_in_bytes").write_text(mem_max, encoding="utf-8") - except PermissionError: - log.warning( - "Failed to set the memory swap limit for the cgroup. " - "This is probably because CONFIG_MEMCG_SWAP or CONFIG_MEMCG_SWAP_ENABLED is unset. " - "Please ensure swap memory is disabled on the system." - ) - - return cgroup - @staticmethod def _parse_log(log_lines: Iterable[str]) -> None: """Parse and log NsJail's log messages.""" @@ -203,10 +146,10 @@ 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 = self._create_dynamic_cgroups() + cgroup = utils.cgroup.create_dynamic(self.config) with NamedTemporaryFile() as nsj_log: - if self._probe_cgroup_version() == 2: + if utils.cgroup.probe_version() == 2: nsjail_args = (["--use_cgroupv2"]).extend(nsjail_args) args = ( diff --git a/snekbox/utils/__init__.py b/snekbox/utils/__init__.py new file mode 100644 index 0000000..22ece40 --- /dev/null +++ b/snekbox/utils/__init__.py @@ -0,0 +1,3 @@ +from . import cgroup + +__all__ = ("cgroup",) diff --git a/snekbox/utils/cgroup.py b/snekbox/utils/cgroup.py new file mode 100644 index 0000000..018fbaa --- /dev/null +++ b/snekbox/utils/cgroup.py @@ -0,0 +1,62 @@ +import logging +import uuid +from pathlib import Path + +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: + """ + Create a PID and memory cgroup for NsJail to use as the parent cgroup. + + Returns the name of the cgroup, located in the cgroup root. + + NsJail doesn't do this automatically because it requires privileges NsJail usually doesn't + have. + + Disables memory swapping. + """ + # Pick a name for the cgroup + cgroup = "snekbox-" + str(uuid.uuid4()) + + pids = Path(config.cgroup_pids_mount, cgroup) + mem = Path(config.cgroup_mem_mount, cgroup) + mem_max = str(config.cgroup_mem_max) + + pids.mkdir(parents=True, exist_ok=True) + mem.mkdir(parents=True, exist_ok=True) + + # Swap limit cannot be set to a value lower than memory.limit_in_bytes. + # Therefore, this must be set before the swap limit. + # + # Since child cgroups are dynamically created, the swap limit has to be set on the parent + # instead so that children inherit it. Given the swap's dependency on the memory limit, + # the memory limit must also be set on the parent. NsJail only sets the memory limit for + # child cgroups, not the parent. + (mem / "memory.limit_in_bytes").write_text(mem_max, encoding="utf-8") + + try: + # Swap limit is specified as the sum of the memory and swap limits. + # Therefore, setting it equal to the memory limit effectively disables swapping. + (mem / "memory.memsw.limit_in_bytes").write_text(mem_max, encoding="utf-8") + except PermissionError: + log.warning( + "Failed to set the memory swap limit for the cgroup. " + "This is probably because CONFIG_MEMCG_SWAP or CONFIG_MEMCG_SWAP_ENABLED is unset. " + "Please ensure swap memory is disabled on the system." + ) + + 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 |