From 821ebfdca62383bf92d87f7372ce51d01a450cd7 Mon Sep 17 00:00:00 2001
From: Gabriel Corado <gabriel.oliveira@goteleport.com>
Date: Tue, 27 Feb 2024 13:43:31 -0300
Subject: [PATCH] [v15] Add option to change cgroup root path (enhanced
 recording) (#38394)

* feat(cgroup): add option to customize root path

* test(cgroup): fix loop issue

* chore(config): add comment on default value
---
 lib/bpf/bpf.go                |  1 +
 lib/cgroup/cgroup.go          |  8 +++++-
 lib/cgroup/cgroup_test.go     | 47 +++++++++++++++++++++++++++++++++++
 lib/config/fileconf.go        |  5 ++++
 lib/service/servicecfg/bpf.go |  3 +++
 5 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/lib/bpf/bpf.go b/lib/bpf/bpf.go
index 7eea97c5cba4e..b725790d60c55 100644
--- a/lib/bpf/bpf.go
+++ b/lib/bpf/bpf.go
@@ -137,6 +137,7 @@ func New(config *servicecfg.BPFConfig) (BPF, error) {
 	// Create a cgroup controller to add/remote cgroups.
 	cgroup, err := controlgroup.New(&controlgroup.Config{
 		MountPath: config.CgroupPath,
+		RootPath:  config.RootPath,
 	})
 	if err != nil {
 		return nil, trace.Wrap(err)
diff --git a/lib/cgroup/cgroup.go b/lib/cgroup/cgroup.go
index d95a77363412f..15197688fc15f 100644
--- a/lib/cgroup/cgroup.go
+++ b/lib/cgroup/cgroup.go
@@ -54,6 +54,9 @@ var log = logrus.WithFields(logrus.Fields{
 type Config struct {
 	// MountPath is where the cgroupv2 hierarchy is mounted.
 	MountPath string
+	// RootPath directory where the Teleport managed cgroups are going to be
+	// placed.
+	RootPath string
 }
 
 // CheckAndSetDefaults checks BPF configuration.
@@ -61,6 +64,9 @@ func (c *Config) CheckAndSetDefaults() error {
 	if c.MountPath == "" {
 		c.MountPath = defaults.CgroupPath
 	}
+	if c.RootPath == "" {
+		c.RootPath = teleportRoot
+	}
 	return nil
 }
 
@@ -82,7 +88,7 @@ func New(config *Config) (*Service, error) {
 
 	s := &Service{
 		Config:       config,
-		teleportRoot: filepath.Join(config.MountPath, teleportRoot, uuid.New().String()),
+		teleportRoot: filepath.Join(config.MountPath, config.RootPath, uuid.New().String()),
 	}
 
 	// Mount the cgroup2 filesystem.
diff --git a/lib/cgroup/cgroup_test.go b/lib/cgroup/cgroup_test.go
index 867b90ab062d8..0ca776c1692f5 100644
--- a/lib/cgroup/cgroup_test.go
+++ b/lib/cgroup/cgroup_test.go
@@ -81,6 +81,53 @@ func TestRootCreate(t *testing.T) {
 	require.NoDirExists(t, service.teleportRoot)
 }
 
+// TestRootCreateCustomRootPath given a service configured with a custom root
+// path, cgroups must be placed on the correct path.
+func TestRootCreateCustomRootPath(t *testing.T) {
+	// This test must be run as root. Only root can create cgroups.
+	if !isRoot() {
+		t.Skip("Tests for package cgroup can only be run as root.")
+	}
+
+	t.Parallel()
+
+	for _, rootPath := range []string{
+		"custom",
+		"/custom",
+		"nested/custom",
+		"/deep/nested/custom",
+	} {
+		rootPath := rootPath
+		t.Run(rootPath, func(t *testing.T) {
+			t.Parallel()
+			dir := t.TempDir()
+			service, err := New(&Config{
+				MountPath: dir,
+				RootPath:  rootPath,
+			})
+			require.NoError(t, err)
+			defer service.Close(false)
+
+			sessionID := uuid.New().String()
+			err = service.Create(sessionID)
+			require.NoError(t, err)
+
+			cgroupPath := path.Join(service.teleportRoot, sessionID)
+			require.DirExists(t, cgroupPath)
+			require.Contains(t, cgroupPath, rootPath)
+
+			err = service.Remove(sessionID)
+			require.NoError(t, err)
+			require.NoDirExists(t, cgroupPath)
+
+			// Teardown
+			err = service.Close(false)
+			require.NoError(t, err)
+			require.NoDirExists(t, service.teleportRoot)
+		})
+	}
+}
+
 // TestRootCleanup tests the ability for Teleport to remove and cleanup all
 // cgroups which is performed upon startup.
 func TestRootCleanup(t *testing.T) {
diff --git a/lib/config/fileconf.go b/lib/config/fileconf.go
index 47309e6dfc320..67b9e3a37dc7b 100644
--- a/lib/config/fileconf.go
+++ b/lib/config/fileconf.go
@@ -1570,6 +1570,10 @@ type BPF struct {
 
 	// CgroupPath controls where cgroupv2 hierarchy is mounted.
 	CgroupPath string `yaml:"cgroup_path"`
+
+	// RootPath root directory for the Teleport cgroups.
+	// Optional, defaults to /teleport
+	RootPath string `yaml:"root_path"`
 }
 
 // Parse will parse the enhanced session recording configuration.
@@ -1581,6 +1585,7 @@ func (b *BPF) Parse() *servicecfg.BPFConfig {
 		DiskBufferSize:    b.DiskBufferSize,
 		NetworkBufferSize: b.NetworkBufferSize,
 		CgroupPath:        b.CgroupPath,
+		RootPath:          b.RootPath,
 	}
 }
 
diff --git a/lib/service/servicecfg/bpf.go b/lib/service/servicecfg/bpf.go
index c7ec8713ba678..38dcbec30d231 100644
--- a/lib/service/servicecfg/bpf.go
+++ b/lib/service/servicecfg/bpf.go
@@ -36,6 +36,9 @@ type BPFConfig struct {
 
 	// CgroupPath is where the cgroupv2 hierarchy is mounted.
 	CgroupPath string
+
+	// RootPath root directory for the Teleport cgroups.
+	RootPath string
 }
 
 // CheckAndSetDefaults checks BPF configuration.