Merge pull request #3369 from zenhack/cgroup2

Put grains in cgroups.

src/sandstorm/backend.c++

@@ -49,9 +49,14 @@ static void tryRecursivelyDelete(kj::StringPtr path) {
 }
 
 BackendImpl::BackendImpl(kj::LowLevelAsyncIoProvider& ioProvider, kj::Network& network,
-  SandstormCoreFactory::Client&& sandstormCoreFactory, kj::Maybe<uid_t> sandboxUid)
+  SandstormCoreFactory::Client&& sandstormCoreFactory,
+  kj::Maybe<Cgroup>&& cgroup,
+  kj::Maybe<uid_t> sandboxUid)
     : ioProvider(ioProvider), network(network), coreFactory(kj::mv(sandstormCoreFactory)),
-      sandboxUid(sandboxUid), tasks(*this) {}
+      sandboxUid(sandboxUid),
+      tasks(*this),
+      cgroup(kj::mv(cgroup))
+    {}
 
 void BackendImpl::taskFailed(kj::Exception&& exception) {
   KJ_LOG(ERROR, exception);
@@ -161,7 +166,8 @@ kj::Promise<Supervisor::Client> BackendImpl::bootGrain(
   auto addressPromise =
       network.parseAddress(kj::str("unix:/var/sandstorm/grains/", grainId, "/socket"));
 
-  // When both of those are done, connect to the address.
+  // When both of those are done, connect to the address, and move the
+  // supervisor into a cgroup.
   auto finalPromise = promise
       .then([KJ_MVCAP(addressPromise)](size_t n) mutable {
     return kj::mv(addressPromise);
@@ -169,6 +175,12 @@ kj::Promise<Supervisor::Client> BackendImpl::bootGrain(
     return address->connect();
   }).then([this,KJ_MVCAP(stdoutPipe),KJ_MVCAP(process),grainId = kj::heapString(grainId)]
           (kj::Own<kj::AsyncIoStream>&& connection) mutable {
+
+    KJ_IF_MAYBE(cg, cgroup) {
+      cg->getOrMakeChild(grainId)
+          .addPid(process.getPid());
+    }
+
     // Connected. Create the RunningGrain and fulfill promises.
     auto ignorePromise = ignoreAll(*stdoutPipe);
     tasks.add(ignorePromise.attach(kj::mv(stdoutPipe)));
@@ -229,6 +241,9 @@ BackendImpl::RunningGrain::RunningGrain(
 
 BackendImpl::RunningGrain::~RunningGrain() noexcept(false) {
   backend.supervisors.erase(grainId);
+  KJ_IF_MAYBE(cg, backend.cgroup) {
+    cg->removeChild(grainId);
+  }
 }
 
 kj::Promise<void> BackendImpl::ping(PingContext context) {

src/sandstorm/backend.h

@@ -23,6 +23,7 @@
 #include <capnp/rpc-twoparty.h>
 #include <kj/one-of.h>
 #include <kj/vector.h>
+#include <sandstorm/cgroup2.h>
 
 namespace kj {
   class InputStream;
@@ -34,6 +35,7 @@ class BackendImpl final: public Backend::Server, private kj::TaskSet::ErrorHandl
 public:
   BackendImpl(kj::LowLevelAsyncIoProvider& ioProvider, kj::Network& network,
               SandstormCoreFactory::Client&& sandstormCoreFactory,
+	      kj::Maybe<Cgroup>&& cgroup,
               kj::Maybe<uid_t> sandboxUid);
 
 protected:
@@ -59,6 +61,7 @@ private:
   SandstormCoreFactory::Client coreFactory;
   kj::Maybe<uid_t> sandboxUid;   // if not using user namespaces
   kj::TaskSet tasks;
+  kj::Maybe<Cgroup> cgroup;
 
   class RunningGrain {
   public:

src/sandstorm/cgroup2.c++

@@ -0,0 +1,36 @@
+#include <sandstorm/cgroup2.h>
+#include <sandstorm/util.h>
+
+#include <kj/debug.h>
+
+namespace sandstorm {
+Cgroup::Cgroup(kj::StringPtr path)
+  : dirfd(raiiOpen(path, O_DIRECTORY|O_CLOEXEC))
+{}
+
+Cgroup::Cgroup(kj::AutoCloseFd&& dirfd)
+  : dirfd(kj::mv(dirfd))
+{}
+
+Cgroup Cgroup::getOrMakeChild(kj::StringPtr path) {
+  KJ_SYSCALL_HANDLE_ERRORS(mkdirat(dirfd.get(), path.cStr(), 0700)) {
+    case EEXIST:
+      break;
+    default:
+      KJ_FAIL_SYSCALL("mkdirat()", error);
+  }
+
+  return Cgroup(raiiOpenAt(dirfd.get(), path, O_DIRECTORY|O_CLOEXEC));
+}
+
+void Cgroup::removeChild(kj::StringPtr path) {
+  KJ_SYSCALL(unlinkat(dirfd.get(), path.cStr(), AT_REMOVEDIR));
+}
+
+void Cgroup::addPid(pid_t pid) {
+  auto procsfd = raiiOpenAt(dirfd.get(), "cgroup.procs", O_WRONLY);
+  auto pidStr = kj::str(pid);
+  auto cStr = pidStr.cStr();
+  KJ_SYSCALL(write(procsfd.get(), cStr, strlen(cStr)));
+}
+};

src/sandstorm/cgroup2.h

@@ -0,0 +1,35 @@
+#ifndef SANDSTORM_CGROUP2_H_
+#define SANDSTORM_CGROUP2_H_
+
+#include <sys/types.h>  // For pid_t
+#include <kj/io.h>
+
+namespace sandstorm {
+class Cgroup {
+  // A Linux control group (version 2).
+
+  public:
+    Cgroup() = delete;
+    KJ_DISALLOW_COPY(Cgroup);
+    Cgroup(Cgroup&&) noexcept = default;
+
+    Cgroup(kj::StringPtr path);
+    // Open the cgroup corresponding to the directory `path`.
+
+    Cgroup getOrMakeChild(kj::StringPtr path);
+    // Open a cgroup that is a child of this one, creating it if it does not
+    // exist.
+
+    void removeChild(kj::StringPtr path);
+    // Delete a child of this cgroup. The child must not contain any
+    // processes.
+
+    void addPid(pid_t pid);
+    // Add the given process to the cgroup.
+  private:
+    Cgroup(kj::AutoCloseFd&& dirfd);
+    kj::AutoCloseFd dirfd;
+};
+};
+
+#endif

src/sandstorm/run-bundle.c++

@@ -799,7 +799,7 @@ public:
     const Config config = readConfig();
 
     // We'll run under the chroot.
-    enterChroot(false);
+    enterChroot(nullptr, false);
 
     // Don't run as root.
     dropPrivs(config.uids);
@@ -1432,7 +1432,7 @@ private:
     }
   }
 
-  void enterChroot(bool inPidNamespace) {
+  void enterChroot(kj::Maybe<const UserIds&> chownCgroupsTo, bool inPidNamespace) {
     KJ_REQUIRE(changedDir);
 
     // Verify ownership is intact.
@@ -1522,6 +1522,53 @@ private:
     // in, copy its contents to /etc/resolv.conf.host-initial so we can use that if needed.
     backupResolvConf();
 
+    // Best effort attempt to set up cgroup2. Depending on system configuration
+    // this may not be possible; if not we continue without cgroup support.
+    KJ_SYSCALL_HANDLE_ERRORS(unshare(CLONE_NEWCGROUP)) {
+      case EINVAL:
+        // Might happen on very old kernels before CLONE_NEWCGROUP was defined.
+        break;
+      default:
+        KJ_FAIL_SYSCALL("unshare(CLONE_NEWCGROUP)", error);
+    } else {
+      KJ_SYSCALL(mkdir("run/cgroup2", 0700));
+      KJ_SYSCALL_HANDLE_ERRORS(mount("none", "run/cgroup2", "cgroup2", MS_NOSUID | MS_NOEXEC, nullptr)) {
+        default:
+          // This could fail if the kernel is new enough to support cgroup namespaces, but
+          // wasn't actually built with cgroup support (CONFIG_CGORUPS) enabled; in this case
+          // unshare() silently ignores the relevant flag, so we'd get an error here.
+          //
+          // Get rid of the mount point, so attempts to open it later will fail:
+          KJ_SYSCALL(rmdir("run/cgroup2"));
+      } else {
+        KJ_IF_MAYBE(uids, chownCgroupsTo) {
+          // Give our unprivileged selves access to manage the cgroup.
+          // See the 'Delegation' section of 'Documentation/admin-guide/cgroup-v2.txt'
+          // in the Linux kernel source tree.
+          //
+          // We only do this if we were given uids to work with; we don't need to do
+          // this if the sandbox is already set up, so commands that expect
+          // an already running sandstorm will pass us nullptr to indicate that
+          // we don't need to do this.
+          //
+          // Additionally, we *can't* do this if we weren't started as root, so in
+          // that case we just skip it. Other parts of the system are built to
+          // handle the case where this isn't available gracefully.
+          if(runningAsRoot) {
+            KJ_SYSCALL(chown("run/cgroup2", uids->uid, uids->gid));
+            KJ_SYSCALL(chown("run/cgroup2/cgroup.procs", uids->uid, uids->gid));
+            KJ_SYSCALL(chown("run/cgroup2/cgroup.subtree_control", uids->uid, uids->gid));
+            // This one doesn't exist on earlier kernels (it appeared some time after 4.6).
+            // If it's absent don't worry about it; we don't use it currently anyway.
+            KJ_SYSCALL_HANDLE_ERRORS(access("run/cgroup2/cgroup.threads", F_OK)) {
+            } else {
+              KJ_SYSCALL(chown("run/cgroup2/cgroup.threads", uids->uid, uids->gid));
+            }
+          }
+        }
+      }
+    }
+
     // OK, change our root directory.
     KJ_SYSCALL(syscall(SYS_pivot_root, ".", "tmp"));
     KJ_SYSCALL(chdir("/"));
@@ -1991,7 +2038,7 @@ private:
 
     setProcessName("montr", "(server monitor)");
 
-    enterChroot(true);
+    enterChroot(config.uids, true);
 
     // Make sure socket directory exists (since the installer doesn't create it).
     if (mkdir("/var/sandstorm/socket", 0770) == 0) {
@@ -2336,8 +2383,19 @@ private:
 
       auto paf = kj::newPromiseAndFulfiller<Backend::Client>();
       TwoPartyServerWithClientBootstrap server(kj::mv(paf.promise));
+
+      // Collect all of the grains' cgroups into a single
+      // container, so they can't collectively starve sandstorm
+      // itself.
+      kj::Maybe<Cgroup> grainsCgroup;
+      kj::runCatchingExceptions([&grainsCgroup]() {
+        grainsCgroup = Cgroup("/run/cgroup2").getOrMakeChild("grains");
+      });
+
       paf.fulfiller->fulfill(kj::heap<BackendImpl>(*io.lowLevelProvider, network,
-        server.getBootstrap().castAs<SandstormCoreFactory>(), sandboxUid));
+        server.getBootstrap().castAs<SandstormCoreFactory>(),
+        kj::mv(grainsCgroup),
+        sandboxUid));
 
       auto gatewayServer = kj::heap<capnp::TwoPartyServer>(kj::refcounted<CapRedirector>([&]() {
         return server.getBootstrap().castAs<SandstormCoreFactory>()