Restore to original cgroup after sandbox and gofer processes are created

The original code assumed that it was safe to join and not restore cgroup,
but Container.Run will not exit after calling start, making cgroup cleanup
fail because there were still processes inside the cgroup.

PiperOrigin-RevId: 228529199
Change-Id: I12a48d9adab4bbb02f20d71ec99598c336cbfe51

2 files changed, 44 insertions, 19 deletions

diff --git a/runsc/container/BUILD b/runsc/container/BUILD
index 28ec81d3f..d9534cbcc 100644
--- a/runsc/container/BUILD
+++ b/runsc/container/BUILD
@@ -19,6 +19,7 @@ go_library(
         "//pkg/log",
         "//pkg/sentry/control",
         "//runsc/boot",
+        "//runsc/cgroup",
         "//runsc/sandbox",
         "//runsc/specutils",
         "@com_github_cenkalti_backoff//:go_default_library",
diff --git a/runsc/container/container.go b/runsc/container/container.go
index 07924d23a..dc9ef86fa 100644
--- a/runsc/container/container.go
+++ b/runsc/container/container.go
@@ -36,6 +36,7 @@ import (
 	"gvisor.googlesource.com/gvisor/pkg/log"
 	"gvisor.googlesource.com/gvisor/pkg/sentry/control"
 	"gvisor.googlesource.com/gvisor/runsc/boot"
+	"gvisor.googlesource.com/gvisor/runsc/cgroup"
 	"gvisor.googlesource.com/gvisor/runsc/sandbox"
 	"gvisor.googlesource.com/gvisor/runsc/specutils"
 )
@@ -286,18 +287,26 @@ func Create(id string, spec *specs.Spec, conf *boot.Config, bundleDir, consoleSo
 			return nil, fmt.Errorf("writing clean spec: %v", err)
 		}
 
-		ioFiles, err := c.createGoferProcess(spec, conf, bundleDir)
-		if err != nil {
-			return nil, err
+		// Create and join cgroup before processes are created to ensure they are
+		// part of the cgroup from the start (and all tneir children processes).
+		cg := cgroup.New(spec)
+		if cg != nil {
+			// If there is cgroup config, install it before creating sandbox process.
+			if err := cg.Install(spec.Linux.Resources); err != nil {
+				return nil, fmt.Errorf("configuring cgroup: %v", err)
+			}
 		}
+		if err := runInCgroup(cg, func() error {
+			ioFiles, err := c.createGoferProcess(spec, conf, bundleDir)
+			if err != nil {
+				return err
+			}
 
-		// Start a new sandbox for this container. Any errors after this point
-		// must destroy the container.
-		c.Sandbox, err = sandbox.Create(id, spec, conf, bundleDir, consoleSocket, userLog, ioFiles)
-		if err != nil {
-			return nil, err
-		}
-		if err := c.Sandbox.AddGoferToCgroup(c.GoferPid); err != nil {
+			// Start a new sandbox for this container. Any errors after this point
+			// must destroy the container.
+			c.Sandbox, err = sandbox.New(id, spec, conf, bundleDir, consoleSocket, userLog, ioFiles, cg)
+			return err
+		}); err != nil {
 			return nil, err
 		}
 	} else {
@@ -381,15 +390,16 @@ func (c *Container) Start(conf *boot.Config) error {
 			return fmt.Errorf("writing clean spec: %v", err)
 		}
 
-		// Create the gofer process.
-		ioFiles, err := c.createGoferProcess(c.Spec, conf, c.BundleDir)
-		if err != nil {
-			return err
-		}
-		if err := c.Sandbox.StartContainer(c.Spec, conf, c.ID, ioFiles); err != nil {
-			return err
-		}
-		if err := c.Sandbox.AddGoferToCgroup(c.GoferPid); err != nil {
+		// Join cgroup to strt gofer process to ensure it's part of the cgroup from
+		// the start (and all tneir children processes).
+		if err := runInCgroup(c.Sandbox.Cgroup, func() error {
+			// Create the gofer process.
+			ioFiles, err := c.createGoferProcess(c.Spec, conf, c.BundleDir)
+			if err != nil {
+				return err
+			}
+			return c.Sandbox.StartContainer(c.Spec, conf, c.ID, ioFiles)
+		}); err != nil {
 			return err
 		}
 	}
@@ -883,3 +893,17 @@ func lockContainerMetadata(containerRootDir string) (func() error, error) {
 	}
 	return l.Unlock, nil
 }
+
+// runInCgroup executes fn inside the specified cgroup. If cg is nil, execute
+// it in the current context.
+func runInCgroup(cg *cgroup.Cgroup, fn func() error) error {
+	if cg == nil {
+		return fn()
+	}
+	restore, err := cg.Join()
+	defer restore()
+	if err != nil {
+		return err
+	}
+	return fn()
+}