Support selecting the cgroup (#25)

This change allows the runner to specify the cgroup path. This is a
forwards-compatibility feature that allows the new version of omegajail
(which is compatible with cgroup v1 and v2) to coexist with this
version.

Author: lhchavez

args.cpp

@@ -107,6 +107,10 @@ ParseArgs(int argc, char* argv[], const std::string_view cwd) throw() {
      cxxopts::value<uint64_t>(), "bytes")
     ("m,memory-limit", "sets the memory limit",
      cxxopts::value<int64_t>()->default_value("-1"), "bytes")
+    ("cgroup-path",
+     "the path under the cgroup hierarchy where the jailed processes should "
+     "be placed",
+     cxxopts::value<std::string>()->default_value("/omegajail"), "cgroup path")
     ("cgroup-memory-limit", "sets the memory limit with cgroups",
      cxxopts::value<ssize_t>(), "bytes")
     ("disable-sandboxing",
@@ -146,6 +150,13 @@ bool Args::Parse(int argc, char* argv[], struct minijail* j) throw() {
     return false;
   }
 
+  cgroup_path = options["cgroup-path"].as<std::string>();
+  if (!cgroup_path.empty() && cgroup_path.front() == '/') {
+    // cgroup paths are normally expressed with a leading /, but that's not
+    // conducive to path joining.
+    cgroup_path = cgroup_path.substr(1);
+  }
+
   if (options.count("disable-sandboxing"))
     disable_sandboxing = options["disable-sandboxing"].as<bool>();
 

args.h

@@ -28,6 +28,7 @@ struct Args {
   std::string stderr_redirect;
   std::string meta;
   std::string script_basename;
+  std::string cgroup_path;
   ssize_t memory_limit_in_bytes = -1;
   size_t vm_memory_size_in_bytes = 0;
   uint64_t wall_time_limit_msec = kMaxWallTimeLimitMsec;

main.cpp

@@ -56,6 +56,7 @@ struct InitPayload {
   ScopedMinijail jail;
   std::string comm;
   std::string cgroup_path;
+  std::string pid_cgroup_path;
   ssize_t memory_limit_in_bytes;
   size_t vm_memory_size_in_bytes;
   std::vector<ResourceLimit> rlimits;
@@ -354,8 +355,8 @@ int MetaInit(void* raw_payload) {
       return -errno;
     }
   } else if (payload->memory_limit_in_bytes >= 0) {
-    memory_cgroup =
-        std::make_unique<ScopedCgroup>("/sys/fs/cgroup/memory/omegajail");
+    memory_cgroup = std::make_unique<ScopedCgroup>(
+        PathJoin("/sys/fs/cgroup/memory", payload->cgroup_path));
     if (!*memory_cgroup) {
       {
         ScopedErrnoPreserver preserve_errno;
@@ -377,8 +378,8 @@ int MetaInit(void* raw_payload) {
   }
 
   std::unique_ptr<ScopedCgroup> pid_cgroup;
-  if (!payload->cgroup_path.empty()) {
-    pid_cgroup = std::make_unique<ScopedCgroup>(payload->cgroup_path);
+  if (!payload->pid_cgroup_path.empty()) {
+    pid_cgroup = std::make_unique<ScopedCgroup>(payload->pid_cgroup_path);
     if (!*pid_cgroup) {
       {
         ScopedErrnoPreserver preserve_errno;
@@ -777,10 +778,12 @@ int main(int argc, char* argv[]) {
                                 "/mnt/stdio/stderr", true);
     }
 
+    std::string memory_cgroup_path =
+        PathJoin("/sys/fs/cgroup/memory", args.cgroup_path);
     if (args.memory_limit_in_bytes >= 0 &&
-        minijail_mount(j.get(), "/sys/fs/cgroup/memory/omegajail",
-                       "/sys/fs/cgroup/memory/omegajail", "", MS_BIND)) {
-      LOG(ERROR) << "Failed to mount /sys/fs/cgroup/memory";
+        minijail_mount(j.get(), memory_cgroup_path.c_str(),
+                       memory_cgroup_path.c_str(), "", MS_BIND)) {
+      LOG(ERROR) << "Failed to mount " << memory_cgroup_path;
       return 1;
     }
   } else {
@@ -811,24 +814,25 @@ int main(int argc, char* argv[]) {
                       MINIJAIL_HOOK_EVENT_PRE_DROP_CAPS);
   }
 
-  std::string cgroup_path;
+  std::string pid_cgroup_path;
   if (!args.script_basename.empty()) {
-    cgroup_path = StringPrintf("/sys/fs/cgroup/pids/omegajail/%s",
-                               args.script_basename.c_str());
-    if (access(cgroup_path.c_str(), W_OK) != 0) {
-      cgroup_path.clear();
+    pid_cgroup_path =
+        PathJoin("/sys/fs/cgroup/pids", args.cgroup_path, args.script_basename);
+    if (access(pid_cgroup_path.c_str(), W_OK) != 0) {
+      pid_cgroup_path.clear();
     } else if (!args.disable_sandboxing &&
-               minijail_mount(j.get(), "/sys/fs/cgroup/pids/omegajail",
-                              "/sys/fs/cgroup/pids/omegajail", "", MS_BIND)) {
-      LOG(ERROR) << "Failed to mount /sys/fs/cgroup/pids";
+               minijail_mount(j.get(), pid_cgroup_path.c_str(),
+                              pid_cgroup_path.c_str(), "", MS_BIND)) {
+      LOG(ERROR) << "Failed to mount " << pid_cgroup_path.c_str();
       return 1;
     }
   }
 
   InitPayload payload{
       .disable_sandboxing = args.disable_sandboxing,
       .comm = args.comm,
-      .cgroup_path = cgroup_path,
+      .cgroup_path = args.cgroup_path,
+      .pid_cgroup_path = pid_cgroup_path,
       .memory_limit_in_bytes = args.memory_limit_in_bytes,
       .vm_memory_size_in_bytes = args.vm_memory_size_in_bytes,
       .rlimits = args.rlimits,

smoketest/test

@@ -73,7 +73,12 @@ def _check_call(args: List[str]) -> bool:
         return False
 
 
-def _omegajail_compile(root: str, lang: str, strace: bool) -> bool:
+def _omegajail_compile(
+    root: str,
+    lang: str,
+    strace: bool,
+    cgroup_path: str,
+) -> bool:
     lang_dir = os.path.join(_PWD, 'run', lang)
     if os.path.isdir(lang_dir):
         shutil.rmtree(lang_dir, True)
@@ -114,15 +119,23 @@ def _omegajail_compile(root: str, lang: str, strace: bool) -> bool:
         source,
         '--compile-target',
         target,
+        '--cgroup-path',
+        cgroup_path,
     ]
     if lang == 'cs':
         os.symlink('/usr/share/dotnet/Main.runtimeconfig.json',
                    os.path.join(lang_dir, 'Main.runtimeconfig.json'))
     return _check_call(args)
 
 
-def _omegajail_run(root: str, lang: str, strace: bool, input_file: str,
-                   output_file: str) -> bool:
+def _omegajail_run(
+    root: str,
+    lang: str,
+    strace: bool,
+    input_path: str,
+    output_path: str,
+    cgroup_path: str,
+) -> bool:
     lang_dir = os.path.join(_PWD, 'run', lang)
     if strace:
         args = [
@@ -137,7 +150,7 @@ def _omegajail_run(root: str, lang: str, strace: bool, input_file: str,
         '--homedir',
         lang_dir,
         '-0',
-        os.path.join(_PWD, input_file),
+        os.path.join(_PWD, input_path),
         '-1',
         os.path.join(lang_dir, 'run.out'),
         '-2',
@@ -156,14 +169,16 @@ def _omegajail_run(root: str, lang: str, strace: bool, input_file: str,
         root,
         '--run',
         lang,
+        '--cgroup-path',
+        cgroup_path,
         '--run-target',
         'Main',
     ]
     if not _check_call(args):
         return False
     with open(os.path.join(lang_dir, 'run.out'), 'r') as run_out:
         got = run_out.read().strip()
-    with open(os.path.join(_PWD, output_file), 'r') as output:
+    with open(os.path.join(_PWD, output_path), 'r') as output:
         expected = output.read().strip()
     if got != expected:
         logging.error('Wrong answer when running \'%s\'. got %r, expected %r',
@@ -177,8 +192,28 @@ def _main() -> None:
     parser.add_argument('--strace', action='store_true')
     parser.add_argument('--verbose', action='store_true')
     parser.add_argument('--root', default='/var/lib/omegajail', type=str)
+    parser.add_argument('--cgroup-path',
+                        default='/omegajail',
+                        type=str)
     args = parser.parse_args()
 
+    # Set this process up for cgroups v2, since it uses slightly different
+    # rules.
+    if os.path.isdir(os.path.join('/sys/fs/cgroup', args.cgroup_path[1:])):
+        smoketest_cgroup = os.path.join('/sys/fs/cgroup', args.cgroup_path[1:],
+                                        'smoketest')
+        if not os.path.isdir(smoketest_cgroup):
+            with open(
+                    os.path.join('/sys/fs/cgroup', args.cgroup_path[1:],
+                                 'cgroup.subtree_control'), 'w') as f:
+                f.write('+memory')
+            os.makedirs(smoketest_cgroup, exist_ok=True)
+            with open(os.path.join(smoketest_cgroup, 'cgroup.subtree_control'),
+                      'w') as f:
+                f.write('+memory')
+        with open(os.path.join(smoketest_cgroup, 'cgroup.procs'), 'w') as f:
+            f.write(str(os.getpid()))
+
     args.root = os.path.abspath(args.root)
 
     languages = _LANGUAGES
@@ -192,15 +227,27 @@ def _main() -> None:
 
     for lang in languages:
         print('%-20s' % lang, end='')
-        if not _omegajail_compile(args.root, lang, args.strace):
+        if not _omegajail_compile(
+                root=args.root,
+                lang=lang,
+                strace=args.strace,
+                cgroup_path=args.cgroup_path,
+        ):
             print('ERROR (COMPILE)')
             passed = False
             continue
         if lang in _KAREL_LANGUAGES:
             input_path, output_path = 'input-karel', 'output-karel'
         else:
             input_path, output_path = 'input', 'output'
-        if _omegajail_run(args.root, lang, args.strace, input_path, output_path):
+        if _omegajail_run(
+                root=args.root,
+                lang=lang,
+                strace=args.strace,
+                input_path=input_path,
+                output_path=output_path,
+                cgroup_path=args.cgroup_path,
+        ):
             print('OK')
         else:
             print('ERROR')