From 90bf8f22fc86f1d6298e7c91534c24237e27be6c Mon Sep 17 00:00:00 2001
From: Shambhavi Srivastava <shambhavisr@google.com>
Date: Thu, 15 Jun 2023 12:02:53 -0700
Subject: [PATCH] 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  | 15 +++++++++++++++
 pkg/sentry/kernel/kernel.go      |  4 ++++
 pkg/sentry/kernel/task_cgroup.go |  9 +++++++--
 pkg/sentry/kernel/task_start.go  | 10 ++++++++--
 4 files changed, 34 insertions(+), 4 deletions(-)

diff --git a/pkg/sentry/control/lifecycle.go b/pkg/sentry/control/lifecycle.go
index 56eb3d2c5..0ad94378b 100644
--- a/pkg/sentry/control/lifecycle.go
+++ b/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
diff --git a/pkg/sentry/kernel/kernel.go b/pkg/sentry/kernel/kernel.go
index e6ce47083..c41cd69f0 100644
--- a/pkg/sentry/kernel/kernel.go
+++ b/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()
diff --git a/pkg/sentry/kernel/task_cgroup.go b/pkg/sentry/kernel/task_cgroup.go
index 37084e7d9..ca577061d 100644
--- a/pkg/sentry/kernel/task_cgroup.go
+++ b/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
diff --git a/pkg/sentry/kernel/task_start.go b/pkg/sentry/kernel/task_start.go
index f8a525bcb..ae8f70ad9 100644
--- a/pkg/sentry/kernel/task_start.go
+++ b/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 {