Enabling container to be initialized to it's initial cgroups

Currently there's no implementation to enable the container to be
initialized to it's cgroups and hence the `EnterInitialCgroups` would
inherit all the root's cgroups when the container would start.

With this implementation we make sure that when the container starts
it enter's into the cgroups that is passed down it to from the sandbox.

PiperOrigin-RevId: 540650868

pkg/sentry/control/lifecycle.go

@@ -114,6 +114,9 @@ type StartContainerArgs struct {
 	// ContainerID is the container for the process being executed.
 	ContainerID string `json:"container_id"`
 
+	// InitialCgroups is the set of cgroup controllers container needs to be initialised to.
+	InitialCgroups map[kernel.CgroupControllerType]string `json:"initial_cgroups"`
+
 	// Limits is the limit set for the process being executed.
 	Limits map[string]limits.Limit `json:"limits"`
 
@@ -294,6 +297,18 @@ func (l *Lifecycle) StartContainer(args *StartContainerArgs, _ *uint32) error {
 		}
 	}()
 
+	initialCgroups := make(map[kernel.Cgroup]struct{}, len(args.InitialCgroups))
+	cgroupRegistry := l.Kernel.CgroupRegistry()
+	// path is relative to the container's cgroup controller of specified type.
+	for initialCgroupController, path := range args.InitialCgroups {
+		cg, err := cgroupRegistry.FindCgroup(ctx, initialCgroupController, path)
+		if err != nil {
+			return fmt.Errorf("FindCgroup can't locate cgroup controller: %v err: %v", initialCgroupController, err)
+		}
+		initialCgroups[cg] = struct{}{}
+	}
+	initArgs.InitialCgroups = initialCgroups
+
 	tg, _, err := l.Kernel.CreateProcess(initArgs)
 	if err != nil {
 		return err

pkg/sentry/kernel/kernel.go

@@ -716,6 +716,9 @@ type CreateProcessArgs struct {
 
 	// ContainerID is the container that the process belongs to.
 	ContainerID string
+
+	// InitialCgroups are the cgroups the container is initialized to.
+	InitialCgroups map[Cgroup]struct{}
 }
 
 // NewContext returns a context.Context that represents the task that will be
@@ -911,6 +914,7 @@ func (k *Kernel) CreateProcess(args CreateProcessArgs) (*ThreadGroup, ThreadID,
 		AbstractSocketNamespace: args.AbstractSocketNamespace,
 		MountNamespace:          mntns,
 		ContainerID:             args.ContainerID,
+		InitialCgroups:          args.InitialCgroups,
 		UserCounters:            k.GetUserCounters(args.Credentials.RealKUID),
 	}
 	config.NetworkNamespace.IncRef()

pkg/sentry/kernel/task_cgroup.go

@@ -25,13 +25,18 @@ import (
 )
 
 // EnterInitialCgroups moves t into an initial set of cgroups.
+// If initCgroups is not nil, the new task will be placed in the specified cgroups.
+// Otherwise, if parent is not nil, the new task will be placed in the parent's cgroups.
+// If neither is specified, the new task will be in the root cgroups.
 //
 // This is analogous to Linux's kernel/cgroup/cgroup.c:cgroup_css_set_fork().
 //
 // Precondition: t isn't in any cgroups yet, t.cgroups is empty.
-func (t *Task) EnterInitialCgroups(parent *Task) {
+func (t *Task) EnterInitialCgroups(parent *Task, initCgroups map[Cgroup]struct{}) {
 	var inherit map[Cgroup]struct{}
-	if parent != nil {
+	if initCgroups != nil {
+		inherit = initCgroups
+	} else if parent != nil {
 		parent.mu.Lock()
 		defer parent.mu.Unlock()
 		inherit = parent.cgroups

pkg/sentry/kernel/task_start.go

@@ -97,6 +97,9 @@ type TaskConfig struct {
 	// ContainerID is the container the new task belongs to.
 	ContainerID string
 
+	// InitialCgroups are the cgroups the container is initialised to.
+	InitialCgroups map[Cgroup]struct{}
+
 	// UserCounters is user resource counters.
 	UserCounters *userCounters
 }
@@ -237,8 +240,11 @@ func (ts *TaskSet) newTask(ctx context.Context, cfg *TaskConfig) (*Task, error)
 		t.parent.children[t] = struct{}{}
 	}
 
-	// srcT may be nil, in which case we default to root cgroups.
-	t.EnterInitialCgroups(srcT)
+	// If InitialCgroups is not nil, the new task will be placed in the
+	// specified cgroups. Otherwise, if srcT is not nil, the new task will
+	// be placed in the srcT's cgroups. If neither is specified, the new task
+	// will be in the root cgroups.
+	t.EnterInitialCgroups(srcT, cfg.InitialCgroups)
 	committed = true
 
 	if tg.leader == nil {